Fix for profiler in TensorBoard 2.5

[wchargin]

Hi!

The profile plugin doesn’t work out of the box with TensorBoard’s fast
data loading mode because of the atypical ways in which it loads data.
Here is a patch that fixes the plugin:

diff --git a/plugin/tensorboard_plugin_profile/profile_plugin.py b/plugin/tensorboard_plugin_profile/profile_plugin.py
index 0773327..45950c2 100644
--- a/plugin/tensorboard_plugin_profile/profile_plugin.py
+++ b/plugin/tensorboard_plugin_profile/profile_plugin.py
@@ -28,6 +28,7 @@ import six
 import tensorflow.compat.v2 as tf
 from werkzeug import wrappers
 
+from tensorboard import context
 from tensorboard.backend.event_processing import plugin_asset_util
 from tensorboard.plugins import base_plugin
 from tensorflow.python.profiler import profiler_client  # pylint: disable=g-direct-tensorflow-import
@@ -266,6 +267,15 @@ def respond(body, content_type, code=200, content_encoding=None):
       body, content_type=content_type, status=code, headers=headers)
 
 
+def _plugin_assets(logdir, runs, plugin_name):
+  result = {}
+  for run in runs:
+    run_path = os.path.join(logdir, run)
+    assets = plugin_asset_util.ListAssets(run_path, plugin_name)
+    result[run] = assets
+  return result
+
+
 class ProfilePlugin(base_plugin.TBPlugin):
   """Profile Plugin for TensorBoard."""
 
@@ -279,7 +289,7 @@ class ProfilePlugin(base_plugin.TBPlugin):
       context: A base_plugin.TBContext instance.
     """
     self.logdir = context.logdir
-    self.multiplexer = context.multiplexer
+    self.data_provider = context.data_provider
     self.stub = None
     self.master_tpu_unsecure_channel = context.flags.master_tpu_unsecure_channel
 
@@ -636,14 +646,9 @@ class ProfilePlugin(base_plugin.TBPlugin):
     tb_run_name, profile_run_name = os.path.split(run)
     if not tb_run_name:
       tb_run_name = '.'
-    tb_run_directory = self.multiplexer.RunPaths().get(tb_run_name)
-    if tb_run_directory is None:
-      # Check if logdir is a directory to handle case where it's actually a
-      # multipart directory spec, which this plugin does not support.
-      if tb_run_name == '.' and tf.io.gfile.isdir(self.logdir):
-        tb_run_directory = self.logdir
-      else:
-        raise RuntimeError('No matching run directory for run %s' % run)
+    tb_run_directory = os.path.join(self.logdir, tb_run_name)
+    if not tf.io.gfile.isdir(tb_run_directory):
+      raise RuntimeError('No matching run directory for run %s' % run)
     plugin_directory = plugin_asset_util.PluginDirectory(
         tb_run_directory, PLUGIN_NAME)
     return os.path.join(plugin_directory, profile_run_name)
@@ -696,8 +701,14 @@ class ProfilePlugin(base_plugin.TBPlugin):
     """
     self.start_grpc_stub_if_necessary()
 
-    plugin_assets = self.multiplexer.PluginAssets(PLUGIN_NAME)
-    tb_run_names_to_dirs = self.multiplexer.RunPaths()
+    # Create a background context; we may not be in a request.
+    ctx = context.RequestContext()
+    tb_run_names_to_dirs = {
+        run.run_name: os.path.join(self.logdir, run.run_name)
+        for run in self.data_provider.list_runs(ctx, experiment_id='')
+    }
+    plugin_assets = _plugin_assets(
+        self.logdir, list(tb_run_names_to_dirs), PLUGIN_NAME)
 
     # Ensure that we also check the root logdir, even if it isn't a recognized
     # TensorBoard run (i.e. has no tfevents file directly under it), to remain

I’ve tested this by manually patching it into the profile plugin, at
version 2.4.0 of the published Pip package. But the build is broken:
running bazel build :build_pip_package fails with Python deps errors,
and running a plain python setup.py develop fails at runtime with
protobuf linking errors.

This patch should work with both TensorBoard 2.4 and 2.5.

Please feel free to take the above patch and integrate it. You can
attribute it to me, or not, as desired. I’ve signed the CLA and all.

[jimicy]

Fixed by #307

[dbl001]

I applied the patch, but I still don't see anything in my tensor board profile tab.
I am on Big Sur 11.5, running tensorflow-macos 2.5, tensorflow-metal, ...

W0725 12:05:40.806071 123145492811776 application.py:556] path /data/index.js not found, sending 404

My environment:

% patch /Users/davidlaxer/tensorflow-metal/lib/python3.8/site-packages/tensorboard_plugin_profile/profile_plugin.py patch 
patching file /Users/davidlaxer/tensorflow-metal/lib/python3.8/site-packages/tensorboard_plugin_profile/profile_plugin.py
Hunk #4 succeeded at 640 (offset -6 lines).
Hunk #5 succeeded at 695 (offset -6 lines).
(tensorflow-metal) (base) davidlaxer@x86_64-apple-darwin13 VariationalDeepSemanticHashing % pip list
Package                    Version
-------------------------- -------------------
absl-py                    0.12.0
anyio                      3.2.1
appnope                    0.1.2
argon2-cffi                20.1.0
astunparse                 1.6.3
async-generator            1.10
attrs                      21.2.0
Babel                      2.9.1
backcall                   0.2.0
bleach                     3.3.1
cachetools                 4.2.2
certifi                    2021.5.30
cffi                       1.14.6
charset-normalizer         2.0.1
cycler                     0.10.0
Cython                     0.29.24
debugpy                    1.3.0
decorator                  5.0.9
defusedxml                 0.7.1
dill                       0.3.4
dotmap                     1.3.23
entrypoints                0.3
flatbuffers                1.12
future                     0.18.2
gast                       0.4.0
gensim                     4.0.1
google-auth                1.32.1
google-auth-oauthlib       0.4.4
google-pasta               0.2.0
googleapis-common-protos   1.53.0
grpcio                     1.34.1
gviz-api                   1.9.0
h5py                       3.1.0
idna                       3.2
importlib-resources        5.2.0
ipykernel                  6.0.1
ipython                    7.25.0
ipython-genutils           0.2.0
ipywidgets                 7.6.3
jedi                       0.18.0
Jinja2                     3.0.1
json5                      0.9.6
jsonschema                 3.2.0
jupyter-client             6.1.12
jupyter-core               4.7.1
jupyter-server             1.9.0
jupyterlab                 3.0.16
jupyterlab-pygments        0.1.2
jupyterlab-server          2.6.1
jupyterlab-widgets         1.0.0
keras-nightly              2.5.0.dev2021032900
Keras-Preprocessing        1.1.2
kiwisolver                 1.3.1
Markdown                   3.3.4
MarkupSafe                 2.0.1
matplotlib                 3.4.2
matplotlib-inline          0.1.2
mistune                    0.8.4
nbclassic                  0.3.1
nbclient                   0.5.3
nbconvert                  6.1.0
nbformat                   5.1.3
nest-asyncio               1.5.1
notebook                   6.4.0
numpy                      1.19.5
oauthlib                   3.1.1
opt-einsum                 3.3.0
packaging                  21.0
pandas                     1.3.0
pandocfilters              1.4.3
parso                      0.8.2
pexpect                    4.8.0
pickleshare                0.7.5
Pillow                     8.3.1
pip                        21.1.3
prometheus-client          0.11.0
promise                    2.3
prompt-toolkit             3.0.19
protobuf                   3.17.3
ptyprocess                 0.7.0
pyasn1                     0.4.8
pyasn1-modules             0.2.8
pybind11                   2.6.2
pycparser                  2.20
Pygments                   2.9.0
pyparsing                  2.4.7
pyrsistent                 0.18.0
python-dateutil            2.8.2
pytz                       2021.1
pyzmq                      22.1.0
requests                   2.26.0
requests-oauthlib          1.3.0
requests-unixsocket        0.2.0
rsa                        4.7.2
scipy                      1.7.0
Send2Trash                 1.7.1
setuptools                 41.2.0
six                        1.15.0
smart-open                 5.1.0
sniffio                    1.2.0
tensorboard                2.5.0
tensorboard-data-server    0.6.1
tensorboard-plugin-profile 2.4.0
tensorboard-plugin-wit     1.8.0
tensorflow-datasets        4.3.0
tensorflow-estimator       2.5.0
tensorflow-hub             0.12.0
tensorflow-macos           2.5.0
tensorflow-metadata        1.1.0
tensorflow-metal           0.1.1
termcolor                  1.1.0
terminado                  0.10.1
testpath                   0.5.0
tornado                    6.1
tqdm                       4.61.2
traitlets                  5.0.5
typing-extensions          3.7.4.3
urllib3                    1.26.6
wcwidth                    0.2.5
webencodings               0.5.1
websocket-client           1.1.0
Werkzeug                   2.0.1
wheel                      0.36.2
widgetsnbextension         3.5.1
wrapt                      1.12.1
zipp                       3.5.0

% ls -lR logs
total 0
drwxr-xr-x  4 davidlaxer  staff  128 Jul 25 11:19 profile

logs/profile:
total 0
drwxr-xr-x  4 davidlaxer  staff  128 Jul 25 10:36 20210725-103510
drwxr-xr-x  4 davidlaxer  staff  128 Jul 25 11:19 20210725-111913

logs/profile/20210725-103510:
total 0
drwxr-xr-x  5 davidlaxer  staff  160 Jul 25 10:36 train
drwxr-xr-x  3 davidlaxer  staff   96 Jul 25 10:36 validation

logs/profile/20210725-103510/train:
total 104
-rw-r--r--  1 davidlaxer  staff  48655 Jul 25 10:39 events.out.tfevents.1627234568.BlueDiamond.local.11275.165.v2
-rw-r--r--  1 davidlaxer  staff     40 Jul 25 10:36 events.out.tfevents.1627234573.BlueDiamond.local.profile-empty
drwxr-xr-x  3 davidlaxer  staff     96 Jul 25 10:36 plugins

logs/profile/20210725-103510/train/plugins:
total 0
drwxr-xr-x  3 davidlaxer  staff  96 Jul 25 10:36 profile

logs/profile/20210725-103510/train/plugins/profile:
total 0
drwxr-xr-x  9 davidlaxer  staff  288 Jul 25 10:36 2021_07_25_10_36_13

logs/profile/20210725-103510/train/plugins/profile/2021_07_25_10_36_13:
total 80
-rw-r--r--  1 davidlaxer  staff  3014 Jul 25 10:36 BlueDiamond.local.input_pipeline.pb
-rw-r--r--  1 davidlaxer  staff     0 Jul 25 10:36 BlueDiamond.local.kernel_stats.pb
-rw-r--r--  1 davidlaxer  staff   564 Jul 25 10:36 BlueDiamond.local.memory_profile.json.gz
-rw-r--r--  1 davidlaxer  staff  4243 Jul 25 10:36 BlueDiamond.local.overview_page.pb
-rw-r--r--  1 davidlaxer  staff  6171 Jul 25 10:36 BlueDiamond.local.tensorflow_stats.pb
-rw-r--r--  1 davidlaxer  staff  2781 Jul 25 10:36 BlueDiamond.local.trace.json.gz
-rw-r--r--  1 davidlaxer  staff  9517 Jul 25 10:36 BlueDiamond.local.xplane.pb

logs/profile/20210725-103510/validation:
total 8
-rw-r--r--  1 davidlaxer  staff  3636 Jul 25 10:39 events.out.tfevents.1627234584.BlueDiamond.local.11275.1666.v2

logs/profile/20210725-111913:
total 0
drwxr-xr-x  5 davidlaxer  staff  160 Jul 25 11:19 train
drwxr-xr-x  3 davidlaxer  staff   96 Jul 25 11:19 validation

logs/profile/20210725-111913/train:
total 568
-rw-r--r--  1 davidlaxer  staff  285873 Jul 25 11:22 events.out.tfevents.1627237166.BlueDiamond.local.11275.16014.v2
-rw-r--r--  1 davidlaxer  staff      40 Jul 25 11:19 events.out.tfevents.1627237166.BlueDiamond.local.profile-empty
drwxr-xr-x  3 davidlaxer  staff      96 Jul 25 11:19 plugins

logs/profile/20210725-111913/train/plugins:
total 0
drwxr-xr-x  3 davidlaxer  staff  96 Jul 25 11:19 profile

logs/profile/20210725-111913/train/plugins/profile:
total 0
drwxr-xr-x  9 davidlaxer  staff  288 Jul 25 11:19 2021_07_25_11_19_26

logs/profile/20210725-111913/train/plugins/profile/2021_07_25_11_19_26:
total 72
-rw-r--r--  1 davidlaxer  staff  2858 Jul 25 11:19 BlueDiamond.local.input_pipeline.pb
-rw-r--r--  1 davidlaxer  staff     0 Jul 25 11:19 BlueDiamond.local.kernel_stats.pb
-rw-r--r--  1 davidlaxer  staff   562 Jul 25 11:19 BlueDiamond.local.memory_profile.json.gz
-rw-r--r--  1 davidlaxer  staff  4087 Jul 25 11:19 BlueDiamond.local.overview_page.pb
-rw-r--r--  1 davidlaxer  staff  5777 Jul 25 11:19 BlueDiamond.local.tensorflow_stats.pb
-rw-r--r--  1 davidlaxer  staff  2747 Jul 25 11:19 BlueDiamond.local.trace.json.gz
-rw-r--r--  1 davidlaxer  staff  9514 Jul 25 11:19 BlueDiamond.local.xplane.pb

logs/profile/20210725-111913/validation:
total 8
-rw-r--r--  1 davidlaxer  staff  3636 Jul 25 11:22 events.out.tfevents.1627237178.BlueDiamond.local.11275.17515.v2