Check Tensorboard Instance (kubeflow#3189)

* Add tensorboard pod status check * print debug log * Need to keep pulling status if the first attemp to ping tb instance returns non-200 * add test * revise comment * use fetch instead of axios to send http head request * remove accidentally checked-in files * format * address comments * format * address comments * comments
Jeffwan · Dec 9, 2020 · 3566760 · 3566760
1 parent d1becef
commit 3566760
Show file tree

Hide file tree

Showing 4 changed files with 117 additions and 1 deletion.
diff --git a/frontend/src/components/viewers/Tensorboard.test.tsx b/frontend/src/components/viewers/Tensorboard.test.tsx
@@ -23,8 +23,14 @@ import { ReactWrapper, ShallowWrapper, shallow, mount } from 'enzyme';
 
 describe('Tensorboard', () => {
   let tree: ReactWrapper | ShallowWrapper;
+  const flushPromisesAndTimers = async () => {
+    jest.runOnlyPendingTimers();
+    await TestUtils.flushPromises();
+  };
+
   beforeEach(() => {
     jest.clearAllMocks();
+    jest.useFakeTimers();
   });
 
   afterEach(async () => {
@@ -102,9 +108,13 @@ 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();
+    await flushPromisesAndTimers();
+    expect(Apis.isTensorboardPodReady).toHaveBeenCalledTimes(1);
+    expect(Apis.isTensorboardPodReady).toHaveBeenCalledWith('apis/v1beta1/_proxy/test/address');
     expect(tree).toMatchSnapshot();
   });
 
@@ -266,4 +276,44 @@ 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();
+    await flushPromisesAndTimers();
+    tree.update();
+    expect(Apis.isTensorboardPodReady).toHaveBeenCalledTimes(1);
+    expect(Apis.isTensorboardPodReady).toHaveBeenCalledWith('apis/v1beta1/_proxy/podaddress');
+    expect(tree.findWhere(el => el.text() === 'Open Tensorboard').exists()).toBeTruthy();
+    expect(
+      tree
+        .findWhere(
+          el =>
+            el.text() === '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));
+    await flushPromisesAndTimers();
+    tree.update();
+    expect(
+      tree
+        .findWhere(
+          el =>
+            el.text() === `Tensorboard is starting, and you may need to wait for a few minutes.`,
+        )
+        .exists(),
+    ).toEqual(false);
+  });
 });
diff --git a/frontend/src/components/viewers/Tensorboard.tsx b/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 => {
@@ -128,6 +143,11 @@ class TensorboardViewer extends Viewer<TensorboardViewerProps, TensorboardViewer
               >
                 Open Tensorboard
               </Button>
+              {this.state.tensorboardReady ? (
+                ``
+              ) : (
+                <div>Tensorboard is starting, and you may need to wait for a few minutes.</div>
+              )}
             </a>
 
             <div>
@@ -235,6 +255,18 @@ 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) {
+      // Remove protocol prefix bofore ":" from pod address if any.
+      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 +285,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 +301,7 @@ class TensorboardViewer extends Viewer<TensorboardViewerProps, TensorboardViewer
         deleteDialogOpen: false,
         podAddress: '',
         tensorflowVersion: DEFAULT_TENSORBOARD_VERSION,
+        tensorboardReady: false,
       });
     });
   };

diff --git a/frontend/src/lib/Apis.test.ts b/frontend/src/lib/Apis.test.ts
@@ -26,6 +26,17 @@ const fetchSpy = (response: string) => {
   return spy;
 };
 
+const failedFetchSpy = (response: string) => {
+  const spy = jest.fn(() =>
+    Promise.resolve({
+      ok: false,
+      text: () => response,
+    }),
+  );
+  window.fetch = spy;
+  return spy;
+};
+
 describe('Apis', () => {
   it('hosts a singleton experimentServiceApi', () => {
     expect(Apis.experimentServiceApi).toBe(Apis.experimentServiceApi);
@@ -182,4 +193,18 @@ describe('Apis', () => {
       },
     );
   });
+
+  it('checks if Tensorboard pod is ready', async () => {
+    const spy = fetchSpy('');
+    const ready = await Apis.isTensorboardPodReady('apis/v1beta1/_proxy/pod_address');
+    expect(ready).toBe(true);
+    expect(spy).toHaveBeenCalledWith('apis/v1beta1/_proxy/pod_address', { method: 'HEAD' });
+  });
+
+  it('checks if Tensorboard pod is not ready', async () => {
+    const spy = failedFetchSpy('');
+    const ready = await Apis.isTensorboardPodReady('apis/v1beta1/_proxy/pod_address');
+    expect(ready).toBe(false);
+    expect(spy).toHaveBeenCalledWith('apis/v1beta1/_proxy/pod_address', { method: 'HEAD' });
+  });
 });
diff --git a/frontend/src/lib/Apis.ts b/frontend/src/lib/Apis.ts
@@ -204,6 +204,14 @@ export class Apis {
     );
   }
 
+  /**
+   * Check if the underlying Tensorboard pod is actually up, given the pod address
+   */
+  public static async isTensorboardPodReady(path: string): Promise<boolean> {
+    const resp = await fetch(path, { method: 'HEAD' });
+    return resp.ok;
+  }
+
   /**
    * Delete a deployment and its service of the Tensorboard given the URL
    */