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 | |
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 | (void)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 | (void)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 | (void)writer->Close(); |
122 | }); |
123 | }; |
124 | |