[UI] Improve TFX artifact visualization speed (#3712)

* Improve TFX artifact visualization speed * Update OutputArtifactLoader.ts * Fix loading * Add nodeId back as an identifier
kubeflow · May 8, 2020 · eab5d31 · eab5d31
1 parent 54c490c
commit eab5d31
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 35 deletions.
diff --git a/frontend/src/lib/OutputArtifactLoader.ts b/frontend/src/lib/OutputArtifactLoader.ts
@@ -239,12 +239,12 @@ export class OutputArtifactLoader {
    * @returns config array, also returns empty array when expected erros happen
    */
   public static async buildTFXArtifactViewerConfig({
-    argoPodName,
     namespace,
+    execution,
     reportProgress = () => null,
   }: {
     namespace: string;
-    argoPodName: string;
+    execution: Execution;
     reportProgress: (progress: number) => void;
   }): Promise<HTMLViewerConfig[]> {
     // Error handling assumptions:
@@ -260,33 +260,15 @@ export class OutputArtifactLoader {
     // Since artifact types don't change per run, this can be optimized further so
     // that we don't fetch them on every page load.
     reportProgress(10);
-    const artifactTypes = await getArtifactTypes();
-    if (artifactTypes.length === 0) {
-      // There are no artifact types data.
+    const [artifactTypes, artifacts] = await Promise.all([
+      getArtifactTypes(),
+      getOutputArtifactsInExecution(execution),
+    ]);
+    if (artifactTypes.length === 0 || artifacts.length === 0) {
+      // There are no artifact types data or no artifacts.
       return [];
     }
-    reportProgress(20);
-
-    const context = await getMlmdContext(argoPodName);
-    if (!context) {
-      // Failed finding corresponding MLMD context.
-      return [];
-    }
-    reportProgress(40);
-
-    const execution = await getExecutionInContextWithPodName(argoPodName, context);
-    if (!execution) {
-      // Failed finding corresponding MLMD execution.
-      return [];
-    }
-    reportProgress(60);
-
-    const artifacts = await getOutputArtifactsInExecution(execution);
-    if (artifacts.length === 0) {
-      // There are no artifacts in this execution.
-      return [];
-    }
-    reportProgress(80);
+    reportProgress(70);
 
     // TODO: Visualize non-TFDV artifacts, such as ModelEvaluation using TFMA
     let viewers: Array<Promise<HTMLViewerConfig>> = [];

diff --git a/frontend/src/pages/RunDetails.tsx b/frontend/src/pages/RunDetails.tsx
@@ -240,7 +240,7 @@ class RunDetails extends Page<RunDetailsInternalProps, RunDetailsState> {
       mlmdExecutions,
     } = this.state;
     const { projectId, clusterName } = this.props.gkeMetadata;
-    const selectedNodeId = selectedNodeDetails ? selectedNodeDetails.id : '';
+    const selectedNodeId = selectedNodeDetails?.id || '';
     const namespace = workflow?.metadata?.namespace;
     let stackdriverK8sLogsUrl = '';
     if (projectId && clusterName && selectedNodeDetails && selectedNodeDetails.id) {
@@ -329,7 +329,8 @@ class RunDetails extends Page<RunDetailsInternalProps, RunDetailsState> {
                                   this.state.selectedNodeDetails &&
                                   this.state.workflow && (
                                     <ArtifactsTabContent
-                                      nodeId={this.state.selectedNodeDetails.id}
+                                      execution={selectedExecution}
+                                      nodeId={selectedNodeId}
                                       nodeStatus={
                                         this.state.workflow && this.state.workflow.status
                                           ? this.state.workflow.status.nodes[
@@ -997,6 +998,7 @@ const COMPLETED_NODE_PHASES: ArgoNodePhase[] = ['Succeeded', 'Failed', 'Error'];
  */
 const ArtifactsTabContent: React.FC<{
   visualizationCreatorConfig: VisualizationCreatorConfig;
+  execution?: Execution;
   nodeId: string;
   nodeStatus?: NodeStatus;
   generatedVisualizations: GeneratedVisualization[];
@@ -1005,6 +1007,7 @@ const ArtifactsTabContent: React.FC<{
 }> = ({
   visualizationCreatorConfig,
   generatedVisualizations,
+  execution,
   nodeId,
   nodeStatus,
   namespace,
@@ -1047,11 +1050,15 @@ const ArtifactsTabContent: React.FC<{
       // Load the viewer configurations from the output paths
       const builtConfigs = (
         await Promise.all([
-          OutputArtifactLoader.buildTFXArtifactViewerConfig({
-            argoPodName: nodeId,
-            reportProgress,
-            namespace: namespace || '',
-          }).catch(reportErrorAndReturnEmpty),
+          ...(!execution
+            ? []
+            : [
+                OutputArtifactLoader.buildTFXArtifactViewerConfig({
+                  reportProgress,
+                  execution,
+                  namespace: namespace || '',
+                }).catch(reportErrorAndReturnEmpty),
+              ]),
           ...outputPaths.map(path =>
             OutputArtifactLoader.load(path, namespace).catch(reportErrorAndReturnEmpty),
           ),
@@ -1077,7 +1084,7 @@ const ArtifactsTabContent: React.FC<{
     // nodeStatus object instance will keep changing after new requests to get
     // workflow status.
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [nodeId, nodeCompleted, onError, namespace]);
+  }, [nodeId, execution?.getId(), nodeCompleted, onError, namespace]);
 
   return (
     <div className={commonCss.page}>