(Frontend): Add commands for generic snapshots

labextension/src/lib/Commands.ts

@@ -20,6 +20,7 @@ import {
   _legacy_executeRpc,
   _legacy_executeRpcAndShowRPCError,
   RPCError,
+  IRPCError,
 } from './RPCUtils';
 import { wait } from './Utils';
 import {
@@ -37,6 +38,7 @@ import {
 } from '../widgets/VolumesPanel';
 import { IDocumentManager } from '@jupyterlab/docmanager';
 import CellUtils from './CellUtils';
+import * as React from 'react';
 
 enum RUN_CELL_STATUS {
   OK = 'ok',
@@ -73,6 +75,8 @@ interface IKatibRunArgs {
 }
 
 export default class Commands {
+  rokError: IRPCError;
+  snapshotError: IRPCError;
   private readonly _notebook: NotebookPanel;
   private readonly _kernel: Kernel.IKernelConnection;
 
@@ -89,6 +93,14 @@ export default class Commands {
     );
   };
 
+  genericsnapshotNotebook = async () => {
+    return await _legacy_executeRpcAndShowRPCError(
+      this._notebook,
+      this._kernel,
+      'snapshot.snapshot_notebook',
+    );
+  };
+
   getSnapshotProgress = async (task_id: string, ms?: number) => {
     const task = await _legacy_executeRpcAndShowRPCError(
       this._notebook,
@@ -104,18 +116,31 @@ export default class Commands {
     return task;
   };
 
+  genericgetSnapshotStatus = async (snapshot_name: string, ms?: number) => {
+    const isReady = await _legacy_executeRpcAndShowRPCError(
+      this._notebook,
+      this._kernel,
+      'snapshot.check_snapshot_status',
+      {
+        snapshot_name,
+      },
+    );
+    if (ms) {
+      await wait(ms);
+    }
+    return isReady;
+  };
+
   runSnapshotProcedure = async (onUpdate: Function) => {
     const showSnapshotProgress = true;
     const snapshot = await this.snapshotNotebook();
     const taskId = snapshot.task.id;
     let task = await this.getSnapshotProgress(taskId);
     onUpdate({ task, showSnapshotProgress });
-
     while (!['success', 'error', 'canceled'].includes(task.status)) {
       task = await this.getSnapshotProgress(taskId, 1000);
       onUpdate({ task });
     }
-
     if (task.status === 'success') {
       console.log('Snapshotting successful!');
       return task;
@@ -130,6 +155,29 @@ export default class Commands {
     return null;
   };
 
+  runGenericSnapshotProcedure = async (onUpdate: Function) => {
+    const showSnapshotProgress = true;
+    const snapshot = await this.genericsnapshotNotebook();
+    let snapshot_names = snapshot;
+    for (let i of snapshot_names) {
+      let isReady = await this.genericgetSnapshotStatus(i);
+      onUpdate({ isReady, showSnapshotProgress });
+      while ((isReady = false)) {
+        isReady = await this.genericgetSnapshotStatus(i, 1000);
+        onUpdate({ isReady });
+      }
+      if ((isReady = true)) {
+        console.log('Snapshotting successful!');
+        return isReady;
+      } else if ((isReady = false)) {
+        console.error('Snapshot not ready');
+        console.error('Stopping the deployment...');
+      }
+    }
+
+    return null;
+  };
+
   replaceClonedVolumes = async (
     bucket: string,
     obj: string,
@@ -149,6 +197,17 @@ export default class Commands {
     );
   };
 
+  replaceGenericClonedVolumes = async (volumes: IVolumeMetadata[]) => {
+    return await _legacy_executeRpcAndShowRPCError(
+      this._notebook,
+      this._kernel,
+      'snapshot.replace_cloned_volumes',
+      {
+        volumes,
+      },
+    );
+  };
+
   getMountedVolumes = async (currentNotebookVolumes: IVolumeMetadata[]) => {
     let notebookVolumes: IVolumeMetadata[] = await _legacy_executeRpcAndShowRPCError(
       this._notebook,

labextension/src/widget.tsx

@@ -84,6 +84,7 @@ async function activate(
   //  env we are in (like Local Laptop, MiniKF, GCP, UI without Kale, ...)
   const backend = await getBackend(kernel);
   let rokError: IRPCError = null;
+  let snapshotError: IRPCError = null;
   if (backend) {
     try {
       await executeRpc(kernel, 'log.setup_logging');
@@ -111,6 +112,26 @@ async function activate(
         throw error;
       }
     }
+
+    try {
+      await executeRpc(kernel, 'snapshot.check_snapshot_availability');
+    } catch (error) {
+      const unexpectedErrorCodes = [
+        RPC_CALL_STATUS.EncodingError,
+        RPC_CALL_STATUS.ImportError,
+        RPC_CALL_STATUS.UnhandledError,
+      ];
+      if (
+        error instanceof RPCError &&
+        !unexpectedErrorCodes.includes(error.error.code)
+      ) {
+        snapshotError = error.error;
+        console.warn('Snapshots are not available', snapshotError);
+      } else {
+        globalUnhandledRejection({ reason: error });
+        throw error;
+      }
+    }
   } else {
     rokError = {
       rpc: 'rok.check_rok_availability',
@@ -175,6 +196,7 @@ async function activate(
         backend={backend}
         kernel={kernel}
         rokError={rokError}
+        snapshotError={snapshotError}
       />,
     );
     widget.id = 'kubeflow-kale/kubeflowDeployment';

labextension/src/widgets/LeftPanel.tsx

@@ -60,6 +60,7 @@ interface IProps {
   backend: boolean;
   kernel: Kernel.IKernelConnection;
   rokError: IRPCError;
+  snapshotError: IRPCError;
 }
 
 interface IState {
@@ -442,6 +443,16 @@ export class KubeflowKaleLeftPanel extends React.Component<IProps, IState> {
             notebookVolumes,
             selectVolumeTypes,
           });
+        } else if (!this.props.snapshotError) {
+          // Get information about volumes currently mounted on the notebook server
+          const {
+            notebookVolumes,
+            selectVolumeTypes,
+          } = await commands.getMountedVolumes(this.state.notebookVolumes);
+          this.setState({
+            notebookVolumes,
+            selectVolumeTypes,
+          });
         } else {
           this.setState((prevState, props) => ({
             selectVolumeTypes: prevState.selectVolumeTypes.map(t => {
@@ -532,8 +543,25 @@ export class KubeflowKaleLeftPanel extends React.Component<IProps, IState> {
               }
               return volume;
             });
+        let snapstateVolumes = this.props.snapshotError
+          ? metadataVolumes
+          : metadataVolumes.map((volume: IVolumeMetadata) => {
+              if (
+                volume.type === 'new_pvc' &&
+                volume.annotations.length > 0 &&
+                volume.annotations[0].key === 'rok/origin'
+              ) {
+                return { ...volume, type: 'snap' };
+              }
+              return volume;
+            });
         if (stateVolumes.length === 0 && metadataVolumes.length === 0) {
           metadataVolumes = stateVolumes = this.state.notebookVolumes;
+        } else if (
+          snapstateVolumes.length === 0 &&
+          metadataVolumes.length === 0
+        ) {
+          metadataVolumes = snapstateVolumes = this.state.notebookVolumes;
         } else {
           metadataVolumes = metadataVolumes.concat(this.state.notebookVolumes);
           stateVolumes = stateVolumes.concat(this.state.notebookVolumes);
@@ -557,11 +585,15 @@ export class KubeflowKaleLeftPanel extends React.Component<IProps, IState> {
           },
           autosnapshot:
             notebookMetadata['autosnapshot'] === undefined
-              ? !this.props.rokError && this.state.notebookVolumes.length > 0
+              ? !this.props.rokError &&
+                !this.props.snapshotError &&
+                this.state.notebookVolumes.length > 0
               : notebookMetadata['autosnapshot'],
           snapshot_volumes:
             notebookMetadata['snapshot_volumes'] === undefined
-              ? !this.props.rokError && this.state.notebookVolumes.length > 0
+              ? !this.props.rokError &&
+                !this.props.snapshotError &&
+                this.state.notebookVolumes.length > 0
               : notebookMetadata['snapshot_volumes'],
           // fixme: for now we are using the 'steps_defaults' field just for poddefaults
           //        so we replace any existing value every time
@@ -577,9 +609,13 @@ export class KubeflowKaleLeftPanel extends React.Component<IProps, IState> {
             ...DefaultState.metadata,
             volumes: prevState.notebookVolumes,
             snapshot_volumes:
-              !this.props.rokError && prevState.notebookVolumes.length > 0,
+              !this.props.rokError &&
+              !this.props.snapshotError &&
+              prevState.notebookVolumes.length > 0,
             autosnapshot:
-              !this.props.rokError && prevState.notebookVolumes.length > 0,
+              !this.props.rokError &&
+              !this.props.snapshotError &&
+              prevState.notebookVolumes.length > 0,
           },
           volumes: prevState.notebookVolumes,
         }));
@@ -647,18 +683,32 @@ export class KubeflowKaleLeftPanel extends React.Component<IProps, IState> {
       metadata.volumes.filter((v: IVolumeMetadata) => v.type === 'clone')
         .length > 0
     ) {
-      const task = await commands.runSnapshotProcedure(_updateDeployProgress);
-      console.log(task);
-      if (!task) {
-        this.setState({ runDeployment: false });
-        return;
+      if (!this.props.rokError) {
+        const task = await commands.runSnapshotProcedure(_updateDeployProgress);
+        console.log(task);
+        if (!task) {
+          this.setState({ runDeployment: false });
+          return;
+        }
+        metadata.volumes = await commands.replaceClonedVolumes(
+          task.bucket,
+          task.result.event.object,
+          task.result.event.version,
+          metadata.volumes,
+        );
+      } else if (!this.props.snapshotError) {
+        const task = await commands.runGenericSnapshotProcedure(
+          _updateDeployProgress,
+        );
+        console.log(task);
+        if (!task) {
+          this.setState({ runDeployment: false });
+          return;
+        }
+        metadata.volumes = await commands.replaceGenericClonedVolumes(
+          metadata.volumes,
+        );
       }
-      metadata.volumes = await commands.replaceClonedVolumes(
-        task.bucket,
-        task.result.event.object,
-        task.result.event.version,
-        metadata.volumes,
-      );
     }
 
     // CREATE PIPELINE
@@ -817,6 +867,7 @@ export class KubeflowKaleLeftPanel extends React.Component<IProps, IState> {
         autosnapshot={this.state.metadata.autosnapshot}
         updateAutosnapshotSwitch={this.updateAutosnapshotSwitch}
         rokError={this.props.rokError}
+        snapshotError={this.props.snapshotError}
         updateVolumes={this.updateVolumes}
         storageClassName={this.state.metadata.storage_class_name}
         updateStorageClassName={this.updateStorageClassName}