@mcmcjs/julia 0.6.1 → 0.7.0

package/dist/index.js

@@ -53,7 +53,7 @@ var juliaEngine = {
 };
 
 // src/fit.ts
-import { existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "fs";
+import { existsSync as existsSync2, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "fs";
 import { join as join3 } from "path";
 import {
   canonicalJson,
@@ -63,6 +63,7 @@ import {
 
 // src/runner-common.ts
 import { createHash } from "crypto";
+import { existsSync } from "fs";
 import { tmpdir } from "os";
 import { dirname, join as join2 } from "path";
 import { fileURLToPath } from "url";
@@ -70,11 +71,19 @@ function sharedTmpParent() {
   const uid = typeof process.getuid === "function" ? `-${process.getuid()}` : "";
   return join2(tmpdir(), `mcmcjs${uid}`);
 }
+function driverFile(name) {
+  const here = dirname(fileURLToPath(import.meta.url));
+  const beside = join2(here, name);
+  return existsSync(beside) ? beside : join2(here, "driver", name);
+}
 function driverPath() {
-  return join2(dirname(fileURLToPath(import.meta.url)), "driver.jl");
+  return driverFile("driver.jl");
 }
 function workerPath() {
-  return join2(dirname(fileURLToPath(import.meta.url)), "worker.jl");
+  return driverFile("worker.jl");
+}
+function pinnedEnvDir() {
+  return join2(dirname(fileURLToPath(import.meta.url)), "julia-env");
 }
 function sha256(text) {
   return createHash("sha256").update(text).digest("hex");
@@ -110,7 +119,7 @@ function makeRequestDir(kind) {
   mkdirSync(parent, { recursive: true });
   return mkdtempSync(join3(parent, `${kind}-`));
 }
-function fitRequest(spec, outPath) {
+function fitRequest(spec, outPath, opts) {
   return {
     schema_version: spec.schema_version,
     backend: { id: spec.backend.id },
@@ -118,12 +127,13 @@ function fitRequest(spec, outPath) {
     data: spec.data,
     sampler: spec.sampler,
     seed: spec.seed,
-    out: outPath
+    out: outPath,
+    ...opts?.streamDraws ? { stream_draws: true } : {}
   };
 }
 function finalizeOkFit(spec, resolved, io, provenance, startedAt, elapsedMs) {
   const runtimeRequested = spec.backend.version;
-  if (!existsSync(io.outPath)) {
+  if (!existsSync2(io.outPath)) {
     return {
       status: "error",
       runtimeRequested,
@@ -152,7 +162,7 @@ function finalizeOkFit(spec, resolved, io, provenance, startedAt, elapsedMs) {
     runtime: { requested: runtimeRequested, actual: julia, path: resolved.command },
     manifest_sha256: provenance?.manifest_sha256,
     packages: provenance?.packages,
-    model_sha256: existsSync(spec.modelPath) ? sha256(readFileSync(spec.modelPath, "utf8")) : void 0,
+    model_sha256: existsSync2(spec.modelPath) ? sha256(readFileSync(spec.modelPath, "utf8")) : void 0,
     data_sha256: io.dataSha256 ?? sha256(canonicalJson(spec.data)),
     ...io.dataFile ? { data_file: io.dataFile } : {},
     samples_file: io.outPath,
@@ -173,10 +183,16 @@ async function runFit(spec, resolved, io) {
   const startedAt = (/* @__PURE__ */ new Date()).toISOString();
   const start = performance.now();
   const runtimeRequested = spec.backend.version;
+  if (io.signal?.aborted) {
+    return { status: "cancelled", runtimeRequested, elapsedMs: 0 };
+  }
   const ownTmp = io.tmpDir === void 0;
   const tmp = io.tmpDir ?? makeRequestDir("fit");
   const requestPath = join3(tmp, "request.json");
-  writeFileSync(requestPath, JSON.stringify(fitRequest(spec, io.outPath)));
+  writeFileSync(
+    requestPath,
+    JSON.stringify(fitRequest(spec, io.outPath, { streamDraws: io.onDraws !== void 0 }))
+  );
   const args = [
     ...resolved.args,
     "--startup-file=no",
@@ -187,15 +203,21 @@ async function runFit(spec, resolved, io) {
   let stdout;
   let stderr;
   let code;
+  let cancelled;
   try {
-    ({ stdout, stderr, code } = await io.spawn(resolved.command, args, {
-      onProgress: io.onProgress
+    ({ stdout, stderr, code, cancelled } = await io.spawn(resolved.command, args, {
+      onProgress: io.onProgress,
+      onDraws: io.onDraws,
+      signal: io.signal
     }));
   } finally {
     if (ownTmp) rmSync(tmp, { recursive: true, force: true });
   }
   const elapsedMs = Math.round(performance.now() - start);
-  if (code !== 0 || !existsSync(io.outPath)) {
+  if (cancelled) {
+    return { status: "cancelled", runtimeRequested, elapsedMs };
+  }
+  if (code !== 0 || !existsSync2(io.outPath)) {
     const failure = lastJsonLine(stderr);
     const stage = toStage(failure?.stage);
     const error = failure?.error ?? stderr.trim().split("\n").pop() ?? `julia exited with code ${code}`;
@@ -205,7 +227,7 @@ async function runFit(spec, resolved, io) {
 }
 
 // src/predict.ts
-import { existsSync as existsSync2, readFileSync as readFileSync2, rmSync as rmSync2, writeFileSync as writeFileSync2 } from "fs";
+import { existsSync as existsSync3, readFileSync as readFileSync2, rmSync as rmSync2, writeFileSync as writeFileSync2 } from "fs";
 import { join as join4 } from "path";
 import {
   canonicalJson as canonicalJson2,
@@ -260,7 +282,7 @@ async function runPredict(spec, resolved, io) {
     if (ownTmp) rmSync2(tmp, { recursive: true, force: true });
   }
   const elapsedMs = Math.round(performance.now() - start);
-  if (code !== 0 || !existsSync2(io.outPath)) {
+  if (code !== 0 || !existsSync3(io.outPath)) {
     const failure = lastJsonLine(stderr);
     const error = failure?.error ?? stderr.trim().split("\n").pop() ?? `julia exited with code ${code}`;
     return { status: "error", runtimeRequested, elapsedMs, stage: toStage(failure?.stage), error };
@@ -286,10 +308,10 @@ async function runPredict(spec, resolved, io) {
     runtime: { requested: runtimeRequested, actual: julia, path: resolved.command },
     manifest_sha256: provenance?.manifest_sha256,
     packages: provenance?.packages,
-    model_sha256: existsSync2(spec.modelPath) ? sha256(readFileSync2(spec.modelPath, "utf8")) : void 0,
+    model_sha256: existsSync3(spec.modelPath) ? sha256(readFileSync2(spec.modelPath, "utf8")) : void 0,
     data_sha256: sha256(canonicalJson2(data)),
     posterior_samples: io.samplesPath,
-    posterior_samples_sha256: existsSync2(io.samplesPath) ? sha256(readFileSync2(io.samplesPath, "utf8")) : void 0,
+    posterior_samples_sha256: existsSync3(io.samplesPath) ? sha256(readFileSync2(io.samplesPath, "utf8")) : void 0,
     samples_file: io.outPath,
     started_at: startedAt,
     elapsed_ms: elapsedMs
@@ -306,10 +328,10 @@ async function runPredict(spec, resolved, io) {
 }
 
 // src/project.ts
-import { existsSync as existsSync3, mkdirSync as mkdirSync2, readFileSync as readFileSync3, writeFileSync as writeFileSync3 } from "fs";
+import { copyFileSync, existsSync as existsSync4, mkdirSync as mkdirSync2, readFileSync as readFileSync3, writeFileSync as writeFileSync3 } from "fs";
 import { homedir as homedir2 } from "os";
 import { join as join5 } from "path";
-import { canonicalJson as canonicalJson3 } from "@mcmcjs/core";
+import { canonicalJson as canonicalJson3, DEFAULT_JULIA_CHANNEL } from "@mcmcjs/core";
 import { createRunner as createRunner3 } from "@mcmcjs/engine";
 var PACKAGES = [
   "Turing",
@@ -354,8 +376,13 @@ function managedProjectDir(version, pins) {
 function sentinelPath(dir) {
   return join5(dir, ".mcmcjs-packages.json");
 }
+var PROVISION_GENERATION = 1;
 function expectedSentinel(pins) {
-  return canonicalJson3({ packages: [...PACKAGES].sort(), pins: pins ?? {} });
+  return canonicalJson3({
+    gen: PROVISION_GENERATION,
+    packages: [...PACKAGES].sort(),
+    pins: pins ?? {}
+  });
 }
 function provisioned(dir, pins) {
   try {
@@ -365,7 +392,7 @@ function provisioned(dir, pins) {
   }
 }
 function managedProjectReady(dir = managedProjectDir(), pins) {
-  return existsSync3(join5(dir, "Project.toml")) && provisioned(dir, pins);
+  return existsSync4(join5(dir, "Project.toml")) && provisioned(dir, pins);
 }
 function packageSpecsCode(pins) {
   const specs = PACKAGES.map((name) => {
@@ -374,12 +401,26 @@ function packageSpecsCode(pins) {
   });
   return `[${specs.join(", ")}]`;
 }
-async function ensureProject(juliaBin, run = createRunner3(30 * 6e4), dir = managedProjectDir(), pins) {
+function usesPinnedEnv(version, pins) {
+  return !hasPins(pins) && version === DEFAULT_JULIA_CHANNEL && existsSync4(pinnedEnvDir());
+}
+async function ensureProject(juliaBin, run = createRunner3(30 * 6e4), dir = managedProjectDir(), pins, opts) {
   validatePins(pins);
   if (managedProjectReady(dir, pins)) return dir;
   mkdirSync2(dir, { recursive: true });
-  const code = `using Pkg; Pkg.add(${packageSpecsCode(pins)}); Pkg.precompile()`;
-  await run(juliaBin, ["--startup-file=no", `--project=${dir}`, "-e", code]);
+  if (usesPinnedEnv(opts?.version, pins)) {
+    copyFileSync(join5(pinnedEnvDir(), "Project.toml"), join5(dir, "Project.toml"));
+    copyFileSync(join5(pinnedEnvDir(), "Manifest.toml"), join5(dir, "Manifest.toml"));
+    await run(juliaBin, [
+      "--startup-file=no",
+      `--project=${dir}`,
+      "-e",
+      "using Pkg; Pkg.instantiate(); Pkg.precompile()"
+    ]);
+  } else {
+    const code = `using Pkg; Pkg.add(${packageSpecsCode(pins)}); Pkg.precompile()`;
+    await run(juliaBin, ["--startup-file=no", `--project=${dir}`, "-e", code]);
+  }
   writeFileSync3(sentinelPath(dir), expectedSentinel(pins));
   return dir;
 }
@@ -414,6 +455,7 @@ async function runMatrix(spec, versions, io) {
       entry = { version, status: "error", elapsedMs: 0, error: error.message };
     }
     entries.push(entry);
+    if (entry.status === "cancelled") break;
     if (entry.status === "error" && !io.keepGoing) break;
   }
   return {
@@ -423,28 +465,38 @@ async function runMatrix(spec, versions, io) {
 }
 
 // src/setup.ts
+import { DEFAULT_JULIA_CHANNEL as DEFAULT_JULIA_CHANNEL2 } from "@mcmcjs/core";
 var JULIAUP_INSTALL_URL = "https://install.julialang.org";
 function juliaupInstallCommand(platform) {
   if (platform === "win32") {
     return null;
   }
-  return { command: "sh", args: ["-c", `curl -fsSL ${JULIAUP_INSTALL_URL} | sh -s -- --yes`] };
+  return {
+    command: "sh",
+    args: [
+      "-c",
+      `curl -fsSL ${JULIAUP_INSTALL_URL} | sh -s -- --yes --default-channel ${DEFAULT_JULIA_CHANNEL2}`
+    ]
+  };
 }
 function planSetup(report, platform) {
-  if (report.julia.found) return [];
+  if (report.julia.found && report.julia.version === DEFAULT_JULIA_CHANNEL2) return [];
   if (report.juliaup.found) {
     return [
       {
         tool: "julia",
-        label: "install the latest stable Julia via juliaup",
-        command: { command: report.juliaup.path ?? "juliaup", args: ["add", "release"] }
+        label: `install the pinned Julia ${DEFAULT_JULIA_CHANNEL2} via juliaup`,
+        command: {
+          command: report.juliaup.path ?? "juliaup",
+          args: ["add", DEFAULT_JULIA_CHANNEL2]
+        }
       }
     ];
   }
   return [
     {
       tool: "juliaup",
-      label: "install juliaup (this also installs Julia)",
+      label: `install juliaup with the pinned Julia ${DEFAULT_JULIA_CHANNEL2}`,
       command: juliaupInstallCommand(platform)
     }
   ];
@@ -457,8 +509,8 @@ async function runSetup(options = {}) {
     dryRun = false
   } = options;
   const before = await runDoctor(runner);
-  if (before.ready) return { ...before, steps: [] };
   const plan = planSetup(before, platform);
+  if (plan.length === 0) return { ...before, steps: [] };
   if (dryRun) {
     return {
       ...before,
@@ -563,10 +615,10 @@ function assertVersionsInstalled(channels, installed) {
 
 // src/worker.ts
 import { spawn } from "child_process";
-import { mkdirSync as mkdirSync3, readdirSync, rmdirSync, rmSync as rmSync3, statSync } from "fs";
+import { mkdirSync as mkdirSync3, readdirSync, readFileSync as readFileSync4, rmdirSync, rmSync as rmSync3, statSync } from "fs";
 import { createConnection } from "net";
 import { join as join7 } from "path";
-import { parseProgressLine } from "@mcmcjs/engine";
+import { parseDrawBatchLine, parseProgressLine } from "@mcmcjs/engine";
 var START_WAIT_MS = 18e4;
 var CONNECT_TIMEOUT_MS = 1e3;
 var STOP_ACK_TIMEOUT_MS = 3e3;
@@ -576,6 +628,27 @@ function fitInactivityTimeoutMs() {
 }
 var WorkerTimeoutError = class extends Error {
 };
+var FitCancelledError = class extends Error {
+};
+function workerPidPath(socket) {
+  return `${socket}.pid`;
+}
+function killWorker(socket) {
+  const pidPath = workerPidPath(socket);
+  try {
+    const pid = Number(readFileSync4(pidPath, "utf8").trim());
+    if (Number.isInteger(pid) && pid > 0) {
+      try {
+        process.kill(-pid, "SIGKILL");
+      } catch {
+        process.kill(pid, "SIGKILL");
+      }
+    }
+  } catch {
+  }
+  rmSync3(socket, { force: true });
+  rmSync3(pidPath, { force: true });
+}
 function workersDir() {
   const runtime = process.env.XDG_RUNTIME_DIR;
   const dir = runtime ? join7(runtime, "mcmcjs") : join7(sharedTmpParent(), "workers");
@@ -627,12 +700,13 @@ function releaseStartLock(socketPath) {
   } catch {
   }
 }
-async function ensureWorker(resolved, projectDir, notify) {
+async function ensureWorker(resolved, projectDir, notify, signal) {
   const path = workerSocketPath(resolved.command, projectDir);
   try {
     return await connectOnce(path);
   } catch {
   }
+  if (signal?.aborted) throw new FitCancelledError();
   let spawnFailure;
   const startedHere = acquireStartLock(path);
   if (startedHere) {
@@ -661,6 +735,7 @@ async function ensureWorker(resolved, projectDir, notify) {
   try {
     const deadline = Date.now() + START_WAIT_MS;
     for (; ; ) {
+      if (signal?.aborted) throw new FitCancelledError();
       if (spawnFailure) throw spawnFailure;
       try {
         return await connectOnce(path);
@@ -673,16 +748,27 @@ async function ensureWorker(resolved, projectDir, notify) {
     if (startedHere) releaseStartLock(path);
   }
 }
-function readWorkerResponse(sock, onProgress, inactivityMs = fitInactivityTimeoutMs()) {
+function readWorkerResponse(sock, onProgress, onDraws, signal, inactivityMs = fitInactivityTimeoutMs()) {
   return new Promise((resolve, reject) => {
     let pending = "";
     let settled = false;
+    const onAbort = () => {
+      settle(() => {
+        sock.destroy();
+        reject(new FitCancelledError());
+      });
+    };
     const settle = (fn) => {
       if (settled) return;
       settled = true;
       sock.setTimeout(0);
+      signal?.removeEventListener("abort", onAbort);
       fn();
     };
+    if (signal) {
+      if (signal.aborted) onAbort();
+      else signal.addEventListener("abort", onAbort, { once: true });
+    }
     sock.setTimeout(inactivityMs, () => {
       const error = new WorkerTimeoutError(
         `the worker sent nothing for ${Math.round(inactivityMs / 6e4)} minutes`
@@ -705,6 +791,11 @@ function readWorkerResponse(sock, onProgress, inactivityMs = fitInactivityTimeou
           onProgress?.(progress);
           continue;
         }
+        const batch = parseDrawBatchLine(line);
+        if (batch) {
+          onDraws?.(batch);
+          continue;
+        }
         try {
           const doc = JSON.parse(line);
           settle(() => {
@@ -723,10 +814,48 @@ function readWorkerResponse(sock, onProgress, inactivityMs = fitInactivityTimeou
 async function runFitViaWorker(spec, resolved, io) {
   const startedAt = (/* @__PURE__ */ new Date()).toISOString();
   const start = performance.now();
-  const sock = await ensureWorker(resolved, io.projectDir, io.notify);
-  sock.write(`${JSON.stringify(fitRequest(spec, io.outPath))}
-`);
-  const final = await readWorkerResponse(sock, io.onProgress);
+  if (io.signal?.aborted) {
+    return { status: "cancelled", runtimeRequested: spec.backend.version, elapsedMs: 0 };
+  }
+  let sock;
+  try {
+    sock = await ensureWorker(resolved, io.projectDir, io.notify, io.signal);
+  } catch (error) {
+    if (error instanceof FitCancelledError) {
+      return {
+        status: "cancelled",
+        runtimeRequested: spec.backend.version,
+        elapsedMs: Math.round(performance.now() - start)
+      };
+    }
+    throw error;
+  }
+  if (io.signal?.aborted) {
+    sock.destroy();
+    return {
+      status: "cancelled",
+      runtimeRequested: spec.backend.version,
+      elapsedMs: Math.round(performance.now() - start)
+    };
+  }
+  sock.write(
+    `${JSON.stringify(fitRequest(spec, io.outPath, { streamDraws: io.onDraws !== void 0 }))}
+`
+  );
+  let final;
+  try {
+    final = await readWorkerResponse(sock, io.onProgress, io.onDraws, io.signal);
+  } catch (error) {
+    if (error instanceof FitCancelledError) {
+      killWorker(workerSocketPath(resolved.command, io.projectDir));
+      return {
+        status: "cancelled",
+        runtimeRequested: spec.backend.version,
+        elapsedMs: Math.round(performance.now() - start)
+      };
+    }
+    throw error;
+  }
   const elapsedMs = Math.round(performance.now() - start);
   if (final.ok !== true) {
     if (final.stage === "worker") throw new Error(String(final.error ?? "worker error"));
@@ -748,6 +877,9 @@ async function runFitViaWorker(spec, resolved, io) {
   );
 }
 async function runFitAuto(spec, resolved, io) {
+  if (io.signal?.aborted) {
+    return { status: "cancelled", runtimeRequested: spec.backend.version, elapsedMs: 0 };
+  }
   if (io.daemon && process.platform !== "win32") {
     const start = performance.now();
     try {
@@ -819,6 +951,7 @@ export {
   gcVersions,
   juliaEngine,
   juliaupInstallCommand,
+  killWorker,
   listVersions,
   listWorkers,
   managedProjectDir,