1/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations 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
26PYBIND11_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