Spaces:

cardef
/

grid-lab

Running

grid-lab / engine_worker.js

Work in a private Space inside the Hub page: pass the page's ?__sign= to the app module, the fonts, the worker and its libraries

dbdc890 verified 4 days ago

Raw History Blame Contribute Delete

5.98 kB

// The engine worker (ES module, ENGINE_CONTRACT.md / engine_api.js): GENCO as ONNX in onnxruntime-web (WASM) and the
// app's scenario builder, "dsq" repair and fast correction (whatif_app's modules, copied verbatim) in Pyodide, both
// self-hosted under assets/engine/ and both in this worker, off the page's thread.
//   messages in:  {id, type: "init" | "preview" | "check" | "compare" | "exact" | "cancel", knobs?, opts?, gen?}
//   messages out: {id, ok: true, result} | {id, ok: false, error};  {type: "status", stage, detail, progress}
// Threads: onnxruntime-web uses 2 WASM threads when the page is cross-origin isolated (COOP + COEP), else 1.
// Requests run one at a time, in arrival order (Pyodide is single-threaded).

import { createEngine, transferables } from "./assets/engine/engine_core.js";

const BASE = new URL("./assets/engine/", import.meta.url).href;
// a private Space: this worker's URL carries the page's signature (engine_api.js); Pyodide and onnxruntime-web import
// their own files with their URL as the Referer, so they are imported signed too
const SIGN = new URLSearchParams(self.location.search).get("__sign");
const signed = (u) => (SIGN ? `${u}?__sign=${encodeURIComponent(SIGN)}` : u);
const STAGES = { start: [0.02, "starting Python (Pyodide)"], pyodide: [0.25, "loading numpy, scipy, pandas"],
  packages: [0.6, "fetching the grid and GENCO"], fetch_rest: [0.65, "loading the engine"],
  engine: [0.8, "loading GENCO (onnxruntime-web)"], model: [0.9, "warming up"], warmup: [1.0, "ready"] };

let engine = null;
let queue = Promise.resolve();
let seq = 0, lastNewer = 0;          // the sequence number of the latest preview / check / compare / cancel received

const status = (stage, detail, progress) => self.postMessage({ type: "status", stage, detail, progress });

async function start(opts = {}) {
  status("loading", "loading the engine", 0.01);
  const [{ loadPyodide }, ortModule] = await Promise.all([import(signed(BASE + "pyodide/pyodide.mjs")), import(signed(BASE + "ort/ort.wasm.min.mjs"))]);
  const ort = ortModule.default ?? ortModule;
  ort.env.wasm.wasmPaths = BASE + "ort/";
  // GENCO's WASM threads: the visitor's cores, up to 8, when the page is cross-origin isolated (else 1). Measured on a
  // 6-core desktop: 2 threads ~550-600 ms per forward, 4-6 threads ~380-420 ms, 8 threads (more than the cores)
  // ~1,080 ms, so never more than navigator.hardwareConcurrency. Outputs are bit-identical for 1-8 threads.
  const cores = Number.isInteger(self.navigator?.hardwareConcurrency) ? self.navigator.hardwareConcurrency : 2;
  const threads = Number.isInteger(opts.threads) ? opts.threads : (self.crossOriginIsolated ? Math.max(1, Math.min(cores, 8)) : 1);
  const fetchBytes = async (url) => {
    const r = await fetch(url);
    if (!r.ok) throw new Error(`${url}: HTTP ${r.status}`);
    return new Uint8Array(await r.arrayBuffer());
  };
  const e = await createEngine({
    base: BASE, loadPyodide, ort, fetchBytes, threads, warmup: opts.warmup !== false,
    onProgress: ({ stage }) => { const s = STAGES[stage]; if (s) status(stage === "warmup" ? "ready" : "loading", s[1], s[0]); },
  });
  status("ready", "engine ready", 1);
  return e;
}

async function handle(m) {
  if (m.type === "init") {
    if (!engine) engine = start(m.opts || {});
    const e = await engine;
    return { ...e.initInfo };
  }
  if (!engine) engine = start({});
  const e = await engine;
  if (m.type === "preview" || m.type === "compare") return e[m.type](m.knobs);
  // progress messages for the page's live timers: {type: "progress", id, stage}, stage "check" (a check's time_ms
  // starts), "setup" (an exact loads casadi or builds its model), "solve" (an exact's interior point starts)
  const notify = (stage) => self.postMessage({ type: "progress", id: m.id, stage });
  if (m.type === "check") {
    const opts = { ...(m.opts || {}), notify };
    if (Number.isInteger(m.gen)) opts.gen = m.gen;
    if (lastNewer > m._seq) {                              // superseded while queued (a newer preview, check, ...)
      const n = e.initInfo.info, nan = (k) => new Float32Array(k).fill(NaN);
      return { status: "cancelled", reason: "superseded before it started", cached: false, converged: false, cost: null,
        max_loading: null, n_viol: null, viol_detail: "cancelled", feasible: false, loops: 0, time_ms: 0,
        line_loading: nan(n.n_branch), bus_vm: nan(n.n_bus), gen_p: nan(n.n_gen), v_out: new Int32Array(0) };
    }
    return e.check(m.knobs, opts);
  }
  if (m.type === "exact") {
    const opts = { ...(m.opts || {}), notify };
    if (Number.isInteger(m.gen)) opts.gen = m.gen;
    await new Promise((r) => setTimeout(r, 20));           // let messages posted right after this one arrive first
    if (lastNewer > m._seq) {                              // superseded while queued: not started
      const start = opts.start ?? "check";
      return { status: "cancelled", start, start_used: start, reason: "superseded before it started", cached: false,
        time_ms: 0, build_ms: 0, check_ms: 0, total_ms: 0 };
    }
    return e.exact(m.knobs, opts);
  }
  if (m.type === "cancel") return { cancelled: true, can_cancel: !!e.initInfo.can_cancel };
  throw new Error(`unknown request ${m.type}`);
}

self.onmessage = (ev) => {
  const m = ev.data;
  m._seq = ++seq;
  if (m.type === "preview" || m.type === "check" || m.type === "compare" || m.type === "cancel") lastNewer = m._seq;
  queue = queue.then(async () => {
    try {
      const result = await handle(m);
      self.postMessage({ id: m.id, ok: true, result }, transferables(result));
    } catch (err) {
      if (m.type === "init") {                              // a failed start can be retried
        engine = null;
        status("error", `the engine failed to load: ${err && err.message ? err.message : err}`, 0);
      }
      self.postMessage({ id: m.id, ok: false, error: String(err && err.stack ? err.stack : err) });
    }
  });
};

	// The engine worker (ES module, ENGINE_CONTRACT.md / engine_api.js): GENCO as ONNX in onnxruntime-web (WASM) and the
	// app's scenario builder, "dsq" repair and fast correction (whatif_app's modules, copied verbatim) in Pyodide, both
	// self-hosted under assets/engine/ and both in this worker, off the page's thread.
	// messages in: {id, type: "init" \| "preview" \| "check" \| "compare" \| "exact" \| "cancel", knobs?, opts?, gen?}
	// messages out: {id, ok: true, result} \| {id, ok: false, error}; {type: "status", stage, detail, progress}
	// Threads: onnxruntime-web uses 2 WASM threads when the page is cross-origin isolated (COOP + COEP), else 1.
	// Requests run one at a time, in arrival order (Pyodide is single-threaded).

	import { createEngine, transferables } from "./assets/engine/engine_core.js";

	const BASE = new URL("./assets/engine/", import.meta.url).href;
	// a private Space: this worker's URL carries the page's signature (engine_api.js); Pyodide and onnxruntime-web import
	// their own files with their URL as the Referer, so they are imported signed too
	const SIGN = new URLSearchParams(self.location.search).get("__sign");
	const signed = (u) => (SIGN ? `${u}?__sign=${encodeURIComponent(SIGN)}` : u);
	const STAGES = { start: [0.02, "starting Python (Pyodide)"], pyodide: [0.25, "loading numpy, scipy, pandas"],
	packages: [0.6, "fetching the grid and GENCO"], fetch_rest: [0.65, "loading the engine"],
	engine: [0.8, "loading GENCO (onnxruntime-web)"], model: [0.9, "warming up"], warmup: [1.0, "ready"] };

	let engine = null;
	let queue = Promise.resolve();
	let seq = 0, lastNewer = 0; // the sequence number of the latest preview / check / compare / cancel received

	const status = (stage, detail, progress) => self.postMessage({ type: "status", stage, detail, progress });

	async function start(opts = {}) {
	status("loading", "loading the engine", 0.01);
	const [{ loadPyodide }, ortModule] = await Promise.all([import(signed(BASE + "pyodide/pyodide.mjs")), import(signed(BASE + "ort/ort.wasm.min.mjs"))]);
	const ort = ortModule.default ?? ortModule;
	ort.env.wasm.wasmPaths = BASE + "ort/";
	// GENCO's WASM threads: the visitor's cores, up to 8, when the page is cross-origin isolated (else 1). Measured on a
	// 6-core desktop: 2 threads ~550-600 ms per forward, 4-6 threads ~380-420 ms, 8 threads (more than the cores)
	// ~1,080 ms, so never more than navigator.hardwareConcurrency. Outputs are bit-identical for 1-8 threads.
	const cores = Number.isInteger(self.navigator?.hardwareConcurrency) ? self.navigator.hardwareConcurrency : 2;
	const threads = Number.isInteger(opts.threads) ? opts.threads : (self.crossOriginIsolated ? Math.max(1, Math.min(cores, 8)) : 1);
	const fetchBytes = async (url) => {
	const r = await fetch(url);
	if (!r.ok) throw new Error(`${url}: HTTP ${r.status}`);
	return new Uint8Array(await r.arrayBuffer());
	};
	const e = await createEngine({
	base: BASE, loadPyodide, ort, fetchBytes, threads, warmup: opts.warmup !== false,
	onProgress: ({ stage }) => { const s = STAGES[stage]; if (s) status(stage === "warmup" ? "ready" : "loading", s[1], s[0]); },
	});
	status("ready", "engine ready", 1);
	return e;
	}

	async function handle(m) {
	if (m.type === "init") {
	if (!engine) engine = start(m.opts \|\| {});
	const e = await engine;
	return { ...e.initInfo };
	}
	if (!engine) engine = start({});
	const e = await engine;
	if (m.type === "preview" \|\| m.type === "compare") return e[m.type](m.knobs);
	// progress messages for the page's live timers: {type: "progress", id, stage}, stage "check" (a check's time_ms
	// starts), "setup" (an exact loads casadi or builds its model), "solve" (an exact's interior point starts)
	const notify = (stage) => self.postMessage({ type: "progress", id: m.id, stage });
	if (m.type === "check") {
	const opts = { ...(m.opts \|\| {}), notify };
	if (Number.isInteger(m.gen)) opts.gen = m.gen;
	if (lastNewer > m._seq) { // superseded while queued (a newer preview, check, ...)
	const n = e.initInfo.info, nan = (k) => new Float32Array(k).fill(NaN);
	return { status: "cancelled", reason: "superseded before it started", cached: false, converged: false, cost: null,
	max_loading: null, n_viol: null, viol_detail: "cancelled", feasible: false, loops: 0, time_ms: 0,
	line_loading: nan(n.n_branch), bus_vm: nan(n.n_bus), gen_p: nan(n.n_gen), v_out: new Int32Array(0) };
	}
	return e.check(m.knobs, opts);
	}
	if (m.type === "exact") {
	const opts = { ...(m.opts \|\| {}), notify };
	if (Number.isInteger(m.gen)) opts.gen = m.gen;
	await new Promise((r) => setTimeout(r, 20)); // let messages posted right after this one arrive first
	if (lastNewer > m._seq) { // superseded while queued: not started
	const start = opts.start ?? "check";
	return { status: "cancelled", start, start_used: start, reason: "superseded before it started", cached: false,
	time_ms: 0, build_ms: 0, check_ms: 0, total_ms: 0 };
	}
	return e.exact(m.knobs, opts);
	}
	if (m.type === "cancel") return { cancelled: true, can_cancel: !!e.initInfo.can_cancel };
	throw new Error(`unknown request ${m.type}`);
	}

	self.onmessage = (ev) => {
	const m = ev.data;
	m._seq = ++seq;
	if (m.type === "preview" \|\| m.type === "check" \|\| m.type === "compare" \|\| m.type === "cancel") lastNewer = m._seq;
	queue = queue.then(async () => {
	try {
	const result = await handle(m);
	self.postMessage({ id: m.id, ok: true, result }, transferables(result));
	} catch (err) {
	if (m.type === "init") { // a failed start can be retried
	engine = null;
	status("error", `the engine failed to load: ${err && err.message ? err.message : err}`, 0);
	}
	self.postMessage({ id: m.id, ok: false, error: String(err && err.stack ? err.stack : err) });
	}
	});
	};