@citeark/agent 0.3.20

package/src/compute/autodl-pro-compiler.mjs

@@ -0,0 +1,347 @@
+import { recordResearchCleanup } from '../workspace/control.mjs';
+import { restoreRunRecovery } from './coordinator-recovery.mjs';
+import { mkdir, mkdtemp, readFile, rm, writeFile } from "node:fs/promises";
+import os from "node:os";
+import path from "node:path";
+
+import { ensureFinalWorkspaceCapture, updateRunState } from "../job.mjs";
+import {
+  acceptCompilerDraft,
+  compilerPrompt,
+  normalizeCompilationTask,
+  prepareCompilationBundle,
+} from "../research/compiler.mjs";
+import { agentInvocationPromptArgument, buildAgentInvocation } from "../runtime/index.mjs";
+import { publicAgentRuntime, resolveAgentRuntime } from "../runtime/config.mjs";
+import {
+  CODEX_FINAL_MESSAGE_PATH,
+  recoverJsonFromFinalMessage,
+} from "../runtime/final-output.mjs";
+import { readAgentUsage } from "../telemetry/usage.mjs";
+import {
+  CiteArkError,
+  createRunId,
+  pathExists,
+  readJson,
+  writeJson,
+} from "../util.mjs";
+import { remoteProviderSession } from "./gcp-batch-job.mjs";
+import {
+  classifyProviderActionRequired,
+  describeRemoteAgentFailure,
+} from "./remote-attempt.mjs";
+import {
+  acquireAutoDlProLease,
+  autoDlInstanceRequest,
+  resolveAutoDlProConfig,
+} from "./autodl-pro-job.mjs";
+import {
+  createAutoDlStageArchive,
+  executeAutoDlRemoteAttempt,
+} from "./autodl-pro-remote.mjs";
+
+export async function executeAutoDlProCompilation({
+  taskPath,
+  runsDirectory,
+  runId = createRunId(),
+  agentRuntime,
+  computeProfile,
+  executorOptions = {},
+  dryRun = false,
+  progress = () => {},
+  sourceReviewer,
+  prepareMarkdown,
+  validateResearch = async () => [],
+}) {
+  assertCompilerProfile(computeProfile);
+  const absoluteTaskPath = path.resolve(taskPath);
+  const taskDirectory = path.dirname(absoluteTaskPath);
+  const task = normalizeCompilationTask(await readJson(absoluteTaskPath));
+  const runtimeAgent = resolveAgentRuntime({
+    taskAgent: task.agent,
+    ...agentRuntime,
+    requireApi: false,
+  });
+  task.environment = {
+    ...task.environment,
+    image: computeProfile.image,
+    cpus: computeProfile.cpuCores,
+    memoryGb: computeProfile.memoryGb,
+    gpu: computeProfile.accelerator === "gpu" ? "all" : "none",
+  };
+  const bundle = await restoreRunRecovery({ runsDirectory, runId }) ?? await prepareCompilationBundle({
+    task,
+    taskDirectory,
+    runsDirectory,
+    runId,
+    prepareMarkdown,
+  });
+  bundle.runtimeAgent = runtimeAgent;
+  bundle.computeDecision = {
+    schemaVersion: "0.1",
+    selectedProfile: structuredClone(computeProfile),
+  };
+  bundle.executionEnvironment = structuredClone(task.environment);
+  bundle.sandboxProvider = "autodl-pro";
+  await updateRunState(bundle, {
+    runtime: runtimeAgent.runtime,
+    agent: publicAgentRuntime(runtimeAgent),
+    sandbox: "autodl-pro",
+    image: task.environment.image,
+    executionEnvironment: bundle.executionEnvironment,
+  });
+
+  const instructions = await readFile(
+    path.join(bundle.directories.input, "agent-instructions.md"),
+    "utf8",
+  );
+  const autoDlOptions = executorOptions.autoDl ?? {};
+  const config = resolveAutoDlProConfig({
+    bundle,
+    options: autoDlOptions,
+    requireToken: !dryRun,
+  });
+  bundle.sourceReviewProviderSecret = config.providerSecret;
+  bundle.paperBudget = autoDlOptions.paperBudget;
+  const instanceRequest = autoDlInstanceRequest(config, bundle.runId);
+  await writeJson(
+    path.join(bundle.directories.execution, "autodl-instance-request.json"),
+    { schemaVersion: "0.1", executor: "autodl-pro", instanceRequest },
+  );
+  if (dryRun) {
+    await updateRunState(bundle, { status: "prepared" });
+    return {
+      dryRun: true,
+      bundle,
+      instanceRequest,
+      cloud: publicAutoDlExecution(config, null, []),
+    };
+  }
+
+  const temporaryDirectory = await mkdtemp(
+    path.join(os.tmpdir(), "citeark-autodl-compiler-"),
+  );
+  const records = [];
+  let validationIssues = ["research.json 尚未生成"];
+  let normalizationWarnings = [];
+  let normalizedResearch = null;
+  let blockingFailure = null;
+  let lease = null;
+  let successful = false;
+  await mkdir(path.join(bundle.directories.execution, "prompts"), { recursive: true });
+  try {
+    await updateRunState(bundle, {
+      status: "running",
+      startedAt: new Date().toISOString(),
+    });
+    await autoDlOptions.assertActive?.();
+    lease = await acquireAutoDlProLease({
+      config,
+      runId: bundle.runId,
+      clients: autoDlOptions.clients,
+      progress,
+    });
+    await autoDlOptions.assertActive?.();
+    const maxAttempts = 1 + runtimeAgent.validationRetries;
+    for (let attempt = 1; attempt <= maxAttempts; attempt += 1) {
+      const invocation = await buildAgentInvocation({
+        bundle,
+        attempt,
+        prompt: compilerPrompt(attempt, validationIssues, {
+          continuation: bundle.runtimeTask?.continuation,
+          repairBaseline: bundle.repairBaselineResearch,
+        }),
+        instructions,
+        providerSession: remoteProviderSession(runtimeAgent),
+      });
+      const resumeInvocation = await buildAgentInvocation({
+        bundle,
+        attempt: 2,
+        prompt: [
+          "Continue from the restored compilation workspace. Inspect preserved artifacts and apply the current task below; a prior research.json may still contain the rejected draft.",
+          compilerPrompt(attempt, validationIssues, {
+            continuation: bundle.runtimeTask?.continuation,
+            repairBaseline: bundle.repairBaselineResearch,
+          }),
+        ].join("\n\n"),
+        instructions,
+        providerSession: remoteProviderSession(runtimeAgent),
+      });
+      await writeFile(
+        path.join(bundle.directories.execution, "prompts", `attempt-${attempt}.txt`),
+        `${agentInvocationPromptArgument(invocation)}\n`,
+      );
+      await writeFile(
+        path.join(bundle.directories.execution, "prompts", `attempt-${attempt}-resume.txt`),
+        `${agentInvocationPromptArgument(resumeInvocation)}\n`,
+      );
+      const stageArchive = path.join(temporaryDirectory, `compile-stage-a${attempt}.tar.gz`);
+      const resultArchive = path.join(temporaryDirectory, `compile-result-a${attempt}.tar.gz`);
+      await createAutoDlStageArchive(bundle, stageArchive);
+      const remote = await executeAutoDlRemoteAttempt({
+        bundle,
+        lease,
+        config,
+        invocation,
+        resumeInvocation,
+        attempt,
+        stageArchive,
+        resultArchive,
+        checkpointPolicy: config.checkpointing,
+        restoreCheckpointOnStart: attempt > 1,
+        reuseCompletedResultOnRestore: false,
+        completionRelativePath: "output/research.json",
+        completionHasEvidence: false,
+        secretManager: autoDlOptions.secretManager,
+        stageTransfer: autoDlOptions.stageTransfer,
+        stageClients: autoDlOptions.clients,
+        assertActive: autoDlOptions.assertActive,
+        progress,
+      });
+      records.push(remote.record);
+      const researchPath = path.join(bundle.directories.output, "research.json");
+      await recoverJsonFromFinalMessage({
+        finalMessagePath: hostRuntimePath(bundle, CODEX_FINAL_MESSAGE_PATH),
+        outputPath: researchPath,
+      });
+      const remoteFailure = await describeRemoteAgentFailure(bundle, attempt);
+      const providerFailure = classifyProviderActionRequired(remoteFailure?.message);
+      if (providerFailure) {
+        blockingFailure = providerFailure;
+        validationIssues = [providerFailure.message];
+        break;
+      }
+      // A syntactically valid compiler file can still be an unfinished draft.
+      // Transport publication errors retain their independent recovery path.
+      if (remoteFailure && remoteFailure.failureCategory !== "publication") {
+        blockingFailure = { code: remoteFailure.failureCode ?? "compiler.agent_incomplete",
+          failureCategory: remoteFailure.failureCategory, retryable: remoteFailure.retryable };
+        validationIssues = [remoteFailure.message];
+        break;
+      }
+      try {
+        if (!(await pathExists(researchPath))) {
+          throw new CiteArkError("缺少 /job/output/research.json");
+        }
+        const accepted = await acceptCompilerDraft({ bundle, draft: await readJson(researchPath), validateResearch, sourceReviewer });
+        normalizedResearch = accepted.research;
+        normalizationWarnings = accepted.normalizationWarnings;
+        validationIssues = [];
+        break;
+      } catch (error) {
+        validationIssues = [
+          ...(remoteFailure ? [remoteFailure.message] : []),
+          error instanceof Error ? error.message : String(error),
+        ];
+        if (error.failureCode === "compiler.source_review_unavailable") { blockingFailure = { code: error.failureCode, retryable: false }; break; }
+        if (attempt < maxAttempts) {
+          progress(`Research Compiler 结果校验失败，将在同一 AutoDL 实例继续：${validationIssues.join("；")}`);
+        }
+      }
+    }
+
+    await ensureFinalWorkspaceCapture(bundle);
+    const usage = await readAgentUsage(
+      path.join(bundle.directories.execution, "trace.jsonl"),
+      { runtimeHome: bundle.directories.runtimeHome, model: runtimeAgent.model },
+    );
+    const status = validationIssues.length ? "invalid_result" : "completed";
+    await updateRunState(bundle, {
+      status,
+      finishedAt: new Date().toISOString(),
+      validationIssues,
+      normalizationWarnings,
+      output: normalizedResearch ? "output/research.normalized.json" : null,
+      sandboxImageIdentity: {
+        id: config.imageUuid,
+        repositoryDigests: [],
+      },
+      cloudExecution: publicAutoDlExecution(config, lease, records),
+      usage,
+    });
+    if (validationIssues.length || !normalizedResearch) {
+      throw new CiteArkError(
+        `研究编译失败：\n- ${validationIssues.join("\n- ")}`,
+        {
+          exitCode: 2,
+          retryable: blockingFailure?.retryable,
+          failureCode: blockingFailure?.code,
+          failureCategory: blockingFailure?.failureCategory,
+          recoveryAction: blockingFailure?.recoveryAction,
+        },
+      );
+    }
+    successful = true;
+    return {
+      dryRun: false,
+      bundle,
+      research: normalizedResearch,
+      normalizationWarnings,
+      runs: records,
+    };
+  } catch (error) {
+    if (bundle.state.status === "running") {
+      await updateRunState(bundle, {
+        status: "failed",
+        finishedAt: new Date().toISOString(),
+        error: error instanceof Error ? error.message : String(error),
+        cloudExecution: publicAutoDlExecution(config, lease, records),
+      });
+    }
+    throw error;
+  } finally {
+    if (lease) {
+      const cleanup = await lease.close({ successful }).catch(error => ({ released: false, error: error.message }));
+      await recordResearchCleanup({ provider: 'autodl', resource: lease.instanceUuid, status: cleanup?.released ? 'released' : 'failed', ...(cleanup?.error ? { error: cleanup.error } : {}) });
+    }
+    await rm(temporaryDirectory, { recursive: true, force: true });
+  }
+}
+
+function publicAutoDlExecution(config, lease, records) {
+  const runSeconds = records.reduce(
+    (total, record) => total + (Number(record.durationSeconds) || 0),
+    0,
+  );
+  const hourlyCny = lease?.observedHourlyCny ?? config.hourlyCny;
+  return {
+    provider: "autodl",
+    service: "container-instance-pro",
+    instance: lease?.publicRecord() ?? null,
+    gpuSpecUuid: lease?.config.gpuSpecUuid ?? config.gpuSpecUuid,
+    imageUuid: config.imageUuid,
+    inputStaging: {
+      preferred: config.stagingBucket && config.stageSigningServiceAccount
+        ? "gcs-signed-url"
+        : "sftp",
+      fallback: "sftp",
+      signedUrlExpiresSeconds: config.stageSignedUrlExpiresSeconds,
+    },
+    credentialInjection: {
+      mode: "secret-manager-to-ephemeral-ssh-environment",
+      valueRecorded: false,
+    },
+    computeUsage: {
+      runSeconds,
+      observedHourlyCny: hourlyCny ?? null,
+      estimatedCostCny: typeof hourlyCny === "number"
+        ? Number(((runSeconds / 3600) * hourlyCny).toFixed(6))
+        : null,
+      invoiceReconciled: false,
+    },
+    runs: structuredClone(records),
+  };
+}
+
+function hostRuntimePath(bundle, containerPath) {
+  return path.join(bundle.directories.run, containerPath.replace(/^\/job\//, ""));
+}
+
+function assertCompilerProfile(profile) {
+  if (!profile || profile.executor !== "autodl-pro") {
+    throw new CiteArkError("AutoDL Research Compiler 必须使用 autodl-pro 资源规格");
+  }
+  if (typeof profile.image !== "string" || !profile.image) {
+    throw new CiteArkError("AutoDL Research Compiler 资源规格必须声明私有镜像标识");
+  }
+}