spb/autoscholar

1import { chmod, mkdir, readFile, writeFile } from "node:fs/promises";2import { hostname } from "node:os";3import { join, resolve } from "node:path";4import { ControlPlaneError, WorkerClient } from "@autoscholar/maclustr-adapter";5import { createLogger, redact } from "@autoscholar/observability";6import { type RunnerHost, runResearch, simulateResearch } from "@autoscholar/research-runner";7import type { ArtifactKind, RunEventInput, StageKey } from "@autoscholar/schemas";8import { collectNodeMetadata, loadAverage, sdkVersion, workerVersion } from "./node-meta.js";9import { runSelfTest } from "./selftest.js";1011const log = createLogger("worker");1213interface Config {14  controlPlaneUrl: string;15  workerName: string;16  workerRoot: string;17  token?: string;18  enrollmentSecret?: string;19  mode: "real" | "simulated";20  maxConcurrent: number;21  profiles: string[];22  apiKey?: string;23  extraDomains: string[];24  tectonicBin?: string;25  latexTemplateDir: string;26  pollMs: number;27}2829function loadConfig(): Config {30  const e = process.env;31  const workerRoot = e.WORKER_ROOT ?? "/var/autoscholar";32  return {33    controlPlaneUrl: e.CONTROL_PLANE_URL ?? "http://127.0.0.1:8470",34    workerName: e.WORKER_NAME ?? hostname(),35    workerRoot,36    token: e.WORKER_TOKEN || undefined,37    enrollmentSecret: e.WORKER_ENROLLMENT_SECRET || undefined,38    mode: e.WORKER_MODE === "simulated" ? "simulated" : "real",39    maxConcurrent: Number(e.WORKER_MAX_CONCURRENT ?? 1),40    profiles: (e.WORKER_PROFILES ?? "standard")41      .split(",")42      .map((s) => s.trim())43      .filter(Boolean),44    apiKey: e.ANTHROPIC_API_KEY || undefined,45    extraDomains: (e.RESEARCH_ALLOWED_DOMAINS ?? "")46      .split(",")47      .map((s) => s.trim())48      .filter(Boolean),49    tectonicBin: e.TECTONIC_BIN || undefined,50    latexTemplateDir: e.LATEX_TEMPLATE_DIR ?? resolve(workerRoot, "templates/latex"),51    pollMs: Number(e.WORKER_POLL_MS ?? 4000),52  };53}5455class Worker {56  private readonly client: WorkerClient;57  private readonly running = new Map<58    string,59    { fencingToken: number; cancel: boolean; restart?: boolean; events: RunEventInput[] }60  >();61  private stopping = false;62  private draining = false;63  private tokenPath: string;6465  constructor(private readonly cfg: Config) {66    this.client = new WorkerClient({ baseUrl: cfg.controlPlaneUrl, token: cfg.token });67    this.tokenPath = join(cfg.workerRoot, "state/worker-token");68  }6970  async start() {71    await mkdir(join(this.cfg.workerRoot, "runs"), { recursive: true });72    await mkdir(join(this.cfg.workerRoot, "state"), { recursive: true });73    if (!this.cfg.token) {74      try {75        const saved = (await readFile(this.tokenPath, "utf8")).trim();76        if (saved) this.client.setToken(saved);77      } catch {78        /* none */79      }80    }81    const node = collectNodeMetadata({82      simulated: this.cfg.mode === "simulated",83      profiles: this.cfg.profiles,84    });85    if (this.cfg.mode === "real" && !node.sandboxAvailable)86      log.error(87        "sandbox-exec not available: this worker will never be dispatched research (sandboxAvailable=false)",88      );89    if (this.cfg.mode === "real" && !this.cfg.apiKey)90      log.error("ANTHROPIC_API_KEY missing: research sessions cannot start");91    // enroll if needed92    let enrolled = false;93    for (let i = 0; i < 30 && !enrolled && !this.stopping; i++) {94      try {95        await this.client.heartbeat({96          status: "online",97          capacity: this.cfg.maxConcurrent,98          running: [],99          load: loadAverage(),100          node,101          leases: [],102        });103        enrolled = true;104      } catch (e) {105        if (e instanceof ControlPlaneError && e.status === 401 && this.cfg.enrollmentSecret) {106          try {107            const r = await this.client.enroll(this.cfg.enrollmentSecret, this.cfg.workerName, node);108            this.client.setToken(r.token);109            await writeFile(this.tokenPath, r.token, { mode: 0o600 });110            await chmod(this.tokenPath, 0o600);111            log.info({ workerId: r.workerId }, "enrolled with control plane");112            enrolled = true;113          } catch (e2) {114            log.error({ err: (e2 as Error).message }, "enrollment failed");115          }116        } else log.warn({ err: (e as Error).message, attempt: i }, "control plane unreachable; retrying");117        if (!enrolled) await new Promise((r) => setTimeout(r, 5000));118      }119    }120    if (!enrolled) throw new Error("could not reach/enroll with the control plane");121    log.info(122      {123        name: this.cfg.workerName,124        mode: this.cfg.mode,125        sdk: sdkVersion(),126        worker: workerVersion(),127        controlPlane: this.cfg.controlPlaneUrl,128        sandbox: node.sandboxAvailable,129      },130      "worker online",131    );132    void this.heartbeatLoop(node);133    void this.leaseLoop();134  }135136  private async heartbeatLoop(node: ReturnType<typeof collectNodeMetadata>) {137    while (!this.stopping) {138      try {139        const leases = [...this.running.entries()].map(([runId, r]) => ({140          runId,141          fencingToken: r.fencingToken,142        }));143        const hb = await this.client.heartbeat({144          status: this.draining ? "draining" : "online",145          capacity: this.cfg.maxConcurrent,146          running: [...this.running.keys()],147          load: loadAverage(),148          leases,149          node: { tools: node.tools, sandboxAvailable: node.sandboxAvailable },150        });151        this.draining = hb.drain;152        for (const id of hb.cancel) {153          const r = this.running.get(id);154          if (r) r.cancel = true;155        }156        for (const id of hb.revokedLeases) {157          const r = this.running.get(id);158          if (r) {159            log.warn({ runId: id }, "lease revoked by control plane (fencing) — aborting local work");160            r.cancel = true;161          }162        }163        await this.flushEvents();164      } catch (e) {165        log.warn({ err: (e as Error).message }, "heartbeat failed");166      }167      await new Promise((r) => setTimeout(r, 20_000));168    }169  }170171  private async flushEvents() {172    for (const [runId, r] of this.running) {173      if (r.events.length === 0) continue;174      const batch = r.events.splice(0, 200);175      try {176        await this.client.events(177          runId,178          r.fencingToken,179          batch.map((e) => ({ ...e, level: e.level ?? "info" })),180        );181      } catch (e) {182        if (e instanceof ControlPlaneError && e.status === 409) {183          r.cancel = true;184          log.warn({ runId }, "lease lost while posting events; cancelling local work");185        } else r.events.unshift(...batch.slice(0, 50));186      }187    }188  }189190  private async leaseLoop() {191    while (!this.stopping) {192      if (!this.draining && this.running.size < this.cfg.maxConcurrent) {193        try {194          const { job } = await this.client.lease(this.cfg.profiles);195          if (job) {196            log.info({ runId: job.runId, attempt: job.attemptNo, resume: job.resumeFromStage }, "leased job");197            void this.execute(job);198          }199        } catch (e) {200          log.warn({ err: (e as Error).message }, "lease request failed");201        }202      }203      await new Promise((r) => setTimeout(r, this.cfg.pollMs));204    }205  }206207  private async execute(job: Awaited<ReturnType<WorkerClient["lease"]>>["job"] & object) {208    const state: { fencingToken: number; cancel: boolean; restart?: boolean; events: RunEventInput[] } = {209      fencingToken: job.fencingToken,210      cancel: job.cancelRequested,211      events: [],212    };213    this.running.set(job.runId, state);214    const renew = setInterval(async () => {215      try {216        const r = await this.client.renew(job.runId, job.fencingToken);217        if (r.cancelRequested) state.cancel = true;218        await this.flushEvents();219      } catch (e) {220        if (e instanceof ControlPlaneError && e.status === 409) {221          state.cancel = true;222          log.warn({ runId: job.runId }, "lease lost (stale fencing token); aborting");223        }224      }225    }, 30_000);226    const host: RunnerHost = {227      workspaceRoot: this.cfg.workerRoot,228      apiKey: this.cfg.apiKey ?? "",229      extraDomains: this.cfg.extraDomains,230      tectonicBin: this.cfg.tectonicBin,231      tectonicCache: join(this.cfg.workerRoot, "cache/tectonic"),232      latexTemplateDir: this.cfg.latexTemplateDir,233      workerName: this.cfg.workerName,234      sdkVersion: sdkVersion(),235      toolchain: collectNodeMetadata({236        simulated: this.cfg.mode === "simulated",237        profiles: this.cfg.profiles,238      }).tools,239      emit: (ev) => {240        state.events.push(redact(ev));241        if (state.events.length >= 50) void this.flushEvents();242      },243      stage: async (report) => {244        await this.flushEvents();245        await this.client.stage(job.runId, { fencingToken: job.fencingToken, ...report });246      },247      upload: async (path: string, kind: ArtifactKind, contentType: string, buf: Buffer) => {248        await this.client.uploadArtifact(249          job.runId,250          job.fencingToken,251          {252            path,253            kind,254            contentType,255            stage: (job.resumeFromStage ?? "preflight") as StageKey,256            synthetic: this.cfg.mode === "simulated",257          },258          buf,259        );260      },261      isCancelled: () => state.cancel || this.stopping,262      fetchArtifacts: async () => {263        const { artifacts } = await this.client.listArtifacts(job.runId, job.fencingToken);264        return artifacts.map((a) => ({265          path: a.path,266          read: () => this.client.downloadArtifact(job.runId, job.fencingToken, a.path),267        }));268      },269    };270    let outcome: Awaited<ReturnType<typeof runResearch>>;271    try {272      const runnerJob = {273        runId: job.runId,274        attemptNo: job.attemptNo,275        manifest: job.manifest,276        instructionText: job.instructionText,277        resumeFromStage: job.resumeFromStage,278        checkpoint: job.checkpoint,279      };280      if (this.cfg.mode === "simulated") {281        const sim = process.env.WORKER_SIM_BEHAVIOUR ?? "";282        outcome = await simulateResearch(runnerJob, host, {283          failAtStage: sim.startsWith("fail:") ? sim.slice(5) : undefined,284          hang: sim === "hang",285          delayMs: Number(process.env.WORKER_SIM_DELAY_MS ?? 0),286        });287      } else {288        if (!this.cfg.apiKey) throw new Error("ANTHROPIC_API_KEY not configured on this worker");289        outcome = await runResearch(runnerJob, host);290      }291      if (state.restart && outcome.outcome === "cancelled") {292        outcome = {293          ...outcome,294          outcome: "failed",295          error: "WORKER_RESTART: worker daemon restarted; resume from checkpoint",296        };297      }298    } catch (e) {299      outcome = {300        outcome: "failed",301        error: `runner crashed: ${(e as Error).message}`.slice(0, 2000),302        costUsd: 0,303        turns: 0,304        stagesCompleted: [],305        resourceUsage: { wallSeconds: 0, diskBytes: 0, peakRssBytes: 0 },306      };307      log.error(308        { runId: job.runId, err: (e as Error).message, stack: (e as Error).stack?.split("\n").slice(0, 5) },309        "runner crashed",310      );311    }312    clearInterval(renew);313    await this.flushEvents();314    try {315      await this.client.complete(job.runId, {316        fencingToken: job.fencingToken,317        outcome: outcome.outcome,318        error: outcome.error,319        costUsd: outcome.costUsd,320        summary: { stagesCompleted: outcome.stagesCompleted, turns: outcome.turns },321        resourceUsage: outcome.resourceUsage,322      });323      log.info({ runId: job.runId, outcome: outcome.outcome, costUsd: outcome.costUsd }, "attempt completed");324    } catch (e) {325      log.error({ runId: job.runId, err: (e as Error).message }, "could not report completion (lease lost?)");326    }327    this.running.delete(job.runId);328  }329330  async stop() {331    this.stopping = true;332    // graceful restart: abort sessions (isCancelled() → true via `stopping`) but report a retryable failure, not a cancel333    for (const r of this.running.values()) r.restart = true;334    const t0 = Date.now();335    while (this.running.size > 0 && Date.now() - t0 < 60_000) await new Promise((r) => setTimeout(r, 500));336  }337}338339const cfg = loadConfig();340if (process.argv[2] === "selftest") {341  runSelfTest(cfg.workerRoot, cfg.apiKey, cfg.latexTemplateDir, cfg.tectonicBin)342    .then((r) => {343      console.log(JSON.stringify(r, null, 2));344      process.exit(r.ok ? 0 : 1);345    })346    .catch((e) => {347      console.error(e);348      process.exit(1);349    });350} else {351  const w = new Worker(cfg);352  process.on("SIGTERM", () => void w.stop().then(() => process.exit(0)));353  process.on("SIGINT", () => void w.stop().then(() => process.exit(0)));354  w.start().catch((e) => {355    log.fatal({ err: (e as Error).message }, "worker failed to start");356    process.exit(1);357  });358}359