[front-end-server] Allow viewer:tensorboard podTemplateSpec to be customizable (#1906)

* Allow front-end server to provide custom viewer podTemplateSpec via path/configmap

* Fix JSON.parse input to string

frontend/server/k8s-helper.ts

@@ -30,47 +30,7 @@ const viewerGroup = 'kubeflow.org';
 const viewerVersion = 'v1beta1';
 const viewerPlural = 'viewers';
 
-export const isInCluster = fs.existsSync(namespaceFilePath);
+export const defaultPodTemplateSpec = {
-
-if (isInCluster) {
-  namespace = fs.readFileSync(namespaceFilePath, 'utf-8');
-  const kc = new KubeConfig();
-  kc.loadFromDefault();
-  k8sV1Client = kc.makeApiClient(Core_v1Api);
-  k8sV1CustomObjectClient = kc.makeApiClient(Custom_objectsApi);
-}
-
-function getNameOfViewerResource(logdir: string): string {
-  // TODO: find some hash function with shorter resulting message.
-  return 'viewer-' + crypto.SHA1(logdir);
-}
-
-/**
- * Create Tensorboard instance via CRD with the given logdir if there is no
- * existing Tensorboard instance.
- */
-export async function newTensorboardInstance(logdir: string): Promise<void> {
-  if (!k8sV1CustomObjectClient) {
-    throw new Error('Cannot access kubernetes Custom Object API');
-  }
-  const currentPod = await getTensorboardInstance(logdir);
-  if (currentPod) {
-    return;
-  }
-
-  const body = {
-    apiVersion: viewerGroup + '/' + viewerVersion,
-    kind: 'Viewer',
-    metadata: {
-      name: getNameOfViewerResource(logdir),
-      namespace: namespace,
-    },
-    spec: {
-      type: 'tensorboard',
-      tensorboardSpec: {
-        logDir: logdir,
-      },
-      podTemplateSpec: {
   spec: {
     containers: [{
       env: [{
@@ -93,6 +53,48 @@ export async function newTensorboardInstance(logdir: string): Promise<void> {
     }]
   }
 }
+
+export const isInCluster = fs.existsSync(namespaceFilePath);
+
+if (isInCluster) {
+  namespace = fs.readFileSync(namespaceFilePath, 'utf-8');
+  const kc = new KubeConfig();
+  kc.loadFromDefault();
+  k8sV1Client = kc.makeApiClient(Core_v1Api);
+  k8sV1CustomObjectClient = kc.makeApiClient(Custom_objectsApi);
+}
+
+function getNameOfViewerResource(logdir: string): string {
+  // TODO: find some hash function with shorter resulting message.
+  return 'viewer-' + crypto.SHA1(logdir);
+}
+
+/**
+ * Create Tensorboard instance via CRD with the given logdir if there is no
+ * existing Tensorboard instance.
+ */
+export async function newTensorboardInstance(logdir: string, podTemplateSpec: Object = defaultPodTemplateSpec): Promise<void> {
+  if (!k8sV1CustomObjectClient) {
+    throw new Error('Cannot access kubernetes Custom Object API');
+  }
+  const currentPod = await getTensorboardInstance(logdir);
+  if (currentPod) {
+    return;
+  }
+
+  const body = {
+    apiVersion: viewerGroup + '/' + viewerVersion,
+    kind: 'Viewer',
+    metadata: {
+      name: getNameOfViewerResource(logdir),
+      namespace: namespace,
+    },
+    spec: {
+      type: 'tensorboard',
+      tensorboardSpec: {
+        logDir: logdir,
+      },
+      podTemplateSpec
     }
   };
   await k8sV1CustomObjectClient.createNamespacedCustomObject(viewerGroup,

frontend/server/server.ts

@@ -27,6 +27,8 @@ import proxyMiddleware from './proxy-middleware';
 import { Storage } from '@google-cloud/storage';
 import {Stream} from 'stream';
 
+import {loadJSON} from './utils';
+
 const BASEPATH = '/pipeline';
 
 /** All configurable environment variables can be found here. */
@@ -50,7 +52,9 @@ const {
   /** API service will listen to this host */
   ML_PIPELINE_SERVICE_HOST = 'localhost',
   /** API service will listen to this port */
-  ML_PIPELINE_SERVICE_PORT = '3001'
+  ML_PIPELINE_SERVICE_PORT = '3001',
+  /** path to viewer:tensorboard pod template spec */
+  VIEWER_TENSORBOARD_POD_TEMPLATE_SPEC_PATH
 } = process.env;
 
 /** construct minio endpoint from host and namespace (optional) */
@@ -75,6 +79,9 @@ const s3Client = new MinioClient({
   secretKey: AWS_SECRET_ACCESS_KEY,
 } as any);
 
+/** pod template spec to use for viewer crd */
+const podTemplateSpec = loadJSON(VIEWER_TENSORBOARD_POD_TEMPLATE_SPEC_PATH, k8sHelper.defaultPodTemplateSpec)
+
 const app = express() as Application;
 
 app.use(function (req, _, next) {
@@ -278,7 +285,7 @@ const createTensorboardHandler = async (req, res) => {
   }
 
   try {
-    await k8sHelper.newTensorboardInstance(logdir);
+    await k8sHelper.newTensorboardInstance(logdir, podTemplateSpec);
     const tensorboardAddress = await k8sHelper.waitForTensorboardInstance(logdir, 60 * 1000);
     res.send(tensorboardAddress);
   } catch (err) {

frontend/server/utils.ts

@@ -11,6 +11,7 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
+import {readFileSync} from 'fs';
 
 export function equalArrays(a1: any[], a2: any[]): boolean {
   if (!Array.isArray(a1) || !Array.isArray(a2) || a1.length !== a2.length) {
@@ -32,3 +33,12 @@ export function generateRandomString(length: number): string {
   }
   return str;
 }
+
+export function loadJSON(filepath: string, defaultValue: Object = {}): Object {
+  if (!filepath) return defaultValue;
+  try {
+    return JSON.parse(readFileSync(filepath, "utf-8"))
+  } catch (error) {
+    return defaultValue;
+  }
+}