Consolidate workers for data.py and analysis.py

flatironinstitute · Jul 3, 2024 · 5e4152d · 5e4152d
1 parent e6d5b7f
commit 5e4152d
Show file tree

Hide file tree

Showing 7 changed files with 444 additions and 312 deletions.
diff --git a/gui/src/app/prototypes/AnalysisPyPrototype/AnalysisPyFileEditor.tsx b/gui/src/app/prototypes/AnalysisPyPrototype/AnalysisPyFileEditor.tsx
@@ -7,9 +7,8 @@ import {
 } from "react";
 import { PlayArrow } from "@mui/icons-material";
 import TextEditor, { ToolbarItem } from "../../FileEditor/TextEditor";
-// https://vitejs.dev/guide/assets#importing-script-as-a-worker
-// https://vitejs.dev/guide/assets#importing-asset-as-url
-import analysisPyWorkerURL from "./analysisPyWorker?worker&url";
+import { PydodideWorkerStatus } from "../pyodideWorker/pyodideWorkerTypes";
+import PyodideWorkerInterface from "../pyodideWorker/pyodideWorkerInterface";
 
 type Props = {
   fileName: string;
@@ -34,64 +33,53 @@ const AnalysisPyFileEditor: FunctionComponent<Props> = ({
   height,
   outputDiv,
 }) => {
-  const [status, setStatus] = useState<
-    "idle" | "loading" | "running" | "completed" | "failed"
-  >("idle");
+  const [status, setStatus] = useState<PydodideWorkerStatus>("idle");
 
-  const [analysisPyWorker, setAnalysisPyWorker] = useState<Worker | undefined>(
-    undefined,
-  );
+  const [analysisPyWorker, setAnalysisPyWorker] = useState<
+    PyodideWorkerInterface | undefined
+  >(undefined);
 
   // worker creation
   useEffect(() => {
-    const worker = new Worker(analysisPyWorkerURL, {
-      name: "dataPyWorker",
-      type: "module",
-    });
-    setAnalysisPyWorker(worker);
-    return () => {
-      console.log("terminating dataPy worker");
-      worker.terminate();
-    };
-  }, []);
-
-  // message handling
-  useEffect(() => {
-    if (!analysisPyWorker) return;
-
-    analysisPyWorker.onmessage = (e: MessageEvent) => {
-      const dd = e.data;
-      if (dd.type === "setStatus") {
-        setStatus(dd.status);
-      } else if (dd.type === "addImage") {
-        const b64 = dd.image;
-        const imageUrl = `data:image/png;base64,${b64}`;
-
-        const img = document.createElement("img");
-        img.src = imageUrl;
-
-        const divElement = document.createElement("div");
-        divElement.appendChild(img);
-        outputDiv?.appendChild(divElement);
-      } else if (dd.type === "stdout") {
-        console.log(dd.data);
+    const worker = PyodideWorkerInterface.create("analysis.py", {
+      onStdout: (x) => {
+        console.log(x);
         const divElement = document.createElement("div");
         divElement.style.color = "blue";
         const preElement = document.createElement("pre");
         divElement.appendChild(preElement);
-        preElement.textContent = dd.data;
+        preElement.textContent = x;
         outputDiv?.appendChild(divElement);
-      } else if (dd.type === "stderr") {
-        console.error(dd.data);
+      },
+      onStderr: (x) => {
+        console.error(x);
         const divElement = document.createElement("div");
         divElement.style.color = "red";
         const preElement = document.createElement("pre");
         divElement.appendChild(preElement);
-        preElement.textContent = dd.data;
+        preElement.textContent = x;
+        outputDiv?.appendChild(divElement);
+      },
+      onStatus: (status) => {
+        setStatus(status);
+      },
+      onImage: (image) => {
+        const b64 = image;
+        const imageUrl = `data:image/png;base64,${b64}`;
+
+        const img = document.createElement("img");
+        img.src = imageUrl;
+
+        const divElement = document.createElement("div");
+        divElement.appendChild(img);
         outputDiv?.appendChild(divElement);
-      }
+      },
+    });
+    setAnalysisPyWorker(worker);
+    return () => {
+      worker.destroy();
     };
-  }, [analysisPyWorker, outputDiv]);
+  }, [outputDiv]);
 
   const handleRun = useCallback(async () => {
     if (status === "running") {
@@ -100,12 +88,11 @@ const AnalysisPyFileEditor: FunctionComponent<Props> = ({
     if (editedFileContent !== fileContent) {
       throw new Error("Cannot run edited code");
     }
-    if (outputDiv) outputDiv.innerHTML = "";
-    analysisPyWorker?.postMessage({
-      type: "run",
-      code: fileContent,
-    });
-  }, [editedFileContent, fileContent, status, analysisPyWorker, outputDiv]);
+    if (!analysisPyWorker) {
+      throw new Error("analysisPyWorker is not defined");
+    }
+    analysisPyWorker.run(fileContent);
+  }, [editedFileContent, fileContent, status, analysisPyWorker]);
   const toolbarItems: ToolbarItem[] = useMemo(() => {
     const ret: ToolbarItem[] = [];
     const runnable = fileContent === editedFileContent && outputDiv;

diff --git a/gui/src/app/prototypes/AnalysisPyPrototype/analysisPyWorker.ts b/gui/src/app/prototypes/AnalysisPyPrototype/analysisPyWorker.ts
diff --git a/gui/src/app/prototypes/DataPyPrototype/DataPyFileEditor.tsx b/gui/src/app/prototypes/DataPyPrototype/DataPyFileEditor.tsx
@@ -7,9 +7,8 @@ import {
 } from "react";
 import { PlayArrow } from "@mui/icons-material";
 import TextEditor, { ToolbarItem } from "../../FileEditor/TextEditor";
-// https://vitejs.dev/guide/assets#importing-script-as-a-worker
-// https://vitejs.dev/guide/assets#importing-asset-as-url
-import dataPyWorkerURL from "./dataPyWorker?worker&url";
+import PyodideWorkerInterface from "../pyodideWorker/pyodideWorkerInterface";
+import { PydodideWorkerStatus } from "../pyodideWorker/pyodideWorkerTypes";
 
 type Props = {
   fileName: string;
@@ -34,44 +33,31 @@ const DataPyFileEditor: FunctionComponent<Props> = ({
   width,
   height,
 }) => {
-  const [status, setStatus] = useState<
-    "idle" | "loading" | "running" | "completed" | "failed"
-  >("idle");
+  const [status, setStatus] = useState<PydodideWorkerStatus>("idle");
 
-  const [dataPyWorker, setDataPyWorker] = useState<Worker | undefined>(
-    undefined,
-  );
+  const [dataPyWorker, setDataPyWorker] = useState<
+    PyodideWorkerInterface | undefined
+  >(undefined);
 
   // worker creation
   useEffect(() => {
-    const worker = new Worker(dataPyWorkerURL, {
-      name: "dataPyWorker",
-      type: "module",
+    const worker = PyodideWorkerInterface.create("data.py", {
+      onStdout: (x) => {
+        console.log(x);
+      },
+      onStderr: (x) => {
+        console.error(x);
+      },
+      onStatus: (status) => {
+        setStatus(status);
+      },
+      onData: setData,
     });
     setDataPyWorker(worker);
     return () => {
-      console.log("terminating dataPy worker");
-      worker.terminate();
+      worker.destroy();
     };
-  }, []);
-
-  // message handling
-  useEffect(() => {
-    if (!dataPyWorker) return;
-
-    dataPyWorker.onmessage = (e: MessageEvent) => {
-      const dd = e.data;
-      if (dd.type === "setStatus") {
-        setStatus(dd.status);
-      } else if (dd.type === "setData") {
-        setData && setData(dd.data);
-      } else if (dd.type === "stdout") {
-        console.log(dd.data);
-      } else if (dd.type === "stderr") {
-        console.error(dd.data);
-      }
-    };
-  }, [dataPyWorker, setData]);
+  }, [setData]);
 
   const handleRun = useCallback(async () => {
     if (status === "running") {
@@ -80,10 +66,10 @@ const DataPyFileEditor: FunctionComponent<Props> = ({
     if (editedFileContent !== fileContent) {
       throw new Error("Cannot run edited code");
     }
-    dataPyWorker?.postMessage({
-      type: "run",
-      code: fileContent,
-    });
+    if (!dataPyWorker) {
+      throw new Error("dataPyWorker is not defined");
+    }
+    dataPyWorker.run(fileContent);
   }, [editedFileContent, fileContent, status, dataPyWorker]);
   const toolbarItems: ToolbarItem[] = useMemo(() => {
     const ret: ToolbarItem[] = [];