Retrieve pod logs from argo archive

kubeflow · Sep 10, 2019 · 1b7fb57 · 1b7fb57
1 parent 30b97dd
commit 1b7fb57
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 0 deletions.
diff --git a/frontend/server/k8s-helper.ts b/frontend/server/k8s-helper.ts
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 // @ts-ignore
+import {Client as MinioClient} from 'minio';
 import {Core_v1Api, Custom_objectsApi, KubeConfig} from '@kubernetes/client-node';
 import * as crypto from 'crypto-js';
 import * as fs from 'fs';
@@ -155,4 +156,19 @@ export function getPodLogs(podName: string): Promise<string> {
       (response: any) => (response && response.body) ? response.body.toString() : '',
       (error: any) => {throw new Error(JSON.stringify(error.body));}
     );
+}
+
+/** Returns a function to retrieve pod logs archived by argo. */
+export function getPodLogsFromArtifactoryHelper(client: MinioClient, bucketName: string, prefix: string) {
+  return async function(podName: string) {
+    let workflowName = workflowNameFromPodName(podName);
+    return client.getObject(bucketName, `${prefix}/${workflowName}/${podName}/main.log`)
+  }
+}
+
+/** Infers workflow name from pod name. */
+function workflowNameFromPodName(podName: string) {
+  let chunks = podName.split("-");
+  chunks.pop();
+  return chunks.join("-");
 }
diff --git a/frontend/server/server.ts b/frontend/server/server.ts
@@ -61,6 +61,14 @@ const {
   METADATA_ENVOY_SERVICE_SERVICE_HOST = 'localhost',
   /** Envoy service will listen to this port */
   METADATA_ENVOY_SERVICE_SERVICE_PORT = '9090',
+  /** Is Argo log archive enabled? */
+  ARGO_ARCHIVE_LOGS = "false",
+  /** Use minio or s3 client to retrieve archives. */
+  ARGO_ARCHIVE_ARTIFACTORY = 'minio',
+  /** Bucket to retrive logs from */
+  ARGO_ARCHIVE_BUCKETNAME = 'mlpipeline',
+  /** Prefix to logs. */
+  ARGO_ARCHIVE_PREFIX = 'logs',
 } = process.env;
 
 /** construct minio endpoint from host and namespace (optional) */
@@ -88,6 +96,13 @@ const s3Client = new MinioClient({
 /** pod template spec to use for viewer crd */
 const podTemplateSpec = loadJSON(VIEWER_TENSORBOARD_POD_TEMPLATE_SPEC_PATH, k8sHelper.defaultPodTemplateSpec)
 
+/** helper function to retrieve pod logs from argo artifactory. */
+const getPodLogsFromArtifactory = _as_bool(ARGO_ARCHIVE_LOGS) ? k8sHelper.getPodLogsFromArtifactoryHelper(
+  ARGO_ARCHIVE_ARTIFACTORY==='minio' ? minioClient : s3Client,
+  ARGO_ARCHIVE_BUCKETNAME,
+  ARGO_ARCHIVE_PREFIX
+) : undefined;
+
 const app = express() as Application;
 
 app.use(function (req, _, next) {
@@ -316,6 +331,18 @@ const logsHandler = async (req, res) => {
   try {
     res.send(await k8sHelper.getPodLogs(podName));
   } catch (err) {
+    // try argo archive if cannot retrieve from k8s pod
+    if (!!getPodLogsFromArtifactory) {
+      try {
+        let stream = await getPodLogsFromArtifactory(podName);
+        stream.pipe(res);
+        return;
+      } catch (archiveErr) {
+        res.status(500).send(
+          `Could not get main container logs: ${err}\nCould not get main container logs from archive: ${archiveErr}`);
+        return;
+      }
+    }
     res.status(500).send('Could not get main container logs: ' + err);
   }
 };