Pause graph during training

src/components/LiveGraph.tsx

@@ -24,7 +24,11 @@ export const smoothenDataPoint = (curr: number, next: number) => {
   return (next / 1000) * 0.25 + curr * 0.75;
 };
 
-const LiveGraph = () => {
+interface LiveGraphProps {
+  paused?: boolean;
+}
+
+const LiveGraph = ({ paused }: LiveGraphProps) => {
   const { isConnected, status } = useConnectionStage();
   const connectActions = useConnectActions();
   const [{ graphColorScheme, graphLineScheme, graphLineWeight }] =
@@ -110,12 +114,15 @@ const LiveGraph = () => {
   ]);
 
   useEffect(() => {
-    if (isConnected || status === ConnectionStatus.ReconnectingAutomatically) {
+    if (
+      (isConnected || status === ConnectionStatus.ReconnectingAutomatically) &&
+      !paused
+    ) {
       chart?.start();
     } else {
       chart?.stop();
     }
-  }, [chart, isConnected, status]);
+  }, [chart, isConnected, paused, status]);
 
   // Draw on graph to display that users are recording.
   const isRecording = useStore((s) => s.isRecording);

src/components/LiveGraphPanel.tsx

@@ -26,6 +26,8 @@ import AlertIcon from "./AlertIcon";
 import InfoToolTip from "./InfoToolTip";
 import LiveGraph from "./LiveGraph";
 import PredictedAction from "./PredictedAction";
+import { TrainModelDialogStage } from "../model";
+import { useStore } from "../store";
 
 interface LiveGraphPanelProps {
   showPredictedAction?: boolean;
@@ -41,6 +43,11 @@ const LiveGraphPanel = ({
   showDisconnectedOverlay = true,
 }: LiveGraphPanelProps) => {
   const { actions, status, isConnected } = useConnectionStage();
+  const isTraining = useStore(
+    (s) =>
+      s.trainModelDialogStage === TrainModelDialogStage.Help ||
+      s.trainModelDialogStage === TrainModelDialogStage.TrainingInProgress
+  );
   const parentPortalRef = useRef(null);
   const logging = useLogging();
   const isReconnecting =
@@ -165,7 +172,7 @@ const LiveGraphPanel = ({
           </HStack>
         </Portal>
         <HStack position="absolute" width="100%" height="100%" spacing={0}>
-          <LiveGraph />
+          <LiveGraph paused={isTraining} />
           {showPredictedAction && <PredictedAction />}
         </HStack>
       </HStack>

src/components/TrainModelFlowDialogs.tsx

@@ -31,9 +31,11 @@ const TrainModelDialogs = ({ finalFocusRef }: TrainModelDialogsProps) => {
       const result = await trainModel();
       if (result) {
         navigate(createTestingModelPageUrl());
+        // Push it onto the event queue after the navigate call
+        setTimeout(() => closeTrainModelDialogs(), 0);
       }
     },
-    [navigate, setSettings, trainModel]
+    [closeTrainModelDialogs, navigate, setSettings, trainModel]
   );
   return (
     <>

src/store.ts

@@ -1156,6 +1156,13 @@ const createMlStore = (logging: Logging) => {
           } else {
             await trainModel();
             callback?.();
+            // Push the trainModelDialogStage change to the back of the event queue so it happens
+            // after navigation changes.
+            setTimeout(
+              () =>
+                set({ trainModelDialogStage: TrainModelDialogStage.Closed }),
+              0
+            );
           }
         },
 
@@ -1195,7 +1202,7 @@ const createMlStore = (logging: Logging) => {
             {
               model,
               trainModelDialogStage: model
-                ? TrainModelDialogStage.Closed
+                ? TrainModelDialogStage.TrainingInProgress
                 : TrainModelDialogStage.TrainingError,
               timestamp,
               ...updatedProject,