Reactive helpers for Jupyter kernels

packages/frontend/src/stdlib/analyses/decapodes.tsx

@@ -1,5 +1,4 @@
-import type { IReplyErrorContent } from "@jupyterlab/services/lib/kernel/messages";
-import { For, Match, Show, Switch, createMemo, createResource, onCleanup } from "solid-js";
+import { For, Match, Show, Switch, createMemo } from "solid-js";
 import { isMatching } from "ts-pattern";
 
 import type { DiagramAnalysisProps } from "../../analysis";
@@ -22,6 +21,7 @@ import type { ModelJudgment, MorphismDecl } from "../../model";
 import type { DiagramAnalysisMeta } from "../../theory";
 import { uniqueIndexArray } from "../../util/indexing";
 import { PDEPlot2D, type PDEPlotData2D } from "../../visualization";
+import { createJuliaKernel, executeAndRetrieve } from "./jupyter";
 
 import Loader from "lucide-solid/icons/loader";
 import RotateCcw from "lucide-solid/icons/rotate-ccw";
@@ -31,19 +31,14 @@ import "./decapodes.css";
 import "./simulation.css";
 
 /** Configuration for a Decapodes analysis of a diagram. */
-export type DecapodesContent = JupyterSettings & {
+export type DecapodesContent = {
     domain: string | null;
     mesh: string | null;
     initialConditions: Record<string, string>;
     plotVariables: Record<string, boolean>;
     scalars: Record<string, number>;
 };
 
-type JupyterSettings = {
-    baseUrl?: string;
-    token?: string;
-};
-
 export function configureDecapodes(options: {
     id?: string;
     name?: string;
@@ -73,88 +68,37 @@ export function configureDecapodes(options: {
  */
 export function Decapodes(props: DiagramAnalysisProps<DecapodesContent>) {
     // Step 1: Start the Julia kernel.
-    const [kernel, { refetch: restartKernel }] = createResource(async () => {
-        const jupyter = await import("@jupyterlab/services");
-
-        const serverSettings = jupyter.ServerConnection.makeSettings({
-            baseUrl: props.content.baseUrl ?? "http://127.0.0.1:8888",
-            token: props.content.token ?? "",
-        });
-
-        const kernelManager = new jupyter.KernelManager({ serverSettings });
-        const kernel = await kernelManager.startNew({ name: "julia-1.11" });
-
-        return kernel;
+    const [kernel, restartKernel] = createJuliaKernel({
+        baseUrl: "http://127.0.0.1:8888",
+        token: "",
     });
 
-    onCleanup(() => kernel()?.shutdown());
-
     // Step 2: Run initialization code in the kernel.
     const startedKernel = () => (kernel.error ? undefined : kernel());
 
-    const [options] = createResource(startedKernel, async (kernel) => {
-        // Request that the kernel run code to initialize the service.
-        const future = kernel.requestExecute({ code: initCode });
-
-        // Look for simulation options as output from the kernel.
-        let options: SimulationOptions | undefined;
-        future.onIOPub = (msg) => {
-            if (msg.header.msg_type === "execute_result") {
-                const content = msg.content as JsonDataContent<SimulationOptions>;
-                options = content["data"]?.["application/json"];
-            }
-        };
-
-        const reply = await future.done;
-        if (reply.content.status === "error") {
-            await kernel.shutdown();
-            throw new Error(formatError(reply.content));
-        }
-        if (!options) {
-            throw new Error("Allowed options not received after initialization");
-        }
-        return {
+    const [options] = executeAndRetrieve(
+        startedKernel,
+        makeInitCode,
+        (options: SimulationOptions) => ({
             domains: uniqueIndexArray(options.domains, (domain) => domain.name),
-        };
-    });
+        }),
+    );
 
     // Step 3: Run the simulation in the kernel!
     const initedKernel = () =>
         kernel.error || options.error || options.loading ? undefined : kernel();
 
-    const [result, { refetch: rerunSimulation }] = createResource(initedKernel, async (kernel) => {
-        // Construct the data to send to kernel.
-        const simulationData = makeSimulationData(props.liveDiagram, props.content);
-        if (!simulationData) {
-            return undefined;
-        }
-        console.log(JSON.parse(JSON.stringify(simulationData)));
-        // Request that the kernel run a simulation with the given data.
-        const future = kernel.requestExecute({
-            code: makeSimulationCode(simulationData),
-        });
-
-        // Look for simulation results as output from the kernel.
-        let result: PDEPlotData2D | undefined;
-        future.onIOPub = (msg) => {
-            if (
-                msg.header.msg_type === "execute_result" &&
-                msg.parent_header.msg_id === future.msg.header.msg_id
-            ) {
-                const content = msg.content as JsonDataContent<PDEPlotData2D>;
-                result = content["data"]?.["application/json"];
+    const [result, rerunSimulation] = executeAndRetrieve(
+        initedKernel,
+        () => {
+            const simulationData = makeSimulationData(props.liveDiagram, props.content);
+            if (!simulationData) {
+                return undefined;
             }
-        };
-
-        const reply = await future.done;
-        if (reply.content.status === "error") {
-            throw new Error(formatError(reply.content));
-        }
-        if (!result) {
-            throw new Error("Result not received from the simulator");
-        }
-        return result;
-    });
+            return makeSimulationCode(simulationData);
+        },
+        (data: PDEPlotData2D) => data,
+    );
 
     const obDecls = createMemo<DiagramObjectDecl[]>(() =>
         props.liveDiagram.formalJudgments().filter((jgmt) => jgmt.tag === "object"),
@@ -346,17 +290,6 @@ export function Decapodes(props: DiagramAnalysisProps<DecapodesContent>) {
     );
 }
 
-const formatError = (content: IReplyErrorContent): string =>
-    // Trackback list already includes `content.evalue`.
-    content.traceback.join("\n");
-
-/** JSON data returned from a Jupyter kernel. */
-type JsonDataContent<T> = {
-    data?: {
-        "application/json"?: T;
-    };
-};
-
 /** Options supported by Decapodes, defined by the Julia service. */
 type SimulationOptions = {
     /** Geometric domains supported by Decapodes. */
@@ -400,14 +333,15 @@ type SimulationData = {
 };
 
 /** Julia code run after kernel is started. */
-const initCode = `
-import IJulia
-IJulia.register_jsonmime(MIME"application/json"())
+const makeInitCode = () =>
+    `
+    import IJulia
+    IJulia.register_jsonmime(MIME"application/json"())
 
-using AlgebraicJuliaService
+    using AlgebraicJuliaService
 
-JsonValue(supported_decapodes_geometries())
-`;
+    JsonValue(supported_decapodes_geometries())
+    `;
 
 /** Julia code run to perform a simulation. */
 const makeSimulationCode = (data: SimulationData) =>

packages/frontend/src/stdlib/analyses/jupyter.ts

@@ -0,0 +1,110 @@
+import type { ServerConnection } from "@jupyterlab/services";
+import type { IKernelConnection, IKernelOptions } from "@jupyterlab/services/lib/kernel/kernel";
+import {
+    type Accessor,
+    type Resource,
+    type ResourceReturn,
+    createResource,
+    onCleanup,
+} from "solid-js";
+
+type ResourceRefetch<T> = ResourceReturn<T>[1]["refetch"];
+
+type ServerSettings = Parameters<typeof ServerConnection.makeSettings>[0];
+
+/** Create a Jupyter kernel in a reactive context.
+
+Returns a kernel as a Solid.js resource and a callback to restart the kernel.
+The kernel is automatically shut down when the component is unmounted.
+ */
+export function createKernel(
+    serverOptions: ServerSettings,
+    kernelOptions: IKernelOptions,
+): [Resource<IKernelConnection>, ResourceRefetch<IKernelConnection>] {
+    const [kernel, { refetch: restartKernel }] = createResource(async () => {
+        const jupyter = await import("@jupyterlab/services");
+
+        const serverSettings = jupyter.ServerConnection.makeSettings(serverOptions);
+
+        const kernelManager = new jupyter.KernelManager({ serverSettings });
+        const kernel = await kernelManager.startNew(kernelOptions);
+
+        return kernel;
+    });
+
+    onCleanup(() => kernel()?.shutdown());
+
+    return [kernel, restartKernel];
+}
+
+/** Create a Julia kernel in a reactive context. */
+export function createJuliaKernel(serverOptions: ServerSettings) {
+    return createKernel(serverOptions, {
+        // XXX: Do I have to specify the Julia version?
+        name: "julia-1.11",
+    });
+}
+
+/** Execute code in a Jupyter kernel and retrieve JSON data reactively.
+
+Assumes that the computation will return JSON data using the "application/json"
+MIME type in Jupyter. Returns the post-processed data as a Solid.js resource and
+a callback to rerun the computation.
+
+The resource depends reactively on the kernel: if the kernel changes, the code
+will be automatically re-executed. It does *not* depend reactively on the code.
+If the code changes, it must be rerun manually.
+ */
+export function executeAndRetrieve<S, T>(
+    kernel: Accessor<IKernelConnection | undefined>,
+    executeCode: Accessor<string | undefined>,
+    postprocess: (data: S) => T,
+): [Resource<T | undefined>, ResourceRefetch<T>] {
+    const [data, { refetch: reexecute }] = createResource(kernel, async (kernel) => {
+        // Request that kernel execute code, if defined.
+        const code = executeCode();
+        if (code === undefined) {
+            return undefined;
+        }
+        const future = kernel.requestExecute({ code });
+
+        // Set up handler for result from kernel.
+        let result: { data: S } | undefined;
+        future.onIOPub = (msg) => {
+            if (
+                msg.header.msg_type === "execute_result" &&
+                msg.parent_header.msg_id === future.msg.header.msg_id
+            ) {
+                const content = msg.content as JsonDataContent<S>;
+                const data = content["data"]?.["application/json"];
+                if (data !== undefined) {
+                    result = { data };
+                }
+            }
+        };
+
+        // Wait for execution to finish and process result.
+        const reply = await future.done;
+        if (reply.content.status === "abort") {
+            throw new Error("Execution was aborted");
+        }
+        if (reply.content.status === "error") {
+            // Trackback list already includes `reply.content.evalue`.
+            const msg = reply.content.traceback.join("\n");
+            throw new Error(msg);
+        }
+        if (result === undefined) {
+            throw new Error("Data was not received from the kernel");
+        }
+        return postprocess(result.data);
+    });
+
+    return [data, reexecute];
+}
+
+/** JSON data returned from a Jupyter kernel. */
+type JsonDataContent<T> = {
+    data?: {
+        "application/json"?: T;
+    };
+};