@citeark/agent 0.3.20

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

@@ -0,0 +1,459 @@
+import { recordResearchCleanup } from '../workspace/control.mjs';
+import { mkdir, mkdtemp, rm, writeFile } from "node:fs/promises";
+import os from "node:os";
+import path from "node:path";
+
+import { readAgentInstructions, updateRunState } from "../job.mjs";
+import { validateResultFile } from "../protocol.mjs";
+import { agentInvocationPromptArgument, buildAgentInvocation } from "../runtime/index.mjs";
+import {
+  CODEX_FINAL_MESSAGE_PATH,
+  recoverJsonFromFinalMessage,
+} from "../runtime/final-output.mjs";
+import { CiteArkError, printableCommand, writeJson } from "../util.mjs";
+import {
+  finalizeReproduction,
+  prepareReproduction,
+  recoverInvalidExecutionResult,
+} from "../reproduction/lifecycle.mjs";
+import {
+  initialReproductionPrompt,
+  retryReproductionPrompt,
+} from "../reproduction/prompts.mjs";
+import { infrastructureResumePrompt } from "./checkpoint-protocol.mjs";
+import { remoteProviderSession } from "./gcp-batch-job.mjs";
+import {
+  describeRemoteAgentFailure,
+  readRemoteAttempt,
+  reconcileRemoteAgentOutcome,
+} from "./remote-attempt.mjs";
+import {
+  acquireAutoDlProLease,
+  autoDlInstanceRequest,
+  resolveAutoDlProConfig,
+} from "./autodl-pro-job.mjs";
+import {
+  createAutoDlStageArchive,
+  executeAutoDlRemoteAttempt,
+} from "./autodl-pro-remote.mjs";
+
+export async function executeAutoDlProReproduction({
+  taskPath,
+  referenceMaterialsPath,
+  runsDirectory,
+  imageOverride,
+  repositoryOverride,
+  paperOverride,
+  computeDecision,
+  runId,
+  agentRuntime,
+  executorOptions = {},
+  dryRun = false,
+  progress = () => {},
+}) {
+  const autoDlOptions = executorOptions.autoDl ?? {};
+  const bundle = await prepareReproduction({
+    taskPath,
+    referenceMaterialsPath,
+    runsDirectory,
+    imageOverride,
+    repositoryOverride,
+    paperOverride,
+    computeDecision,
+    runId,
+    agentRuntime,
+    requireApi: false,
+    sandboxProvider: "autodl-pro",
+  });
+  const config = resolveAutoDlProConfig({
+    bundle,
+    options: autoDlOptions,
+    requireToken: !dryRun,
+  });
+  const instructions = await readAgentInstructions(bundle.runtimeTask);
+  await mkdir(path.join(bundle.directories.execution, "prompts"), { recursive: true });
+
+  let validation = { result: null, issues: ["尚未运行 Agent"], protocolVerification: null };
+  const firstInvocation = await buildRemoteInvocation({
+    bundle,
+    attempt: 1,
+    validation,
+    instructions,
+  });
+  await writeFile(
+    path.join(bundle.directories.execution, "prompts", "attempt-1.txt"),
+    `${agentInvocationPromptArgument(firstInvocation)}\n`,
+  );
+  const preview = {
+    schemaVersion: "0.1",
+    executor: "autodl-pro",
+    instanceRequest: autoDlInstanceRequest(config, bundle.runId),
+    ssh: {
+      authentication: "instance-password-from-autodl-api",
+      hostKeyPolicy: "sha256-tofu-per-lease",
+    },
+  };
+  await writeJson(
+    path.join(bundle.directories.execution, "autodl-instance-request.json"),
+    preview,
+  );
+  if (dryRun) {
+    const cloudExecution = publicAutoDlExecution(config, null, []);
+    await updateRunState(bundle, { status: "prepared", cloudExecution });
+    return {
+      dryRun: true,
+      bundle,
+      invocation: {
+        printable: printableCommand("POST", [
+          `${config.apiBaseUrl}/api/v1/dev/instance/pro/create`,
+        ]),
+        executor: "autodl-pro",
+        instanceRequest: preview.instanceRequest,
+      },
+      cloud: cloudExecution,
+    };
+  }
+
+  let lease = null;
+  const attempts = [];
+  const records = [];
+  const temporaryDirectory = await mkdtemp(
+    path.join(os.tmpdir(), "citeark-autodl-pro-"),
+  );
+  let timedOut = false;
+  let successful = false;
+  let leaseReusable = true;
+  const leasePool = autoDlOptions.leasePool ?? null;
+  try {
+    await updateRunState(bundle, {
+      status: "running",
+      startedAt: new Date().toISOString(),
+      cloudExecution: publicAutoDlExecution(config, null, records),
+    });
+    await autoDlOptions.assertActive?.();
+    lease = leasePool
+      ? await leasePool.acquire({
+          config,
+          runId: bundle.runId,
+          clients: autoDlOptions.clients,
+          progress,
+        })
+      : await acquireAutoDlProLease({
+          config,
+          runId: bundle.runId,
+          clients: autoDlOptions.clients,
+          progress,
+        });
+    try {
+      await autoDlOptions.assertActive?.();
+    } catch (error) {
+      if (error?.processingLeaseLost === true) leaseReusable = false;
+      throw error;
+    }
+    await updateRunState(bundle, {
+      cloudExecution: publicAutoDlExecution(config, lease, records),
+    });
+    const maxAttempts = 1 + bundle.runtimeAgent.validationRetries;
+    for (let attempt = 1; attempt <= maxAttempts; attempt += 1) {
+      const invocation = attempt === 1
+        ? firstInvocation
+        : await buildRemoteInvocation({ bundle, attempt, validation, instructions });
+      const resumeInvocation = await buildAgentInvocation({
+        bundle,
+        attempt: 2,
+        prompt: infrastructureResumePrompt(bundle),
+        instructions,
+        providerSession: remoteProviderSession(bundle.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, `stage-a${attempt}.tar.gz`);
+      const resultArchive = path.join(temporaryDirectory, `result-a${attempt}.tar.gz`);
+      await createAutoDlStageArchive(bundle, stageArchive);
+      const liveRecord = {
+        attempt,
+        instanceUuid: lease.instanceUuid,
+        state: "RUNNING",
+        startedAt: new Date().toISOString(),
+        checkpoint: null,
+      };
+      records.push(liveRecord);
+      let remote;
+      try {
+        const infrastructureExecution = await executeAutoDlRemoteAttemptWithLeaseRetry({
+          lease,
+          runId: bundle.runId,
+          maximumRetries: config.remoteInfrastructureRetries,
+          progress,
+          execute: async (infrastructureAttempt) => {
+            if (infrastructureAttempt > 1) await rm(resultArchive, { force: true });
+            return await executeAutoDlRemoteAttempt({
+              bundle,
+              lease,
+              config,
+              invocation,
+              resumeInvocation,
+              attempt,
+              stageArchive,
+              resultArchive,
+              checkpointPolicy: config.checkpointing,
+              restoreCheckpointOnStart: attempt > 1 || infrastructureAttempt > 1,
+              reuseCompletedResultOnRestore: false,
+              secretManager: autoDlOptions.secretManager,
+              stageTransfer: autoDlOptions.stageTransfer,
+              stageClients: autoDlOptions.clients,
+              assertActive: autoDlOptions.assertActive,
+              progress,
+              onCheckpoint: async (checkpoint) => {
+                liveRecord.checkpoint = publicCheckpoint(checkpoint);
+                await updateRunState(bundle, {
+                  cloudExecution: publicAutoDlExecution(config, lease, records),
+                });
+              },
+            });
+          },
+        });
+        remote = infrastructureExecution.remote;
+        liveRecord.infrastructureRetryCount = infrastructureExecution.retryCount;
+      } catch (error) {
+        leaseReusable = false;
+        liveRecord.state = "FAILED";
+        liveRecord.finishedAt = new Date().toISOString();
+        liveRecord.error = error instanceof Error ? error.message : String(error);
+        if (error?.processingLeaseLost === true) throw error;
+        throw new CiteArkError(
+          `AutoDL 第 ${attempt} 次执行没有返回可恢复的运行包：${error instanceof Error ? error.message : String(error)}`,
+          { cause: error },
+        );
+      }
+      Object.assign(liveRecord, remote.record);
+      timedOut ||= remote.record.timedOut;
+      const remoteAttempt = await readRemoteAttempt(bundle, attempt);
+      attempts.push({
+        attempt,
+        startedAt: remote.record.startedAt,
+        finishedAt: remote.record.finishedAt,
+        exitCode: remoteAttempt?.exitCode ?? remote.record.exitCode,
+        signal: remote.record.signal,
+        timedOut: remote.record.timedOut,
+        executor: "autodl-pro",
+        instanceUuid: lease.instanceUuid,
+        checkpoint: remote.record.checkpoint,
+      });
+      await recoverJsonFromFinalMessage({
+        finalMessagePath: hostRuntimePath(bundle, CODEX_FINAL_MESSAGE_PATH),
+        outputPath: path.join(bundle.directories.output, "result.json"),
+      });
+      const submittedValidation = await validateResultFile(
+        path.join(bundle.directories.output, "result.json"),
+        bundle.runtimeTask,
+        bundle.directories.output,
+      );
+      const remoteFailure = await describeRemoteAgentFailure(bundle, attempt);
+      const reconciled = reconcileRemoteAgentOutcome({
+        validation: submittedValidation,
+        remoteFailure,
+      });
+      if (reconciled.blockingFailure) throw new CiteArkError(reconciled.blockingFailure.message);
+      validation = reconciled.validation;
+      await updateRunState(bundle, {
+        cloudExecution: publicAutoDlExecution(config, lease, records),
+      });
+      if (validation.issues.length === 0) break;
+      if (attempt < maxAttempts) {
+        progress(`结果协议校验失败，将在同一 AutoDL 实例继续：${validation.issues.join("；")}`);
+      }
+    }
+
+    const cloudExecution = publicAutoDlExecution(config, lease, records);
+    await writeJson(path.join(bundle.directories.execution, "autodl-runs.json"), {
+      schemaVersion: "0.1",
+      executor: "autodl-pro",
+      ...cloudExecution,
+    });
+    await writeInvocation(bundle, cloudExecution, records);
+    validation = await recoverInvalidExecutionResult({
+      bundle,
+      validation,
+      attempts,
+      timedOut,
+    });
+    const result = await finalizeReproduction({
+      bundle,
+      validation,
+      attempts,
+      timedOut,
+      sandboxImageIdentity: {
+        id: config.imageUuid,
+        repositoryDigests: [],
+      },
+      cloudExecution,
+    });
+    successful = validation.issues.length === 0;
+    return result;
+  } 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 (
+        leasePool
+          ? leasePool.release(lease, { reusable: leaseReusable, runId: bundle.runId })
+          : lease.close({ successful })
+      ).then((value) => value ?? { released: false }).catch((error) => ({
+          released: false,
+          error: error instanceof Error ? error.message : String(error),
+        }));
+      for (const record of records) record.instanceReleased = cleanup.released;
+      await recordResearchCleanup({ provider: 'autodl', resource: lease.instanceUuid, status: cleanup.released ? 'released' : leasePool && leaseReusable ? 'pooled' : 'failed', ...(cleanup.error ? { error: cleanup.error } : {}) });
+    }
+    await rm(temporaryDirectory, { recursive: true, force: true });
+  }
+}
+
+export async function executeAutoDlRemoteAttemptWithLeaseRetry({
+  lease,
+  runId,
+  maximumRetries = 2,
+  execute,
+  progress = () => {},
+}) {
+  const retries = Math.max(0, Number.parseInt(String(maximumRetries), 10) || 0);
+  for (let infrastructureAttempt = 1; ; infrastructureAttempt += 1) {
+    try {
+      return {
+        remote: await execute(infrastructureAttempt),
+        retryCount: infrastructureAttempt - 1,
+      };
+    } catch (error) {
+      if (
+        infrastructureAttempt > retries
+        || !isRetryableAutoDlRemoteInfrastructureError(error)
+        || !(await lease.isHealthy())
+      ) {
+        throw error;
+      }
+      progress(
+        `AutoDL ${runId}: 瞬时网络或传输故障，实例 ${lease.instanceUuid} 仍健康；保留已拉取镜像并在同一租约进行第 ${infrastructureAttempt + 1}/${retries + 1} 次基础设施尝试`,
+      );
+    }
+  }
+}
+
+export function isRetryableAutoDlRemoteInfrastructureError(error) {
+  const message = error instanceof Error ? error.message : String(error);
+  if (/远程运行.*已超时|执行超时|timed out after/i.test(message)) return false;
+  return /网络预检|network.?preflight|HTTP\s*5\d\d|\b50[234]\b|SSH|SFTP|GCS|连接|connect|ECONN|EPIPE|socket|上传|下载|transfer/i.test(message);
+}
+
+async function buildRemoteInvocation({ bundle, attempt, validation, instructions }) {
+  return await buildAgentInvocation({
+    bundle,
+    attempt,
+    prompt: attempt === 1
+      ? initialReproductionPrompt(bundle)
+      : retryReproductionPrompt(validation.issues, bundle.runtimeTask),
+    instructions,
+    providerSession: remoteProviderSession(bundle.runtimeAgent),
+  });
+}
+
+function hostRuntimePath(bundle, containerPath) {
+  return path.join(bundle.directories.run, containerPath.replace(/^\/job\//, ""));
+}
+
+function publicAutoDlExecution(config, lease, records) {
+  const runSeconds = records.reduce(
+    (total, record) => total + (Number(record.durationSeconds) || 0),
+    0,
+  );
+  const observedHourlyCny = 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,
+    },
+    checkpointing: {
+      ...config.checkpointing,
+      storage: "instance-system-disk",
+      durabilityBoundary: "instance-lease",
+    },
+    temporaryResources: {
+      cleanupPolicy: {
+        releaseOnSuccess: !config.keepInstance,
+        keepFailedInstance: config.keepFailedInstance,
+      },
+    },
+    credentialInjection: {
+      mode: config.credentialSource === "task-key" ? "authorized-task-key-to-ephemeral-ssh-environment" : "secret-manager-to-ephemeral-ssh-environment",
+      valueRecorded: false,
+    },
+    computeUsage: {
+      runSeconds,
+      observedHourlyCny: observedHourlyCny ?? null,
+      estimatedCostCny: typeof observedHourlyCny === "number"
+        ? Number(((runSeconds / 3600) * observedHourlyCny).toFixed(6))
+        : null,
+      estimatedCostUsd: typeof config.hourlyUsd === "number"
+        ? Number(((runSeconds / 3600) * config.hourlyUsd).toFixed(6))
+        : null,
+      invoiceReconciled: false,
+    },
+    runs: structuredClone(records),
+  };
+}
+
+async function writeInvocation(bundle, cloudExecution, records) {
+  await writeJson(path.join(bundle.directories.execution, "invocation.json"), {
+    schemaVersion: "0.1",
+    sandbox: "autodl-pro",
+    credentials: {
+      agent: bundle.runtimeAgent.runtime,
+      protocol: bundle.runtimeAgent.api.protocol,
+      source: "gcp-secret-manager",
+      mode: "ephemeral-ssh-environment",
+      injectedIntoSandbox: true,
+      valueRecorded: false,
+    },
+    networkEgress: {
+      mode: "direct",
+      allowedHosts: [new URL(bundle.runtimeAgent.api.baseUrl).hostname],
+    },
+    cloud: cloudExecution,
+    invocations: records.map((record) => ({
+      attempt: record.attempt,
+      instanceUuid: record.instanceUuid,
+      state: record.state,
+    })),
+  });
+}
+
+function publicCheckpoint(checkpoint) {
+  return {
+    schemaVersion: checkpoint.schemaVersion,
+    sequence: checkpoint.sequence,
+    kind: checkpoint.kind,
+    createdAt: checkpoint.createdAt,
+  };
+}