badgr-cli 1.0.43 → 1.0.45

package/src/commands/serve.js

@@ -1,9 +1,10 @@
 import { requireApiKey } from '../config.js';
 import { callApi, listDeployments } from '../api.js';
 import { addDeployment, addReceipt, updateReceipt, generateReceiptId } from '../store.js';
-import { normalizeTier, callWithFallback, HIGH_RATE_THRESHOLD } from '../fallback.js';
+import { normalizeTier, callWithFallback } from '../fallback.js';
 import { formatCliError } from '../errors.js';
-import { TEMPLATE_MAP, buildTemplateFlags, parseTemplateOverrides, BLESSED_VLLM_MODELS } from '../catalog.js';
+import { TEMPLATE_MAP, buildTemplateFlags, parseTemplateOverrides, BLESSED_VLLM_MODELS, isLikelyGatedModel } from '../catalog.js';
+import { stage as _stage, stageDone as _stageDone, writeBlock as _writeBlock, clearBlock as _clearBlock, renderLiveBlock as _renderLiveBlock, printFailureClass as _printFailureClass } from '../progress.js';
 
 const LLAMA_CPP_IMAGE = 'michaelmanleyx/llama-cpp:server-cuda';
 
@@ -64,18 +65,36 @@ function parseEnvFlag(envList) {
   return obj;
 }
 
+function envObjHasHfToken(envList) {
+  return (envList || []).some(kv => kv.startsWith('HF_TOKEN='));
+}
+
+// Extract a parameter-count-in-billions hint from a model/file name, e.g.
+// "Qwen2.5-0.5B-Instruct" → 0.5, "Llama-3.1-8B-Instruct" → 8, "Mixtral-8x7B" → 56.
+// Splits into delimiter-bounded segments first so a version number like "2.5"
+// in "Qwen2.5-0.5B" is never mistaken for the param count — only a segment that
+// IS entirely "<digits>b" or "<digits>x<digits>b" counts as a size hint.
+function _extractParamsB(name) {
+  const s = name.toLowerCase();
+  const segments = s.split(/[^a-z0-9.]+/).filter(Boolean);
+  for (const seg of segments) {
+    const moe = seg.match(/^(\d+)x(\d+)b$/);
+    if (moe) return parseInt(moe[1], 10) * parseInt(moe[2], 10);
+  }
+  for (const seg of segments) {
+    const m = seg.match(/^(\d+(?:\.\d+)?)b$/);
+    if (m) return parseFloat(m[1]);
+  }
+  return null;
+}
+
 // Mirror of backend workload_profile.py infer_profile_from_model — for pre-flight display.
+// Sizing is only asserted when a param-count hint is found in the name; unknown
+// sizing falls back to the 7B–8B/24GB+ default rather than guessing small or large.
 function _inferServeProfile(modelName) {
-  const s = modelName.toLowerCase();
-  const moe = s.match(/(\d+)x(\d+)b/);
-  let paramsB;
-  if (moe) {
-    paramsB = parseInt(moe[1]) * parseInt(moe[2]);
-  } else {
-    const m = s.match(/(\d+)b/);
-    paramsB = m ? parseInt(m[1]) : null;
-  }
+  const paramsB = _extractParamsB(modelName);
   if (paramsB === null) return { label: 'inference (7B–8B model)', vram: '24+ GB', gpus: ['RTX 4090', 'A6000', 'L40S'] };
+  if (paramsB <= 3)    return { label: 'inference (≤3B model)',    vram: '8+ GB',  gpus: ['RTX 3090', 'RTX 4090', 'L4'] };
   if (paramsB <= 9)    return { label: 'inference (7B–8B model)', vram: '24+ GB', gpus: ['RTX 4090', 'A6000', 'L40S'] };
   if (paramsB <= 35)   return { label: 'inference (30B–34B model)', vram: '40+ GB', gpus: ['A6000', 'L40S', 'A100'] };
   return                      { label: 'inference (70B+ model)',    vram: '80+ GB', gpus: ['H100', 'A100'] };
@@ -84,32 +103,16 @@ function _inferServeProfile(modelName) {
 // Mirror of backend workload_profile.py infer_profile_from_gguf.
 // Accepts the --hf-file filename; looks for param-count hints like "35B" or "8x7B".
 function _inferGgufProfile(ggufPath) {
-  const s = ggufPath.toLowerCase();
-  const moe = s.match(/(\d+)x(\d+)b/);
-  let paramsB;
-  if (moe) {
-    paramsB = parseInt(moe[1]) * parseInt(moe[2]);
-  } else {
-    const m = s.match(/(\d+)b/);
-    paramsB = m ? parseInt(m[1]) : null;
-  }
+  const paramsB = _extractParamsB(ggufPath);
   if (paramsB === null || paramsB <= 9) return { label: 'GGUF inference (≤9B, llama.cpp)',      vram: '8+ GB',  gpus: ['RTX 4090', 'RTX 3090', 'A6000'] };
   if (paramsB <= 35)                   return { label: 'GGUF inference (10B–35B, llama.cpp)',   vram: '24+ GB', gpus: ['RTX 4090', 'A6000', 'L40S'] };
   return                                      { label: 'GGUF inference (36B+, llama.cpp)',      vram: '48+ GB', gpus: ['A6000', 'L40S', 'A100'] };
 }
 
-// Label derived from the backend's real readiness signal (readiness_reason),
-// not a fixed timer — the timer only breaks ties when we have no signal yet.
-function _serveStageLabel(elapsedSec, healthPath, readinessReason) {
-  switch (readinessReason) {
-    case 'port_unreachable': return 'Waiting for container to start…';
-    case 'http_404':         return 'Model loading — waiting for the API to come up…';
-    case 'http_error':       return `Waiting for ${healthPath || 'endpoint'} (non-200 response)…`;
-    default: break;
-  }
-  if (elapsedSec < 30) return 'Starting container…';
-  return `Waiting for ${healthPath || 'endpoint'}…`;
-}
+// If the readiness reason hasn't changed for this long, the status word
+// flips from "starting" to "stuck" so a silent stall never looks identical
+// to normal progress.
+const SERVE_STUCK_THRESHOLD_MS = 90_000;
 
 function _detectHealthPath(image) {
   if (!image) return null;
@@ -129,10 +132,15 @@ function _detectHealthPath(image) {
 // vLLM cold start (model download + load) often exceeds 5 min on first boot.
 const VLLM_SERVE_WAIT_MS = 15 * 60 * 1000;
 
-async function waitForEndpoint(deploymentId, config, timeoutMs = VLLM_SERVE_WAIT_MS, chalk, healthPath = '/models') {
+async function waitForEndpoint(deploymentId, config, timeoutMs = VLLM_SERVE_WAIT_MS, chalk, healthPath = '/models', costCtx = {}) {
+  const { costPerHour = 0, maxCost = null, stageLine = '' } = costCtx;
   const startMs  = Date.now();
   const deadline = startMs + timeoutMs;
 
+  let lastReason     = null;
+  let reasonSinceMs  = startMs;
+  let blockLines     = 0;
+
   while (Date.now() < deadline) {
     try {
       const dep = await callApi(`/deployments/${deploymentId}`, {
@@ -141,24 +149,35 @@ async function waitForEndpoint(deploymentId, config, timeoutMs = VLLM_SERVE_WAIT
         timeoutMs: 10_000,
       });
       if (['failed', 'terminated', 'error', 'stopped'].includes(dep.status)) {
-        process.stdout.write('\n');
-        return { ready: false, timedOut: false, depFailed: true, failReason: dep.fix_hint || dep.error || dep.status };
+        if (blockLines > 0) { _clearBlock(blockLines); blockLines = 0; }
+        return {
+          ready: false, timedOut: false, depFailed: true,
+          failReason: dep.fix_hint || dep.error || dep.status,
+          failureClass: dep.failure_class ?? null,
+          nextAction: dep.next_action ?? null,
+        };
       }
       if (dep.endpoint_ready) {
-        process.stdout.write('\n');
+        if (blockLines > 0) { _clearBlock(blockLines); blockLines = 0; }
         return { ready: true, timedOut: false, depFailed: false };
       }
-      const elapsed = Math.round((Date.now() - startMs) / 1000);
-      process.stdout.write(
-        `\r  ${chalk.dim(_serveStageLabel(elapsed, dep.health_path || healthPath, dep.readiness_reason) + ` (${elapsed}s)`)}   `
-      );
+
+      const now     = Date.now();
+      const elapsed = Math.round((now - startMs) / 1000);
+      const reason  = dep.readiness_reason || 'starting';
+      if (reason !== lastReason) { lastReason = reason; reasonSinceMs = now; }
+      const statusWord = (now - reasonSinceMs) >= SERVE_STUCK_THRESHOLD_MS ? 'stuck' : 'still starting';
+      const spend   = costPerHour * (elapsed / 3600);
+
+      blockLines = _writeBlock(blockLines, _renderLiveBlock(chalk, {
+        stageLine, elapsedSec: elapsed, statusWord, spend, id: deploymentId,
+      }));
     } catch {
       // status check failed (transient network/API issue) — retry next tick
     }
     await new Promise(r => setTimeout(r, 8000));
   }
 
-  process.stdout.write('\n');
   return { ready: false, timedOut: true, depFailed: false };
 }
 
@@ -194,6 +213,101 @@ const _KNOWN_SERVE_FLAGS = new Set([
   '--persistent', '--yes', '-y', '--runtime', '--hf-repo', '--hf-file',
 ]);
 
+function extractFlag(args, flagName) {
+  const rest = [];
+  let value = null;
+  for (let i = 0; i < args.length; i++) {
+    if (args[i] === flagName) { value = args[++i]; continue; }
+    rest.push(args[i]);
+  }
+  return { value, rest };
+}
+
+// Returns { id, endpointUrl } for a running vLLM endpoint matching modelId, or
+// null. The id is needed so a caller that auto-launches a new one can print a
+// precise `badgr down <id>` hint if something later fails.
+async function findRunningVllmEndpoint(config, modelId) {
+  try {
+    const existing = await listDeployments(config);
+    const ACTIVE = new Set(['running', 'provisioning', 'starting', 'queued']);
+    const match = (existing.deployments ?? []).find(d =>
+      ACTIVE.has(d.status) && d.workload_type === 'endpoint' && d.model === modelId
+    );
+    if (!match) return null;
+    const endpointUrl = match.endpoint_url || match.openai_base_url || null;
+    return endpointUrl ? { id: match.deployment_id, endpointUrl } : null;
+  } catch {
+    return null;
+  }
+}
+
+/**
+ * `badgr serve openwebui [--model <alias-or-id>] [--connect <url>] [flags]`
+ *
+ * Open WebUI serves a chat UI that connects to a model endpoint — usually
+ * vLLM. It is a separate served app from the model itself: this reuses a
+ * running vLLM endpoint for the model if one exists, or launches one via
+ * `serveCommand` first, then launches Open WebUI wired to it via
+ * OPENAI_API_BASE_URL/OPENAI_API_KEY. `--connect` skips vLLM discovery/launch
+ * entirely and points Open WebUI at any existing OpenAI-compatible endpoint —
+ * no model is ever touched in that mode.
+ */
+async function serveOpenWebUICommand(config, args, chalk) {
+  requireApiKey(config);
+  const { value: connect, rest: afterConnect } = extractFlag(args, '--connect');
+  const { value: model, rest } = extractFlag(afterConnect, '--model');
+
+  let endpointUrl = connect || null;
+  let launchedVllmDepId = null; // set only if we launched a new vLLM (for cleanup hints)
+
+  if (!endpointUrl) {
+    const alias = model || 'qwen-7b';
+    const vllmAlias = BLESSED_VLLM_MODELS[alias];
+    const modelId = vllmAlias ? vllmAlias.model_id : alias;
+
+    const existing = await findRunningVllmEndpoint(config, modelId);
+    if (existing) {
+      endpointUrl = existing.endpointUrl;
+      console.log(chalk.dim(`  Reusing running vLLM endpoint for ${modelId}: ${endpointUrl}\n`));
+    } else {
+      // Auto-launching vLLM reuses rest's --max-cost/--persistent, which also
+      // apply to Open WebUI below — that's two separately billed deployments,
+      // not a shared budget, so say so instead of a silent double spend.
+      const { value: webuiMaxCost } = extractFlag(rest, '--max-cost');
+      console.log(chalk.bold(`\n  Open WebUI needs a model endpoint behind it — no running vLLM endpoint for ${modelId} found.\n`));
+      if (webuiMaxCost) {
+        console.log(chalk.yellow(`  Launching vLLM first with its own --max-cost $${webuiMaxCost} cap (separate from Open WebUI's).`));
+        console.log(chalk.yellow(`  Total possible spend across both deployments: ~$${(parseFloat(webuiMaxCost) * 2).toFixed(2)}.\n`));
+      }
+      await serveCommand(config, [alias, ...rest], chalk);
+      const launched = await findRunningVllmEndpoint(config, modelId);
+      if (!launched) {
+        console.error(chalk.red(`\n  ✗ vLLM endpoint for ${modelId} did not come up — see output above.\n`));
+        console.error(chalk.dim(`  Once it's ready, run: badgr serve openwebui --connect <endpoint-url>\n`));
+        process.exitCode = 1;
+        return;
+      }
+      endpointUrl = launched.endpointUrl;
+      launchedVllmDepId = launched.id;
+    }
+  }
+
+  const overrides = parseTemplateOverrides(rest);
+  // User-supplied --env intentionally wins over the auto-wired connection env.
+  overrides.env = {
+    OPENAI_API_BASE_URL: endpointUrl,
+    OPENAI_API_KEY:      config.apiKey || 'sk-local',
+    ...overrides.env,
+  };
+  const flags = buildTemplateFlags(TEMPLATE_MAP.openwebui, overrides);
+  await serveCommand(config, flags, chalk);
+
+  if (process.exitCode && launchedVllmDepId) {
+    console.error(chalk.yellow(`\n  Open WebUI failed to start. The vLLM endpoint it was connecting to is still running and billing separately.`));
+    console.error(chalk.dim(`  Stop it with: badgr down ${launchedVllmDepId}\n`));
+  }
+}
+
 export async function serveCommand(config, args, chalk) {
   // `badgr serve template <name> [flags]` — expand template defaults then re-dispatch
   if (args[0] === 'template') {
@@ -216,14 +330,21 @@ export async function serveCommand(config, args, chalk) {
     return serveCommand(config, expandedArgs, chalk);
   }
 
+  // `badgr serve openwebui|open-webui [flags]` — a chat UI that connects to a
+  // model endpoint (usually vLLM). See serveOpenWebUICommand.
+  if (args[0] && ['openwebui', 'open-webui'].includes(args[0].toLowerCase())) {
+    return serveOpenWebUICommand(config, args.slice(1), chalk);
+  }
+
   if (args.includes('--list-aliases')) {
-    console.log(chalk.bold('\n⚡ Blessed vLLM aliases\n'));
-    console.log(chalk.dim('  Use any of these in place of a full model ID: badgr serve <alias>\n'));
-    for (const [alias, spec] of Object.entries(BLESSED_VLLM_MODELS)) {
-      console.log(`  ${chalk.cyan(alias.padEnd(16))}  ${spec.model_id}`);
-      console.log(`  ${''.padEnd(16)}  ${chalk.dim(spec.description)}`);
+    console.log(chalk.bold('\nTested model routes:\n'));
+    for (const alias of Object.keys(BLESSED_VLLM_MODELS)) {
+      console.log(`  - ${chalk.cyan(alias)}`);
     }
     console.log();
+    console.log(chalk.dim('You can also try a Hugging Face model ID:'));
+    console.log(chalk.dim('  badgr serve Qwen/Qwen2.5-7B-Instruct --max-cost 10'));
+    console.log();
     return;
   }
 
@@ -314,50 +435,60 @@ export async function serveCommand(config, args, chalk) {
 
   const effectiveTier = normalizeTier(flags.tier);
 
+  // Header fields + trailing note (shown after Max cost) per serve mode.
+  // Built as [label, value] pairs so every mode renders through one aligned
+  // printer instead of four hand-spaced copies.
+  let title;
+  const headerLines = [];
+  let trailingNote = null;
+  let sizeProfile = null; // { label, vram } — only shown when GPU sizing was inferred, not chosen
+
   if (isLlamaCpp) {
-    console.log(chalk.bold('\n⚡ Serving HF GGUF (llama.cpp)\n'));
-    console.log(`  ${chalk.bold('HF Repo:')}  ${flags.hfRepo}`);
-    console.log(`  ${chalk.bold('HF File:')}  ${flags.hfFile}`);
-    console.log(`  ${chalk.bold('Runtime:')}  llama.cpp`);
-    console.log(`  ${chalk.bold('GPU:')}      ${gpuLabel}`);
-    if (flags.env?.length) console.log(`  ${chalk.bold('Env:')}      ${flags.env.join(', ')}`);
-    if (gpu === 'AUTO') {
-      const prof = _inferGgufProfile(flags.hfFile);
-      console.log();
-      console.log(`  ${chalk.bold('Estimated workload:')}      ${prof.label}`);
-      console.log(`  ${chalk.bold('Estimated minimum VRAM:')}  ${prof.vram}`);
-    }
+    title = flags.hfRepo;
+    headerLines.push(['Route', 'best-effort Hugging Face model']);
+    headerLines.push(['Mode', 'OpenAI-compatible endpoint (llama.cpp)']);
+    headerLines.push(['File', flags.hfFile]);
+    trailingNote = 'Badgr will try a compatible route.';
+    if (gpu === 'AUTO') sizeProfile = _inferGgufProfile(flags.hfFile);
   } else if (customImage) {
-    console.log(chalk.bold('\n⚡ Serving custom container\n'));
-    console.log(`  ${chalk.bold('Image:')}  ${customImage}`);
-    console.log(`  ${chalk.bold('GPU:')}    ${gpuLabel}`);
-    if (flags.task) console.log(`  ${chalk.bold('Task:')}   ${flags.task}`);
-    if (flags.env?.length) console.log(`  ${chalk.bold('Env:')}    ${flags.env.join(', ')}`);
+    title = 'custom container';
+    headerLines.push(['Mode', 'custom server']);
+    trailingNote = 'Badgr manages runtime, logs, caps, teardown, and receipts.\n  Your container owns the app behavior.';
+  } else if (vllmAlias) {
+    title = model;
+    headerLines.push(['Route', 'tested']);
+    headerLines.push(['Mode', 'OpenAI-compatible endpoint']);
   } else {
-    console.log(chalk.bold('\n⚡ Serving model\n'));
-    if (vllmAlias) {
-      console.log(`  ${chalk.bold('Alias:')}  ${model}  ${chalk.dim(`→ ${effectiveModel}`)}`);
-    } else {
-      console.log(`  ${chalk.bold('Model:')}  ${effectiveModel}`);
-    }
-    console.log(`  ${chalk.bold('GPU:')}    ${gpuLabel}`);
-    if (flags.task) console.log(`  ${chalk.bold('Task:')}   ${flags.task}`);
-    if (flags.env?.length) console.log(`  ${chalk.bold('Env:')}    ${flags.env.join(', ')}`);
-
-    if (gpu === 'AUTO') {
-      const prof = _inferServeProfile(effectiveModel);
-      console.log();
-      console.log(`  ${chalk.bold('Estimated workload:')}      ${prof.label}`);
-      console.log(`  ${chalk.bold('Estimated minimum VRAM:')}  ${prof.vram}`);
-    }
+    title = effectiveModel;
+    headerLines.push(['Route', 'best-effort Hugging Face model']);
+    headerLines.push(['Mode', 'OpenAI-compatible endpoint']);
+    trailingNote = 'Badgr will try a compatible route.';
+    if (gpu === 'AUTO') sizeProfile = _inferServeProfile(effectiveModel);
   }
-  if (flags.maxCost) {
-    console.log(`  ${chalk.bold('Max cost:')}    $${flags.maxCost.toFixed(2)} (auto-stop)`);
-  } else {
-    console.log(chalk.yellow(`  ⚠  Persistent — billing until: badgr down <id>  or  badgr down --all`));
+  if (flags.gpu) headerLines.push(['GPU', gpuLabel]);
+  if (flags.task) headerLines.push(['Task', flags.task]);
+  if (flags.env?.length) headerLines.push(['Env', flags.env.join(', ')]);
+
+  console.log(chalk.bold(`\n⚡ Serving ${title}\n`));
+  const labelWidth = Math.max(...headerLines.map(([label]) => label.length)) + 1;
+  for (const [label, value] of headerLines) {
+    console.log(`  ${chalk.bold(`${label}:`.padEnd(labelWidth + 1))}${value}`);
+  }
+  if (sizeProfile) {
+    console.log();
+    console.log(`  ${chalk.bold('Estimated workload:')}      ${sizeProfile.label}`);
+    console.log(`  ${chalk.bold('Estimated minimum VRAM:')}  ${sizeProfile.vram}`);
+  }
+  console.log(`  ${chalk.bold('Max cost:')}    ${flags.maxCost ? `$${flags.maxCost.toFixed(2)}` : chalk.dim('none')}`);
+  console.log(`  ${chalk.bold('Auto-stop:')}   ${flags.maxCost ? 'enabled' : chalk.yellow('disabled — stop manually with badgr down')}`);
+  if (trailingNote) {
+    console.log();
+    console.log(chalk.dim(`  ${trailingNote}`));
   }
   console.log();
-  process.stdout.write(chalk.dim('  Finding suitable capacity...\n'));
+
+  const STAGE_TOTAL = 5;
+  let stageN = 1;
 
   // ── Duplicate check ────────────────────────────────────────────────────────
   if (config.apiKey) {
@@ -469,6 +600,11 @@ export async function serveCommand(config, args, chalk) {
     createdAt:    new Date().toISOString(),
   });
 
+  console.log(chalk.dim(_stageDone(stageN, STAGE_TOTAL, 'Finding a working route...')));
+  stageN++;
+  console.log(chalk.dim(_stageDone(stageN, STAGE_TOTAL, 'Starting runtime...')));
+  stageN++;
+
   // ── Fix 7: never fall back to config.baseUrl for endpoint health check ─────
   const endpointUrl = dep.endpoint_url || dep.openai_base_url;
   if (!endpointUrl) {
@@ -499,18 +635,48 @@ export async function serveCommand(config, args, chalk) {
     resolvedHealthPath = _detectHealthPath(customImage);  // '/system_stats' for comfyui, null otherwise
   }
 
+  // Gated-model guidance is only shown when it's actually needed — on failure —
+  // not up front, so common launches stay short and uncluttered.
+  const gatedModelId = isLlamaCpp ? flags.hfRepo : effectiveModel;
+  const gatedHintNeeded = isLikelyGatedModel(gatedModelId) && !envObjHasHfToken(flags.env);
+
+  // Shared reporting for "deployment failed to start" — hit from both the
+  // pre-poll status check and the waitForEndpoint poll loop below.
+  function reportDeployFailure(failReason, failureInfo = {}) {
+    console.error(formatCliError('HEALTH_CHECK_DEPLOY_FAILED', {
+      deploymentId: dep.deployment_id,
+      failReason,
+    }, chalk));
+    _printFailureClass(chalk, { failure_class: failureInfo.failureClass, next_action: failureInfo.nextAction });
+    if (gatedHintNeeded) {
+      const rerun = args.join(' ');
+      const retryCmd = rerun.includes('HF_TOKEN=') ? rerun : `${rerun} --env HF_TOKEN=$HF_TOKEN`;
+      console.error();
+      console.error(chalk.yellow('  This model may require Hugging Face access.'));
+      console.error(chalk.dim('  Retry:'));
+      console.error(chalk.dim(`    badgr serve ${retryCmd}`));
+    }
+    updateReceipt(rcptId, { status: 'failed', failReason });
+    process.exitCode = 1;
+  }
+
   // ── Health check ──────────────────────────────────────────────────────────
+  const loadingStageN = stageN;      // "Loading model..."
+  const healthStageN  = stageN + 1;  // "Checking endpoint health..."
+  const readyStageN   = stageN + 2;  // "Ready"
+
   let endpointReady = false;
   if (flags.noWait) {
-    console.log(chalk.yellow('\n  Skipped health check (--no-wait)\n'));
+    console.log(chalk.dim(_stage(loadingStageN, STAGE_TOTAL, 'Loading model...')) + chalk.yellow('  (skipped — --no-wait)'));
+    console.log(chalk.dim(_stage(healthStageN, STAGE_TOTAL, 'Checking endpoint health...')) + chalk.yellow('  (skipped — --no-wait)'));
+    console.log(chalk.yellow(_stage(readyStageN, STAGE_TOTAL, 'Not confirmed ready — check badgr logs') + `\n`));
   } else if (resolvedHealthPath === null) {
+    console.log(chalk.dim(_stage(loadingStageN, STAGE_TOTAL, 'Loading model...')));
     console.log(chalk.yellow(
-      '\n  Skipping health check (custom image — add --health-path /your-readiness-path to enable)\n'
+      _stage(healthStageN, STAGE_TOTAL, 'Checking endpoint health...') +
+      '  (skipped — custom image; add --health-path /your-readiness-path to enable)\n'
     ));
   } else {
-    if (resolvedHealthPath !== '/models') {
-      process.stdout.write(chalk.dim(`  Checking ${resolvedHealthPath} for readiness…\n`));
-    }
     // Check deployment status once before starting the 5-min wait
     try {
       const latest = await callApi(`/deployments/${dep.deployment_id}`, {
@@ -519,33 +685,35 @@ export async function serveCommand(config, args, chalk) {
         timeoutMs: 10_000,
       });
       if (['failed', 'terminated', 'error'].includes(latest.status)) {
-        console.error(formatCliError('HEALTH_CHECK_DEPLOY_FAILED', {
-          deploymentId: dep.deployment_id,
-          failReason:   latest.error || latest.status,
-        }, chalk));
-        updateReceipt(rcptId, { status: 'failed', failReason: latest.error || latest.status });
-        process.exitCode = 1;
+        reportDeployFailure(latest.error || latest.status, {
+          failureClass: latest.failure_class ?? null, nextAction: latest.next_action ?? null,
+        });
         return;
       }
     } catch {
       // status check failed — proceed with endpoint poll anyway
     }
 
-    const healthResult = await waitForEndpoint(dep.deployment_id, config, VLLM_SERVE_WAIT_MS, chalk, resolvedHealthPath);
-    process.stdout.write('\n');
+    const loadingStageLine = _stage(loadingStageN, STAGE_TOTAL, 'Loading model...');
+    const healthResult = await waitForEndpoint(
+      dep.deployment_id, config, VLLM_SERVE_WAIT_MS, chalk, resolvedHealthPath,
+      { costPerHour: dep.cost_per_hour || 0, maxCost: flags.maxCost || null, stageLine: loadingStageLine },
+    );
 
     if (healthResult.depFailed) {
-      console.error(formatCliError('HEALTH_CHECK_DEPLOY_FAILED', {
-        deploymentId: dep.deployment_id,
-        failReason:   healthResult.failReason,
-      }, chalk));
-      updateReceipt(rcptId, { status: 'failed', failReason: healthResult.failReason });
-      process.exitCode = 1;
+      reportDeployFailure(healthResult.failReason, {
+        failureClass: healthResult.failureClass, nextAction: healthResult.nextAction,
+      });
       return;
     }
 
     endpointReady = healthResult.ready;
-    if (!endpointReady) updateReceipt(rcptId, { status: 'health_check_timeout' });
+    if (endpointReady) {
+      console.log(chalk.dim(_stage(healthStageN, STAGE_TOTAL, 'Checking endpoint health...')));
+      console.log(chalk.green(_stage(readyStageN, STAGE_TOTAL, 'Ready')));
+    } else {
+      updateReceipt(rcptId, { status: 'health_check_timeout' });
+    }
   }
 
   // ── Custom-node validation (ComfyUI) ─────────────────────────────────────
@@ -570,8 +738,6 @@ export async function serveCommand(config, args, chalk) {
     console.log();
   }
 
-  const serveRate = dep.cost_per_hour || 0;
-
   console.log(`  ${chalk.bold('Base URL:')}     ${chalk.cyan(endpointUrl)}`);
   if (isLlamaCpp) {
     console.log(`  ${chalk.bold('HF Repo:')}      ${flags.hfRepo}`);
@@ -579,23 +745,20 @@ export async function serveCommand(config, args, chalk) {
   }
   else if (dep.model || effectiveModel) console.log(`  ${chalk.bold('Model:')}        ${dep.model || effectiveModel}`);
   if (customImage)        console.log(`  ${chalk.bold('Image:')}        ${customImage}`);
-  console.log(`  ${chalk.bold('GPU:')}          ${dep.gpu_type} × ${dep.gpu_count}`);
-  if (serveRate > 0) console.log(`  ${chalk.bold('Rate:')}         $${serveRate.toFixed(2)}/hr`);
   if (flags.maxCost) console.log(`  ${chalk.bold('Max cost:')}     $${flags.maxCost.toFixed(2)}`);
-  console.log(`  ${chalk.bold('Receipt:')}      ${chalk.dim(rcptId)}`);
-  console.log(`  ${chalk.bold('Logs:')}         ${chalk.dim(`badgr logs ${dep.deployment_id}`)}`);
+  console.log(`  ${chalk.bold('Logs:')}         badgr logs ${dep.deployment_id}`);
   console.log(`  ${chalk.bold('Stop billing:')} ${chalk.cyan(`badgr down ${dep.deployment_id}`)}`);
+  console.log(`  ${chalk.bold('Receipt:')}      badgr receipts ${rcptId}`);
   console.log();
-
-  if (serveRate > HIGH_RATE_THRESHOLD && !flags.maxCost) {
-    console.log(chalk.yellow(`  Selected capacity rate: $${serveRate.toFixed(2)}/hr`));
-    console.log(chalk.dim('  Tip: use --max-cost to enforce a hard ceiling.\n'));
-  }
   console.log(chalk.dim('  Billing continues until you run: ') + chalk.cyan(`badgr down ${dep.deployment_id}`));
 
   if (endpointReady && !customImage) {
     const keySnip = config.apiKey?.slice(0, 4) || 'sk-...';
     const sdkModel = isLlamaCpp ? 'default' : (dep.model || effectiveModel);
+    console.log(`  ${chalk.bold('Test with curl:')}`);
+    console.log(chalk.dim(`    curl ${endpointUrl}/chat/completions \\`));
+    console.log(chalk.dim(`      -H "Authorization: Bearer ${keySnip}..." -H "Content-Type: application/json" \\`));
+    console.log(chalk.dim(`      -d '{"model":"${sdkModel}","messages":[{"role":"user","content":"Hello"}]}'`));
     console.log(`  ${chalk.bold('Use with OpenAI SDK:')}`);
     console.log(chalk.dim(`    from openai import OpenAI`));
     console.log(chalk.dim(`    client = OpenAI(base_url="${endpointUrl}", api_key="${keySnip}...")`));

package/src/commands/train.js

@@ -10,6 +10,7 @@ import { requireApiKey } from '../config.js';
 import { addReceipt, updateReceipt, generateReceiptId } from '../store.js';
 import { normalizeTier, callWithFallback } from '../fallback.js';
 import { monitorBatchJob, fmtRuntime } from '../batch.js';
+import { pollJobUntilTerminal, renderJobClosingBlock } from '../progress.js';
 
 const MAX_CONFIG_B = 512 * 1024; // 512 KB config limit
 
@@ -239,35 +240,29 @@ export async function trainLoraCommand(config, args, chalk) {
   console.log(`  ${chalk.bold('Receipt:')} ${chalk.dim(rcptId)}`);
   console.log(chalk.dim('\n  Polling for completion — Ctrl+C to detach (GPU keeps running)\n'));
 
-  // Poll until complete
-  const startMs = Date.now();
   const maxMs = maxRuntime * 60 * 1000;
-  while (Date.now() - startMs < maxMs) {
-    await new Promise(r => setTimeout(r, 15_000));
-    let detail;
-    try {
-      detail = await callApi(`/jobs/${job.job_id}`, {
-        apiKey: config.apiKey,
-        baseUrl: config.baseUrl,
-      });
-    } catch { continue; }
-    process.stdout.write(`\r  Status: ${detail.status}  elapsed: ${Math.floor((Date.now() - startMs) / 1000)}s  `);
-    if (detail.status === 'completed') {
-      const out = detail.output || {};
-      console.log(chalk.green('\n\n  ✓ Training complete\n'));
-      if (out.adapter_url) console.log(`  ${chalk.bold('Adapter:')} ${out.adapter_url}`);
-      if (out.checkpoint_url) console.log(`  ${chalk.bold('Checkpoint:')} ${out.checkpoint_url}`);
-      console.log(`  ${chalk.bold('Receipt:')} ${chalk.dim(rcptId)}\n`);
-      return;
-    }
-    if (detail.status === 'failed') {
-      console.error(chalk.red(`\n\n  ✗ Training failed: ${detail.error_code || ''} — ${detail.error_message || ''}\n`));
-      process.exitCode = 1;
-      return;
-    }
+  const { outcome, detail } = await pollJobUntilTerminal(callApi, config, job.job_id, { chalk, maxMs });
+
+  if (outcome === 'polling_failed') {
+    console.error(chalk.red(`\n\n  ✗ Lost contact with Badgr — could not confirm job status.\n`));
+    console.error(chalk.dim(`    badgr status\n    badgr logs ${job.job_id}\n`));
+    process.exitCode = 1;
+    return;
+  }
+  if (outcome === 'timed_out') {
+    console.error(chalk.yellow('\n\n  Training still running — detached. Check status:\n'));
+    console.error(chalk.dim(`    badgr status\n`));
+    return;
+  }
+
+  if (outcome === 'completed') {
+    const out = detail.output || {};
+    if (out.adapter_url) console.log(`\n  ${chalk.bold('Adapter:')} ${out.adapter_url}`);
+    if (out.checkpoint_url) console.log(`  ${chalk.bold('Checkpoint:')} ${out.checkpoint_url}`);
+    if (out.warning) console.log(chalk.yellow(`  ${out.warning}`));
   }
-  console.error(chalk.yellow('\n\n  Training still running — detached. Check status:\n'));
-  console.error(chalk.dim(`    badgr status\n`));
+  console.log(renderJobClosingBlock(chalk, detail, rcptId));
+  if (outcome === 'failed') process.exitCode = 1;
 }
 
 export async function trainCommand(config, args, chalk) {