MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / RunSharding

Function RunSharding

tensorflow/core/util/work_sharder_test.cc:30–62  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

28namespace {
29
30void RunSharding(int64 num_workers, int64 total, int64 cost_per_unit,
31 int64 per_thread_max_parallelism,
32 thread::ThreadPool* threads) {
33 mutex mu;
34 int64 num_shards = 0;
35 int64 num_done_work = 0;
36 std::vector<bool> work(total, false);
37 Shard(num_workers, threads, total, cost_per_unit,
38 [=, &mu, &num_shards, &num_done_work, &work](int64 start, int64 limit) {
39 VLOG(1) << "Shard [" << start << "," << limit << ")";
40 EXPECT_GE(start, 0);
41 EXPECT_LE(limit, total);
42 mutex_lock l(mu);
43 ++num_shards;
44 for (; start < limit; ++start) {
45 EXPECT_FALSE(work[start]); // No duplicate
46 ++num_done_work;
47 work[start] = true;
48 }
49 });
50 LOG(INFO) << num_workers << " " << total << " " << cost_per_unit << " "
51 << num_shards;
52 EXPECT_EQ(num_done_work, total);
53 if (std::min(num_workers, per_thread_max_parallelism) <
54 threads->NumThreads()) {
55 // If the intention is to limit the parallelism explicitly, we'd
56 // better honor it. Ideally, even if per_thread_max_parallelism >
57 // num_workers, we should expect that Shard() implementation do
58 // not over-shard. Unfortunately, ThreadPoolDevice::parallelFor
59 // tends to over-shard.
60 EXPECT_LE(num_shards, 1 + per_thread_max_parallelism);
61 }
62}
63
64TEST(Shard, Basic) {
65 thread::ThreadPool threads(Env::Default(), "test", 16);

Callers 1

TESTFunction · 0.70

Calls 3

ShardFunction · 0.70
minFunction · 0.50
NumThreadsMethod · 0.45

Tested by

no test coverage detected