Skip to content

Commit

Permalink
Merge pull request #285 from ToposInstitute/jupyter-refactor
Browse files Browse the repository at this point in the history
Reactive helpers for Jupyter kernels
  • Loading branch information
epatters authored Dec 6, 2024
2 parents f2727c1 + 8e4bb24 commit bc27dcf
Show file tree
Hide file tree
Showing 2 changed files with 139 additions and 95 deletions.
124 changes: 29 additions & 95 deletions packages/frontend/src/stdlib/analyses/decapodes.tsx
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, onCleanup } 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,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";
Expand All @@ -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;
Expand Down Expand Up @@ -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"),
Expand Down Expand Up @@ -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. */
Expand Down Expand Up @@ -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) =>
Expand Down
110 changes: 110 additions & 0 deletions packages/frontend/src/stdlib/analyses/jupyter.ts
Original file line number Diff line number Diff line change
@@ -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;
};
};

0 comments on commit bc27dcf

Please sign in to comment.