Merge pull request #317 from nvkevihu/cpp-manual-profiler-export

[C++ Benchmark]: Replace ExportToTensorBoard

Author: meena-at-work

tftrt/benchmarking-cpp/BUILD

@@ -26,24 +26,6 @@ python3 convert_model.py --model-dir /path/to/model/dir --output-dir /path/to/de
 
 ## Building
 
-The binary relies on a modified Tensorflow, which will need to be rebuilt. Internal users can use a container with Tensorflow already modified and built, instead of building with Bazel, which will take much longer.
-
-### Bazel
-
-The `setup.sh` script applies the Tensorflow patch and prepares the container for the Bazel build.
-
-```
-/workspace/tensorrt/tftrt/benchmarking-cpp/build-scripts/setup.sh
-cd /opt/tensorflow
-./tftrt-build.sh
-```
-
-The binary will be located at `/opt/tensorflow/tensorflow-source/bazel-bin/tensorflow/examples/benchmarking-cpp/tftrt_benchmark_runner`.
-
-### Prebuilt
-
-For internal NVIDIA users, a container with a prebuilt modified Tensorflow is available. In the container, use CMake to build the binary without needing to rebuild Tensorflow:
-
 ```
 cd /workspace/tensorrt/tftrt/benchmarking-cpp
 mkdir build && cd build
@@ -61,4 +43,6 @@ The binary will be located at `/workspace/tensorrt/tftrt/benchmarking-cpp/tftrt_
 
 ### Profiling
 
-To profile, set the `--out_dir` flag. Run `tensorboard --logdir [out_dir]` to view results.
+To profile, set the `--out_dir` flag. This creates a log directory and serializes the `XSpace` to a location that TensorBoard expects (i.e. `[out_dir]/plugins/profile/[run_id]/[host_id].xplane.pb`).
+
+Run `tensorboard --logdir [out_dir]` to view results. TensorBoard will generate the available dashboards using the `XSpace` directly.

tftrt/benchmarking-cpp/README.md

@@ -1,6 +1,10 @@
 #include <chrono>
+#include <ctime>
+#include <fstream>
+#include <iomanip>
 #include <iostream>
 #include <numeric>
+#include <sstream>
 #include <vector>
 
 #include "tensorflow/cc/ops/array_ops.h"
@@ -14,9 +18,11 @@
 #include "tensorflow/core/platform/init_main.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/host_info.h"
 #include "tensorflow/core/profiler/lib/profiler_session.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
-#include "tensorflow/core/profiler/rpc/client/capture_profile.h"
 #include "tensorflow/core/public/session.h"
 #include "tensorflow/core/util/command_line_flags.h"
 
@@ -35,6 +41,10 @@ using tensorflow::Tensor;
     }                                                                      \
   } while (0)
 
+// Directory required by Tensorboard to load the .xplane.pb properly
+const string kProfilePath = "plugins/profile";
+const string kProfileSuffix = ".xplane.pb";
+
 // Get the GPU and its default context.
 Status GetDevice(std::unique_ptr<tensorflow::Session>& session,
                  tensorflow::Device** device,
@@ -144,6 +154,26 @@ Status SetupCallable(std::unique_ptr<tensorflow::Session>& session,
   return session->MakeCallable(opts, handle);
 }
 
+// Get the current time as a string.
+string GetTime() {
+  std::stringstream ss;
+  std::time_t now = std::time(nullptr);
+  ss << std::put_time(std::localtime(&now), "%EY_%m_%d_%H_%M_%S");
+  return ss.str();
+}
+
+// Create the directory structure expected by TensorBoard.
+Status GetAndCreateProfilePath(const string& out_dir,
+                               string* profile_path) {
+  // Use date for run ID
+  string run_id = GetTime();
+  string host_id = tensorflow::port::Hostname();
+  string run_dir = tensorflow::io::JoinPath(out_dir, kProfilePath, run_id);
+  *profile_path = tensorflow::io::JoinPath(run_dir, host_id + kProfileSuffix);
+  TF_RETURN_IF_ERROR(tensorflow::Env::Default()->RecursivelyCreateDir(run_dir));
+  return Status::OK();
+}
+
 // Start the profiling session.
 Status StartProfiling(std::unique_ptr<tensorflow::ProfilerSession>& profiler) {
   profiler = tensorflow::ProfilerSession::Create(
@@ -157,7 +187,18 @@ Status StopProfiling(std::unique_ptr<tensorflow::ProfilerSession>& profiler,
                      const string& out_dir) {
   tensorflow::profiler::XSpace xspace;
   TF_RETURN_IF_ERROR(profiler->CollectData(&xspace));
-  tensorflow::profiler::ExportToTensorBoard(xspace, out_dir);
+
+  // Get export path
+  string profile_path;
+  TF_RETURN_IF_ERROR(GetAndCreateProfilePath(out_dir, &profile_path));
+
+  // Serialize and write to disk
+  string xspace_serialized;
+  xspace.SerializeToString(&xspace_serialized);
+  std::ofstream ofs(profile_path);
+  ofs << xspace_serialized;
+  ofs.close();
+
   profiler.reset();
   return Status::OK();
 }