From 183ecd577fb655b1c61cc46028b1bd13561bffeb Mon Sep 17 00:00:00 2001
From: Clive Verghese <cliveverghese@gmail.com>
Date: Fri, 13 Dec 2024 23:57:43 +0000
Subject: [PATCH] Update Tensorflow Dependency and include
 _pywrap_profiler_plugin.so in built wheel

---
 .bazelrc                                      |  1 +
 WORKSPACE                                     | 22 +++++--------------
 plugin/build_pip_package.sh                   |  2 ++
 plugin/setup.py                               |  1 +
 plugin/tensorboard_plugin_profile/BUILD       |  1 +
 .../convert/raw_to_tool_data.py               |  5 ++++-
 6 files changed, 14 insertions(+), 18 deletions(-)

diff --git a/.bazelrc b/.bazelrc
index e4661bff6..d11a5cb9d 100644
--- a/.bazelrc
+++ b/.bazelrc
@@ -7,3 +7,4 @@ build --host_cxxopt=-std=c++17
 build --copt=-Wno-gnu-offsetof-extensions
 build --copt=-Wno-error=array-parameter
 build --copt=-Wno-error=unused-command-line-argument
+build --copt=-w
diff --git a/WORKSPACE b/WORKSPACE
index 141edb471..4d070f6b8 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -31,7 +31,7 @@ python_register_toolchains(
     ignore_root_user_error = True,
     # Available versions are listed in @rules_python//python:versions.bzl.
     # We recommend using the same version your team is already standardized on.
-    python_version = "3.10",
+    python_version = "3.11",
 )
 
 load("@python//:defs.bzl", "interpreter")
@@ -160,29 +160,16 @@ load("@io_bazel_rules_sass//:defs.bzl", "sass_repositories")
 
 sass_repositories()
 
-http_archive(
-    name = "com_google_protobuf",
-    sha256 = "f66073dee0bc159157b0bd7f502d7d1ee0bc76b3c1eac9836927511bdc4b3fc1",
-    strip_prefix = "protobuf-3.21.9",
-    urls = [
-        "https://github.com/protocolbuffers/protobuf/archive/v3.21.9.zip",
-    ],
-)
-
-load("@com_google_protobuf//:protobuf_deps.bzl", "protobuf_deps")
-
-protobuf_deps()
-
 http_archive(
     name = "org_tensorflow",
     patch_args = ["-p1"],
     patches = ["//third_party:tensorflow.patch"],
     # NOTE: when updating this, MAKE SURE to also update the protobuf_js runtime version
     # in third_party/workspace.bzl to >= the protobuf/protoc version provided by TF.
-    sha256 = "288f840d70639d54a087b9ba5310ffae3dbc45ede90ada168a0943cb6e43dabe",
-    strip_prefix = "tensorflow-a568daa404774cc90c3729ad7af2c9f7a0a558d2",
+    sha256 = "fe5261186952ec13652e48f219048d1f5e7825d6b83c91ebee71bbc41060baba",
+    strip_prefix = "tensorflow-0c1df46af2b5010ebdeac51f9391649753d3f1d2",
     urls = [
-        "https://github.com/tensorflow/tensorflow/archive/a568daa404774cc90c3729ad7af2c9f7a0a558d2.zip",
+    	"https://github.com/tensorflow/tensorflow/archive/0c1df46af2b5010ebdeac51f9391649753d3f1d2.zip"
     ],
 )
 
@@ -210,6 +197,7 @@ load("@org_tensorflow//tensorflow:workspace0.bzl", "tf_workspace0")
 
 tf_workspace0()
 
+
 load(
     "@local_tsl//third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl",
     "cuda_json_init_repository",
diff --git a/plugin/build_pip_package.sh b/plugin/build_pip_package.sh
index 3761b2233..d6128836d 100755
--- a/plugin/build_pip_package.sh
+++ b/plugin/build_pip_package.sh
@@ -45,6 +45,8 @@ find tensorboard_plugin_profile/protobuf -name \*.py -exec $sedi -e '
     s/^from plugin.tensorboard_plugin_profile/from tensorboard_plugin_profile/
   ' {} +
 
+cp ${BUILD_WORKSPACE_DIRECTORY}/bazel-bin/external/org_tensorflow/tensorflow/python/profiler/internal/_pywrap_profiler_plugin.so tensorboard_plugin_profile/convert/
+
 # Copy static files.
 cd tensorboard_plugin_profile
 mkdir -p static
diff --git a/plugin/setup.py b/plugin/setup.py
index ce6cf75c4..dc3ce1bb8 100644
--- a/plugin/setup.py
+++ b/plugin/setup.py
@@ -54,6 +54,7 @@ def get_readme():
     ),
     package_data={
         'tensorboard_plugin_profile': ['static/**'],
+        '': ['_pywrap_profiler_plugin.so'],
     },
     entry_points={
         'tensorboard_plugins': [
diff --git a/plugin/tensorboard_plugin_profile/BUILD b/plugin/tensorboard_plugin_profile/BUILD
index 44268f2c1..77b1ccddc 100644
--- a/plugin/tensorboard_plugin_profile/BUILD
+++ b/plugin/tensorboard_plugin_profile/BUILD
@@ -14,6 +14,7 @@ py_library(
         "//:expect_tensorflow_installed",
         "@org_pocoo_werkzeug",
         "@org_pythonhosted_six",
+	"@org_tensorflow//tensorflow/python/profiler/internal:_pywrap_profiler_plugin",
         "@org_tensorflow_tensorboard//tensorboard/backend/event_processing:plugin_asset_util",
         "@org_tensorflow_tensorboard//tensorboard/plugins:base_plugin",
         "@org_xprof//plugin/tensorboard_plugin_profile/convert:all_libs",
diff --git a/plugin/tensorboard_plugin_profile/convert/raw_to_tool_data.py b/plugin/tensorboard_plugin_profile/convert/raw_to_tool_data.py
index 57c0f7909..2403b8c9a 100644
--- a/plugin/tensorboard_plugin_profile/convert/raw_to_tool_data.py
+++ b/plugin/tensorboard_plugin_profile/convert/raw_to_tool_data.py
@@ -25,7 +25,10 @@
 
 import logging
 
-from tensorflow.python.profiler.internal import _pywrap_profiler_plugin  # pylint: disable=g-direct-tensorflow-import
+try:
+  from tensorboard_plugin_profile.convert import _pywrap_profiler_plugin
+except:
+  from tensorflow.python.profiler.internal import _pywrap_profiler_plugin  # pylint: disable=g-direct-tensorflow-import
 from tensorboard_plugin_profile.convert import dcn_collective_stats_proto_to_gviz
 from tensorboard_plugin_profile.convert import hlo_stats_proto_to_gviz
 from tensorboard_plugin_profile.convert import inference_stats_proto_to_gviz