Parametrize the image used for the Tensorboard viewer (#3235)

* Parametrize the image used for the Tensorboard viewer

Right now, the image is hard-coded and it's not possible to override it.
This is a problem when running Kubeflow Pipelines in a cluster without
access to the internet or public repositories (e.g.
https://github.com/kubeflow/pipelines/issues/3232). This solves the
issue by parametrizing the image name through an environment
variable.

* Fix formatting issues

* Address review comments

backend/src/crd/controller/viewer/reconciler/reconciler.go

@@ -177,12 +177,11 @@ func setPodSpecForTensorboard(view *viewerV1beta1.Viewer, s *corev1.PodSpec) {
 		"tensorboard",
 		fmt.Sprintf("--logdir=%s", view.Spec.TensorboardSpec.LogDir),
 		fmt.Sprintf("--path_prefix=/tensorboard/%s/", view.Name),
-		// This is needed for tf 2.0. We need to optionally add it
-		// when https://github.com/kubeflow/pipelines/issues/2514 is done
-		// "--bind_all",
 	}
 
-	if !strings.HasPrefix(view.Spec.TensorboardSpec.TensorflowImage, `tensorflow/tensorflow:1.`) {
+	tfImageVersion := strings.Split(view.Spec.TensorboardSpec.TensorflowImage, ":")[1]
+	// This is needed for tf 2.0
+	if !strings.HasPrefix(tfImageVersion, `1.`) {
 		c.Args = append(c.Args, "--bind_all")
 	}
 

backend/src/crd/controller/viewer/reconciler/reconciler_test.go

@@ -38,7 +38,7 @@ import (
 
 var viewer *Reconciler
 
-const tensorflowImage = "tensorflow/tensorflow:dummy"
+const tensorflowImage = "potentially_custom_tensorflow:dummy"
 
 func TestMain(m *testing.M) {
 	viewerV1beta1.AddToScheme(scheme.Scheme)

frontend/server/app.ts

@@ -121,7 +121,7 @@ function createUIServer(options: UIConfigs) {
   registerHandler(
     app.post,
     '/apps/tensorboard',
-    getCreateTensorboardHandler(options.viewer.tensorboard.podTemplateSpec),
+    getCreateTensorboardHandler(options.viewer.tensorboard),
   );
 
   /** Pod logs */

frontend/server/configs.ts

@@ -70,6 +70,8 @@ export function loadConfigs(
     ML_PIPELINE_SERVICE_PORT = '3001',
     /** path to viewer:tensorboard pod template spec */
     VIEWER_TENSORBOARD_POD_TEMPLATE_SPEC_PATH,
+    /** Tensorflow image used for tensorboard viewer */
+    VIEWER_TENSORBOARD_TF_IMAGE_NAME = 'tensorflow/tensorflow',
     /** Whether custom visualizations are allowed to be generated by the frontend */
     ALLOW_CUSTOM_VISUALIZATIONS = 'false',
     /** Envoy service will listen to this host */
@@ -146,6 +148,7 @@ export function loadConfigs(
     viewer: {
       tensorboard: {
         podTemplateSpec: loadJSON<object>(VIEWER_TENSORBOARD_POD_TEMPLATE_SPEC_PATH),
+        tfImageName: VIEWER_TENSORBOARD_TF_IMAGE_NAME,
       },
     },
     visualizations: {
@@ -180,10 +183,12 @@ export interface PipelineConfigs {
   host: string;
   port: string | number;
 }
+export interface ViewerTensorboardConfig {
+  podTemplateSpec?: object;
+  tfImageName: string;
+}
 export interface ViewerConfigs {
-  tensorboard: {
-    podTemplateSpec?: object;
-  };
+  tensorboard: ViewerTensorboardConfig;
 }
 export interface VisualizationsConfigs {
   allowCustomVisualizations: boolean;

frontend/server/handlers/tensorboard.ts

@@ -13,6 +13,7 @@
 // limitations under the License.
 import { Handler } from 'express';
 import * as k8sHelper from '../k8s-helper';
+import { ViewerTensorboardConfig } from '../configs';
 
 /**
  * A handler which retrieve the endpoint for a tensorboard instance. The
@@ -44,10 +45,9 @@ export const getTensorboardHandler: Handler = async (req, res) => {
  * The handler expects the following query strings in the request:
  * - `logdir`
  * - `tfversion`
- * @param podTemplateSpec Custom pod template specification to be applied on the
- * tensorboard pod.
+ * @param tensorboardConfig The configuration for Tensorboard.
  */
-export function getCreateTensorboardHandler(podTemplateSpec?: object): Handler {
+export function getCreateTensorboardHandler(tensorboardConfig: ViewerTensorboardConfig): Handler {
   return async (req, res) => {
     if (!k8sHelper.isInCluster) {
       res.status(500).send('Cannot talk to Kubernetes master');
@@ -68,7 +68,12 @@ export function getCreateTensorboardHandler(podTemplateSpec?: object): Handler {
     const tfversion = decodeURIComponent(req.query.tfversion);
 
     try {
-      await k8sHelper.newTensorboardInstance(logdir, tfversion, podTemplateSpec);
+      await k8sHelper.newTensorboardInstance(
+        logdir,
+        tensorboardConfig.tfImageName,
+        tfversion,
+        tensorboardConfig.podTemplateSpec,
+      );
       const tensorboardAddress = await k8sHelper.waitForTensorboardInstance(logdir, 60 * 1000);
       res.send(tensorboardAddress);
     } catch (err) {

frontend/server/k8s-helper.ts

@@ -73,6 +73,7 @@ function getNameOfViewerResource(logdir: string): string {
  */
 export async function newTensorboardInstance(
   logdir: string,
+  tfImageName: string,
   tfversion: string,
   podTemplateSpec: object = defaultPodTemplateSpec,
 ): Promise<void> {
@@ -101,7 +102,7 @@ export async function newTensorboardInstance(
       podTemplateSpec,
       tensorboardSpec: {
         logDir: logdir,
-        tensorflowImage: 'tensorflow/tensorflow:' + tfversion,
+        tensorflowImage: tfImageName + ':' + tfversion,
       },
       type: 'tensorboard',
     },
@@ -146,10 +147,9 @@ export async function getTensorboardInstance(
           viewer.body.spec.type === 'tensorboard'
         ) {
           const address = `http://${viewer.body.metadata.name}-service.${namespace}.svc.cluster.local:80/tensorboard/${viewer.body.metadata.name}/`;
-          const version = viewer.body.spec.tensorboardSpec.tensorflowImage
-            ? viewer.body.spec.tensorboardSpec.tensorflowImage.replace('tensorflow/tensorflow:', '')
-            : '';
-          return { podAddress: address, tfVersion: version };
+          const tfImageParts = viewer.body.spec.tensorboardSpec.tensorflowImage.split(':', 1);
+          const tfVersion = tfImageParts.length == 2 ? tfImageParts[1] : '';
+          return { podAddress: address, tfVersion: tfVersion };
         } else {
           return { podAddress: '', tfVersion: '' };
         }