Skip to content

Commit 2a542e4

Browse files
mag1c-hmag1c-h
committed
unittest for dump task
1 parent e220065 commit 2a542e4

File tree

4 files changed

+203
-2
lines changed

4 files changed

+203
-2
lines changed

‎ucm/store/cache/cc/dump_queue.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ class DumpQueue {
5252
TransBuffer* buffer_{nullptr};
5353
Store* backend_{nullptr};
5454
SpscRingQueue<TaskPair, 1024> waiting_;
55-
SpscRingQueue<ShardTask, 65536> dumping_;
55+
SpscRingQueue<ShardTask, 32768> dumping_;
5656
std::thread dispatcher_;
5757
std::thread dumper_;
5858

‎ucm/store/cache/cc/load_queue.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ class LoadQueue {
5555
TransBuffer* buffer_{nullptr};
5656
Store* backend_{nullptr};
5757
SpscRingQueue<TaskPair, 1024> waiting_;
58-
SpscRingQueue<ShardTask, 65536> running_;
58+
SpscRingQueue<ShardTask, 32768> running_;
5959
std::thread dispatcher_;
6060
std::thread transfer_;
6161

‎ucm/store/test/case/cache/cache_dump_queue_test.cc‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
/**
2+
* MIT License
3+
*
4+
* Copyright (c) 2025 Huawei Technologies Co., Ltd. All rights reserved.
5+
*
6+
* Permission is hereby granted, free of charge, to any person obtaining a copy
7+
* of this software and associated documentation files (the "Software"), to deal
8+
* in the Software without restriction, including without limitation the rights
9+
* to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
10+
* copies of the Software, and to permit persons to whom the Software is
11+
* furnished to do so, subject to the following conditions:
12+
*
13+
* The above copyright notice and this permission notice shall be included in all
14+
* copies or substantial portions of the Software.
15+
*
16+
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17+
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18+
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
19+
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20+
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
21+
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
22+
* SOFTWARE.
23+
* */
24+
#include <gtest/gtest.h>
25+
#include "cache/cc/dump_queue.h"
26+
#include "detail/data_generator.h"
27+
#include "detail/mock_store.h"
28+
#include "detail/random.h"
29+
#include "detail/types_helper.h"
30+
31+
class UCCacheDumpQueueTest : public testing::Test {
32+
public:
33+
UC::Test::Detail::Random rd;
34+
static UC::Detail::TaskHandle NextId()
35+
{
36+
static std::atomic<size_t> id{1};
37+
return id.fetch_add(1, std::memory_order_relaxed);
38+
}
39+
};
40+
41+
TEST_F(UCCacheDumpQueueTest, DumpOneBlock)
42+
{
43+
using namespace UC::CacheStore;
44+
UC::Test::Detail::MockStore backend;
45+
EXPECT_CALL(backend, Dump).WillOnce(testing::Invoke(NextId));
46+
UC::Latch finish{};
47+
finish.Up();
48+
EXPECT_CALL(backend, Wait).WillOnce(testing::Invoke([&finish]() {
49+
finish.Done();
50+
return UC::Status::OK();
51+
}));
52+
UC::HashSet<UC::Detail::TaskHandle> failureSet;
53+
Config config;
54+
config.backend = (uintptr_t)(void*)&backend;
55+
config.tensorSize = 32768;
56+
config.shardSize = config.tensorSize;
57+
config.blockSize = config.shardSize;
58+
config.deviceId = 0;
59+
config.bufferSize = config.blockSize * 2048;
60+
config.engineId = rd.RandomString(10);
61+
config.shareBufferEnable = true;
62+
TransBuffer buffer;
63+
DumpQueue dumpQ;
64+
auto s = buffer.Setup(config);
65+
ASSERT_EQ(s, UC::Status::OK());
66+
s = dumpQ.Setup(config, &failureSet, &buffer);
67+
ASSERT_EQ(s, UC::Status::OK());
68+
auto blockId = UC::Test::Detail::TypesHelper::MakeBlockId("a1b2c3d4e5f6789012345678901234ab");
69+
constexpr size_t shardIdx = 0;
70+
UC::Test::Detail::DataGenerator data{1, config.blockSize};
71+
data.Generate();
72+
UC::Detail::TaskDesc desc{
73+
{blockId, shardIdx, {data.Buffer()}}
74+
};
75+
auto task = std::make_shared<TransTask>(TransTask::Type::DUMP, desc);
76+
auto waiter = std::make_shared<UC::Latch>();
77+
dumpQ.Submit(task, waiter);
78+
waiter->Wait();
79+
ASSERT_FALSE(failureSet.Contains(task->id));
80+
finish.Wait();
81+
}
82+
83+
TEST_F(UCCacheDumpQueueTest, DumpBlockWhileBackendSubmitFailed)
84+
{
85+
using namespace UC::CacheStore;
86+
UC::Test::Detail::MockStore backend;
87+
EXPECT_CALL(backend, Dump).WillOnce(testing::Return(UC::Status::Error()));
88+
UC::HashSet<UC::Detail::TaskHandle> failureSet;
89+
Config config;
90+
config.backend = (uintptr_t)(void*)&backend;
91+
config.tensorSize = 32768;
92+
config.shardSize = config.tensorSize;
93+
config.blockSize = config.shardSize;
94+
config.deviceId = 0;
95+
config.bufferSize = config.blockSize * 2048;
96+
config.engineId = rd.RandomString(10);
97+
config.shareBufferEnable = true;
98+
TransBuffer buffer;
99+
DumpQueue dumpQ;
100+
auto s = buffer.Setup(config);
101+
ASSERT_EQ(s, UC::Status::OK());
102+
s = dumpQ.Setup(config, &failureSet, &buffer);
103+
ASSERT_EQ(s, UC::Status::OK());
104+
auto blockId = UC::Test::Detail::TypesHelper::MakeBlockId("a1b2c3d4e5f6789012345678901234ab");
105+
constexpr size_t shardIdx = 0;
106+
UC::Test::Detail::DataGenerator data{1, config.blockSize};
107+
data.Generate();
108+
UC::Detail::TaskDesc desc{
109+
{blockId, shardIdx, {data.Buffer()}}
110+
};
111+
auto task = std::make_shared<TransTask>(TransTask::Type::DUMP, desc);
112+
auto waiter = std::make_shared<UC::Latch>();
113+
dumpQ.Submit(task, waiter);
114+
waiter->Wait();
115+
ASSERT_TRUE(failureSet.Contains(task->id));
116+
}

‎ucm/store/test/case/cache/cache_trans_manager_test.cc‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
/**
2+
* MIT License
3+
*
4+
* Copyright (c) 2025 Huawei Technologies Co., Ltd. All rights reserved.
5+
*
6+
* Permission is hereby granted, free of charge, to any person obtaining a copy
7+
* of this software and associated documentation files (the "Software"), to deal
8+
* in the Software without restriction, including without limitation the rights
9+
* to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
10+
* copies of the Software, and to permit persons to whom the Software is
11+
* furnished to do so, subject to the following conditions:
12+
*
13+
* The above copyright notice and this permission notice shall be included in all
14+
* copies or substantial portions of the Software.
15+
*
16+
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17+
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18+
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
19+
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20+
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
21+
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
22+
* SOFTWARE.
23+
* */
24+
#include "cache/cc/trans_manager.h"
25+
#include "detail/data_generator.h"
26+
#include "detail/mock_store.h"
27+
#include "detail/random.h"
28+
#include "detail/types_helper.h"
29+
30+
class UCCacheTransManagerTest : public ::testing::Test {
31+
public:
32+
UC::Test::Detail::Random rd;
33+
static UC::Detail::TaskHandle NextId()
34+
{
35+
static std::atomic<size_t> id{1};
36+
return id.fetch_add(1, std::memory_order_relaxed);
37+
}
38+
};
39+
40+
TEST_F(UCCacheTransManagerTest, DumpThenLoad)
41+
{
42+
using namespace UC::CacheStore;
43+
UC::Test::Detail::MockStore backend;
44+
EXPECT_CALL(backend, Dump).WillOnce(testing::Invoke(NextId));
45+
UC::Latch finish{};
46+
finish.Up();
47+
EXPECT_CALL(backend, Wait).WillOnce(testing::Invoke([&finish]() {
48+
finish.Done();
49+
return UC::Status::OK();
50+
}));
51+
Config config;
52+
config.backend = (uintptr_t)(void*)&backend;
53+
config.tensorSize = 32768;
54+
config.shardSize = config.tensorSize;
55+
config.blockSize = config.shardSize;
56+
config.deviceId = 0;
57+
config.bufferSize = config.blockSize * 2048;
58+
config.engineId = rd.RandomString(10);
59+
config.shareBufferEnable = true;
60+
TransManager transMgr;
61+
auto s = transMgr.Setup(config);
62+
ASSERT_EQ(s, UC::Status::OK());
63+
auto block = UC::Test::Detail::TypesHelper::MakeBlockId("a1b2c3d4e5f6789012345678901234ab");
64+
constexpr size_t nBlocks = 1;
65+
UC::Test::Detail::DataGenerator data1{nBlocks, config.blockSize};
66+
data1.GenerateRandom();
67+
UC::Detail::TaskDesc desc1;
68+
desc1.brief = "Dump";
69+
desc1.push_back(UC::Detail::Shard{block, 0, {data1.Buffer()}});
70+
auto handle1 = transMgr.Submit({TransTask::Type::DUMP, desc1});
71+
ASSERT_TRUE(handle1.HasValue());
72+
s = transMgr.Wait(handle1.Value());
73+
ASSERT_EQ(s, UC::Status::OK());
74+
UC::Test::Detail::DataGenerator data2{nBlocks, config.blockSize};
75+
data2.Generate();
76+
UC::Detail::TaskDesc desc2;
77+
desc2.brief = "Load";
78+
desc2.push_back(UC::Detail::Shard{block, 0, {data2.Buffer()}});
79+
auto handle2 = transMgr.Submit({TransTask::Type::LOAD, desc2});
80+
ASSERT_TRUE(handle2.HasValue());
81+
s = transMgr.Wait(handle2.Value());
82+
ASSERT_EQ(s, UC::Status::OK());
83+
ASSERT_EQ(data1.Compare(data2), 0);
84+
finish.Wait();
85+
}

0 commit comments

Comments
 (0)