Skip to content

Commit e71dec1

Browse files
shanshanptshanshanpt
authored
[Runtime] Add UT for session group. (#806)
Signed-off-by: Tao Peng <jiankeng.pt@alibaba-inc.com>
1 parent 7636cd7 commit e71dec1

File tree

2 files changed

+231
-0
lines changed

2 files changed

+231
-0
lines changed

‎tensorflow/core/BUILD‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5162,6 +5162,45 @@ tf_cuda_cc_test(
51625162
] + if_cuda([":cuda"]),
51635163
)
51645164

5165+
tf_cuda_cc_test(
5166+
name = "common_runtime_direct_session_group_test",
5167+
size = "small",
5168+
srcs = ["common_runtime/direct_session_group_test.cc"],
5169+
args = [] + if_cuda(["--heap_check=local"]), # The GPU tracer leaks memory
5170+
linkstatic = tf_kernel_tests_linkstatic(),
5171+
deps = [
5172+
":core_cpu",
5173+
":core_cpu_internal",
5174+
":direct_session_internal",
5175+
":framework",
5176+
":framework_internal",
5177+
":lib",
5178+
":lib_internal",
5179+
":ops",
5180+
":protos_all_cc",
5181+
":test",
5182+
":test_main",
5183+
":testlib",
5184+
"@com_google_absl//absl/memory",
5185+
"@com_google_absl//absl/strings",
5186+
"//third_party/eigen3",
5187+
"//tensorflow/cc:cc_ops",
5188+
"//tensorflow/core/kernels:collective_ops",
5189+
"//tensorflow/core/kernels:control_flow_ops",
5190+
"//tensorflow/core/kernels:cwise_op",
5191+
"//tensorflow/core/kernels:dense_update_ops",
5192+
"//tensorflow/core/kernels:fifo_queue_op",
5193+
"//tensorflow/core/kernels:function_ops",
5194+
"//tensorflow/core/kernels:identity_n_op",
5195+
"//tensorflow/core/kernels:identity_op",
5196+
"//tensorflow/core/kernels:matmul_op",
5197+
"//tensorflow/core/kernels:ops_util",
5198+
"//tensorflow/core/kernels:queue_ops",
5199+
"//tensorflow/core/kernels:session_ops",
5200+
"//tensorflow/core/kernels:variable_ops",
5201+
] + if_cuda([":cuda"]),
5202+
)
5203+
51655204
tf_cuda_cc_test(
51665205
name = "common_runtime_gpu_cuda_graph_mode_session_test",
51675206
size = "small",

‎tensorflow/core/common_runtime/direct_session_group_test.cc‎

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,192 @@
1+
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
16+
#include "tensorflow/core/common_runtime/direct_session.h"
17+
18+
#include <map>
19+
#include <memory>
20+
#include <string>
21+
#include <thread> // NOLINT
22+
#include <unistd.h>
23+
#include <unordered_map>
24+
#include <vector>
25+
26+
#include "absl/memory/memory.h"
27+
#include "absl/strings/match.h"
28+
#include "tensorflow/core/common_runtime/device_factory.h"
29+
#include "tensorflow/core/common_runtime/device_mgr.h"
30+
#include "tensorflow/core/common_runtime/function_testlib.h"
31+
#include "tensorflow/core/framework/allocator.h"
32+
#include "tensorflow/core/framework/graph.pb.h"
33+
#include "tensorflow/core/framework/op_kernel.h"
34+
#include "tensorflow/core/framework/tensor.h"
35+
#include "tensorflow/core/framework/tensor_testutil.h"
36+
#include "tensorflow/core/framework/types.pb.h"
37+
#include "tensorflow/core/graph/costmodel.h"
38+
#include "tensorflow/core/graph/graph.h"
39+
#include "tensorflow/core/graph/node_builder.h"
40+
#include "tensorflow/core/graph/testlib.h"
41+
#include "tensorflow/core/kernels/ops_util.h"
42+
#include "tensorflow/core/lib/core/errors.h"
43+
#include "tensorflow/core/lib/core/status.h"
44+
#include "tensorflow/core/lib/core/status_test_util.h"
45+
#include "tensorflow/core/lib/core/threadpool.h"
46+
#include "tensorflow/core/lib/strings/str_util.h"
47+
#include "tensorflow/core/platform/protobuf.h"
48+
#include "tensorflow/core/platform/test.h"
49+
#include "tensorflow/core/platform/test_benchmark.h"
50+
#include "tensorflow/core/protobuf/rewriter_config.pb.h"
51+
#include "tensorflow/core/public/session.h"
52+
#include "tensorflow/core/public/session_options.h"
53+
#include "tensorflow/core/util/device_name_utils.h"
54+
55+
#if GOOGLE_CUDA
56+
#include "third_party/gpus/cuda/include/cuda.h"
57+
#include "third_party/gpus/cuda/include/cuda_runtime_api.h"
58+
#elif TENSORFLOW_USE_ROCM
59+
#include "rocm/include/hip/hip_runtime.h"
60+
#endif // GOOGLE_CUDA
61+
62+
namespace tensorflow {
63+
namespace {
64+
65+
CallableOptions MakeCallableOptions(gtl::ArraySlice<string> feeds,
66+
gtl::ArraySlice<string> fetches,
67+
gtl::ArraySlice<string> targets) {
68+
CallableOptions ret;
69+
for (const string& feed : feeds) {
70+
ret.add_feed(feed);
71+
}
72+
for (const string& fetch : fetches) {
73+
ret.add_fetch(fetch);
74+
}
75+
for (const string& target : targets) {
76+
ret.add_target(target);
77+
}
78+
return ret;
79+
}
80+
81+
SessionOptions DefaultSessionOptions() {
82+
SessionOptions options;
83+
(*options.config.mutable_device_count())["CPU"] = 2;
84+
return options;
85+
}
86+
87+
std::unique_ptr<SessionGroup> CreateSessionGroup() {
88+
SessionGroup* sg = nullptr;
89+
SessionGroupMetadata metadata;
90+
NewSessionGroup(DefaultSessionOptions(), &sg, metadata);
91+
return std::unique_ptr<SessionGroup>(sg);
92+
}
93+
94+
std::unique_ptr<Session> CreateSession() {
95+
return std::unique_ptr<Session>(NewSession(DefaultSessionOptions()));
96+
}
97+
98+
class DirectSessionMinusAXTest : public ::testing::Test {
99+
public:
100+
void Initialize(std::initializer_list<float> a_values) {
101+
Graph graph(OpRegistry::Global());
102+
103+
Tensor a_tensor(DT_FLOAT, TensorShape({2, 2}));
104+
test::FillValues<float>(&a_tensor, a_values);
105+
Node* a = test::graph::Constant(&graph, a_tensor);
106+
a->set_assigned_device_name("/job:localhost/replica:0/task:0/cpu:0");
107+
a_ = a->name();
108+
109+
Tensor x_tensor(DT_FLOAT, TensorShape({2, 1}));
110+
test::FillValues<float>(&x_tensor, {1, 1});
111+
Node* x = test::graph::Constant(&graph, x_tensor);
112+
x->set_assigned_device_name("/job:localhost/replica:0/task:0/cpu:1");
113+
x_ = x->name();
114+
115+
// y = A * x
116+
Node* y = test::graph::Matmul(&graph, a, x, false, false);
117+
y->set_assigned_device_name("/job:localhost/replica:0/task:0/cpu:0");
118+
y_ = y->name();
119+
120+
Node* y_neg = test::graph::Unary(&graph, "Neg", y);
121+
y_neg_ = y_neg->name();
122+
y_neg->set_assigned_device_name("/job:localhost/replica:0/task:0/cpu:1");
123+
124+
Node* z = test::graph::Unary(&graph, "Identity", y_neg);
125+
z_ = z->name();
126+
z->set_assigned_device_name("/job:localhost/replica:0/task:0/cpu:1");
127+
128+
graph.ToGraphDef(&def_);
129+
}
130+
131+
string a_;
132+
string x_;
133+
string y_;
134+
string y_neg_;
135+
string z_;
136+
GraphDef def_;
137+
};
138+
139+
TEST_F(DirectSessionMinusAXTest, TestDirectSessionGroup) {
140+
for (int i = 0; i < 1000; ++i) {
141+
Initialize({3, 2, -1, 0});
142+
auto sg = CreateSessionGroup();
143+
ASSERT_TRUE(sg != nullptr);
144+
TF_ASSERT_OK(sg->Create(def_));
145+
std::vector<std::pair<string, Tensor>> inputs;
146+
147+
// Request two targets: one fetch output and one non-fetched output.
148+
std::vector<string> output_names = {y_ + ":0"};
149+
std::vector<string> target_nodes = {y_neg_};
150+
std::vector<Tensor> outputs;
151+
Status s = sg->Run(inputs, output_names, target_nodes, &outputs);
152+
TF_ASSERT_OK(s);
153+
154+
ASSERT_EQ(1, outputs.size());
155+
// The first output should be initialized and have the correct
156+
// output.
157+
auto mat = outputs[0].matrix<float>();
158+
ASSERT_TRUE(outputs[0].IsInitialized());
159+
EXPECT_FLOAT_EQ(5.0, mat(0, 0));
160+
161+
//usleep(10000);
162+
}
163+
}
164+
165+
TEST_F(DirectSessionMinusAXTest, TestDirectSession) {
166+
for (int i = 0; i < 1000; ++i) {
167+
Initialize({3, 2, -1, 0});
168+
auto sg = CreateSession();
169+
ASSERT_TRUE(sg != nullptr);
170+
TF_ASSERT_OK(sg->Create(def_));
171+
std::vector<std::pair<string, Tensor>> inputs;
172+
173+
// Request two targets: one fetch output and one non-fetched output.
174+
std::vector<string> output_names = {y_ + ":0"};
175+
std::vector<string> target_nodes = {y_neg_};
176+
std::vector<Tensor> outputs;
177+
Status s = sg->Run(inputs, output_names, target_nodes, &outputs);
178+
TF_ASSERT_OK(s);
179+
180+
ASSERT_EQ(1, outputs.size());
181+
// The first output should be initialized and have the correct
182+
// output.
183+
auto mat = outputs[0].matrix<float>();
184+
ASSERT_TRUE(outputs[0].IsInitialized());
185+
EXPECT_FLOAT_EQ(5.0, mat(0, 0));
186+
187+
//usleep(10000);
188+
}
189+
}
190+
191+
}
192+
} // namespace tensorflow

0 commit comments

Comments
 (0)