@mcmcjs/julia 0.3.0 → 0.5.0

package/dist/index.js

@@ -52,8 +52,7 @@ var juliaEngine = {
 };
 
 // src/fit.ts
-import { existsSync, mkdtempSync, readFileSync, writeFileSync } from "fs";
-import { tmpdir } from "os";
+import { existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "fs";
 import { join as join3 } from "path";
 import {
   canonicalJson,
@@ -63,11 +62,19 @@ import {
 
 // src/runner-common.ts
 import { createHash } from "crypto";
+import { tmpdir } from "os";
 import { dirname, join as join2 } from "path";
 import { fileURLToPath } from "url";
+function sharedTmpParent() {
+  const uid = typeof process.getuid === "function" ? `-${process.getuid()}` : "";
+  return join2(tmpdir(), `mcmcjs${uid}`);
+}
 function driverPath() {
   return join2(dirname(fileURLToPath(import.meta.url)), "driver.jl");
 }
+function workerPath() {
+  return join2(dirname(fileURLToPath(import.meta.url)), "worker.jl");
+}
 function sha256(text) {
   return createHash("sha256").update(text).digest("hex");
 }
@@ -89,45 +96,40 @@ var STAGES = /* @__PURE__ */ new Set([
   "load_samples",
   "predict",
   "write",
-  "spawn"
+  "spawn",
+  "worker"
 ]);
 function toStage(value) {
   return typeof value === "string" && STAGES.has(value) ? value : void 0;
 }
 
 // src/fit.ts
-async function runFit(spec, resolved, io) {
-  const startedAt = (/* @__PURE__ */ new Date()).toISOString();
-  const start = performance.now();
+function makeRequestDir(kind) {
+  const parent = sharedTmpParent();
+  mkdirSync(parent, { recursive: true });
+  return mkdtempSync(join3(parent, `${kind}-`));
+}
+function fitRequest(spec, outPath) {
+  return {
+    schema_version: spec.schema_version,
+    backend: { id: spec.backend.id },
+    model: { file: spec.modelPath, entry: spec.model.entry },
+    data: spec.data,
+    sampler: spec.sampler,
+    seed: spec.seed,
+    out: outPath
+  };
+}
+function finalizeOkFit(spec, resolved, io, provenance, startedAt, elapsedMs) {
   const runtimeRequested = spec.backend.version;
-  const tmp = io.tmpDir ?? mkdtempSync(join3(tmpdir(), "mcmcjs-fit-"));
-  const requestPath = join3(tmp, "request.json");
-  writeFileSync(
-    requestPath,
-    JSON.stringify({
-      schema_version: spec.schema_version,
-      backend: { id: spec.backend.id },
-      model: { file: spec.modelPath, entry: spec.model.entry },
-      data: spec.data,
-      sampler: spec.sampler,
-      seed: spec.seed,
-      out: io.outPath
-    })
-  );
-  const args = [
-    ...resolved.args,
-    "--startup-file=no",
-    `--project=${io.projectDir}`,
-    driverPath(),
-    requestPath
-  ];
-  const { stdout, stderr, code } = await io.spawn(resolved.command, args);
-  const elapsedMs = Math.round(performance.now() - start);
-  if (code !== 0 || !existsSync(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}`;
-    return { status: "error", runtimeRequested, elapsedMs, stage, error };
+  if (!existsSync(io.outPath)) {
+    return {
+      status: "error",
+      runtimeRequested,
+      elapsedMs,
+      stage: "write",
+      error: "the runtime reported success but wrote no samples file"
+    };
   }
   try {
     parseSamples(readFileSync(io.outPath, "utf8"));
@@ -140,7 +142,6 @@ async function runFit(spec, resolved, io) {
       error: `samples file did not parse: ${error.message}`
     };
   }
-  const provenance = lastJsonLine(stdout);
   const julia = provenance?.julia_version;
   const record = {
     schema_version: RUN_RECORD_SCHEMA_VERSION,
@@ -156,7 +157,7 @@ async function runFit(spec, resolved, io) {
     started_at: startedAt,
     elapsed_ms: elapsedMs
   };
-  writeFileSync(`${io.outPath}.run.json`, `${JSON.stringify(record, null, 2)}
+  writeFileSync(io.recordPath ?? `${io.outPath}.run.json`, `${JSON.stringify(record, null, 2)}
 `);
   return {
     status: "ok",
@@ -166,10 +167,43 @@ async function runFit(spec, resolved, io) {
     elapsedMs
   };
 }
+async function runFit(spec, resolved, io) {
+  const startedAt = (/* @__PURE__ */ new Date()).toISOString();
+  const start = performance.now();
+  const runtimeRequested = spec.backend.version;
+  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)));
+  const args = [
+    ...resolved.args,
+    "--startup-file=no",
+    `--project=${io.projectDir}`,
+    driverPath(),
+    requestPath
+  ];
+  let stdout;
+  let stderr;
+  let code;
+  try {
+    ({ stdout, stderr, code } = await io.spawn(resolved.command, args, {
+      onProgress: io.onProgress
+    }));
+  } finally {
+    if (ownTmp) rmSync(tmp, { recursive: true, force: true });
+  }
+  const elapsedMs = Math.round(performance.now() - start);
+  if (code !== 0 || !existsSync(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}`;
+    return { status: "error", runtimeRequested, elapsedMs, stage, error };
+  }
+  return finalizeOkFit(spec, resolved, io, lastJsonLine(stdout), startedAt, elapsedMs);
+}
 
 // src/predict.ts
-import { existsSync as existsSync2, mkdtempSync as mkdtempSync2, readFileSync as readFileSync2, writeFileSync as writeFileSync2 } from "fs";
-import { tmpdir as tmpdir2 } from "os";
+import { existsSync as existsSync2, readFileSync as readFileSync2, rmSync as rmSync2, writeFileSync as writeFileSync2 } from "fs";
 import { join as join4 } from "path";
 import {
   canonicalJson as canonicalJson2,
@@ -191,7 +225,8 @@ async function runPredict(spec, resolved, io) {
   const start = performance.now();
   const runtimeRequested = spec.backend.version;
   const data = predictData(spec);
-  const tmp = io.tmpDir ?? mkdtempSync2(join4(tmpdir2(), "mcmcjs-predict-"));
+  const ownTmp = io.tmpDir === void 0;
+  const tmp = io.tmpDir ?? makeRequestDir("predict");
   const requestPath = join4(tmp, "request.json");
   writeFileSync2(
     requestPath,
@@ -214,7 +249,14 @@ async function runPredict(spec, resolved, io) {
     driverPath(),
     requestPath
   ];
-  const { stdout, stderr, code } = await io.spawn(resolved.command, args);
+  let stdout;
+  let stderr;
+  let code;
+  try {
+    ({ stdout, stderr, code } = await io.spawn(resolved.command, args));
+  } finally {
+    if (ownTmp) rmSync2(tmp, { recursive: true, force: true });
+  }
   const elapsedMs = Math.round(performance.now() - start);
   if (code !== 0 || !existsSync2(io.outPath)) {
     const failure = lastJsonLine(stderr);
@@ -262,12 +304,14 @@ async function runPredict(spec, resolved, io) {
 }
 
 // src/project.ts
-import { existsSync as existsSync3, mkdirSync, readFileSync as readFileSync3, writeFileSync as writeFileSync3 } from "fs";
+import { existsSync as existsSync3, mkdirSync as mkdirSync2, readFileSync as readFileSync3, writeFileSync as writeFileSync3 } from "fs";
 import { homedir as homedir2 } from "os";
 import { join as join5 } from "path";
 import { createRunner as createRunner3 } from "@mcmcjs/engine";
 var PACKAGES = [
   "Turing",
+  "FlexiChains",
+  "DimensionalData",
   "MCMCChains",
   "JuliaBUGS",
   "AdvancedHMC",
@@ -294,7 +338,7 @@ function managedProjectReady(dir = managedProjectDir()) {
 }
 async function ensureProject(juliaBin, run = createRunner3(30 * 6e4), dir = managedProjectDir()) {
   if (managedProjectReady(dir)) return dir;
-  mkdirSync(dir, { recursive: true });
+  mkdirSync2(dir, { recursive: true });
   const code = `using Pkg; Pkg.add(${JSON.stringify(PACKAGES)}); Pkg.precompile()`;
   await run(juliaBin, ["--startup-file=no", `--project=${dir}`, "-e", code]);
   writeFileSync3(sentinelPath(dir), JSON.stringify([...PACKAGES].sort()));
@@ -462,7 +506,7 @@ async function resolveVersion(bin, channel, run) {
       `Julia version "${channel}" is not installed. Add it with: mcmc julia version add ${channel}`
     );
   }
-  return { command: match.File, args: match.Args ?? [] };
+  return { command: match.File, args: match.Args ?? [], version: match.Version };
 }
 function assertVersionsInstalled(channels, installed) {
   const have = new Set(installed.map((v) => v.id));
@@ -474,6 +518,256 @@ 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 { createConnection } from "net";
+import { join as join7 } from "path";
+import { parseProgressLine } from "@mcmcjs/engine";
+var START_WAIT_MS = 18e4;
+var CONNECT_TIMEOUT_MS = 1e3;
+var STOP_ACK_TIMEOUT_MS = 3e3;
+function fitInactivityTimeoutMs() {
+  const fromEnv = Number(process.env.MCMC_WORKER_TIMEOUT_MS);
+  return Number.isFinite(fromEnv) && fromEnv > 0 ? fromEnv : 30 * 6e4;
+}
+var WorkerTimeoutError = class extends Error {
+};
+function workersDir() {
+  const runtime = process.env.XDG_RUNTIME_DIR;
+  const dir = runtime ? join7(runtime, "mcmcjs") : join7(sharedTmpParent(), "workers");
+  mkdirSync3(dir, { recursive: true, mode: 448 });
+  return dir;
+}
+function workerSocketPath(command, projectDir) {
+  return join7(workersDir(), `worker-${sha256(`${command}
+${projectDir}`).slice(0, 12)}.sock`);
+}
+function connectOnce(path) {
+  return new Promise((resolve, reject) => {
+    const sock = createConnection(path);
+    const timer = setTimeout(() => {
+      sock.destroy();
+      reject(new Error("connect timeout"));
+    }, CONNECT_TIMEOUT_MS);
+    sock.once("connect", () => {
+      clearTimeout(timer);
+      resolve(sock);
+    });
+    sock.once("error", (error) => {
+      clearTimeout(timer);
+      reject(error);
+    });
+  });
+}
+var delay = (ms) => new Promise((resolve) => setTimeout(resolve, ms));
+function acquireStartLock(socketPath) {
+  const lock = `${socketPath}.start`;
+  try {
+    mkdirSync3(lock);
+    return true;
+  } catch {
+    try {
+      if (Date.now() - statSync(lock).mtimeMs > START_WAIT_MS) {
+        rmdirSync(lock);
+        mkdirSync3(lock);
+        return true;
+      }
+    } catch {
+    }
+    return false;
+  }
+}
+function releaseStartLock(socketPath) {
+  try {
+    rmdirSync(`${socketPath}.start`);
+  } catch {
+  }
+}
+async function ensureWorker(resolved, projectDir, notify) {
+  const path = workerSocketPath(resolved.command, projectDir);
+  try {
+    return await connectOnce(path);
+  } catch {
+  }
+  let spawnFailure;
+  const startedHere = acquireStartLock(path);
+  if (startedHere) {
+    try {
+      rmSync3(path, { force: true });
+      notify?.("starting the Julia worker (first use loads Turing, can take a minute)...");
+      const child = spawn(
+        resolved.command,
+        [...resolved.args, "--startup-file=no", `--project=${projectDir}`, workerPath(), path],
+        { detached: true, stdio: "ignore" }
+      );
+      child.on("error", (error) => {
+        spawnFailure = error;
+      });
+      child.on("exit", (code) => {
+        if (code !== 0 && code !== null) {
+          spawnFailure = new Error(`the worker exited with code ${code} before serving`);
+        }
+      });
+      child.unref();
+    } catch (error) {
+      releaseStartLock(path);
+      throw error;
+    }
+  }
+  try {
+    const deadline = Date.now() + START_WAIT_MS;
+    for (; ; ) {
+      if (spawnFailure) throw spawnFailure;
+      try {
+        return await connectOnce(path);
+      } catch {
+      }
+      if (Date.now() > deadline) throw new Error("the Julia worker did not start in time");
+      await delay(250);
+    }
+  } finally {
+    if (startedHere) releaseStartLock(path);
+  }
+}
+function readWorkerResponse(sock, onProgress, inactivityMs = fitInactivityTimeoutMs()) {
+  return new Promise((resolve, reject) => {
+    let pending = "";
+    let settled = false;
+    const settle = (fn) => {
+      if (settled) return;
+      settled = true;
+      sock.setTimeout(0);
+      fn();
+    };
+    sock.setTimeout(inactivityMs, () => {
+      const error = new WorkerTimeoutError(
+        `the worker sent nothing for ${Math.round(inactivityMs / 6e4)} minutes`
+      );
+      settle(() => {
+        sock.destroy();
+        reject(error);
+      });
+    });
+    sock.on("data", (chunk) => {
+      pending += chunk.toString("utf8");
+      for (; ; ) {
+        const at = pending.indexOf("\n");
+        if (at === -1) break;
+        const line = pending.slice(0, at).trim();
+        pending = pending.slice(at + 1);
+        if (!line) continue;
+        const progress = parseProgressLine(line);
+        if (progress) {
+          onProgress?.(progress);
+          continue;
+        }
+        try {
+          const doc = JSON.parse(line);
+          settle(() => {
+            sock.end();
+            resolve(doc);
+          });
+          return;
+        } catch {
+        }
+      }
+    });
+    sock.on("error", (error) => settle(() => reject(error)));
+    sock.on("close", () => settle(() => reject(new Error("worker closed the connection"))));
+  });
+}
+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);
+  const elapsedMs = Math.round(performance.now() - start);
+  if (final.ok !== true) {
+    if (final.stage === "worker") throw new Error(String(final.error ?? "worker error"));
+    return {
+      status: "error",
+      runtimeRequested: spec.backend.version,
+      elapsedMs,
+      stage: toStage(final.stage),
+      error: String(final.error ?? "worker error")
+    };
+  }
+  return finalizeOkFit(
+    spec,
+    resolved,
+    io,
+    final.provenance,
+    startedAt,
+    elapsedMs
+  );
+}
+async function runFitAuto(spec, resolved, io) {
+  if (io.daemon && process.platform !== "win32") {
+    const start = performance.now();
+    try {
+      return await runFitViaWorker(spec, resolved, io);
+    } catch (error) {
+      if (error instanceof WorkerTimeoutError) {
+        return {
+          status: "error",
+          runtimeRequested: spec.backend.version,
+          elapsedMs: Math.round(performance.now() - start),
+          stage: "worker",
+          error: `${error.message}; the worker may be wedged, try mcmc daemon stop`
+        };
+      }
+      io.notify?.(
+        `worker unavailable (${error.message}); falling back to a one-shot run`
+      );
+    }
+  }
+  return runFit(spec, resolved, io);
+}
+async function listWorkers() {
+  const dir = workersDir();
+  const sockets = readdirSync(dir).filter((f) => f.startsWith("worker-") && f.endsWith(".sock"));
+  const statuses = [];
+  for (const name of sockets) {
+    const socket = join7(dir, name);
+    let alive = false;
+    try {
+      const sock = await connectOnce(socket);
+      sock.destroy();
+      alive = true;
+    } catch {
+    }
+    statuses.push({ socket, alive });
+  }
+  return statuses;
+}
+async function stopWorker(socket) {
+  let sock;
+  try {
+    sock = await connectOnce(socket);
+  } catch {
+    rmSync3(socket, { force: true });
+    return "stale";
+  }
+  const acked = await new Promise((resolve) => {
+    const timer = setTimeout(() => resolve(false), STOP_ACK_TIMEOUT_MS);
+    const finish = (ok) => {
+      clearTimeout(timer);
+      resolve(ok);
+    };
+    sock.resume();
+    sock.on("close", () => finish(true));
+    sock.on("error", () => finish(false));
+    sock.write('{"mcmcjs":"stop"}\n');
+  });
+  sock.destroy();
+  if (!acked) return "pending";
+  rmSync3(socket, { force: true });
+  return "stopped";
+}
 export {
   addVersion,
   assertVersionsInstalled,
@@ -484,6 +778,7 @@ export {
   juliaEngine,
   juliaupInstallCommand,
   listVersions,
+  listWorkers,
   managedProjectDir,
   managedProjectReady,
   planSetup,
@@ -492,10 +787,13 @@ export {
   resolveVersion,
   runDoctor,
   runFit,
+  runFitAuto,
   runMatrix,
   runPredict,
   runSetup,
   setDefaultVersion,
+  sharedTmpParent,
+  stopWorker,
   updateVersion
 };
 //# sourceMappingURL=index.js.map