Add frontend support for Python based visualizations

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,27 @@ 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 {
+      // This should never be thrown as the html property of a generated
+      // visualization is always set for successful visualization generations.
+      throw new Error('Visualization was generated successfully but generated HTML was not found.');
+    }
+  }
+
   /**
    * 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(() => []);
@@ -645,6 +647,19 @@ describe('RunDetails', () => {
     });
   });
 
+  it('shows an error banner if the custom visualizations state API fails', async () => {
+    TestUtils.makeErrorResponseOnce(isCustomVisualizationsAllowedSpy, 'woops');
+    tree = shallow(<RunDetails {...generateProps()} />);
+    await isCustomVisualizationsAllowedSpy;
+    await TestUtils.flushPromises();
+    expect(updateBannerSpy).toHaveBeenCalledTimes(2); // Once initially to clear
+    expect(updateBannerSpy).toHaveBeenLastCalledWith(expect.objectContaining({
+      additionalInfo: 'woops',
+      message:'Error: Unable to enable custom visualizations. Click Details for more information.',
+      mode: 'error',
+    }));
+  });
+
   describe('logs tab', () => {
     it('switches to logs tab in side pane', async () => {
       testRun.pipeline_runtime!.workflow_manifest = JSON.stringify({

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,6 +239,13 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
                         <div className={commonCss.page}>
                           {sidepanelSelectedTab === SidePaneTab.ARTIFACTS && (
                             <div className={commonCss.page}>
+                              <div className={padding(20, 'lrt')}>
+                                <PlotCard
+                                  configs={[visualizationCreatorConfig]}
+                                  title={VisualizationCreator.prototype.getDisplayName()}
+                                  maxDimension={500} />
+                                <Hr />
+                              </div>
                               {(selectedNodeDetails.viewerConfigs || []).map((config, i) => {
                                 const title = componentMap[config.type].prototype.getDisplayName();
                                 return (
@@ -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) {
+      this.showPageError('Error: Unable to enable custom visualizations.', err);
+    }
+
     try {
       const runDetail = await Apis.runServiceApi.getRun(runId);
 
@@ -585,7 +623,7 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
   }
 
   private async _loadSelectedNodeOutputs(): Promise<void> {
-    const selectedNodeDetails = this.state.selectedNodeDetails;
+    const { generatedVisualizations, selectedNodeDetails } = this.state;
     if (!selectedNodeDetails) {
       return;
     }
@@ -594,12 +632,15 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
     if (workflow && workflow.status && workflow.status.nodes) {
       // Load runtime outputs from the selected Node
       const outputPaths = WorkflowParser.loadNodeOutputPaths(workflow.status.nodes[selectedNodeDetails.id]);
-
       // Load the viewer configurations from the output paths
       let viewerConfigs: ViewerConfig[] = [];
       for (const path of outputPaths) {
         viewerConfigs = viewerConfigs.concat(await OutputArtifactLoader.load(path));
       }
+      const generatedConfigs = generatedVisualizations
+        .filter(visualization => visualization.nodeId === selectedNodeDetails.id)
+        .map(visualization => visualization.config);
+      viewerConfigs = viewerConfigs.concat(generatedConfigs);
 
       selectedNodeDetails.viewerConfigs = viewerConfigs;
       this.setStateSafe({ selectedNodeDetails });
@@ -641,6 +682,49 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
       this.setStateSafe({ sidepanelBusy: false });
     }
   }
+
+  private async _onGenerate(visualizationArguments: string, source: string, type: ApiVisualizationType): Promise<void> {
+    const nodeId = this.state.selectedNodeDetails ? this.state.selectedNodeDetails.id : '';
+    if (nodeId.length === 0) {
+      this.showPageError('Unable to generate visualization, no component selected.');
+      return;
+    }
+
+    if (visualizationArguments.length) {
+      try {
+        // Attempts to validate JSON, if attempt fails an error is displayed.
+        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, selectedNodeDetails } = this.state;
+      const generatedVisualization: GeneratedVisualization = {
+        config,
+        nodeId,
+      };
+      generatedVisualizations.push(generatedVisualization);
+      if (selectedNodeDetails) {
+        const viewerConfigs = selectedNodeDetails.viewerConfigs || []; 
+        viewerConfigs.push(generatedVisualization.config);
+        selectedNodeDetails.viewerConfigs = viewerConfigs;
+      }
+      this.setState({ generatedVisualizations, selectedNodeDetails });
+    } catch (err) {
+      this.showPageError('Unable to generate visualization, an unexpected error was encountered.', err);
+    } finally {
+      this.setState({ isGeneratingVisualization: false });
+    }
+  }
 }
 
 export default RunDetails;

frontend/src/pages/__snapshots__/RunDetails.test.tsx.snap

@@ -1367,7 +1367,27 @@ exports[`RunDetails opens side panel when graph node is clicked 1`] = `
               >
                 <div
                   className="page"
-                />
+                >
+                  <div
+                    className=""
+                  >
+                    <PlotCard
+                      configs={
+                        Array [
+                          Object {
+                            "allowCustomVisualizations": false,
+                            "isBusy": false,
+                            "onGenerate": [Function],
+                            "type": "visualization-creator",
+                          },
+                        ]
+                      }
+                      maxDimension={500}
+                      title="Visualization Creator"
+                    />
+                    <Component />
+                  </div>
+                </div>
               </div>
             </div>
           </SidePanel>
@@ -1591,7 +1611,27 @@ exports[`RunDetails shows clicked node message in side panel 1`] = `
               >
                 <div
                   className="page"
-                />
+                >
+                  <div
+                    className=""
+                  >
+                    <PlotCard
+                      configs={
+                        Array [
+                          Object {
+                            "allowCustomVisualizations": false,
+                            "isBusy": false,
+                            "onGenerate": [Function],
+                            "type": "visualization-creator",
+                          },
+                        ]
+                      }
+                      maxDimension={500}
+                      title="Visualization Creator"
+                    />
+                    <Component />
+                  </div>
+                </div>
               </div>
             </div>
           </SidePanel>
@@ -1694,6 +1734,25 @@ exports[`RunDetails shows clicked node output in side pane 1`] = `
                 <div
                   className="page"
                 >
+                  <div
+                    className=""
+                  >
+                    <PlotCard
+                      configs={
+                        Array [
+                          Object {
+                            "allowCustomVisualizations": false,
+                            "isBusy": false,
+                            "onGenerate": [Function],
+                            "type": "visualization-creator",
+                          },
+                        ]
+                      }
+                      maxDimension={500}
+                      title="Visualization Creator"
+                    />
+                    <Component />
+                  </div>
                   <div
                     className=""
                     key="0"