@elizaos/plugin-training 2.0.3-beta.5 → 2.0.3-beta.7

package/dist/services/training-vast-service.js

@@ -0,0 +1,648 @@
+import { spawn } from "node:child_process";
+import { existsSync, promises as fs } from "node:fs";
+import { join, resolve } from "node:path";
+import { fileURLToPath } from "node:url";
+import { logger } from "@elizaos/core";
+import {
+  aggregateInferenceStats,
+  emptyInferenceStatsAggregate,
+  parseStatRow
+} from "./vast-inference-stats.js";
+import {
+  inferenceStatsPath,
+  readInferenceEndpoints,
+  readJobLogTail,
+  VastJobStore,
+  writeInferenceEndpoints
+} from "./vast-job-store.js";
+import {
+  runCapture,
+  runDetachedToLog,
+  VastServiceError
+} from "./vast-subprocess.js";
+import { VastServiceError as VastServiceError2 } from "./vast-subprocess.js";
+const DEFAULT_PYTHON = {
+  command: "uv",
+  preArgs: ["run", "--quiet", "python"]
+};
+const HERE = fileURLToPath(import.meta.url);
+const TRAINING_ROOT_DEFAULT = resolve(
+  HERE,
+  "..",
+  "..",
+  "..",
+  "..",
+  "..",
+  "packages",
+  "training"
+);
+const RUN_NAME_PATTERN = /^[A-Za-z0-9._-]{1,64}$/;
+const LABEL_PATTERN = /^[A-Za-z0-9._-]{1,64}$/;
+const REGISTRY_KEY_PATTERN = /^[A-Za-z0-9._-]{1,64}$/;
+class VastTrainingService {
+  store;
+  trainingRoot;
+  python;
+  spawnImpl;
+  registryCache = null;
+  registryLoadedAt = null;
+  constructor(options = {}) {
+    this.store = options.store ?? new VastJobStore();
+    this.trainingRoot = options.trainingRoot ?? TRAINING_ROOT_DEFAULT;
+    this.python = options.pythonLauncher ?? DEFAULT_PYTHON;
+    this.spawnImpl = options.spawnImpl ?? spawn;
+  }
+  // ── Registry ──────────────────────────────────────────────────────────
+  async getRegistry(refresh = false) {
+    if (this.registryCache && !refresh) return this.registryCache;
+    const dumpScript = join(
+      this.trainingRoot,
+      "scripts",
+      "dump_registry_json.py"
+    );
+    if (!existsSync(dumpScript)) {
+      throw new VastServiceError(
+        `Registry dump script not found at ${dumpScript}`,
+        503
+      );
+    }
+    const stdout = await this.runCapture(
+      this.python.command,
+      [...this.python.preArgs, dumpScript],
+      { cwd: this.trainingRoot }
+    );
+    let parsed;
+    try {
+      parsed = JSON.parse(stdout);
+    } catch (err) {
+      throw new VastServiceError(
+        `Registry dump produced non-JSON output: ${err instanceof Error ? err.message : String(err)}`,
+        500
+      );
+    }
+    if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) {
+      throw new VastServiceError(
+        "Registry dump returned unexpected shape",
+        500
+      );
+    }
+    const out = {};
+    for (const [key, value] of Object.entries(
+      parsed
+    )) {
+      const entry = narrowRegistryEntry(value);
+      if (entry) out[key] = entry;
+    }
+    this.registryCache = out;
+    this.registryLoadedAt = (/* @__PURE__ */ new Date()).toISOString();
+    return out;
+  }
+  async listRegistry(refresh = false) {
+    const registry = await this.getRegistry(refresh);
+    const entries = Object.entries(registry).map(([short_name, entry]) => ({
+      short_name,
+      entry
+    }));
+    return { loaded_at: this.registryLoadedAt, entries };
+  }
+  async ensureRegistryKey(key) {
+    if (!REGISTRY_KEY_PATTERN.test(key)) {
+      throw new VastServiceError(
+        "registry_key must match [A-Za-z0-9._-]{1,64}",
+        400
+      );
+    }
+    const registry = await this.getRegistry(false);
+    const entry = registry[key];
+    if (!entry) {
+      throw new VastServiceError(
+        `Unknown registry_key '${key}'. Refresh /api/training/vast/models to update the cache.`,
+        400
+      );
+    }
+    return entry;
+  }
+  // ── Jobs ──────────────────────────────────────────────────────────────
+  async listJobs() {
+    return this.store.list();
+  }
+  async getJob(jobId) {
+    return this.store.get(jobId);
+  }
+  async createJob(input) {
+    await this.ensureRegistryKey(input.registry_key);
+    if (!Number.isInteger(input.epochs) || input.epochs < 1 || input.epochs > 64) {
+      throw new VastServiceError(
+        "epochs must be an integer between 1 and 64",
+        400
+      );
+    }
+    if (input.run_name !== void 0 && !RUN_NAME_PATTERN.test(input.run_name)) {
+      throw new VastServiceError(
+        "run_name must match [A-Za-z0-9._-]{1,64}",
+        400
+      );
+    }
+    const job_id = makeJobId();
+    const run_name = input.run_name ?? defaultRunName(input.registry_key);
+    const now = (/* @__PURE__ */ new Date()).toISOString();
+    const record = {
+      job_id,
+      run_name,
+      registry_key: input.registry_key,
+      status: "queued",
+      epochs: input.epochs,
+      created_at: now,
+      updated_at: now,
+      started_at: null,
+      ended_at: null,
+      vast_instance_id: null,
+      exit_code: null,
+      error: null
+    };
+    await this.store.insert(record);
+    void this.dispatchJob(record);
+    return record;
+  }
+  async cancelJob(jobId) {
+    const job = await this.store.get(jobId);
+    if (!job) throw new VastServiceError("Job not found", 404);
+    if (job.status === "completed" || job.status === "failed") {
+      throw new VastServiceError(
+        `Cannot cancel job in terminal state '${job.status}'`,
+        409
+      );
+    }
+    if (job.status === "cancelled") return job;
+    return this.store.update(jobId, {
+      status: "cancelled",
+      ended_at: (/* @__PURE__ */ new Date()).toISOString()
+    });
+  }
+  async readJobLog(jobId, tailLines) {
+    const job = await this.store.get(jobId);
+    if (!job) throw new VastServiceError("Job not found", 404);
+    return readJobLogTail(jobId, tailLines);
+  }
+  async runEval(jobId, input) {
+    const job = await this.store.get(jobId);
+    if (!job) throw new VastServiceError("Job not found", 404);
+    const evalScript = join(this.trainingRoot, "scripts", "eval_checkpoint.py");
+    if (!existsSync(evalScript)) {
+      throw new VastServiceError(
+        `eval_checkpoint.py not found at ${evalScript} \u2014 CheckpointSyncAgent has not landed it yet`,
+        503
+      );
+    }
+    const checkpointDir = input.checkpoint_dir ?? join(this.trainingRoot, "checkpoints", job.run_name, "final");
+    if (!isSafeCheckpointPath(checkpointDir, this.trainingRoot)) {
+      throw new VastServiceError(
+        "checkpoint_dir must resolve under the training root",
+        400
+      );
+    }
+    const valJsonl = input.val_jsonl ?? "data/smoke/val.jsonl";
+    if (!isSafeRelativePath(valJsonl)) {
+      throw new VastServiceError(
+        "val_jsonl must be a relative path without traversal",
+        400
+      );
+    }
+    const maxExamples = input.max_examples ?? 50;
+    if (!Number.isInteger(maxExamples) || maxExamples < 1 || maxExamples > 5e3) {
+      throw new VastServiceError(
+        "max_examples must be an integer between 1 and 5000",
+        400
+      );
+    }
+    const outPath = join(
+      this.trainingRoot,
+      "checkpoints",
+      job.run_name,
+      `eval-${Date.now()}.json`
+    );
+    const args = [
+      ...this.python.preArgs,
+      evalScript,
+      "--checkpoint",
+      checkpointDir,
+      "--registry-key",
+      job.registry_key,
+      "--val-jsonl",
+      valJsonl,
+      "--max-examples",
+      String(maxExamples),
+      "--out",
+      outPath
+    ];
+    const exitCode = await this.runDetachedToLog(
+      jobId,
+      this.python.command,
+      args,
+      this.trainingRoot
+    );
+    let summary = null;
+    if (existsSync(outPath)) {
+      try {
+        const raw = await fs.readFile(outPath, "utf8");
+        const parsed = JSON.parse(raw);
+        if (parsed && typeof parsed === "object" && !Array.isArray(parsed)) {
+          summary = parsed;
+        }
+      } catch (err) {
+        logger.warn(
+          `[VastTrainingService] failed to read eval output ${outPath}: ${err instanceof Error ? err.message : String(err)}`
+        );
+      }
+    }
+    return { job_id: jobId, exit_code: exitCode, out_path: outPath, summary };
+  }
+  // ── Checkpoints ───────────────────────────────────────────────────────
+  async listCheckpointsForRegistryKey(registryKey) {
+    await this.ensureRegistryKey(registryKey);
+    const checkpointsRoot = join(this.trainingRoot, "checkpoints");
+    if (!existsSync(checkpointsRoot)) return [];
+    const out = [];
+    const runDirs = await fs.readdir(checkpointsRoot, { withFileTypes: true });
+    for (const runDir of runDirs) {
+      if (!runDir.isDirectory()) continue;
+      if (!runDir.name.startsWith(registryKey.replace(/\./g, "-"))) continue;
+      const runPath = join(checkpointsRoot, runDir.name);
+      const status = join(runPath, "STATUS.md");
+      if (existsSync(status)) {
+        try {
+          const body = await fs.readFile(status, "utf8");
+          if (body.includes("FAILED RUN")) continue;
+        } catch {
+          continue;
+        }
+      }
+      const ckptEntries = await fs.readdir(runPath, { withFileTypes: true });
+      for (const ckpt of ckptEntries) {
+        if (!ckpt.isDirectory()) continue;
+        if (!/^checkpoint-\d+$|^final$/.test(ckpt.name)) continue;
+        const ckptPath = join(runPath, ckpt.name);
+        const stepMatch = /^checkpoint-(\d+)$/.exec(ckpt.name);
+        const evalDoneFile = join(ckptPath, "_eval.json");
+        let evalSummary = null;
+        const evaluated = existsSync(evalDoneFile);
+        if (evaluated) {
+          try {
+            const raw = await fs.readFile(evalDoneFile, "utf8");
+            const parsed = JSON.parse(raw);
+            if (parsed && typeof parsed === "object" && !Array.isArray(parsed)) {
+              evalSummary = parsed;
+            }
+          } catch {
+            evalSummary = null;
+          }
+        }
+        out.push({
+          name: `${runDir.name}/${ckpt.name}`,
+          path: ckptPath,
+          step: stepMatch ? Number(stepMatch[1]) : null,
+          evaluated,
+          eval_summary: evalSummary
+        });
+      }
+    }
+    out.sort((a, b) => (b.step ?? -1) - (a.step ?? -1));
+    return out;
+  }
+  // ── Inference endpoints ───────────────────────────────────────────────
+  async listInferenceEndpoints() {
+    return readInferenceEndpoints();
+  }
+  async createInferenceEndpoint(input) {
+    if (!LABEL_PATTERN.test(input.label)) {
+      throw new VastServiceError("label must match [A-Za-z0-9._-]{1,64}", 400);
+    }
+    let url;
+    try {
+      url = new URL(input.base_url);
+    } catch {
+      throw new VastServiceError("base_url must be a valid URL", 400);
+    }
+    if (url.protocol !== "http:" && url.protocol !== "https:") {
+      throw new VastServiceError("base_url must use http:// or https://", 400);
+    }
+    await this.ensureRegistryKey(input.registry_key);
+    const endpoints = await readInferenceEndpoints();
+    if (endpoints.some((e) => e.label === input.label)) {
+      throw new VastServiceError(
+        `Inference endpoint with label '${input.label}' already exists`,
+        409
+      );
+    }
+    const record = {
+      id: `ep_${makeShortId()}`,
+      label: input.label,
+      base_url: input.base_url,
+      registry_key: input.registry_key,
+      created_at: (/* @__PURE__ */ new Date()).toISOString()
+    };
+    endpoints.push(record);
+    await writeInferenceEndpoints(endpoints);
+    return record;
+  }
+  async deleteInferenceEndpoint(id) {
+    const endpoints = await readInferenceEndpoints();
+    const next = endpoints.filter((e) => e.id !== id);
+    if (next.length === endpoints.length) return false;
+    await writeInferenceEndpoints(next);
+    return true;
+  }
+  /**
+   * Read a per-job budget snapshot via `scripts.lib.vast_budget snapshot
+   * --json`. Returns `null` when the job has no provisioned instance
+   * yet (the UI shows a "not provisioned" state in that case).
+   *
+   * The python module is the single source of truth — it both renders
+   * the watcher's status line and answers this endpoint, so the UI and
+   * the watcher never disagree about the budget state.
+   */
+  async getJobBudget(jobId) {
+    const job = await this.store.get(jobId);
+    if (!job) throw new VastServiceError("Job not found", 404);
+    const instanceIdRaw = job.vast_instance_id;
+    if (!instanceIdRaw) return null;
+    const instanceId = Number(instanceIdRaw);
+    if (!Number.isFinite(instanceId) || instanceId <= 0) return null;
+    const dumpScript = "-m";
+    let stdout;
+    try {
+      stdout = await this.runCapture(
+        this.python.command,
+        [
+          ...this.python.preArgs,
+          dumpScript,
+          "scripts.lib.vast_budget",
+          "snapshot",
+          String(instanceId),
+          "--pipeline",
+          job.registry_key,
+          "--run-name",
+          job.run_name,
+          "--json"
+        ],
+        { cwd: this.trainingRoot }
+      );
+    } catch (err) {
+      logger.warn(
+        `[VastTrainingService] budget snapshot failed for ${jobId}: ${err instanceof Error ? err.message : String(err)}`
+      );
+      return null;
+    }
+    const trimmed = stdout.trim();
+    if (!trimmed) return null;
+    let parsed;
+    try {
+      parsed = JSON.parse(trimmed);
+    } catch {
+      return null;
+    }
+    return narrowBudgetSnapshot(parsed, jobId);
+  }
+  async getInferenceStats(label, lastMinutes) {
+    if (!Number.isFinite(lastMinutes) || lastMinutes <= 0 || lastMinutes > 24 * 60) {
+      throw new VastServiceError(
+        "last_minutes must be a positive number \u2264 1440",
+        400
+      );
+    }
+    if (label !== null && !LABEL_PATTERN.test(label)) {
+      throw new VastServiceError("label must match [A-Za-z0-9._-]{1,64}", 400);
+    }
+    const path = inferenceStatsPath();
+    if (!existsSync(path)) {
+      return emptyInferenceStatsAggregate(label, lastMinutes);
+    }
+    const cutoffMs = Date.now() - lastMinutes * 6e4;
+    const raw = await fs.readFile(path, "utf8");
+    const samples = [];
+    for (const line of raw.split("\n")) {
+      const trimmed = line.trim();
+      if (!trimmed) continue;
+      const row = parseStatRow(trimmed);
+      if (!row) continue;
+      if (label !== null && row.label !== label) continue;
+      const ts = Date.parse(row.ts);
+      if (!Number.isFinite(ts) || ts < cutoffMs) continue;
+      samples.push(row);
+    }
+    return aggregateInferenceStats(samples, label, lastMinutes);
+  }
+  // ── Internals ─────────────────────────────────────────────────────────
+  async dispatchJob(record) {
+    const trainScript = join(this.trainingRoot, "scripts", "train_vast.sh");
+    if (!existsSync(trainScript)) {
+      logger.error(
+        `[VastTrainingService] train_vast.sh not found at ${trainScript}`
+      );
+      await this.store.update(record.job_id, {
+        status: "failed",
+        ended_at: (/* @__PURE__ */ new Date()).toISOString(),
+        error: "train_vast.sh not found"
+      });
+      return;
+    }
+    await this.store.update(record.job_id, {
+      status: "provisioning",
+      started_at: (/* @__PURE__ */ new Date()).toISOString()
+    });
+    const args = [
+      trainScript,
+      "provision-and-train",
+      "--registry-key",
+      record.registry_key,
+      "--epochs",
+      String(record.epochs)
+    ];
+    try {
+      const exitCode = await this.runDetachedToLog(
+        record.job_id,
+        "bash",
+        args,
+        this.trainingRoot,
+        {
+          RUN_NAME: record.run_name
+        }
+      );
+      const instanceId = await this.discoverInstanceIdForRun(record);
+      if (exitCode === 0) {
+        await this.store.update(record.job_id, {
+          status: "completed",
+          ended_at: (/* @__PURE__ */ new Date()).toISOString(),
+          exit_code: exitCode,
+          vast_instance_id: instanceId
+        });
+      } else {
+        await this.store.update(record.job_id, {
+          status: "failed",
+          ended_at: (/* @__PURE__ */ new Date()).toISOString(),
+          exit_code: exitCode,
+          vast_instance_id: instanceId,
+          error: `train_vast.sh exited with code ${exitCode}`
+        });
+      }
+    } catch (err) {
+      logger.error(
+        `[VastTrainingService] dispatch failure for ${record.job_id}: ${err instanceof Error ? err.stack ?? err.message : String(err)}`
+      );
+      await this.store.update(record.job_id, {
+        status: "failed",
+        ended_at: (/* @__PURE__ */ new Date()).toISOString(),
+        error: err instanceof Error ? err.message : String(err)
+      });
+    }
+  }
+  async discoverInstanceIdForRun(record) {
+    const stateFile = join(
+      this.trainingRoot,
+      ".eliza",
+      "vast-state",
+      `${record.run_name}.json`
+    );
+    if (existsSync(stateFile)) {
+      try {
+        const raw = await fs.readFile(stateFile, "utf8");
+        const parsed = JSON.parse(raw);
+        if (parsed && typeof parsed === "object") {
+          const obj = parsed;
+          const id = obj.instance_id ?? obj.vast_instance_id;
+          if (typeof id === "string" && id.trim()) return id.trim();
+          if (typeof id === "number") return String(id);
+        }
+      } catch (err) {
+        logger.warn(
+          `[VastTrainingService] could not parse vast-state file ${stateFile}: ${err instanceof Error ? err.message : String(err)}`
+        );
+      }
+    }
+    const logTail = await readJobLogTail(record.job_id, 200);
+    for (const line of logTail) {
+      const match = /ELIZA_VAST_INSTANCE_ID=([A-Za-z0-9_-]+)/.exec(line);
+      if (match) return match[1];
+    }
+    return null;
+  }
+  runCapture(command, args, options) {
+    return runCapture(this.spawnImpl, command, args, options);
+  }
+  runDetachedToLog(jobId, command, args, cwd, extraEnv = {}) {
+    return runDetachedToLog(
+      this.spawnImpl,
+      jobId,
+      command,
+      args,
+      cwd,
+      extraEnv
+    );
+  }
+}
+function defaultRunName(registryKey) {
+  return `${registryKey.replace(/\./g, "-")}-apollo`;
+}
+function makeJobId() {
+  return `vjob_${makeShortId()}`;
+}
+function makeShortId() {
+  const buf = new Uint8Array(8);
+  if (typeof globalThis.crypto.getRandomValues === "function") {
+    globalThis.crypto.getRandomValues(buf);
+  } else {
+    for (let i = 0; i < buf.length; i += 1)
+      buf[i] = Math.floor(Math.random() * 256);
+  }
+  let n = 0n;
+  for (const byte of buf) n = n << 8n | BigInt(byte);
+  return n.toString(36).padStart(12, "0").slice(0, 12);
+}
+function narrowRegistryEntry(value) {
+  if (!value || typeof value !== "object") return null;
+  const obj = value;
+  if (typeof obj.eliza_short_name === "string" && typeof obj.eliza_repo_id === "string" && typeof obj.gguf_repo_id === "string" && typeof obj.base_hf_id === "string" && typeof obj.tier === "string" && typeof obj.inference_max_context === "number") {
+    return {
+      eliza_short_name: obj.eliza_short_name,
+      eliza_repo_id: obj.eliza_repo_id,
+      gguf_repo_id: obj.gguf_repo_id,
+      base_hf_id: obj.base_hf_id,
+      tier: obj.tier,
+      inference_max_context: obj.inference_max_context
+    };
+  }
+  return null;
+}
+function narrowBudgetSnapshot(raw, jobId) {
+  if (!raw || typeof raw !== "object" || Array.isArray(raw)) return null;
+  const obj = raw;
+  const num = (k) => {
+    const v = obj[k];
+    if (typeof v !== "number" || !Number.isFinite(v)) return null;
+    return v;
+  };
+  const str = (k) => {
+    const v = obj[k];
+    return typeof v === "string" ? v : null;
+  };
+  const bool = (k) => {
+    const v = obj[k];
+    return typeof v === "boolean" ? v : null;
+  };
+  const instanceId = num("instance_id");
+  const pipeline = str("pipeline");
+  const runName = str("run_name");
+  const gpuName = str("gpu_name");
+  const gpuSku = str("gpu_sku");
+  const state = str("state");
+  const uptime = num("uptime_seconds");
+  const uptimePretty = str("uptime_pretty");
+  const dph = num("dph_total");
+  const total = num("total_so_far_usd");
+  const fetchedAt = num("fetched_at");
+  const numGpus = num("num_gpus");
+  const overSoft = bool("over_soft");
+  const overHard = bool("over_hard");
+  if (instanceId === null || !pipeline || runName === null || !gpuName || !gpuSku || !state || uptime === null || !uptimePretty || dph === null || total === null || fetchedAt === null || numGpus === null || overSoft === null || overHard === null) {
+    return null;
+  }
+  const softRaw = obj.soft_cap_usd;
+  const hardRaw = obj.hard_cap_usd;
+  const soft = typeof softRaw === "number" && Number.isFinite(softRaw) ? softRaw : null;
+  const hard = typeof hardRaw === "number" && Number.isFinite(hardRaw) ? hardRaw : null;
+  return {
+    job_id: jobId,
+    instance_id: instanceId,
+    pipeline,
+    run_name: runName,
+    gpu_name: gpuName,
+    num_gpus: numGpus,
+    gpu_sku: gpuSku,
+    state,
+    uptime_seconds: uptime,
+    uptime_pretty: uptimePretty,
+    dph_total: dph,
+    total_so_far_usd: total,
+    soft_cap_usd: soft,
+    hard_cap_usd: hard,
+    over_soft: overSoft,
+    over_hard: overHard,
+    fetched_at: fetchedAt
+  };
+}
+function isSafeRelativePath(p) {
+  if (typeof p !== "string" || !p.trim()) return false;
+  if (p.startsWith("/") || p.includes("..")) return false;
+  return true;
+}
+function isSafeCheckpointPath(checkpointDir, trainingRoot) {
+  if (typeof checkpointDir !== "string" || !checkpointDir.trim()) return false;
+  const root = resolve(trainingRoot);
+  const resolved = resolve(checkpointDir);
+  return resolved === root || resolved.startsWith(`${root}/`);
+}
+export {
+  VastServiceError2 as VastServiceError,
+  VastTrainingService
+};
+//# sourceMappingURL=training-vast-service.js.map