feat: Use feature flag for streaming updates - manually update projec…

…t store (#9170)
determined-ai · Apr 18, 2024 · 72344e0 · 72344e0
1 parent dd7f4b5
commit 72344e0
Show file tree

Hide file tree

Showing 8 changed files with 105 additions and 33 deletions.
diff --git a/webui/react/src/App.tsx b/webui/react/src/App.tsx
@@ -18,6 +18,7 @@ import PageMessage from 'components/PageMessage';
 import Router from 'components/Router';
 import useUI, { Mode, ThemeProvider } from 'components/ThemeProvider';
 import useAuthCheck from 'hooks/useAuthCheck';
+import useFeature from 'hooks/useFeature';
 import useKeyTracker from 'hooks/useKeyTracker';
 import usePageVisibility from 'hooks/usePageVisibility';
 import usePermissions from 'hooks/usePermissions';
@@ -61,6 +62,7 @@ const AppView: React.FC = () => {
   const settings = useObservable(themeSetting);
   const [isSettingsReady, setIsSettingsReady] = useState(false);
   const { ui, actions: uiActions, theme, isDarkMode } = useUI();
+  const streamingUpdatesOn = useFeature().isOn('streaming_updates');
 
   useEffect(() => {
     if (isServerReachable) checkAuth();
@@ -71,8 +73,8 @@ const AppView: React.FC = () => {
   useRouteTracker();
 
   useEffect(() => {
-    return streamStore.on(projectStore);
-  }, []);
+    return streamingUpdatesOn ? streamStore.on(projectStore) : streamStore.off('projects');
+  }, [streamingUpdatesOn]);
 
   useEffect(() => (isAuthenticated ? userStore.fetchCurrentUser() : undefined), [isAuthenticated]);
   useEffect(() => (isAuthenticated ? clusterStore.startPolling() : undefined), [isAuthenticated]);

diff --git a/webui/react/src/components/ProjectActionDropdown.tsx b/webui/react/src/components/ProjectActionDropdown.tsx
@@ -20,7 +20,7 @@ interface Props {
   direction?: 'vertical' | 'horizontal';
   isContextMenu?: boolean;
   onDelete?: () => void;
-  onEdit?: (name: string, archived: boolean) => void;
+  onEdit?: (name: string, archived: boolean, description?: string) => void;
   onMove?: () => void;
   project: Project;
   showChildrenIfEmpty?: boolean;
@@ -29,7 +29,7 @@ interface Props {
 
 interface ProjectMenuPropsIn {
   onDelete?: () => void;
-  onEdit?: (name: string, archived: boolean) => void;
+  onEdit?: (name: string, archived: boolean, description?: string) => void;
   onMove?: () => void;
   project?: Project;
   workspaceArchived?: boolean;

diff --git a/webui/react/src/components/ProjectCard.tsx b/webui/react/src/components/ProjectCard.tsx
@@ -20,7 +20,7 @@ import css from './ProjectCard.module.scss';
 
 interface Props {
   hideActionMenu?: boolean;
-  onEdit?: (name: string, archived: boolean) => void;
+  onEdit?: (name: string, archived: boolean, description?: string) => void;
   onRemove?: () => void;
   project: Project;
   showWorkspace?: boolean;

diff --git a/webui/react/src/components/ProjectEditModal.tsx b/webui/react/src/components/ProjectEditModal.tsx
@@ -15,7 +15,7 @@ interface FormInputs {
 }
 
 interface Props {
-  onEdit?: (name: string, archived: boolean) => void;
+  onEdit?: (name: string, archived: boolean, description?: string) => void;
   project: Project;
 }
 
@@ -31,7 +31,7 @@ const ProjectEditModalComponent: React.FC<Props> = ({ onEdit, project }: Props)
 
     try {
       await patchProject({ description, id: project.id, name });
-      onEdit?.(name, project.archived);
+      onEdit?.(name, project.archived, description);
     } catch (e) {
       handleError(e, {
         level: ErrorLevel.Error,

diff --git a/webui/react/src/hooks/useFeature.ts b/webui/react/src/hooks/useFeature.ts
@@ -7,7 +7,12 @@ import determinedStore, { DeterminedInfo } from 'stores/determinedInfo';
 import userSettings from 'stores/userSettings';
 
 // add new feature switches here
-export type ValidFeature = 'explist_v2' | 'rp_binding' | 'genai' | 'flat_runs';
+export type ValidFeature =
+  | 'explist_v2'
+  | 'rp_binding'
+  | 'genai'
+  | 'flat_runs'
+  | 'streaming_updates';
 
 type FeatureDescription = {
   friendlyName: string;
@@ -37,6 +42,11 @@ export const FEATURES: Record<ValidFeature, FeatureDescription> = {
     description: 'Allow resource pools to be assigned to workspaces',
     friendlyName: 'Resource Pool Binding',
   },
+  streaming_updates: {
+    defaultValue: false,
+    description: 'Allow streaming updates through websockets for better performance',
+    friendlyName: 'Streaming Updates',
+  },
 } as const;
 export const FEATURE_SETTINGS_PATH = 'global-features';
 

diff --git a/webui/react/src/pages/WorkspaceDetails/WorkspaceProjects.tsx b/webui/react/src/pages/WorkspaceDetails/WorkspaceProjects.tsx
@@ -31,6 +31,7 @@ import {
   stateRenderer,
   userRenderer,
 } from 'components/Table/Table';
+import useFeature from 'hooks/useFeature';
 import usePermissions from 'hooks/usePermissions';
 import usePrevious from 'hooks/usePrevious';
 import { useSettings } from 'hooks/useSettings';
@@ -67,6 +68,7 @@ const WorkspaceProjects: React.FC<Props> = ({ workspace, id, pageRef }) => {
   const ProjectCreateModal = useModal(ProjectCreateModalComponent);
   const config = useMemo(() => configForWorkspace(id), [id]);
   const { settings, updateSettings } = useSettings<WorkspaceDetailsSettings>(config);
+  const streamingUpdatesOn = useFeature().isOn('streaming_updates');
 
   const loadableProjects: Loadable<List<Project>> = useObservable(
     projectStore.getProjectsByWorkspace(id),
@@ -158,27 +160,65 @@ const WorkspaceProjects: React.FC<Props> = ({ workspace, id, pageRef }) => {
     }
   }, [currentUser, loadableUsers, prevWhose, settings.whose, updateSettings, users]);
 
-  const saveProjectDescription = useCallback(async (newDescription: string, projectId: number) => {
-    try {
-      await patchProject({ description: newDescription, id: projectId });
-    } catch (e) {
-      handleError(e, {
-        level: ErrorLevel.Error,
-        publicMessage: 'Please try again later.',
-        publicSubject: 'Unable to edit project.',
-        silent: false,
-        type: ErrorType.Server,
-      });
-    }
-  }, []);
+  const onEdit = useCallback(
+    (projectId: number, name: string, archived: boolean, description?: string) => {
+      if (!streamingUpdatesOn) {
+        const project = projects.find((p) => p.id === projectId);
+        project &&
+          projectStore.upsertProject({
+            ...project,
+            archived,
+            description: description || project?.description,
+            name,
+          } as Project);
+      }
+    },
+    [streamingUpdatesOn, projects],
+  );
+
+  const onRemove = useCallback(
+    (projectId: number) => {
+      if (!streamingUpdatesOn) {
+        projectStore.delete(projectId);
+      }
+    },
+    [streamingUpdatesOn],
+  );
+
+  const saveProjectDescription = useCallback(
+    async (newDescription: string, projectId: number) => {
+      try {
+        await patchProject({ description: newDescription, id: projectId });
+        const project = projects.find((p) => p.id === projectId);
+        project && projectStore.upsertProject({ ...project, description: newDescription });
+      } catch (e) {
+        handleError(e, {
+          level: ErrorLevel.Error,
+          publicMessage: 'Please try again later.',
+          publicSubject: 'Unable to edit project.',
+          silent: false,
+          type: ErrorType.Server,
+        });
+      }
+    },
+    [projects],
+  );
 
   const columns = useMemo(() => {
     const projectNameRenderer = (value: string, record: Project) => (
       <Link path={paths.projectDetails(record.id)}>{value}</Link>
     );
 
     const actionRenderer: GenericRenderer<Project> = (_, record) => (
-      <ProjectActionDropdown project={record} workspaceArchived={workspace?.archived} />
+      <ProjectActionDropdown
+        project={record}
+        workspaceArchived={workspace?.archived}
+        onDelete={() => onRemove(record.id)}
+        onEdit={(name: string, archived: boolean, description?: string) =>
+          onEdit(record.id, name, archived, description)
+        }
+        onMove={() => onRemove(record.id)}
+      />
     );
 
     const descriptionRenderer = (value: string, record: Project) => (
@@ -262,7 +302,7 @@ const WorkspaceProjects: React.FC<Props> = ({ workspace, id, pageRef }) => {
         title: '',
       },
     ] as ColumnDef<Project>[];
-  }, [saveProjectDescription, workspace?.archived, users]);
+  }, [saveProjectDescription, workspace?.archived, users, onEdit, onRemove]);
 
   const switchShowArchived = useCallback(
     (showArchived: boolean) => {
@@ -303,12 +343,20 @@ const WorkspaceProjects: React.FC<Props> = ({ workspace, id, pageRef }) => {
 
   const actionDropdown = useCallback(
     ({ record, children }: { children: React.ReactNode; record: Project }) => (
-      <ProjectActionDropdown isContextMenu project={record} workspaceArchived={workspace?.archived}>
+      <ProjectActionDropdown
+        isContextMenu
+        project={record}
+        workspaceArchived={workspace?.archived}
+        onDelete={() => onRemove(record.id)}
+        onEdit={(name: string, archived: boolean, description?: string) =>
+          onEdit(record.id, name, archived, description)
+        }
+        onMove={() => onRemove(record.id)}>
         {children}
       </ProjectActionDropdown>
     ),
     // eslint-disable-next-line react-hooks/exhaustive-deps
-    [workspace?.archived],
+    [workspace?.archived, onEdit, onRemove],
   );
 
   const projectsList = useMemo(() => {
@@ -323,6 +371,10 @@ const WorkspaceProjects: React.FC<Props> = ({ workspace, id, pageRef }) => {
                 key={project.id}
                 project={project}
                 workspaceArchived={workspace?.archived}
+                onEdit={(name: string, archived: boolean, description?: string) =>
+                  onEdit(project.id, name, archived, description)
+                }
+                onRemove={() => onRemove(project.id)}
               />
             ))}
           </Card.Group>
@@ -359,11 +411,13 @@ const WorkspaceProjects: React.FC<Props> = ({ workspace, id, pageRef }) => {
     settings,
     updateSettings,
     workspace?.archived,
+    onEdit,
+    onRemove,
   ]);
 
   useEffect(() => {
-    projectStore.fetch(id);
-  }, [id]);
+    projectStore.fetch(id, canceler.signal, true);
+  }, [id, canceler]);
 
   useEffect(() => {
     return () => canceler.abort();

diff --git a/webui/react/src/stores/projects.tsx b/webui/react/src/stores/projects.tsx
@@ -89,6 +89,10 @@ export class ProjectStore implements StreamSubscriber {
     }
   }
 
+  public upsert(content: StreamContent): void {
+    this.upsertProject(mapStreamProject(content));
+  }
+
   #upsert(p: Project, np: Project): Project {
     p.name = np.name;
     p.description = np.description;
@@ -98,8 +102,7 @@ export class ProjectStore implements StreamSubscriber {
     return { ...p };
   }
 
-  public upsert(content: StreamContent): void {
-    const p = mapStreamProject(content);
+  public upsertProject(p: Project): void {
     let prevProjectWorkspaceId: number | undefined;
 
     this.#projects.update((prev) =>

diff --git a/webui/react/src/stores/stream.ts b/webui/react/src/stores/stream.ts
@@ -12,10 +12,10 @@ export interface StreamSubscriber {
 }
 
 class StreamStore {
-  #stream: Stream;
+  #stream?: Stream = undefined;
   #subscribers: Partial<Record<Streamable, StreamSubscriber>> = {};
 
-  constructor() {
+  #opensocket() {
     const socketUrl = new URL(serverAddress());
     socketUrl.pathname = 'stream';
     socketUrl.protocol = socketUrl.protocol.replace('http', 'ws');
@@ -41,22 +41,25 @@ class StreamStore {
   }
 
   public on(sub: StreamSubscriber): () => void {
+    if (!this.#stream) this.#opensocket();
+
     this.#subscribers[sub.id()] = sub;
 
     return () => this.#closeSocket.bind(this);
   }
 
   public off(key: Streamable) {
+    this.#closeSocket();
     delete this.#subscribers[key];
   }
 
   public emit(spec: StreamSpec, id?: string) {
-    this.#stream.subscribe(spec, id);
+    this.#stream?.subscribe(spec, id);
   }
 
   #closeSocket() {
     this.#subscribers = {};
-    this.#stream.close();
+    this.#stream?.close();
   }
 }