From af456bcc61f00fad4cb44bca09a088551040d243 Mon Sep 17 00:00:00 2001
From: Eterna2 <eterna2@hotmail.com>
Date: Fri, 23 Aug 2019 02:08:33 +0800
Subject: [PATCH] [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 | 50 ++++++++++++++++++-----------------
 frontend/server/server.ts     | 11 ++++++--
 frontend/server/utils.ts      | 10 +++++++
 3 files changed, 45 insertions(+), 26 deletions(-)

diff --git a/frontend/server/k8s-helper.ts b/frontend/server/k8s-helper.ts
index af8a791017c..7260c4cb34c 100644
--- a/frontend/server/k8s-helper.ts
+++ b/frontend/server/k8s-helper.ts
@@ -30,6 +30,30 @@ const viewerGroup = 'kubeflow.org';
 const viewerVersion = 'v1beta1';
 const viewerPlural = 'viewers';
 
+export const defaultPodTemplateSpec = {
+  spec: {
+    containers: [{
+      env: [{
+        name: "GOOGLE_APPLICATION_CREDENTIALS",
+        value: "/secret/gcp-credentials/user-gcp-sa.json"
+      }],
+      volumeMounts: [{
+        name: "gcp-credentials",
+        mountPath: "/secret/gcp-credentials/user-gcp-sa.json",
+        readOnly: true
+      }]
+    }],
+    volumes: [{
+      name: "gcp-credentials",
+      volumeSource: {
+        secret: {
+          secretName: "user-gcp-sa"
+        }
+      }
+    }]
+  }
+}
+
 export const isInCluster = fs.existsSync(namespaceFilePath);
 
 if (isInCluster) {
@@ -49,7 +73,7 @@ function getNameOfViewerResource(logdir: string): string {
  * Create Tensorboard instance via CRD with the given logdir if there is no
  * existing Tensorboard instance.
  */
-export async function newTensorboardInstance(logdir: string): Promise<void> {
+export async function newTensorboardInstance(logdir: string, podTemplateSpec: Object = defaultPodTemplateSpec): Promise<void> {
   if (!k8sV1CustomObjectClient) {
     throw new Error('Cannot access kubernetes Custom Object API');
   }
@@ -70,29 +94,7 @@ export async function newTensorboardInstance(logdir: string): Promise<void> {
       tensorboardSpec: {
         logDir: logdir,
       },
-      podTemplateSpec: {
-        spec: {
-          containers: [{
-            env: [{
-              name: "GOOGLE_APPLICATION_CREDENTIALS",
-              value: "/secret/gcp-credentials/user-gcp-sa.json"
-            }],
-            volumeMounts: [{
-              name: "gcp-credentials",
-              mountPath: "/secret/gcp-credentials/user-gcp-sa.json",
-              readOnly: true
-            }]
-          }],
-          volumes: [{
-            name: "gcp-credentials",
-            volumeSource: {
-              secret: {
-                secretName: "user-gcp-sa"
-              }
-            }
-          }]
-        }
-      }
+      podTemplateSpec
     }
   };
   await k8sV1CustomObjectClient.createNamespacedCustomObject(viewerGroup,
diff --git a/frontend/server/server.ts b/frontend/server/server.ts
index 4bb4f6784fc..4304e4aefc5 100644
--- a/frontend/server/server.ts
+++ b/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) {
diff --git a/frontend/server/utils.ts b/frontend/server/utils.ts
index 4233fa0ad41..d6bbdb37984 100644
--- a/frontend/server/utils.ts
+++ b/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;
+  }
+}
\ No newline at end of file