Skip to content

Commit

Permalink
REFACTOR: Extract Jupyter code execution into independent function.
Browse files Browse the repository at this point in the history
  • Loading branch information
@epatters
epatters committed Dec 6, 2024
1 parent 0c0c1d7 commit f2a80ff
Show file tree
Hide file tree
Showing 2 changed files with 98 additions and 80 deletions.
101 changes: 25 additions & 76 deletions packages/frontend/src/stdlib/analyses/decapodes.tsx
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
import type { IReplyErrorContent } from "@jupyterlab/services/lib/kernel/messages";
import { For, Match, Show, Switch, createMemo, createResource } from "solid-js";
import { For, Match, Show, Switch, createMemo } from "solid-js";
import { isMatching } from "ts-pattern";

import type { DiagramAnalysisProps } from "../../analysis";
Expand All @@ -22,7 +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 { createKernel } from "./jupyter";
import { createKernel, executeAndRetrieve } from "./jupyter";

import Loader from "lucide-solid/icons/loader";
import RotateCcw from "lucide-solid/icons/rotate-ccw";
Expand Down Expand Up @@ -82,69 +81,29 @@ export function Decapodes(props: DiagramAnalysisProps<DecapodesContent>) {
// 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"),
Expand Down Expand Up @@ -336,17 +295,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. */
Expand Down Expand Up @@ -390,14 +338,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) =>
Expand Down
77 changes: 73 additions & 4 deletions packages/frontend/src/stdlib/analyses/jupyter.ts
View file
Original file line number Diff line number Diff line change
@@ -1,15 +1,21 @@
import type { ServerConnection } from "@jupyterlab/services";
import type { IKernelConnection, IKernelOptions } from "@jupyterlab/services/lib/kernel/kernel";
import { type Resource, type ResourceReturn, createResource, onCleanup } from "solid-js";
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 as a Solid resource.
/** Create a Jupyter kernel in a reactive context.
Returns a handle to the resource and a callback to restart the kernel. The
kernel is automatically shut down when the component is unmounted.
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,
Expand All @@ -30,3 +36,66 @@ export function createKernel(

return [kernel, restartKernel];
}

/** Execute code in a Jupyter kernel and retrieve JSON data reactively.
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;
};
};

0 comments on commit f2a80ff

Please sign in to comment.