Merge pull request #278 from flatironinstitute/revert-277-revert-274-…

…cancellable-scripts Allow scripts to be cancelable
flatironinstitute · Feb 13, 2025 · f864de8 · f864de8
2 parents 9a23544 + 481d89e
commit f864de8
Show file tree

Hide file tree

Showing 13 changed files with 172 additions and 51 deletions.
diff --git a/backend/stan-wasm-server/src/app/main.py b/backend/stan-wasm-server/src/app/main.py
@@ -50,6 +50,7 @@ def setup_logger() -> None:
     allow_origins=[
         "https://stan-playground.flatironinstitute.org",
         "https://stan-playground.vercel.app",
+        "https://stan-playground-staging.vercel.app",
         "http://127.0.0.1:3000",  # yarn dev
         "http://127.0.0.1:4173",  # yarn preview
     ],
@@ -106,7 +107,13 @@ async def download_file(
     file_path = model_dir / filename
     if not file_path.is_file():
         raise FileNotFoundError(f"File not found: {file_path}")
-    return FileResponse(file_path)
+    return FileResponse(
+        file_path,
+        headers={
+            "Cross-Origin-Embedder-Policy": "require-corp",
+            "Cross-Origin-Opener-Policy": "same-origin",
+        },
+    )
 
 
 @app.post("/compile")

diff --git a/gui/src/app/areas/ControlArea/DataGenerationArea/DataPyPanel.tsx b/gui/src/app/areas/ControlArea/DataGenerationArea/DataPyPanel.tsx
@@ -31,7 +31,7 @@ const DataPyPanel: FunctionComponent = () => {
     [consoleRef, onData, onStatus],
   );
 
-  const { run } = usePyodideWorker(callbacks);
+  const { run, cancel } = usePyodideWorker(callbacks);
 
   const handleRun = useCallback(
     (code: string) => {
@@ -62,6 +62,7 @@ const DataPyPanel: FunctionComponent = () => {
       language="python"
       status={status}
       onRun={handleRun}
+      onCancel={cancel}
       runnable={true}
       notRunnableReason=""
       onHelp={handleHelp}

diff --git a/gui/src/app/areas/ControlArea/DataGenerationArea/DataRPanel.tsx b/gui/src/app/areas/ControlArea/DataGenerationArea/DataRPanel.tsx
@@ -18,7 +18,8 @@ const handleHelp = () =>
 const DataRPanel: FunctionComponent = () => {
   const { consoleRef, status, onStatus, onData } = useDataGenState("r");
 
-  const { run } = useWebR({ consoleRef, onStatus, onData });
+  const { run, cancel } = useWebR({ consoleRef, onStatus, onData });
+
   const handleRun = useCallback(
     async (code: string) => {
       clearOutputDivs(consoleRef);
@@ -40,6 +41,7 @@ const DataRPanel: FunctionComponent = () => {
       language="r"
       status={status}
       onRun={handleRun}
+      onCancel={cancel}
       runnable={true}
       notRunnableReason=""
       onHelp={handleHelp}

diff --git a/gui/src/app/areas/ControlArea/SamplingArea/ResultsArea/AnalysisArea/AnalysisPyPanel.tsx b/gui/src/app/areas/ControlArea/SamplingArea/ResultsArea/AnalysisArea/AnalysisPyPanel.tsx
@@ -36,7 +36,7 @@ const AnalysisPyPanel: FunctionComponent<NeedsLatestRun> = ({ latestRun }) => {
     [consoleRef, imagesRef, onStatus],
   );
 
-  const { run } = usePyodideWorker(callbacks);
+  const { run, cancel } = usePyodideWorker(callbacks);
 
   const handleRun = useCallback(
     (code: string) => {
@@ -69,6 +69,7 @@ const AnalysisPyPanel: FunctionComponent<NeedsLatestRun> = ({ latestRun }) => {
       language="python"
       status={status}
       onRun={handleRun}
+      onCancel={cancel}
       runnable={runnable}
       notRunnableReason={notRunnableReason}
       imagesRef={imagesRef}

diff --git a/gui/src/app/areas/ControlArea/SamplingArea/ResultsArea/AnalysisArea/AnalysisRPanel.tsx b/gui/src/app/areas/ControlArea/SamplingArea/ResultsArea/AnalysisArea/AnalysisRPanel.tsx
@@ -24,7 +24,7 @@ const AnalysisRPanel: FunctionComponent<NeedsLatestRun> = ({ latestRun }) => {
     files,
   } = useAnalysisState(latestRun);
 
-  const { run } = useWebR({ consoleRef, imagesRef, onStatus });
+  const { run, cancel } = useWebR({ consoleRef, imagesRef, onStatus });
 
   const handleRun = useCallback(
     async (userCode: string) => {
@@ -48,6 +48,7 @@ const AnalysisRPanel: FunctionComponent<NeedsLatestRun> = ({ latestRun }) => {
       language="r"
       status={status}
       onRun={handleRun}
+      onCancel={cancel}
       runnable={runnable}
       notRunnableReason={notRunnableReason}
       imagesRef={imagesRef}

diff --git a/gui/src/app/areas/ControlArea/SamplingArea/ResultsArea/SamplerOutputArea/DrawsTablePanel.tsx b/gui/src/app/areas/ControlArea/SamplingArea/ResultsArea/SamplerOutputArea/DrawsTablePanel.tsx
@@ -95,7 +95,9 @@ const DrawsTablePanel: FunctionComponent<DrawsTableProps> = ({
               <TableCell key="chain">Chain</TableCell>
               <TableCell key="draw">Draw</TableCell>
               {paramNames.map((name, i) => (
-                <TableCell key={i}>{name}</TableCell>
+                <TableCell padding="checkbox" key={i}>
+                  {name}
+                </TableCell>
               ))}
             </TableRow>
           </SuccessColoredTableHead>
@@ -105,7 +107,9 @@ const DrawsTablePanel: FunctionComponent<DrawsTableProps> = ({
                 <TableCell>{drawChainIds[i]}</TableCell>
                 <TableCell>{drawNumbers[i]}</TableCell>
                 {formattedDraws.map((draw, j) => (
-                  <TableCell key={j}>{draw[i]}</TableCell>
+                  <TableCell padding="checkbox" key={j}>
+                    {draw[i]}
+                  </TableCell>
                 ))}
               </SuccessBorderedTableRow>
             ))}

diff --git a/gui/src/app/components/FileEditor/ScriptEditor.tsx b/gui/src/app/components/FileEditor/ScriptEditor.tsx
@@ -1,6 +1,6 @@
 import { FunctionComponent, RefObject, useCallback, use, useMemo } from "react";
 
-import { Help, PlayArrow } from "@mui/icons-material";
+import { Close, Help, PlayArrow } from "@mui/icons-material";
 import Box from "@mui/material/Box";
 import { Split } from "@geoffcox/react-splitter";
 import { useMonaco } from "@monaco-editor/react";
@@ -22,6 +22,7 @@ export type ScriptEditorProps = {
   filename: FileNames;
   dataKey: ProjectKnownFiles;
   onRun: (code: string) => void;
+  onCancel?: () => void;
   runnable: boolean;
   notRunnableReason?: string;
   onHelp?: () => void;
@@ -35,6 +36,7 @@ const ScriptEditor: FunctionComponent<ScriptEditorProps> = ({
   filename,
   dataKey,
   onRun,
+  onCancel,
   runnable,
   notRunnableReason,
   onHelp,
@@ -74,11 +76,12 @@ const ScriptEditor: FunctionComponent<ScriptEditorProps> = ({
 
   const monacoInstance = useMonaco();
 
-  const runCtrlEnter: editor.IActionDescriptor[] = useMemo(() => {
+  const scriptShortcuts: editor.IActionDescriptor[] = useMemo(() => {
     if (!monacoInstance) {
       return [];
     }
-    return [
+    const actions = [
+      // Ctrl-Enter to run
       {
         id: "run-script",
         label: "Run Script",
@@ -92,7 +95,19 @@ const ScriptEditor: FunctionComponent<ScriptEditorProps> = ({
         },
       },
     ];
-  }, [monacoInstance, runCode, runnable, unsavedChanges]);
+    if (onCancel) {
+      // Ctrl-C to cancel
+      actions.push({
+        id: "cancel-script",
+        label: "Cancel Script",
+        keybindings: [
+          monacoInstance.KeyMod.CtrlCmd | monacoInstance.KeyCode.KeyC,
+        ],
+        run: onCancel,
+      });
+    }
+    return actions;
+  }, [monacoInstance, onCancel, runCode, runnable, unsavedChanges]);
 
   const toolbarItems: ToolbarItem[] = useMemo(() => {
     return makeToolbar({
@@ -101,11 +116,13 @@ const ScriptEditor: FunctionComponent<ScriptEditorProps> = ({
       runnable: runnable && !unsavedChanges,
       notRunnableReason,
       onRun: runCode,
+      onCancel,
       onHelp,
     });
   }, [
     language,
     notRunnableReason,
+    onCancel,
     onHelp,
     runCode,
     runnable,
@@ -124,7 +141,7 @@ const ScriptEditor: FunctionComponent<ScriptEditorProps> = ({
         onSaveText={onSaveText}
         toolbarItems={toolbarItems}
         contentOnEmpty={contentOnEmpty}
-        actions={runCtrlEnter}
+        actions={scriptShortcuts}
       />
       <ConsoleOutputPanel consoleRef={consoleRef} />
     </Split>
@@ -138,8 +155,9 @@ const makeToolbar = (o: {
   notRunnableReason?: string;
   onRun: () => void;
   onHelp?: () => void;
+  onCancel?: () => void;
 }): ToolbarItem[] => {
-  const { status, onRun, runnable, onHelp, name } = o;
+  const { status, onRun, runnable, onHelp, name, onCancel } = o;
   const ret: ToolbarItem[] = [];
   if (onHelp !== undefined) {
     ret.push({
@@ -166,6 +184,17 @@ const makeToolbar = (o: {
     });
   }
 
+  if (onCancel && status === "running") {
+    ret.push({
+      type: "button",
+      tooltip: "Cancel",
+      label: "Cancel",
+      icon: <Close />,
+      onClick: onCancel,
+      color: "error",
+    });
+  }
+
   let label: string;
   let color: ColorOptions;
   if (status === "loading") {

diff --git a/gui/src/app/core/Scripting/pyodide/pyodideWorker.ts b/gui/src/app/core/Scripting/pyodide/pyodideWorker.ts
@@ -3,33 +3,29 @@ import { isMonacoWorkerNoise } from "@SpUtil/isMonacoWorkerNoise";
 import { InterpreterStatus } from "@SpCore/Scripting/InterpreterTypes";
 import {
   MessageFromPyodideWorker,
+  MessageToPyodideWorker,
   PyodideRunSettings,
 } from "./pyodideWorkerTypes";
 import spDrawsScript from "./sp_load_draws.py?raw";
 import spMPLScript from "./sp_patch_matplotlib.py?raw";
 
-let pyodide: PyodideInterface | null = null;
 const loadPyodideInstance = async () => {
-  if (pyodide === null) {
-    pyodide = await loadPyodide({
-      indexURL: "https://cdn.jsdelivr.net/pyodide/v0.27.2/full",
-      stdout: (x: string) => {
-        sendStdout(x);
-      },
-      stderr: (x: string) => {
-        sendStderr(x);
-      },
-      packages: ["numpy", "micropip", "pandas"],
-    });
-    setStatus("installing");
-
-    pyodide.FS.writeFile("sp_load_draws.py", spDrawsScript);
-    pyodide.FS.writeFile("sp_patch_matplotlib.py", spMPLScript);
-
-    return pyodide;
-  } else {
-    return pyodide;
-  }
+  const pyodide = await loadPyodide({
+    indexURL: "https://cdn.jsdelivr.net/pyodide/v0.27.2/full",
+    stdout: (x: string) => {
+      sendStdout(x);
+    },
+    stderr: (x: string) => {
+      sendStderr(x);
+    },
+    packages: ["numpy", "micropip", "pandas"],
+  });
+  console.log("pyodide loaded");
+
+  pyodide.FS.writeFile("sp_load_draws.py", spDrawsScript);
+  pyodide.FS.writeFile("sp_patch_matplotlib.py", spMPLScript);
+
+  return pyodide;
 };
 
 const sendMessageToMain = (message: MessageFromPyodideWorker) => {
@@ -56,25 +52,38 @@ const addImage = (image: any) => {
   sendMessageToMain({ type: "addImage", image });
 };
 
-console.log("pyodide worker loaded");
-
-self.onmessage = async (e) => {
+self.onmessage = async (e: MessageEvent<MessageToPyodideWorker>) => {
   if (isMonacoWorkerNoise(e.data)) {
     return;
   }
   const message = e.data;
-  await run(message.code, message.spData, message.spRunSettings, message.files);
+  await run(
+    message.code,
+    message.spData,
+    message.spRunSettings,
+    message.files,
+    message.interruptBuffer,
+  );
 };
+console.log("pyodide worker initialized");
+
+console.log("opportunistically loading pyodide");
+const pyodidePromise: Promise<PyodideInterface> = loadPyodideInstance();
 
 const run = async (
   code: string,
   spData: Record<string, any> | undefined,
   spPySettings: PyodideRunSettings,
   files: Record<string, string> | undefined,
+  interruptBuffer: Uint8Array | undefined,
 ) => {
   setStatus("loading");
   try {
-    const pyodide = await loadPyodideInstance();
+    const pyodide = await pyodidePromise;
+    if (interruptBuffer) {
+      pyodide.setInterruptBuffer(interruptBuffer);
+    }
+    setStatus("installing");
 
     const [scriptPreamble, scriptPostamble] = getScriptParts(spPySettings);
 
@@ -95,6 +104,7 @@ const run = async (
     let succeeded = false;
     try {
       const packageFutures = [];
+      let patch_http = false;
       const micropip = pyodide.pyimport("micropip");
 
       if (spPySettings.showsPlots) {
@@ -104,10 +114,20 @@ const run = async (
           packageFutures.push(micropip.install("arviz"));
         }
       }
+      if (script.includes("requests")) {
+        patch_http = true;
+        packageFutures.push(
+          micropip.install(["requests", "lzma", "pyodide-http"]),
+        );
+      }
       packageFutures.push(micropip.install("stanio"));
       packageFutures.push(pyodide.loadPackagesFromImports(script));
-      for (const f of packageFutures) {
-        await f;
+      await Promise.all(packageFutures);
+      if (patch_http) {
+        await pyodide.runPythonAsync(`
+        from pyodide_http import patch_all
+        patch_all()
+        `);
       }
 
       if (files) {

diff --git a/gui/src/app/core/Scripting/pyodide/pyodideWorkerTypes.ts b/gui/src/app/core/Scripting/pyodide/pyodideWorkerTypes.ts
@@ -17,6 +17,7 @@ export type MessageToPyodideWorker = {
   spData: Record<string, any> | undefined;
   spRunSettings: PyodideRunSettings;
   files: Record<string, string> | undefined;
+  interruptBuffer: Uint8Array | undefined;
 };
 
 export const isMessageToPyodideWorker = (