From 1ad688e4213cc8de3607bfc01b5d050a13764019 Mon Sep 17 00:00:00 2001
From: Thijs Walcarius <thijs.walcarius@ugent.be>
Date: Thu, 1 Aug 2024 10:43:35 +0200
Subject: [PATCH] Improve handling of tensorboard (#2126)

* Improve handling of tensorboard

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Improve style

* Move bash file to cuda-subdir

* Update tensorflow-notebook info

* Fix redirect link

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Ayaz Salikhov <mathbunnyru@users.noreply.github.com>
---
 docs/using/selecting.md                                |  3 ++-
 images/tensorflow-notebook/Dockerfile                  |  6 +++++-
 .../cuda/20tensorboard-proxy-env.sh                    | 10 ++++++++++
 images/tensorflow-notebook/cuda/Dockerfile             |  6 +++++-
 4 files changed, 22 insertions(+), 3 deletions(-)
 create mode 100644 images/tensorflow-notebook/cuda/20tensorboard-proxy-env.sh

diff --git a/docs/using/selecting.md b/docs/using/selecting.md
index fb2ec53bd9..d4cd71dc54 100644
--- a/docs/using/selecting.md
+++ b/docs/using/selecting.md
@@ -193,7 +193,8 @@ It contains:
 `jupyter/tensorflow-notebook` includes popular Python deep learning libraries.
 
 - Everything in `jupyter/scipy-notebook` and its ancestor images
-- [tensorflow](https://www.tensorflow.org/) machine learning library
+- [TensorFlow](https://www.tensorflow.org/) machine learning library
+- [Jupyter Server Proxy](https://jupyter-server-proxy.readthedocs.io/en/latest/) to support [TensorBoard](https://www.tensorflow.org/tensorboard)
 
 ### jupyter/pytorch-notebook
 
diff --git a/images/tensorflow-notebook/Dockerfile b/images/tensorflow-notebook/Dockerfile
index be366d4959..01ace0f1a1 100644
--- a/images/tensorflow-notebook/Dockerfile
+++ b/images/tensorflow-notebook/Dockerfile
@@ -13,6 +13,10 @@ SHELL ["/bin/bash", "-o", "pipefail", "-c"]
 
 # Install tensorflow with pip, on x86_64 tensorflow-cpu
 RUN [[ $(uname -m) = x86_64 ]] && TF_POSTFIX="-cpu" || TF_POSTFIX="" && \
-    pip install --no-cache-dir "tensorflow${TF_POSTFIX}" && \
+    pip install --no-cache-dir \
+    "tensorflow${TF_POSTFIX}" \
+    "jupyter-server-proxy" && \
     fix-permissions "${CONDA_DIR}" && \
     fix-permissions "/home/${NB_USER}"
+
+COPY --chown="${NB_UID}:${NB_GID}" cuda/20tensorboard-proxy-env.sh /usr/local/bin/before-notebook.d/
diff --git a/images/tensorflow-notebook/cuda/20tensorboard-proxy-env.sh b/images/tensorflow-notebook/cuda/20tensorboard-proxy-env.sh
new file mode 100644
index 0000000000..f31c7488dd
--- /dev/null
+++ b/images/tensorflow-notebook/cuda/20tensorboard-proxy-env.sh
@@ -0,0 +1,10 @@
+#!/bin/bash
+# Copyright (c) Jupyter Development Team.
+# Distributed under the terms of the Modified BSD License.
+
+set -e
+
+# Initialize the TENSORBOARD_PROXY_URL with the appropriate path
+# to use jupyter-server-proxy.
+
+export TENSORBOARD_PROXY_URL="${JUPYTERHUB_SERVICE_PREFIX:-/}proxy/%PORT%/"
diff --git a/images/tensorflow-notebook/cuda/Dockerfile b/images/tensorflow-notebook/cuda/Dockerfile
index a21d7e941c..8854276184 100644
--- a/images/tensorflow-notebook/cuda/Dockerfile
+++ b/images/tensorflow-notebook/cuda/Dockerfile
@@ -12,10 +12,14 @@ LABEL maintainer="Jupyter Project <jupyter@googlegroups.com>"
 SHELL ["/bin/bash", "-o", "pipefail", "-c"]
 
 # Install TensorFlow, CUDA and cuDNN with pip
-RUN pip install --no-cache-dir "tensorflow[and-cuda]" && \
+RUN pip install --no-cache-dir \
+    "tensorflow[and-cuda]" \
+    "jupyter-server-proxy" && \
     fix-permissions "${CONDA_DIR}" && \
     fix-permissions "/home/${NB_USER}"
 
+COPY --chown="${NB_UID}:${NB_GID}" 20tensorboard-proxy-env.sh /usr/local/bin/before-notebook.d/
+
 # workaround for https://github.com/tensorflow/tensorflow/issues/63362
 RUN mkdir -p "${CONDA_DIR}/etc/conda/activate.d/" && \
     fix-permissions "${CONDA_DIR}"