Check Tensorboard Instance

frontend/src/components/viewers/Tensorboard.test.tsx

@@ -25,6 +25,7 @@ describe('Tensorboard', () => {
   let tree: ReactWrapper | ShallowWrapper;
   beforeEach(() => {
     jest.clearAllMocks();
+    jest.useFakeTimers();
   });
 
   afterEach(async () => {
@@ -102,9 +103,14 @@ describe('Tensorboard', () => {
     const config = { type: PlotType.TENSORBOARD, url: 'http://test/url' };
     const getAppMock = () => Promise.resolve({ podAddress: 'test/address', tfVersion: '1.14.0' });
     jest.spyOn(Apis, 'getTensorboardApp').mockImplementation(getAppMock);
+    jest.spyOn(Apis, 'isTensorboardPodReady').mockImplementation(() => Promise.resolve(true));
     tree = shallow(<TensorboardViewer configs={[config]} />);
 
     await TestUtils.flushPromises();
+    jest.runOnlyPendingTimers();
+    await TestUtils.flushPromises();
+    expect(setInterval).toHaveBeenCalledTimes(1);
+    expect(setInterval).toHaveBeenCalledWith(expect.any(Function), 5000);
     expect(tree).toMatchSnapshot();
   });
 
@@ -266,4 +272,48 @@ describe('Tensorboard', () => {
     expect(tree.findWhere(el => el.text() === 'Open Tensorboard').exists()).toBeTruthy();
     expect(tree.findWhere(el => el.text() === 'Delete Tensorboard').exists()).toBeTruthy();
   });
+
+  it('asks user to wait when Tensorboard status is not ready', async () => {
+    const getAppMock = jest.fn(() =>
+      Promise.resolve({ podAddress: 'podaddress', tfVersion: '1.14.0' }),
+    );
+    jest.spyOn(Apis, 'getTensorboardApp').mockImplementation(getAppMock);
+    jest.spyOn(Apis, 'isTensorboardPodReady').mockImplementation(() => Promise.resolve(false));
+    jest.spyOn(Apis, 'deleteTensorboardApp').mockImplementation(jest.fn(() => Promise.resolve('')));
+    const config = { type: PlotType.TENSORBOARD, url: 'http://test/url' };
+    tree = mount(<TensorboardViewer configs={[config]} />);
+
+    await TestUtils.flushPromises();
+    jest.runOnlyPendingTimers();
+    await TestUtils.flushPromises();
+    tree.update();
+    expect(setInterval).toHaveBeenCalledTimes(1);
+    expect(setInterval).toHaveBeenCalledWith(expect.any(Function), 5000);
+    expect(
+      tree
+        .findWhere(
+          el =>
+            el.text() === 'Open Tensorboard' &&
+            el.prop('title') ===
+              `Tensorboard is starting, and you may need to wait for a few minutes.`,
+        )
+        .exists(),
+    ).toBeTruthy();
+    expect(tree.findWhere(el => el.text() === 'Delete Tensorboard').exists()).toBeTruthy();
+
+    // After a while, it is ready and wait message is not shwon any more
+    jest.spyOn(Apis, 'isTensorboardPodReady').mockImplementation(() => Promise.resolve(true));
+    jest.runOnlyPendingTimers();
+    await TestUtils.flushPromises();
+    tree.update();
+    expect(
+      tree
+        .findWhere(
+          el =>
+            el.prop('title') ===
+            `Tensorboard is starting, and you may need to wait for a few minutes.`,
+        )
+        .exists(),
+    ).toEqual(false);
+  });
 });

frontend/src/components/viewers/Tensorboard.tsx

@@ -55,19 +55,25 @@ export interface TensorboardViewerConfig extends ViewerConfig {
 
 interface TensorboardViewerProps {
   configs: TensorboardViewerConfig[];
+  // Interval in ms. If not specified, default to 5000.
+  intervalOfCheckingTensorboardPodStatus?: number;
 }
 
 interface TensorboardViewerState {
   busy: boolean;
   deleteDialogOpen: boolean;
   podAddress: string;
   tensorflowVersion: string;
+  // When podAddress is not null, we need to further tell whether the TensorBoard pod is accessible or not
+  tensorboardReady: boolean;
 }
 
 // TODO(jingzhang36): we'll later parse Tensorboard version from mlpipeline-ui-metadata.json file.
 const DEFAULT_TENSORBOARD_VERSION = '2.0.0';
 
 class TensorboardViewer extends Viewer<TensorboardViewerProps, TensorboardViewerState> {
+  timerID: NodeJS.Timeout;
+
   constructor(props: any) {
     super(props);
 
@@ -76,6 +82,7 @@ class TensorboardViewer extends Viewer<TensorboardViewerProps, TensorboardViewer
       deleteDialogOpen: false,
       podAddress: '',
       tensorflowVersion: DEFAULT_TENSORBOARD_VERSION,
+      tensorboardReady: false,
     };
   }
 
@@ -89,6 +96,14 @@ class TensorboardViewer extends Viewer<TensorboardViewerProps, TensorboardViewer
 
   public componentDidMount(): void {
     this._checkTensorboardApp();
+    this.timerID = setInterval(
+      () => this._checkTensorboardPodStatus(),
+      this.props.intervalOfCheckingTensorboardPodStatus || 5000,
+    );
+  }
+
+  public componentWillUnmount(): void {
+    clearInterval(this.timerID);
   }
 
   public handleVersionSelect = (e: React.ChangeEvent<{ name?: string; value: unknown }>): void => {
@@ -125,6 +140,11 @@ class TensorboardViewer extends Viewer<TensorboardViewerProps, TensorboardViewer
                 className={classes(commonCss.buttonAction, css.button)}
                 disabled={this.state.busy}
                 color={'primary'}
+                title={
+                  this.state.tensorboardReady
+                    ? ``
+                    : `Tensorboard is starting, and you may need to wait for a few minutes.`
+                }
               >
                 Open Tensorboard
               </Button>
@@ -235,6 +255,17 @@ class TensorboardViewer extends Viewer<TensorboardViewerProps, TensorboardViewer
     return urls.length === 1 ? urls[0] : urls.map((c, i) => `Series${i + 1}:` + c).join(',');
   }
 
+  private async _checkTensorboardPodStatus(): Promise<void> {
+    // If pod address is not null and tensorboard pod doesn't seem to be read, pull status again
+    if (this.state.podAddress && !this.state.tensorboardReady) {
+      Apis.isTensorboardPodReady(
+        'apis/v1beta1/_proxy/' + this.state.podAddress.replace(/(^\w+:|^)\/\//, ''),
+      ).then(ready => {
+        this.setState(({ tensorboardReady }) => ({ tensorboardReady: tensorboardReady || ready }));
+      });
+    }
+  }
+
   private async _checkTensorboardApp(): Promise<void> {
     this.setState({ busy: true }, async () => {
       const { podAddress, tfVersion } = await Apis.getTensorboardApp(this._buildUrl());
@@ -253,7 +284,7 @@ class TensorboardViewer extends Viewer<TensorboardViewerProps, TensorboardViewer
         encodeURIComponent(this._buildUrl()),
         encodeURIComponent(this.state.tensorflowVersion),
       );
-      this.setState({ busy: false }, () => {
+      this.setState({ busy: false, tensorboardReady: false }, () => {
         this._checkTensorboardApp();
       });
     });
@@ -269,6 +300,7 @@ class TensorboardViewer extends Viewer<TensorboardViewerProps, TensorboardViewer
         deleteDialogOpen: false,
         podAddress: '',
         tensorflowVersion: DEFAULT_TENSORBOARD_VERSION,
+        tensorboardReady: false,
       });
     });
   };

frontend/src/components/viewers/__snapshots__/Tensorboard.test.tsx.snap

@@ -129,6 +129,7 @@ exports[`Tensorboard shows a link to the tensorboard instance if exists 1`] = `
         className="buttonAction button"
         color="primary"
         disabled={false}
+        title=""
       >
         Open Tensorboard
       </WithStyles(Button)>

frontend/src/lib/Apis.test.ts

@@ -182,4 +182,16 @@ describe('Apis', () => {
       },
     );
   });
+
+  it('checks if Tensorboard pod is ready', async () => {
+    const spy = jest.fn(() =>
+      Promise.resolve({
+        status: 200,
+      }),
+    );
+    window.fetch = spy;
+    const ready = await Apis.isTensorboardPodReady('apis/v1beta1/_proxy/pod_address');
+    expect(ready).toBe(true);
+    expect(spy).toHaveBeenCalledWith('apis/v1beta1/_proxy/pod_address', { method: 'HEAD' });
+  });
 });

frontend/src/lib/Apis.ts

@@ -204,6 +204,17 @@ export class Apis {
     );
   }
 
+  /**
+   * Check if the underlying Tensorboard pod is actually up, given
+   */
+  public static async isTensorboardPodReady(path: string): Promise<boolean> {
+    const resp = await fetch(path, { method: 'HEAD' });
+    if (resp.status === 200) {
+      return true;
+    }
+    return false;
+  }
+
   /**
    * Delete a deployment and its service of the Tensorboard given the URL
    */