1 /* Copyright 2020 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 <memory>
17 #include <string>
18 #include <variant>
19 #include <vector>
20
21 #include "pybind11/pybind11.h"
22 #include "tensorflow/core/platform/env.h"
23 #include "tensorflow/core/profiler/convert/tool_options.h"
24 #include "tensorflow/core/profiler/convert/xplane_to_tools_data.h"
25 #include "tensorflow/core/profiler/rpc/profiler_server.h"
26 #include "tensorflow/python/lib/core/pybind11_status.h"
27 #include "tensorflow/python/profiler/internal/profiler_pywrap_impl.h"
28
29 namespace py = ::pybind11;
30
31 namespace {
32
33 using ::tensorflow::profiler::ToolOptions;
34 using ::tensorflow::profiler::pywrap::ProfilerSessionWrapper;
35
36 // These must be called under GIL because it reads Python objects. Reading
37 // Python objects require GIL because the objects can be mutated by other Python
38 // threads. In addition, Python objects are reference counted; reading py::dict
39 // will increase its reference count.
ToolOptionsFromPythonDict(const py::dict & dictionary)40 ToolOptions ToolOptionsFromPythonDict(const py::dict& dictionary) {
41 ToolOptions map;
42 for (const auto& item : dictionary) {
43 std::variant<int, std::string> value;
44 try {
45 value = item.second.cast<int>();
46 } catch (...) {
47 try {
48 value = item.second.cast<std::string>();
49 } catch (...) {
50 continue;
51 }
52 }
53 map.emplace(item.first.cast<std::string>(), value);
54 }
55 return map;
56 }
57
58 } // namespace
59
PYBIND11_MODULE(_pywrap_profiler,m)60 PYBIND11_MODULE(_pywrap_profiler, m) {
61 py::class_<ProfilerSessionWrapper> profiler_session_class(m,
62 "ProfilerSession");
63 profiler_session_class.def(py::init<>())
64 .def("start",
65 [](ProfilerSessionWrapper& wrapper, const char* logdir,
66 const py::dict& options) {
67 tensorflow::Status status;
68 ToolOptions tool_options = ToolOptionsFromPythonDict(options);
69 {
70 py::gil_scoped_release release;
71 status = wrapper.Start(logdir, tool_options);
72 }
73 // Py_INCREF and Py_DECREF must be called holding the GIL.
74 tensorflow::MaybeRaiseRegisteredFromStatus(status);
75 })
76 .def("stop",
77 [](ProfilerSessionWrapper& wrapper) {
78 tensorflow::string content;
79 tensorflow::Status status;
80 {
81 py::gil_scoped_release release;
82 status = wrapper.Stop(&content);
83 }
84 // Py_INCREF and Py_DECREF must be called holding the GIL.
85 tensorflow::MaybeRaiseRegisteredFromStatus(status);
86 // The content is not valid UTF-8. It must be converted to bytes.
87 return py::bytes(content);
88 })
89 .def("export_to_tb", [](ProfilerSessionWrapper& wrapper) {
90 tensorflow::Status status;
91 {
92 py::gil_scoped_release release;
93 status = wrapper.ExportToTensorBoard();
94 }
95 // Py_INCREF and Py_DECREF must be called holding the GIL.
96 tensorflow::MaybeRaiseRegisteredFromStatus(status);
97 });
98
99 m.def("start_server", [](int port) {
100 auto profiler_server =
101 std::make_unique<tensorflow::profiler::ProfilerServer>();
102 profiler_server->StartProfilerServer(port);
103 // Intentionally release profiler server. Should transfer ownership to
104 // caller instead.
105 profiler_server.release();
106 });
107
108 m.def("trace",
109 [](const char* service_addr, const char* logdir,
110 const char* worker_list, bool include_dataset_ops, int duration_ms,
111 int num_tracing_attempts, py::dict options) {
112 tensorflow::Status status;
113 ToolOptions tool_options = ToolOptionsFromPythonDict(options);
114 {
115 py::gil_scoped_release release;
116 status = tensorflow::profiler::pywrap::Trace(
117 service_addr, logdir, worker_list, include_dataset_ops,
118 duration_ms, num_tracing_attempts, tool_options);
119 }
120 // Py_INCREF and Py_DECREF must be called holding the GIL.
121 tensorflow::MaybeRaiseRegisteredFromStatus(status);
122 });
123
124 m.def("monitor", [](const char* service_addr, int duration_ms,
125 int monitoring_level, bool display_timestamp) {
126 tensorflow::string content;
127 tensorflow::Status status;
128 {
129 py::gil_scoped_release release;
130 status = tensorflow::profiler::pywrap::Monitor(
131 service_addr, duration_ms, monitoring_level, display_timestamp,
132 &content);
133 }
134 // Py_INCREF and Py_DECREF must be called holding the GIL.
135 tensorflow::MaybeRaiseRegisteredFromStatus(status);
136 return content;
137 });
138
139 m.def(
140 "xspace_to_tools_data",
141 [](const py::list& xspace_path_list, const py::str& py_tool_name,
142 const py::dict options = py::dict()) {
143 std::vector<tensorflow::profiler::XSpace> xspaces;
144 xspaces.reserve(xspace_path_list.size());
145 std::vector<std::string> filenames;
146 filenames.reserve(xspace_path_list.size());
147 for (py::handle obj : xspace_path_list) {
148 std::string filename = std::string(py::cast<py::str>(obj));
149
150 tensorflow::profiler::XSpace xspace;
151 tensorflow::Status status;
152
153 status = tensorflow::ReadBinaryProto(tensorflow::Env::Default(),
154 filename, &xspace);
155
156 if (!status.ok()) {
157 return py::make_tuple(py::bytes(""), py::bool_(false));
158 }
159
160 xspaces.push_back(xspace);
161 filenames.push_back(filename);
162 }
163 std::string tool_name = std::string(py_tool_name);
164 ToolOptions tool_options = ToolOptionsFromPythonDict(options);
165 auto status_or_tool_data =
166 tensorflow::profiler::ConvertMultiXSpacesToToolData(
167 xspaces, filenames, tool_name, tool_options);
168 if (!status_or_tool_data.ok()) {
169 LOG(ERROR) << status_or_tool_data.status().error_message();
170 return py::make_tuple(py::bytes(""), py::bool_(false));
171 }
172 return py::make_tuple(py::bytes(status_or_tool_data.value()),
173 py::bool_(true));
174 },
175 // TODO: consider defaulting `xspace_path_list` to empty list, since
176 // this parameter is only used for two of the tools...
177 py::arg(), py::arg(), py::arg() = py::dict());
178
179 m.def("xspace_to_tools_data_from_byte_string",
180 [](const py::list& xspace_string_list, const py::list& filenames_list,
181 const py::str& py_tool_name) {
182 std::vector<tensorflow::profiler::XSpace> xspaces;
183 xspaces.reserve(xspace_string_list.size());
184 std::vector<std::string> filenames;
185 filenames.reserve(filenames_list.size());
186
187 // XSpace string inputs
188 for (py::handle obj : xspace_string_list) {
189 std::string xspace_string = std::string(py::cast<py::bytes>(obj));
190
191 tensorflow::profiler::XSpace xspace;
192
193 if (!xspace.ParseFromString(xspace_string)) {
194 return py::make_tuple(py::bytes(""), py::bool_(false));
195 }
196
197 xspaces.push_back(xspace);
198 }
199
200 // Filenames
201 for (py::handle obj : filenames_list) {
202 filenames.push_back(std::string(py::cast<py::str>(obj)));
203 }
204 std::string tool_name = std::string(py_tool_name);
205 auto status_or_tool_data =
206 tensorflow::profiler::ConvertMultiXSpacesToToolData(
207 xspaces, filenames, tool_name, {});
208 if (!status_or_tool_data.ok()) {
209 LOG(ERROR) << status_or_tool_data.status().error_message();
210 return py::make_tuple(py::bytes(""), py::bool_(false));
211 }
212 return py::make_tuple(py::bytes(status_or_tool_data.value()),
213 py::bool_(true));
214 });
215 };
216