EOD Push

Most work on feature is completed. Currently attempting to fix bug where UI does not properly render visualizations for selected component.

frontend/src/components/viewers/VisualizationCreator.tsx

@@ -151,7 +151,15 @@ class VisualizationCreator extends Viewer<VisualizationCreatorProps, Visualizati
       <BusyButton title='Generate Visualization' busy={isBusy} disabled={!canGenerate}
         onClick={() => {
           if (onGenerate && selectedType) {
-            onGenerate(_arguments, source, selectedType);
+            const specifiedArguments: any = JSON.parse(_arguments || '{}');
+            if (selectedType === ApiVisualizationType.CUSTOM) {
+              specifiedArguments.code = code.split('\n');
+            }
+            onGenerate(
+              JSON.stringify(specifiedArguments),
+              source,
+              selectedType
+            );
           }
         }} />
     </div>;

frontend/src/lib/Apis.ts

@@ -18,7 +18,9 @@ import { JobServiceApi } from '../apis/job';
 import { RunServiceApi } from '../apis/run';
 import { PipelineServiceApi, ApiPipeline } from '../apis/pipeline';
 import { StoragePath } from './WorkflowParser';
-import { VisualizationServiceApi } from '../apis/visualization';
+import { VisualizationServiceApi, ApiVisualization } from '../apis/visualization';
+import { HTMLViewerConfig } from 'src/components/viewers/HTMLViewer';
+import { PlotType } from '../components/viewers/Viewer';
 
 const v1beta1Prefix = 'apis/v1beta1';
 
@@ -52,6 +54,25 @@ export class Apis {
     return customVisualizationsAllowed;
   }
 
+  public static async buildPythonVisualizationConfig(visualizationData: ApiVisualization): Promise<HTMLViewerConfig> {
+    const visualization = await Apis.visualizationServiceApi.createVisualization(visualizationData);
+    if (visualization.html) {
+      const htmlContent = visualization.html
+        // Fixes issue with TFX components (and other iframe based
+        // visualizations), where the method in which javascript interacts
+        // with embedded iframes is not allowed when embedded in an additional
+        // iframe. This is resolved by setting the srcdoc value rather that
+        // manipulating the document directly.
+        .replace('contentWindow.document.write', 'srcdoc=');
+      return {
+        htmlContent,
+        type: PlotType.WEB_APP,
+      } as HTMLViewerConfig;
+    } else {
+      throw new Error('Unable to generate visualization!');
+    }
+  }
+
   /**
    * Get pod logs
    */

frontend/src/pages/RunDetails.test.tsx

@@ -40,12 +40,13 @@ describe('RunDetails', () => {
   const getClusterNameSpy = jest.spyOn(Apis, 'getClusterName');
   const getRunSpy = jest.spyOn(Apis.runServiceApi, 'getRun');
   const getExperimentSpy = jest.spyOn(Apis.experimentServiceApi, 'getExperiment');
+  const isCustomVisualizationsAllowedSpy = jest.spyOn(Apis, 'areCustomVisualizationsAllowed');
   const getPodLogsSpy = jest.spyOn(Apis, 'getPodLogs');
   const pathsParser = jest.spyOn(WorkflowParser, 'loadNodeOutputPaths');
   const pathsWithStepsParser = jest.spyOn(WorkflowParser, 'loadAllOutputPathsWithStepNames');
   const loaderSpy = jest.spyOn(OutputArtifactLoader, 'load');
   // We mock this because it uses toLocaleDateString, which causes mismatches between local and CI
-  // test enviroments
+  // test environments
   const formatDateStringSpy = jest.spyOn(Utils, 'formatDateString');
 
   let testRun: ApiRunDetail = {};
@@ -93,6 +94,7 @@ describe('RunDetails', () => {
     getClusterNameSpy.mockImplementation(() => Promise.resolve('some-cluster'));
     getRunSpy.mockImplementation(() => Promise.resolve(testRun));
     getExperimentSpy.mockImplementation(() => Promise.resolve({ id: 'some-experiment-id', name: 'some experiment' }));
+    isCustomVisualizationsAllowedSpy.mockImplementation(() => Promise.resolve(false));
     getPodLogsSpy.mockImplementation(() => 'test logs');
     pathsParser.mockImplementation(() => []);
     pathsWithStepsParser.mockImplementation(() => []);

frontend/src/pages/RunDetails.tsx

@@ -39,14 +39,17 @@ import { OutputArtifactLoader } from '../lib/OutputArtifactLoader';
 import { Page } from './Page';
 import { RoutePage, RouteParams } from '../components/Router';
 import { ToolbarProps } from '../components/Toolbar';
-import { ViewerConfig } from '../components/viewers/Viewer';
+import { ViewerConfig, PlotType } from '../components/viewers/Viewer';
 import { Workflow } from '../../third_party/argo-ui/argo_template';
 import { classes, stylesheet } from 'typestyle';
 import { commonCss, padding, color, fonts, fontsize } from '../Css';
 import { componentMap } from '../components/viewers/ViewerContainer';
 import { flatten } from 'lodash';
 import { formatDateString, getRunDurationFromWorkflow, logger, errorToMessage } from '../lib/Utils';
 import { statusToIcon } from './Status';
+import VisualizationCreator, { VisualizationCreatorConfig } from '../components/viewers/VisualizationCreator';
+import { ApiVisualization, ApiVisualizationType } from '../apis/visualization';
+import { HTMLViewerConfig } from '../components/viewers/HTMLViewer';
 
 enum SidePaneTab {
   ARTIFACTS,
@@ -72,9 +75,17 @@ interface AnnotatedConfig {
   stepName: string;
 }
 
+interface GeneratedVisualization {
+  config: HTMLViewerConfig;
+  nodeId: string;
+}
+
 interface RunDetailsState {
   allArtifactConfigs: AnnotatedConfig[];
+  allowCustomVisualizations: boolean;
   experiment?: ApiExperiment;
+  generatedVisualizations: GeneratedVisualization[];
+  isGeneratingVisualization: boolean;
   legacyStackdriverUrl: string;
   logsBannerAdditionalInfo: string;
   logsBannerMessage: string;
@@ -135,6 +146,9 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
 
     this.state = {
       allArtifactConfigs: [],
+      allowCustomVisualizations: false,
+      generatedVisualizations: [],
+      isGeneratingVisualization: false,
       legacyStackdriverUrl: '',
       logsBannerAdditionalInfo: '',
       logsBannerMessage: '',
@@ -170,7 +184,9 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
   public render(): JSX.Element {
     const {
       allArtifactConfigs,
+      allowCustomVisualizations,
       graph,
+      isGeneratingVisualization,
       legacyStackdriverUrl,
       runFinished,
       runMetadata,
@@ -185,6 +201,14 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
     const workflowParameters = WorkflowParser.getParameters(workflow);
     const nodeInputOutputParams = WorkflowParser.getNodeInputOutputParams(workflow, selectedNodeId);
     const hasMetrics = runMetadata && runMetadata.metrics && runMetadata.metrics.length > 0;
+    const visualizationCreatorConfig: VisualizationCreatorConfig = {
+      allowCustomVisualizations,
+      isBusy: isGeneratingVisualization,
+      onGenerate: (visualizationArguments: string, source: string, type: ApiVisualizationType) => {
+        this._onGenerate(visualizationArguments, source, type);
+      },
+      type: PlotType.VISUALIZATION_CREATOR,
+    };
 
     return (
       <div className={classes(commonCss.page, padding(20, 't'))}>
@@ -215,7 +239,14 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
                         <div className={commonCss.page}>
                           {sidepanelSelectedTab === SidePaneTab.ARTIFACTS && (
                             <div className={commonCss.page}>
-                              {(selectedNodeDetails.viewerConfigs || []).map((config, i) => {
+                              <div className={padding(20, 'lrt')}>
+                                <PlotCard
+                                  configs={[visualizationCreatorConfig]}
+                                  title={VisualizationCreator.prototype.getDisplayName()}
+                                  maxDimension={500} />
+                                <Hr />
+                              </div>
+                              {this._getSelectedNodeConfigs().map((config, i) => {
                                 const title = componentMap[config.type].prototype.getDisplayName();
                                 return (
                                   <div key={i} className={padding(20, 'lrt')}>
@@ -395,6 +426,13 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
     this.clearBanner();
     const runId = this.props.match.params[RouteParams.runId];
 
+    try {
+      const allowCustomVisualizations = await Apis.areCustomVisualizationsAllowed();
+      this.setState({ allowCustomVisualizations });
+    } catch (err) {
+      // Ignore error
+    }
+
     try {
       const runDetail = await Apis.runServiceApi.getRun(runId);
 
@@ -641,6 +679,56 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
       this.setStateSafe({ sidepanelBusy: false });
     }
   }
+
+  private _getSelectedNodeConfigs(): ViewerConfig[] {
+    const { generatedVisualizations, selectedNodeDetails } = this.state;
+    let selectedNodeConfigs: ViewerConfig[] = [];
+    if (selectedNodeDetails !== null) {
+      const nodeConfigs = selectedNodeDetails.viewerConfigs || [];
+      const generatedConfigs = generatedVisualizations
+        .filter(visualization => visualization.nodeId === selectedNodeDetails.id)
+        .map(visualization => visualization.config);
+        selectedNodeConfigs = nodeConfigs.concat(generatedConfigs);
+    }
+    return selectedNodeConfigs;
+  }
+
+  private async _onGenerate(visualizationArguments: string, source: string, type: ApiVisualizationType): Promise<void> {
+    const { selectedNodeDetails } = this.state;
+    const nodeId = selectedNodeDetails ? selectedNodeDetails.id : '';
+    if (nodeId.length === 0) {
+      this.showPageError('Unable to generate visualization, no component selected.');
+      return;
+    }
+
+    if (visualizationArguments.length) {
+      try {
+        JSON.parse(visualizationArguments);
+      } catch (err) {
+        this.showPageError('Unable to generate visualization, invalid JSON provided.', err);
+        return;
+      }
+    }
+    this.setState({ isGeneratingVisualization: true });
+    const visualizationData: ApiVisualization = {
+      arguments: visualizationArguments,
+      source,
+      type,
+    };
+    try {
+      const config = await Apis.buildPythonVisualizationConfig(visualizationData);
+      const { generatedVisualizations } = this.state;
+      generatedVisualizations.push({
+        config,
+        nodeId,
+      });
+      this.setState({ generatedVisualizations });
+    } catch (err) {
+      this.showPageError('Unable to generate visualization, an unexpected error was encountered.', err);
+    } finally {
+      this.setState({ isGeneratingVisualization: false });
+    }
+  }
 }
 
 export default RunDetails;