1 /* Copyright 2019 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 "absl/strings/string_view.h"
17 #include "pybind11/pybind11.h"
18 #include "pybind11/pytypes.h"
19 #include "tensorflow/core/lib/core/status.h"
20 #include "tensorflow/core/lib/strings/stringprintf.h"
21 #include "tensorflow/core/util/debug_events_writer.h"
22 #include "tensorflow/python/lib/core/pybind11_absl.h"
23 #include "tensorflow/python/lib/core/pybind11_proto.h"
24 #include "tensorflow/python/lib/core/pybind11_status.h"
25
PYBIND11_MODULE(_pywrap_debug_events_writer,m)26 PYBIND11_MODULE(_pywrap_debug_events_writer, m) {
27 namespace py = pybind11;
28 using namespace tensorflow; // NOLINT(build/namespaces)
29 using namespace tensorflow::tfdbg; // NOLINT(build/namespaces)
30
31 m.def("Init",
32 [](const std::string& dump_root, const std::string& tfdbg_run_id,
33 const int64 circular_buffer_size) {
34 DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
35 dump_root, tfdbg_run_id, circular_buffer_size);
36 if (!writer->Init().ok()) {
37 throw py::value_error(tensorflow::strings::Printf(
38 "Failed to initialize debug events writer at: %s",
39 dump_root.c_str()));
40 }
41 });
42 m.def("WriteSourceFile",
43 [](const std::string& dump_root, const py::object obj) {
44 CheckProtoType(obj, "tensorflow.DebugEvent");
45 DebugEventsWriter* writer = nullptr;
46 TF_CHECK_OK(
47 DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
48 writer->WriteSerializedNonExecutionDebugEvent(
49 obj.attr("SerializeToString")().cast<std::string>(),
50 tfdbg::DebugEventFileType::SOURCE_FILES);
51 });
52 m.def("WriteStackFrameWithId",
53 [](const std::string& dump_root, const py::object& obj) {
54 CheckProtoType(obj, "tensorflow.DebugEvent");
55 DebugEventsWriter* writer = nullptr;
56 TF_CHECK_OK(
57 DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
58 writer->WriteSerializedNonExecutionDebugEvent(
59 obj.attr("SerializeToString")().cast<std::string>(),
60 tfdbg::DebugEventFileType::STACK_FRAMES);
61 });
62 m.def("WriteGraphOpCreation",
63 [](const std::string& dump_root, const py::object& obj) {
64 CheckProtoType(obj, "tensorflow.DebugEvent");
65 DebugEventsWriter* writer = nullptr;
66 TF_CHECK_OK(
67 DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
68 writer->WriteSerializedNonExecutionDebugEvent(
69 obj.attr("SerializeToString")().cast<std::string>(),
70 tfdbg::DebugEventFileType::GRAPHS);
71 });
72 m.def("WriteDebuggedGraph",
73 [](const std::string& dump_root, const py::object& obj) {
74 CheckProtoType(obj, "tensorflow.DebugEvent");
75 DebugEventsWriter* writer = nullptr;
76 TF_CHECK_OK(
77 DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
78 writer->WriteSerializedNonExecutionDebugEvent(
79 obj.attr("SerializeToString")().cast<std::string>(),
80 tfdbg::DebugEventFileType::GRAPHS);
81 });
82 m.def("WriteExecution",
83 [](const std::string& dump_root, const py::object& obj) {
84 CheckProtoType(obj, "tensorflow.DebugEvent");
85 DebugEventsWriter* writer = nullptr;
86 TF_CHECK_OK(
87 DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
88 writer->WriteSerializedExecutionDebugEvent(
89 obj.attr("SerializeToString")().cast<std::string>(),
90 tfdbg::DebugEventFileType::EXECUTION);
91 });
92 m.def("WriteGraphExecutionTrace",
93 [](const std::string& dump_root, const py::object& obj) {
94 CheckProtoType(obj, "tensorflow.DebugEvent");
95 DebugEventsWriter* writer = nullptr;
96 TF_CHECK_OK(
97 DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
98 writer->WriteSerializedExecutionDebugEvent(
99 obj.attr("SerializeToString")().cast<std::string>(),
100 tfdbg::DebugEventFileType::GRAPH_EXECUTION_TRACES);
101 });
102 m.def("RegisterDeviceAndGetId", [](const std::string& dump_root,
103 const std::string& device_name) {
104 DebugEventsWriter* writer = nullptr;
105 TF_CHECK_OK(DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
106 return writer->RegisterDeviceAndGetId(device_name);
107 });
108 m.def("FlushNonExecutionFiles", [](const std::string& dump_root) {
109 DebugEventsWriter* writer = nullptr;
110 TF_CHECK_OK(DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
111 writer->FlushNonExecutionFiles();
112 });
113 m.def("FlushExecutionFiles", [](const std::string& dump_root) {
114 DebugEventsWriter* writer = nullptr;
115 TF_CHECK_OK(DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
116 writer->FlushExecutionFiles();
117 });
118 m.def("Close", [](const std::string& dump_root) {
119 DebugEventsWriter* writer = nullptr;
120 TF_CHECK_OK(DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
121 writer->Close();
122 });
123 };
124