@tuned-tensor/cli 0.15.1 → 0.16.0

package/dist/local-runtime.js

@@ -14,6 +14,7 @@ import {
   buildModelPrefetchPayload,
   buildSystemMessage,
   canonicalLocalPipeline,
+  checkAwsGpu,
   comparisonReportSchema,
   compileSpecToJsonl,
   createLocalStore,
@@ -64,6 +65,7 @@ import {
   resolveModelLoader,
   resolveRequestedBaseModelRevision,
   resolveTrainingModel,
+  runGpuProcess,
   runLocalFineTune,
   runLocalPipeline,
   runReportSchema,
@@ -81,13 +83,13 @@ import {
   withHuggingFaceCacheEnvironment,
   withOfflineHuggingFaceCacheEnvironment,
   writePrivateJsonAtomic
-} from "./chunk-VVKC563R.js";
+} from "./chunk-6RLNBTM6.js";
 
 // src/local-runtime/index.ts
-import { readFileSync as readFileSync2, realpathSync } from "fs";
+import { readFileSync as readFileSync3, realpathSync } from "fs";
 import { createHash as createHash2 } from "crypto";
 import { readFile as readFile2, stat } from "fs/promises";
-import { basename, dirname, join as join4, resolve as resolve3 } from "path";
+import { basename, dirname, join as join5, resolve as resolve3 } from "path";
 import { cwd } from "process";
 import { fileURLToPath as fileURLToPath2 } from "url";
 
@@ -290,6 +292,22 @@ function foundationPythonProbePlan(requireBf16) {
     env: withFoundationPythonEnvironment(process.env)
   };
 }
+function localPreparationProbePlan(foundation) {
+  const source = [
+    "import json",
+    foundation ? "import numpy, safetensors, tokenizers, torch" : "import torch, transformers, peft, huggingface_hub",
+    "assert torch.ones(1, device='cpu').item() == 1",
+    "print(json.dumps({'python_ok': True, 'device': 'cpu', 'purpose': 'local preparation and CPU work'}))"
+  ].join("; ");
+  const entrypoint = foundation ? buildFoundationPythonCommand("-c", [source]) : buildBundledPythonCommand("-c", [source]);
+  const env = minimalMachineLearningEnvironment(process.env);
+  return {
+    name: "python-runtime",
+    command: entrypoint.command,
+    args: entrypoint.commandArgs,
+    env: foundation ? withFoundationPythonEnvironment(env) : withBundledPythonEnvironment(env)
+  };
+}
 function pythonProbeSource(device) {
   return [
     "import json",
@@ -430,7 +448,7 @@ async function runDoctor(config, request, foundationSpec) {
     checks.push(placeholderSpecCheck(request));
     resolveTrainingModel(request.spec_snapshot.base_model);
   }
-  const pythonPlans = foundationSpec && !config.dryRun ? [foundationPythonProbePlan(foundationSpec.foundation.bf16 !== false)] : buildDoctorPythonPlans(config);
+  const pythonPlans = config.dryRun ? [] : config.gpu ? [localPreparationProbePlan(Boolean(foundationSpec))] : foundationSpec ? [foundationPythonProbePlan(foundationSpec.foundation.bf16 !== false)] : buildDoctorPythonPlans(config);
   if (pythonPlans.length > 0) {
     const uvVersion = await runCommand("uv", ["--version"]);
     checks.push({
@@ -465,7 +483,28 @@ async function runDoctor(config, request, foundationSpec) {
     });
   }
   const device = config.evaluation.inference.device;
-  if (!config.dryRun) {
+  if (config.gpu && !config.dryRun) {
+    try {
+      await checkAwsGpu(config.gpu);
+      const plans = foundationSpec ? [foundationPythonProbePlan(foundationSpec.foundation.bf16 !== false)] : buildDoctorPythonPlans(config);
+      for (const plan of plans) {
+        const result = await runGpuProcess({
+          gpu: config.gpu,
+          runtime: foundationSpec ? "foundation" : "adapter",
+          files: [],
+          command: plan.command,
+          commandArgs: plan.args,
+          env: plan.env,
+          stage: "gpu_check",
+          timeoutMs: 18e5
+        });
+        if (result.exitCode !== 0) throw new Error(`Remote Python/CUDA check failed: ${result.stderr}`);
+      }
+      checks.push({ name: "aws-gpu", ok: true, message: `GPU execution on ${config.gpu.instanceId}; orchestration stays local.` });
+    } catch (error) {
+      checks.push({ name: "aws-gpu", ok: false, message: error instanceof Error ? error.message : String(error) });
+    }
+  } else if (!config.dryRun) {
     const queried = await runCommand("nvidia-smi", [
       "--query-gpu=name,driver_version",
       "--format=csv,noheader"
@@ -503,8 +542,9 @@ async function runDoctor(config, request, foundationSpec) {
 }
 
 // src/local-runtime/model-server.ts
+import { readFileSync as readFileSync2, statSync } from "fs";
 import { spawn as spawn2 } from "child_process";
-import { resolve as resolve2 } from "path";
+import { resolve as resolve2, join as join3 } from "path";
 import { fileURLToPath } from "url";
 
 // src/local-runtime/serving-runtime.ts
@@ -562,7 +602,16 @@ function buildModelServerLaunch(args) {
   const options = args.options ?? {};
   const device = options.device ?? args.config.evaluation.inference.device;
   if (device !== "cuda") throw new Error("The pinned vLLM serving runtime requires CUDA; CPU serving is not supported.");
-  const contextLength = boundedNumber("contextLength", options.contextLength ?? 16384, 512, 131072, true);
+  const contextLength = boundedNumber(
+    "contextLength",
+    options.contextLength ?? args.foundation?.contextLength ?? 16384,
+    args.foundation ? 1 : 512,
+    131072,
+    true
+  );
+  if (args.foundation && contextLength > args.foundation.contextLength) {
+    throw new Error("Context length exceeds the foundation checkpoint sequence_length.");
+  }
   const gpuMemoryUtilization = boundedNumber("gpuMemoryUtilization", options.gpuMemoryUtilization ?? 0.8, 0.05, 0.95);
   const host = options.host ?? "127.0.0.1";
   const remote = !isLoopbackHost(host);
@@ -575,7 +624,7 @@ function buildModelServerLaunch(args) {
     throw new Error("Non-loopback serving requires --api-key-env with a populated environment variable.");
   }
   const port = boundedNumber("port", options.port ?? 8e3, 1, 65535, true);
-  const maxTokens = boundedNumber("maxTokens", options.maxTokens ?? 512, 1, 8192, true);
+  const maxTokens = boundedNumber("maxTokens", options.maxTokens ?? Math.min(512, contextLength), 1, 8192, true);
   const temperature = boundedNumber("temperature", options.temperature ?? 0, 0, 5);
   const topP = boundedNumber("topP", options.topP ?? 1, 0, 1);
   const maxConcurrentRequests = boundedNumber(
@@ -587,14 +636,14 @@ function buildModelServerLaunch(args) {
   );
   const artifactPath = args.artifactUri ? localArtifactPath(args.artifactUri) : void 0;
   const recordedBaseModelPath = options.baseModelArtifactUri ? localArtifactPath(options.baseModelArtifactUri) : void 0;
-  const configuredBaseModelPath = args.config.paths.baseModel ? resolve2(args.config.paths.baseModel) : void 0;
+  const configuredBaseModelPath = !args.foundation && args.config.paths.baseModel ? resolve2(args.config.paths.baseModel) : void 0;
   if (recordedBaseModelPath && configuredBaseModelPath && resolve2(recordedBaseModelPath) !== configuredBaseModelPath) {
     throw new Error(
       `Configured paths.baseModel ${configuredBaseModelPath} does not match the base model recorded for this artifact: ` + recordedBaseModelPath
     );
   }
   const localBaseModelPath = recordedBaseModelPath ?? configuredBaseModelPath;
-  const requestedBaseModelRevision = resolveRequestedBaseModelRevision(
+  const requestedBaseModelRevision = args.foundation ? void 0 : resolveRequestedBaseModelRevision(
     args.baseModel,
     options.baseModelRevision
   );
@@ -610,7 +659,11 @@ function buildModelServerLaunch(args) {
       TT_MODEL_NAME: args.modelName,
       TT_CONTEXT_LENGTH: String(contextLength),
       TT_GPU_MEMORY_UTILIZATION: String(gpuMemoryUtilization),
-      TT_MODEL_LOADER: resolveModelLoader(args.baseModel),
+      TT_MODEL_LOADER: args.foundation ? "foundation" : resolveModelLoader(args.baseModel),
+      ...args.foundation ? {
+        TT_FOUNDATION_CHECKPOINT: args.foundation.checkpoint,
+        TT_FOUNDATION_TOKENIZER: args.foundation.tokenizer
+      } : {},
       TT_HOST: host,
       TT_PORT: String(port),
       TT_DEVICE: device,
@@ -710,16 +763,32 @@ async function serveLocalModel(launch) {
     });
   });
 }
+function buildFoundationModelServerLaunch(args) {
+  const checkpoint = resolve2(args.checkpoint);
+  const tokenizer = resolve2(args.tokenizer);
+  for (const path of [join3(checkpoint, "model.safetensors"), tokenizer]) {
+    if (!statSync(path).isFile()) throw new Error(`Expected foundation artifact file: ${path}`);
+  }
+  const metadata = JSON.parse(readFileSync2(join3(checkpoint, "config.json"), "utf8"));
+  const contextLength = boundedNumber("sequence_length", metadata.sequence_length, 1, 131072, true);
+  return buildModelServerLaunch({
+    modelName: "foundation",
+    baseModel: "foundation",
+    config: args.config,
+    options: args.options,
+    foundation: { checkpoint, tokenizer, contextLength }
+  });
+}
 
 // src/local-runtime/active-model.ts
 import { randomUUID } from "crypto";
 import { appendFile, chmod, readFile, rename, writeFile as writeFile2 } from "fs/promises";
-import { join as join3 } from "path";
+import { join as join4 } from "path";
 function pointerPath(store) {
-  return join3(store.root, "active-model.json");
+  return join4(store.root, "active-model.json");
 }
 function historyPath(store) {
-  return join3(store.root, "activation-history.jsonl");
+  return join4(store.root, "activation-history.jsonl");
 }
 async function readPointer(store) {
   if (!await hardenExistingLocalStore(store.root)) return null;
@@ -759,7 +828,7 @@ async function assertModelCanActivate(store, modelId) {
     throw new Error(`Model ${model.id} belongs to run ${run.id}, which is not completed.`);
   }
   const report = runReportSchema.parse(
-    JSON.parse(await readFile(join3(model.artifact_dir, "run-report.json"), "utf8"))
+    JSON.parse(await readFile(join4(model.artifact_dir, "run-report.json"), "utf8"))
   );
   if (!report.general_regression) {
     throw new Error(`Model ${model.id} has no general regression evaluation.`);
@@ -816,7 +885,7 @@ async function rollbackActiveModel(store) {
 function packageVersion() {
   try {
     const packageJson = JSON.parse(
-      readFileSync2(resolve3(localRuntimePackageRoot(import.meta.url), "package.json"), "utf8")
+      readFileSync3(resolve3(localRuntimePackageRoot(import.meta.url), "package.json"), "utf8")
     );
     return typeof packageJson.version === "string" ? packageJson.version : "unknown";
   } catch {
@@ -851,7 +920,7 @@ Commands:
   hardware [--config local-runner.json] [--quick]
   validate [tunedtensor.json] [--config local-runner.json]
   run [tunedtensor.json] [--config local-runner.json] [--dry-run] [--verbose] [--quiet]
-  serve <model-id|active|base> [--config local-runner.json] [--host 127.0.0.1] [--port 8000]
+  serve <model-id|active|base|foundation> [--config local-runner.json] [--host 127.0.0.1] [--port 8000]
   runs list|get|events|report|compare [args] [--config local-runner.json]
   models list|get|verify|prefetch|verify-base|active|activate|rollback|serve [args] [--config local-runner.json]
 
@@ -869,6 +938,8 @@ var VERBOSE_OPTION = { name: "--verbose", description: "Stream subprocess output
 var QUIET_OPTION = { name: "--quiet", description: "Suppress progress output on stderr" };
 var MODEL_SERVE_OPTIONS = [
   CONFIG_OPTION,
+  { name: "--checkpoint", value: "path", description: "Foundation model directory (target: foundation)" },
+  { name: "--tokenizer", value: "path", description: "Foundation tokenizer.json (target: foundation)" },
   { name: "--host", value: "host", description: "Bind host (localhost by default)" },
   { name: "--port", value: "port", description: "Bind port" },
   { name: "--device", value: "device", description: "cuda (pinned vLLM runtime)" },
@@ -876,7 +947,7 @@ var MODEL_SERVE_OPTIONS = [
   { name: "--temperature", value: "number", description: "Default sampling temperature" },
   { name: "--top-p", value: "number", description: "Default nucleus sampling threshold" },
   { name: "--max-concurrent-requests", value: "count", description: "Maximum sequences in an upstream generation batch" },
-  { name: "--context-length", value: "count", description: "Total context token budget (default 16384)" },
+  { name: "--context-length", value: "count", description: "Total context token budget (foundation: checkpoint limit; otherwise 16384)" },
   { name: "--gpu-memory-utilization", value: "fraction", description: "vLLM GPU memory budget fraction (default 0.8)" },
   { name: "--spec", value: "path", description: "Behavior spec whose instructions are enforced" },
   { name: "--no-spec-prompt", description: "Do not enforce the stored behavior-spec prompt" },
@@ -939,8 +1010,8 @@ var COMMAND_DEFINITIONS = {
     maxPositionals: 1
   },
   serve: {
-    usage: "tt serve <model-id|active|base> [options]",
-    description: "Serve a verified adapter, the active model, or the protected base.",
+    usage: "tt serve <model-id|active|base|foundation> [options]",
+    description: "Serve an adapter, pretrained base, or foundation checkpoint with vLLM.",
     options: MODEL_SERVE_OPTIONS,
     minPositionals: 1,
     maxPositionals: 1,
@@ -1006,7 +1077,7 @@ var COMMAND_GROUPS = {
         maxPositionals: 0
       },
       serve: {
-        usage: "tt models serve <model-id|active|base> [options]",
+        usage: "tt models serve <model-id|active|base|foundation> [options]",
         description: "Alias for `tt serve`.",
         options: MODEL_SERVE_OPTIONS,
         minPositionals: 1,
@@ -1146,7 +1217,7 @@ function readNumberOption(argv, name) {
 async function selectedConfigPath(argv, adjacentTo) {
   const explicitPath = readOption(argv, "--config");
   if (explicitPath) return resolve3(explicitPath);
-  const candidate = join4(adjacentTo ? dirname(resolve3(adjacentTo)) : cwd(), "local-runner.json");
+  const candidate = join5(adjacentTo ? dirname(resolve3(adjacentTo)) : cwd(), "local-runner.json");
   const metadata = await stat(candidate).catch(() => null);
   return metadata?.isFile() ? candidate : void 0;
 }
@@ -1211,7 +1282,7 @@ function createConsoleReporter(options) {
   };
 }
 async function verifyStoredModel(model, config) {
-  const manifestPath = join4(model.artifact_dir, "artifact-manifest.json");
+  const manifestPath = join5(model.artifact_dir, "artifact-manifest.json");
   const integrity = await assertArtifactManifest(manifestPath, {
     requiredPaths: ["stage-metadata.json", "training-report.json"],
     scopeToRequired: true,
@@ -1246,8 +1317,8 @@ async function verifyStoredModel(model, config) {
   };
 }
 async function verifyActivationEvidence(model) {
-  const generalBaseline = join4(model.artifact_dir, "general-baseline-eval.json");
-  const generalCandidate = join4(model.artifact_dir, "general-candidate-eval.json");
+  const generalBaseline = join5(model.artifact_dir, "general-baseline-eval.json");
+  const generalCandidate = join5(model.artifact_dir, "general-candidate-eval.json");
   const [hasBaseline, hasCandidate] = await Promise.all([
     stat(generalBaseline).then((metadata) => metadata.isFile(), () => false),
     stat(generalCandidate).then((metadata) => metadata.isFile(), () => false)
@@ -1257,7 +1328,7 @@ async function verifyActivationEvidence(model) {
       "This model cannot be activated because the run did not include a general-regression suite. Add evaluation.generalRegression to local-runner.json with a held-out dataset, re-run training, then activate. Until then, serve the adapter with `tt serve local-<run-id>`."
     );
   }
-  await assertArtifactManifest(join4(model.artifact_dir, "artifact-manifest.json"), {
+  await assertArtifactManifest(join5(model.artifact_dir, "artifact-manifest.json"), {
     requiredPaths: [
       "run-report.json",
       "general-baseline-eval.json",
@@ -1283,7 +1354,7 @@ async function verifyModelArtifactPath(input, config) {
     if (!metadata) throw new Error(`Model not found and artifact path does not exist: ${input}`);
     let current = metadata.isDirectory() ? inputPath : dirname(inputPath);
     for (let depth = 0; depth < 8; depth += 1) {
-      const candidate = join4(current, "artifact-manifest.json");
+      const candidate = join5(current, "artifact-manifest.json");
       const raw = await readFile2(candidate, "utf8").catch(() => null);
       if (raw) {
         const parsed = JSON.parse(raw);
@@ -1337,11 +1408,11 @@ async function modelSystemPrompt(args) {
       throw new Error(`--spec must contain a tunedtensor.json behavior spec: ${resolve3(specPath)}`);
     }
     if (isFoundationSpecFile(local.data)) {
-      throw new Error("Foundation specs have no Hugging Face base model. `tt serve` cannot host them yet.");
+      throw new Error("Foundation specs have no Hugging Face base model. Use `tt serve foundation --checkpoint <model-dir> --tokenizer <tokenizer.json>`.");
     }
     spec = local.data;
   } else {
-    const runRequestPath = join4(args.store.paths.runsDir, args.model.run_id, "request.json");
+    const runRequestPath = join5(args.store.paths.runsDir, args.model.run_id, "request.json");
     let persistedRequest = null;
     try {
       persistedRequest = fineTuneRunRequestSchema.parse(JSON.parse(await readFile2(runRequestPath, "utf8")));
@@ -1357,7 +1428,7 @@ async function modelSystemPrompt(args) {
   }
   const prompt = buildSystemMessage(spec);
   const metadata = JSON.parse(
-    await readFile2(join4(args.model.artifact_dir, "stage-metadata.json"), "utf8")
+    await readFile2(join5(args.model.artifact_dir, "stage-metadata.json"), "utf8")
   );
   const promptHash = createHash2("sha256").update(prompt).digest("hex");
   if (metadata.system_prompt_sha256 !== promptHash) {
@@ -1400,7 +1471,7 @@ function printServeClientConfig(argv, launch) {
           supportsUsageInStreaming: true,
           maxTokensField: "max_tokens"
         },
-        models: [.../* @__PURE__ */ new Set([launch.modelName, `base:${launch.env.TT_BASE_MODEL}`])].map((id) => ({
+        models: [...new Set(launch.env.TT_FOUNDATION_CHECKPOINT ? [launch.modelName] : [launch.modelName, `base:${launch.env.TT_BASE_MODEL}`])].map((id) => ({
           id,
           name: id,
           // This endpoint exports ordinary text/tool calls, not reasoning-effort controls.
@@ -1488,7 +1559,7 @@ async function serveBaseModelFromCli(args) {
     JSON.parse(await readFile2(specPath, "utf8"))
   ) : void 0;
   if (spec && isFoundationSpecFile(spec)) {
-    throw new Error("Foundation specs have no Hugging Face base model. `tt serve` cannot host them yet.");
+    throw new Error("Foundation specs have no Hugging Face base model. Use `tt serve foundation --checkpoint <model-dir> --tokenizer <tokenizer.json>`.");
   }
   let baseModel = spec?.base_model;
   if (!baseModel) {
@@ -1525,6 +1596,45 @@ async function serveBaseModelFromCli(args) {
   await serveLocalModel(launch);
 }
 async function serveModelTargetFromCli(args) {
+  if (args.target === "foundation") {
+    const checkpoint = readOption(args.argv, "--checkpoint");
+    const tokenizer = readOption(args.argv, "--tokenizer");
+    if (!checkpoint || !tokenizer) {
+      throw new Error("Foundation serving requires --checkpoint <model-dir> and --tokenizer <tokenizer.json>.");
+    }
+    const specPath = readOption(args.argv, "--spec");
+    if (specPath && hasFlag(args.argv, "--no-spec-prompt")) {
+      throw new Error("Use only one of --spec or --no-spec-prompt.");
+    }
+    const spec = specPath ? localBehaviorSpecFileSchema.parse(JSON.parse(await readFile2(resolve3(specPath), "utf8"))) : void 0;
+    if (spec && !isFoundationSpecFile(spec)) throw new Error("Foundation serving requires a foundation spec.");
+    const launch = buildFoundationModelServerLaunch({
+      checkpoint,
+      tokenizer,
+      config: args.config,
+      options: {
+        ...serveOptionsFromArgv(args.argv),
+        systemPrompt: spec ? buildSystemMessage(spec) : void 0
+      }
+    });
+    if (printServeClientConfig(args.argv, launch)) return;
+    if (hasFlag(args.argv, "--print-command")) {
+      printJson({
+        ok: true,
+        model_id: launch.modelName,
+        url: launch.url,
+        command: launch.displayCommand,
+        checkpoint: launch.env.TT_FOUNDATION_CHECKPOINT,
+        tokenizer: launch.env.TT_FOUNDATION_TOKENIZER
+      });
+      return;
+    }
+    await serveLocalModel(launch);
+    return;
+  }
+  if (readOption(args.argv, "--checkpoint") || readOption(args.argv, "--tokenizer")) {
+    throw new Error("--checkpoint and --tokenizer require the foundation target.");
+  }
   if (args.target === "base") {
     return serveBaseModelFromCli(args);
   }
@@ -1903,6 +2013,7 @@ export {
   baseModelRevisionSchema,
   behaviorSpecExampleSchema,
   behaviorSpecSchema,
+  buildFoundationModelServerLaunch,
   buildLocalBaseModelServerLaunch,
   buildLocalModelServerLaunch,
   buildModelPrefetchPayload,