Skip to content

Commit bc27dcf

Browse files
epattersepatters
authored
Merge pull request #285 from ToposInstitute/jupyter-refactor
Reactive helpers for Jupyter kernels
2 parents f2727c1 + 8e4bb24 commit bc27dcf

File tree

2 files changed

+139
-95
lines changed

2 files changed

+139
-95
lines changed

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

Lines changed: 29 additions & 95 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
1-
import type { IReplyErrorContent } from "@jupyterlab/services/lib/kernel/messages";
2-
import { For, Match, Show, Switch, createMemo, createResource, onCleanup } from "solid-js";
1+
import { For, Match, Show, Switch, createMemo } from "solid-js";
32
import { isMatching } from "ts-pattern";
43

54
import type { DiagramAnalysisProps } from "../../analysis";
@@ -22,6 +21,7 @@ import type { ModelJudgment, MorphismDecl } from "../../model";
2221
import type { DiagramAnalysisMeta } from "../../theory";
2322
import { uniqueIndexArray } from "../../util/indexing";
2423
import { PDEPlot2D, type PDEPlotData2D } from "../../visualization";
24+
import { createJuliaKernel, executeAndRetrieve } from "./jupyter";
2525

2626
import Loader from "lucide-solid/icons/loader";
2727
import RotateCcw from "lucide-solid/icons/rotate-ccw";
@@ -31,19 +31,14 @@ import "./decapodes.css";
3131
import "./simulation.css";
3232

3333
/** Configuration for a Decapodes analysis of a diagram. */
34-
export type DecapodesContent = JupyterSettings & {
34+
export type DecapodesContent = {
3535
domain: string | null;
3636
mesh: string | null;
3737
initialConditions: Record<string, string>;
3838
plotVariables: Record<string, boolean>;
3939
scalars: Record<string, number>;
4040
};
4141

42-
type JupyterSettings = {
43-
baseUrl?: string;
44-
token?: string;
45-
};
46-
4742
export function configureDecapodes(options: {
4843
id?: string;
4944
name?: string;
@@ -73,88 +68,37 @@ export function configureDecapodes(options: {
7368
*/
7469
export function Decapodes(props: DiagramAnalysisProps<DecapodesContent>) {
7570
// Step 1: Start the Julia kernel.
76-
const [kernel, { refetch: restartKernel }] = createResource(async () => {
77-
const jupyter = await import("@jupyterlab/services");
78-
79-
const serverSettings = jupyter.ServerConnection.makeSettings({
80-
baseUrl: props.content.baseUrl ?? "http://127.0.0.1:8888",
81-
token: props.content.token ?? "",
82-
});
83-
84-
const kernelManager = new jupyter.KernelManager({ serverSettings });
85-
const kernel = await kernelManager.startNew({ name: "julia-1.11" });
86-
87-
return kernel;
71+
const [kernel, restartKernel] = createJuliaKernel({
72+
baseUrl: "http://127.0.0.1:8888",
73+
token: "",
8874
});
8975

90-
onCleanup(() => kernel()?.shutdown());
91-
9276
// Step 2: Run initialization code in the kernel.
9377
const startedKernel = () => (kernel.error ? undefined : kernel());
9478

95-
const [options] = createResource(startedKernel, async (kernel) => {
96-
// Request that the kernel run code to initialize the service.
97-
const future = kernel.requestExecute({ code: initCode });
98-
99-
// Look for simulation options as output from the kernel.
100-
let options: SimulationOptions | undefined;
101-
future.onIOPub = (msg) => {
102-
if (msg.header.msg_type === "execute_result") {
103-
const content = msg.content as JsonDataContent<SimulationOptions>;
104-
options = content["data"]?.["application/json"];
105-
}
106-
};
107-
108-
const reply = await future.done;
109-
if (reply.content.status === "error") {
110-
await kernel.shutdown();
111-
throw new Error(formatError(reply.content));
112-
}
113-
if (!options) {
114-
throw new Error("Allowed options not received after initialization");
115-
}
116-
return {
79+
const [options] = executeAndRetrieve(
80+
startedKernel,
81+
makeInitCode,
82+
(options: SimulationOptions) => ({
11783
domains: uniqueIndexArray(options.domains, (domain) => domain.name),
118-
};
119-
});
84+
}),
85+
);
12086

12187
// Step 3: Run the simulation in the kernel!
12288
const initedKernel = () =>
12389
kernel.error || options.error || options.loading ? undefined : kernel();
12490

125-
const [result, { refetch: rerunSimulation }] = createResource(initedKernel, async (kernel) => {
126-
// Construct the data to send to kernel.
127-
const simulationData = makeSimulationData(props.liveDiagram, props.content);
128-
if (!simulationData) {
129-
return undefined;
130-
}
131-
console.log(JSON.parse(JSON.stringify(simulationData)));
132-
// Request that the kernel run a simulation with the given data.
133-
const future = kernel.requestExecute({
134-
code: makeSimulationCode(simulationData),
135-
});
136-
137-
// Look for simulation results as output from the kernel.
138-
let result: PDEPlotData2D | undefined;
139-
future.onIOPub = (msg) => {
140-
if (
141-
msg.header.msg_type === "execute_result" &&
142-
msg.parent_header.msg_id === future.msg.header.msg_id
143-
) {
144-
const content = msg.content as JsonDataContent<PDEPlotData2D>;
145-
result = content["data"]?.["application/json"];
91+
const [result, rerunSimulation] = executeAndRetrieve(
92+
initedKernel,
93+
() => {
94+
const simulationData = makeSimulationData(props.liveDiagram, props.content);
95+
if (!simulationData) {
96+
return undefined;
14697
}
147-
};
148-
149-
const reply = await future.done;
150-
if (reply.content.status === "error") {
151-
throw new Error(formatError(reply.content));
152-
}
153-
if (!result) {
154-
throw new Error("Result not received from the simulator");
155-
}
156-
return result;
157-
});
98+
return makeSimulationCode(simulationData);
99+
},
100+
(data: PDEPlotData2D) => data,
101+
);
158102

159103
const obDecls = createMemo<DiagramObjectDecl[]>(() =>
160104
props.liveDiagram.formalJudgments().filter((jgmt) => jgmt.tag === "object"),
@@ -346,17 +290,6 @@ export function Decapodes(props: DiagramAnalysisProps<DecapodesContent>) {
346290
);
347291
}
348292

349-
const formatError = (content: IReplyErrorContent): string =>
350-
// Trackback list already includes `content.evalue`.
351-
content.traceback.join("\n");
352-
353-
/** JSON data returned from a Jupyter kernel. */
354-
type JsonDataContent<T> = {
355-
data?: {
356-
"application/json"?: T;
357-
};
358-
};
359-
360293
/** Options supported by Decapodes, defined by the Julia service. */
361294
type SimulationOptions = {
362295
/** Geometric domains supported by Decapodes. */
@@ -400,14 +333,15 @@ type SimulationData = {
400333
};
401334

402335
/** Julia code run after kernel is started. */
403-
const initCode = `
404-
import IJulia
405-
IJulia.register_jsonmime(MIME"application/json"())
336+
const makeInitCode = () =>
337+
`
338+
import IJulia
339+
IJulia.register_jsonmime(MIME"application/json"())
406340
407-
using AlgebraicJuliaService
341+
using AlgebraicJuliaService
408342
409-
JsonValue(supported_decapodes_geometries())
410-
`;
343+
JsonValue(supported_decapodes_geometries())
344+
`;
411345

412346
/** Julia code run to perform a simulation. */
413347
const makeSimulationCode = (data: SimulationData) =>

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

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
import type { ServerConnection } from "@jupyterlab/services";
2+
import type { IKernelConnection, IKernelOptions } from "@jupyterlab/services/lib/kernel/kernel";
3+
import {
4+
type Accessor,
5+
type Resource,
6+
type ResourceReturn,
7+
createResource,
8+
onCleanup,
9+
} from "solid-js";
10+
11+
type ResourceRefetch<T> = ResourceReturn<T>[1]["refetch"];
12+
13+
type ServerSettings = Parameters<typeof ServerConnection.makeSettings>[0];
14+
15+
/** Create a Jupyter kernel in a reactive context.
16+
17+
Returns a kernel as a Solid.js resource and a callback to restart the kernel.
18+
The kernel is automatically shut down when the component is unmounted.
19+
*/
20+
export function createKernel(
21+
serverOptions: ServerSettings,
22+
kernelOptions: IKernelOptions,
23+
): [Resource<IKernelConnection>, ResourceRefetch<IKernelConnection>] {
24+
const [kernel, { refetch: restartKernel }] = createResource(async () => {
25+
const jupyter = await import("@jupyterlab/services");
26+
27+
const serverSettings = jupyter.ServerConnection.makeSettings(serverOptions);
28+
29+
const kernelManager = new jupyter.KernelManager({ serverSettings });
30+
const kernel = await kernelManager.startNew(kernelOptions);
31+
32+
return kernel;
33+
});
34+
35+
onCleanup(() => kernel()?.shutdown());
36+
37+
return [kernel, restartKernel];
38+
}
39+
40+
/** Create a Julia kernel in a reactive context. */
41+
export function createJuliaKernel(serverOptions: ServerSettings) {
42+
return createKernel(serverOptions, {
43+
// XXX: Do I have to specify the Julia version?
44+
name: "julia-1.11",
45+
});
46+
}
47+
48+
/** Execute code in a Jupyter kernel and retrieve JSON data reactively.
49+
50+
Assumes that the computation will return JSON data using the "application/json"
51+
MIME type in Jupyter. Returns the post-processed data as a Solid.js resource and
52+
a callback to rerun the computation.
53+
54+
The resource depends reactively on the kernel: if the kernel changes, the code
55+
will be automatically re-executed. It does *not* depend reactively on the code.
56+
If the code changes, it must be rerun manually.
57+
*/
58+
export function executeAndRetrieve<S, T>(
59+
kernel: Accessor<IKernelConnection | undefined>,
60+
executeCode: Accessor<string | undefined>,
61+
postprocess: (data: S) => T,
62+
): [Resource<T | undefined>, ResourceRefetch<T>] {
63+
const [data, { refetch: reexecute }] = createResource(kernel, async (kernel) => {
64+
// Request that kernel execute code, if defined.
65+
const code = executeCode();
66+
if (code === undefined) {
67+
return undefined;
68+
}
69+
const future = kernel.requestExecute({ code });
70+
71+
// Set up handler for result from kernel.
72+
let result: { data: S } | undefined;
73+
future.onIOPub = (msg) => {
74+
if (
75+
msg.header.msg_type === "execute_result" &&
76+
msg.parent_header.msg_id === future.msg.header.msg_id
77+
) {
78+
const content = msg.content as JsonDataContent<S>;
79+
const data = content["data"]?.["application/json"];
80+
if (data !== undefined) {
81+
result = { data };
82+
}
83+
}
84+
};
85+
86+
// Wait for execution to finish and process result.
87+
const reply = await future.done;
88+
if (reply.content.status === "abort") {
89+
throw new Error("Execution was aborted");
90+
}
91+
if (reply.content.status === "error") {
92+
// Trackback list already includes `reply.content.evalue`.
93+
const msg = reply.content.traceback.join("\n");
94+
throw new Error(msg);
95+
}
96+
if (result === undefined) {
97+
throw new Error("Data was not received from the kernel");
98+
}
99+
return postprocess(result.data);
100+
});
101+
102+
return [data, reexecute];
103+
}
104+
105+
/** JSON data returned from a Jupyter kernel. */
106+
type JsonDataContent<T> = {
107+
data?: {
108+
"application/json"?: T;
109+
};
110+
};

0 commit comments

Comments
 (0)