@mcmcjs/julia 0.3.0 → 0.5.0

package/dist/index.cjs

@@ -29,6 +29,7 @@ __export(index_exports, {
   juliaEngine: () => juliaEngine,
   juliaupInstallCommand: () => juliaupInstallCommand,
   listVersions: () => listVersions,
+  listWorkers: () => listWorkers,
   managedProjectDir: () => managedProjectDir,
   managedProjectReady: () => managedProjectReady,
   planSetup: () => planSetup,
@@ -37,10 +38,13 @@ __export(index_exports, {
   resolveVersion: () => resolveVersion,
   runDoctor: () => runDoctor,
   runFit: () => runFit,
+  runFitAuto: () => runFitAuto,
   runMatrix: () => runMatrix,
   runPredict: () => runPredict,
   runSetup: () => runSetup,
   setDefaultVersion: () => setDefaultVersion,
+  sharedTmpParent: () => sharedTmpParent,
+  stopWorker: () => stopWorker,
   updateVersion: () => updateVersion
 });
 module.exports = __toCommonJS(index_exports);
@@ -104,17 +108,24 @@ var juliaEngine = {
 
 // src/fit.ts
 var import_node_fs = require("fs");
-var import_node_os2 = require("os");
 var import_node_path3 = require("path");
 var import_core = require("@mcmcjs/core");
 
 // src/runner-common.ts
 var import_node_crypto = require("crypto");
+var import_node_os2 = require("os");
 var import_node_path2 = require("path");
 var import_node_url = require("url");
+function sharedTmpParent() {
+  const uid = typeof process.getuid === "function" ? `-${process.getuid()}` : "";
+  return (0, import_node_path2.join)((0, import_node_os2.tmpdir)(), `mcmcjs${uid}`);
+}
 function driverPath() {
   return (0, import_node_path2.join)((0, import_node_path2.dirname)((0, import_node_url.fileURLToPath)(importMetaUrl)), "driver.jl");
 }
+function workerPath() {
+  return (0, import_node_path2.join)((0, import_node_path2.dirname)((0, import_node_url.fileURLToPath)(importMetaUrl)), "worker.jl");
+}
 function sha256(text) {
   return (0, import_node_crypto.createHash)("sha256").update(text).digest("hex");
 }
@@ -136,45 +147,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();
+  (0, import_node_fs.mkdirSync)(parent, { recursive: true });
+  return (0, import_node_fs.mkdtempSync)((0, import_node_path3.join)(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 ?? (0, import_node_fs.mkdtempSync)((0, import_node_path3.join)((0, import_node_os2.tmpdir)(), "mcmcjs-fit-"));
-  const requestPath = (0, import_node_path3.join)(tmp, "request.json");
-  (0, import_node_fs.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 || !(0, import_node_fs.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 (!(0, import_node_fs.existsSync)(io.outPath)) {
+    return {
+      status: "error",
+      runtimeRequested,
+      elapsedMs,
+      stage: "write",
+      error: "the runtime reported success but wrote no samples file"
+    };
   }
   try {
     (0, import_core.parseSamples)((0, import_node_fs.readFileSync)(io.outPath, "utf8"));
@@ -187,7 +193,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: import_core.RUN_RECORD_SCHEMA_VERSION,
@@ -203,7 +208,7 @@ async function runFit(spec, resolved, io) {
     started_at: startedAt,
     elapsed_ms: elapsedMs
   };
-  (0, import_node_fs.writeFileSync)(`${io.outPath}.run.json`, `${JSON.stringify(record, null, 2)}
+  (0, import_node_fs.writeFileSync)(io.recordPath ?? `${io.outPath}.run.json`, `${JSON.stringify(record, null, 2)}
 `);
   return {
     status: "ok",
@@ -213,10 +218,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 = (0, import_node_path3.join)(tmp, "request.json");
+  (0, import_node_fs.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) (0, import_node_fs.rmSync)(tmp, { recursive: true, force: true });
+  }
+  const elapsedMs = Math.round(performance.now() - start);
+  if (code !== 0 || !(0, import_node_fs.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
 var import_node_fs2 = require("fs");
-var import_node_os3 = require("os");
 var import_node_path4 = require("path");
 var import_core2 = require("@mcmcjs/core");
 function predictData(spec) {
@@ -234,7 +272,8 @@ async function runPredict(spec, resolved, io) {
   const start = performance.now();
   const runtimeRequested = spec.backend.version;
   const data = predictData(spec);
-  const tmp = io.tmpDir ?? (0, import_node_fs2.mkdtempSync)((0, import_node_path4.join)((0, import_node_os3.tmpdir)(), "mcmcjs-predict-"));
+  const ownTmp = io.tmpDir === void 0;
+  const tmp = io.tmpDir ?? makeRequestDir("predict");
   const requestPath = (0, import_node_path4.join)(tmp, "request.json");
   (0, import_node_fs2.writeFileSync)(
     requestPath,
@@ -257,7 +296,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) (0, import_node_fs2.rmSync)(tmp, { recursive: true, force: true });
+  }
   const elapsedMs = Math.round(performance.now() - start);
   if (code !== 0 || !(0, import_node_fs2.existsSync)(io.outPath)) {
     const failure = lastJsonLine(stderr);
@@ -306,11 +352,13 @@ async function runPredict(spec, resolved, io) {
 
 // src/project.ts
 var import_node_fs3 = require("fs");
-var import_node_os4 = require("os");
+var import_node_os3 = require("os");
 var import_node_path5 = require("path");
 var import_engine3 = require("@mcmcjs/engine");
 var PACKAGES = [
   "Turing",
+  "FlexiChains",
+  "DimensionalData",
   "MCMCChains",
   "JuliaBUGS",
   "AdvancedHMC",
@@ -319,7 +367,7 @@ var PACKAGES = [
   "StableRNGs"
 ];
 function managedProjectDir() {
-  const base = process.env.XDG_DATA_HOME || (0, import_node_path5.join)((0, import_node_os4.homedir)(), ".local", "share");
+  const base = process.env.XDG_DATA_HOME || (0, import_node_path5.join)((0, import_node_os3.homedir)(), ".local", "share");
   return (0, import_node_path5.join)(base, "mcmcjs", "julia", "env");
 }
 function sentinelPath(dir) {
@@ -505,7 +553,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));
@@ -517,6 +565,256 @@ function assertVersionsInstalled(channels, installed) {
     );
   }
 }
+
+// src/worker.ts
+var import_node_child_process = require("child_process");
+var import_node_fs4 = require("fs");
+var import_node_net = require("net");
+var import_node_path7 = require("path");
+var import_engine4 = require("@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 ? (0, import_node_path7.join)(runtime, "mcmcjs") : (0, import_node_path7.join)(sharedTmpParent(), "workers");
+  (0, import_node_fs4.mkdirSync)(dir, { recursive: true, mode: 448 });
+  return dir;
+}
+function workerSocketPath(command, projectDir) {
+  return (0, import_node_path7.join)(workersDir(), `worker-${sha256(`${command}
+${projectDir}`).slice(0, 12)}.sock`);
+}
+function connectOnce(path) {
+  return new Promise((resolve, reject) => {
+    const sock = (0, import_node_net.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 {
+    (0, import_node_fs4.mkdirSync)(lock);
+    return true;
+  } catch {
+    try {
+      if (Date.now() - (0, import_node_fs4.statSync)(lock).mtimeMs > START_WAIT_MS) {
+        (0, import_node_fs4.rmdirSync)(lock);
+        (0, import_node_fs4.mkdirSync)(lock);
+        return true;
+      }
+    } catch {
+    }
+    return false;
+  }
+}
+function releaseStartLock(socketPath) {
+  try {
+    (0, import_node_fs4.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 {
+      (0, import_node_fs4.rmSync)(path, { force: true });
+      notify?.("starting the Julia worker (first use loads Turing, can take a minute)...");
+      const child = (0, import_node_child_process.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 = (0, import_engine4.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 = (0, import_node_fs4.readdirSync)(dir).filter((f) => f.startsWith("worker-") && f.endsWith(".sock"));
+  const statuses = [];
+  for (const name of sockets) {
+    const socket = (0, import_node_path7.join)(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 {
+    (0, import_node_fs4.rmSync)(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";
+  (0, import_node_fs4.rmSync)(socket, { force: true });
+  return "stopped";
+}
 // Annotate the CommonJS export names for ESM import in node:
 0 && (module.exports = {
   addVersion,
@@ -528,6 +826,7 @@ function assertVersionsInstalled(channels, installed) {
   juliaEngine,
   juliaupInstallCommand,
   listVersions,
+  listWorkers,
   managedProjectDir,
   managedProjectReady,
   planSetup,
@@ -536,10 +835,13 @@ function assertVersionsInstalled(channels, installed) {
   resolveVersion,
   runDoctor,
   runFit,
+  runFitAuto,
   runMatrix,
   runPredict,
   runSetup,
   setDefaultVersion,
+  sharedTmpParent,
+  stopWorker,
   updateVersion
 });
 //# sourceMappingURL=index.cjs.map