badgr-cli 1.0.30 → 1.0.32

package/src/commands/run.js

@@ -7,6 +7,7 @@ import { normalizeTier, callWithFallback, HIGH_RATE_THRESHOLD } from '../fallbac
  * badgr run python train.py                    # gpu=auto, attached
  * badgr run python train.py --gpu A100         # specific GPU
  * badgr run --image my/image:latest --gpu L40S --detach
+ * badgr run python train.py --env HF_TOKEN=abc --env DATASET=my/data
  */
 export function parseRunArgs(args) {
   const flags = {};
@@ -27,11 +28,26 @@ export function parseRunArgs(args) {
     if (args[i] === '--no-expanded-search') { flags.noFallback = true; i++; continue; }
     if (args[i] === '--max-runtime')     { flags.maxRuntime = parseFloat(args[++i]); i++; continue; }
     if (args[i] === '--max-cost')        { flags.maxCost = parseFloat(args[++i]); i++; continue; }
+    if (args[i] === '--env') {
+      const kv = args[++i]; i++;
+      if (!flags.env) flags.env = [];
+      flags.env.push(kv);
+      continue;
+    }
     positional.push(args[i++]);
   }
   return { flags, positional };
 }
 
+function parseEnvFlag(envList) {
+  const obj = {};
+  for (const kv of (envList || [])) {
+    const idx = kv.indexOf('=');
+    if (idx > 0) obj[kv.slice(0, idx)] = kv.slice(idx + 1);
+  }
+  return obj;
+}
+
 // Mirror of backend workload_profile.py — kept in sync for pre-flight display.
 const _PROFILES = {
   smoke_test:       { label: 'smoke test',            vram: '4 GB',    gpus: ['RTX 3080', 'RTX 3090', 'RTX 4090'] },
@@ -60,6 +76,7 @@ function fmtRuntime(ms) {
 
 const HEARTBEAT_WARN_POLLS  = 3;
 const HEARTBEAT_KILL_POLLS  = 15;
+const STARTUP_STATES = new Set(['queued', 'provisioning', 'starting']);
 
 export function classifyFailure(finalStatus, exitCode) {
   if (finalStatus === 'failed' && (exitCode === null || exitCode === undefined)) return 'infrastructure';
@@ -68,7 +85,7 @@ export function classifyFailure(finalStatus, exitCode) {
 }
 
 // Lines the log stream never needs to print — we surface them in the status bar instead.
-const LOG_META_RE = /^\[dep-[^\]]+\] (status|gpu|region|endpoint|cost|receipt|provider_status|uptime)=/;
+const LOG_META_RE = /^\[dep-[^\]]+\] (status|gpu|region|cost|receipt|provider_status|uptime)=/;
 
 // Extract structured values from provider status lines so we can show them nicely.
 function parseProviderLine(line) {
@@ -103,11 +120,10 @@ function renderStatusBar(chalk, { elapsedMs, ratePerHour, gpuUtil, cpuUtil, maxR
 }
 
 // Wait for status to leave 'starting'/'queued'/'provisioning'.
-// The runtime limit does NOT start until this returns.
-// Returns the final dep object with status 'running' or 'failed'.
+// Returns the dep once it leaves startup states (or the last known state on timeout).
 async function waitForRunning(config, depId, chalk) {
   const POLL_MS     = 3000;
-  const TIMEOUT_MS  = 5 * 60 * 1000; // 5 min startup grace (image pull + container init)
+  const TIMEOUT_MS  = 5 * 60 * 1000;
   const startMs     = Date.now();
   const PHASES = [
     { afterMs:     0, label: '  Starting container'    },
@@ -117,11 +133,9 @@ async function waitForRunning(config, depId, chalk) {
   ];
 
   let lastPhaseIdx = -1;
-  const STARTUP_STATES = new Set(['queued', 'provisioning', 'starting']);
 
   const ticker = setInterval(() => {
     const elapsed = Date.now() - startMs;
-    // Find the latest phase whose afterMs has been passed
     let phaseIdx = 0;
     for (let i = 0; i < PHASES.length; i++) {
       if (elapsed >= PHASES[i].afterMs) phaseIdx = i;
@@ -140,6 +154,7 @@ async function waitForRunning(config, depId, chalk) {
       const dep = await callApi(`/deployments/${depId}`, {
         apiKey: config.apiKey,
         baseUrl: config.baseUrl,
+        timeoutMs: 10_000,
       });
       if (!STARTUP_STATES.has(dep.status)) {
         process.stdout.write('\n');
@@ -151,14 +166,17 @@ async function waitForRunning(config, depId, chalk) {
     process.stdout.write('\n');
   }
 
-  // Timed out — return last known state
-  return await callApi(`/deployments/${depId}`, { apiKey: config.apiKey, baseUrl: config.baseUrl });
+  return await callApi(`/deployments/${depId}`, {
+    apiKey: config.apiKey,
+    baseUrl: config.baseUrl,
+    timeoutMs: 10_000,
+  });
 }
 
-async function attachToJob(config, depId, { chalk, maxRuntimeMs = null, maxCost = null, ratePerHour = 0, onTeardown }) {
+async function attachToJob(config, depId, { chalk, maxRuntimeMs = null, maxCost = null, ratePerHour = 0, onTeardown, isShuttingDown }) {
   const TERMINAL = new Set(['stopped', 'failed', 'completed', 'succeeded']);
   const POLL_MS  = 4000;
-  let seenContent     = new Set();  // track by content, not index, to avoid reprinting stale lines
+  let seenContent     = new Set();
   let lastStatus      = '';
   let consecutiveErrs = 0;
   let gpuUtil         = null;
@@ -167,7 +185,9 @@ async function attachToJob(config, depId, { chalk, maxRuntimeMs = null, maxCost
   let statusBarActive = false;
   const startMs       = Date.now();
 
-  // Ticker: update status bar in place every second between polls.
+  // tearing: guards against double-teardown for cap/heartbeat paths within this function.
+  let tearing = false;
+
   let tickerInterval = null;
   const startTicker = () => {
     if (tickerInterval) return;
@@ -181,22 +201,20 @@ async function attachToJob(config, depId, { chalk, maxRuntimeMs = null, maxCost
   };
   const stopTicker = () => {
     if (tickerInterval) { clearInterval(tickerInterval); tickerInterval = null; }
-    if (statusBarActive) { process.stdout.write('\r\x1b[2K'); statusBarActive = false; } // clear line
-  };
-
-  let tearing = false;
-  const sigintHandler = () => {
-    if (tearing) return;
-    tearing = true;
-    stopTicker();
-    onTeardown('interrupted');
+    if (statusBarActive) { process.stdout.write('\r\x1b[2K'); statusBarActive = false; }
   };
-  process.once('SIGINT', sigintHandler);
 
   try {
     while (true) {
+      // Exit loop if SIGINT handler has started shutdown externally.
+      if (isShuttingDown()) break;
+      if (tearing) break;
+
       await new Promise(r => setTimeout(r, POLL_MS));
 
+      if (isShuttingDown()) break;
+      if (tearing) break;
+
       const elapsedMs  = Date.now() - startMs;
       const spentSoFar = ratePerHour * (elapsedMs / 3_600_000);
 
@@ -204,14 +222,14 @@ async function attachToJob(config, depId, { chalk, maxRuntimeMs = null, maxCost
         tearing = true;
         stopTicker();
         await onTeardown('max-cost');
-        return { status: 'timeout', exitCode: null, runtimeMs: elapsedMs, failureType: null };
+        return { status: 'capped', reason: 'max-cost', exitCode: null, runtimeMs: elapsedMs, failureType: null };
       }
 
       if (maxRuntimeMs !== null && elapsedMs >= maxRuntimeMs) {
         tearing = true;
         stopTicker();
         await onTeardown('max-runtime');
-        return { status: 'timeout', exitCode: null, runtimeMs: elapsedMs, failureType: null };
+        return { status: 'capped', reason: 'max-runtime', exitCode: null, runtimeMs: elapsedMs, failureType: null };
       }
 
       let dep;
@@ -219,6 +237,7 @@ async function attachToJob(config, depId, { chalk, maxRuntimeMs = null, maxCost
         dep = await callApi(`/deployments/${depId}`, {
           apiKey: config.apiKey,
           baseUrl: config.baseUrl,
+          timeoutMs: 10_000,
         });
         consecutiveErrs = 0;
       } catch {
@@ -233,7 +252,7 @@ async function attachToJob(config, depId, { chalk, maxRuntimeMs = null, maxCost
             tearing = true;
             stopTicker();
             await onTeardown('heartbeat-lost');
-            return { status: 'failed', exitCode: null, runtimeMs: elapsedMs, failureType: 'infrastructure' };
+            return { status: 'failed', reason: 'heartbeat-lost', exitCode: null, runtimeMs: elapsedMs, failureType: 'infrastructure' };
           }
         }
         continue;
@@ -252,16 +271,15 @@ async function attachToJob(config, depId, { chalk, maxRuntimeMs = null, maxCost
         const logData = await callApi(`/deployments/${depId}/logs`, {
           apiKey: config.apiKey,
           baseUrl: config.baseUrl,
+          timeoutMs: 10_000,
         });
         const lines = logData?.logs ?? [];
 
         for (const line of lines) {
-          // Extract structured provider values (gpu_util, ssh, etc.) from any line.
           const parsed = parseProviderLine(line);
           if (parsed.gpuUtil  !== null) gpuUtil = parsed.gpuUtil;
           if (parsed.cpuUtil  !== null) cpuUtil = parsed.cpuUtil;
 
-          // Show SSH address once, prominently.
           if (parsed.ssh && !sshShown) {
             stopTicker();
             console.log(`  ${chalk.bold('SSH:')}      ${chalk.cyan(parsed.ssh)}`);
@@ -269,11 +287,9 @@ async function attachToJob(config, depId, { chalk, maxRuntimeMs = null, maxCost
             startTicker();
           }
 
-          // Skip lines we've already printed and pure metadata lines.
           if (seenContent.has(line)) continue;
           seenContent.add(line);
           if (LOG_META_RE.test(line)) continue;
-          // Skip provider util lines — they're shown in the status bar instead.
           if (/\b(gpu_util|cpu_util|provider_status|uptime)=/.test(line)) continue;
 
           const isErrorLine = /^error\b/i.test(line) || /Error response from daemon/i.test(line);
@@ -285,15 +301,14 @@ async function attachToJob(config, depId, { chalk, maxRuntimeMs = null, maxCost
         // logs not ready yet
       }
 
-      // Start the ticker once the job is confirmed running.
       if (status === 'running') startTicker();
 
       if (TERMINAL.has(status)) {
         stopTicker();
-        process.removeListener('SIGINT', sigintHandler);
         const exitCode = dep.exit_code ?? null;
         return {
           status,
+          reason: null,
           exitCode,
           runtimeMs: Date.now() - startMs,
           failureType: classifyFailure(status, exitCode),
@@ -302,8 +317,11 @@ async function attachToJob(config, depId, { chalk, maxRuntimeMs = null, maxCost
     }
   } finally {
     stopTicker();
-    process.removeListener('SIGINT', sigintHandler);
   }
+
+  // Reached when SIGINT (isShuttingDown) or duplicate tearing flag breaks the loop.
+  // The SIGINT handleShutdown() is managing teardown + exit.
+  return { status: 'interrupted', reason: 'signal', exitCode: null, runtimeMs: Date.now() - startMs, failureType: null };
 }
 
 export async function runCommand(config, args, chalk) {
@@ -317,27 +335,45 @@ export async function runCommand(config, args, chalk) {
 
   requireApiKey(config);
 
+  // ── Validate flags early ───────────────────────────────────────────────────
+  if (flags.count !== undefined && (!Number.isFinite(flags.count) || flags.count < 1)) {
+    console.error(chalk.red('  ✗ --count must be an integer greater than 0'));
+    process.exitCode = 1;
+    return;
+  }
+  if (flags.maxCost !== undefined && (!Number.isFinite(flags.maxCost) || flags.maxCost <= 0)) {
+    console.error(chalk.red('  ✗ --max-cost must be a number greater than 0'));
+    process.exitCode = 1;
+    return;
+  }
+  if (flags.maxPrice !== undefined && (!Number.isFinite(flags.maxPrice) || flags.maxPrice <= 0)) {
+    console.error(chalk.red('  ✗ --max-price must be a number greater than 0'));
+    process.exitCode = 1;
+    return;
+  }
+  if (flags.region !== undefined && !['US', 'EU', 'AU'].includes(flags.region.toUpperCase())) {
+    console.error(chalk.red('  ✗ --region must be US, EU, or AU'));
+    process.exitCode = 1;
+    return;
+  }
+
   const command      = positional.length > 0 ? positional : undefined;
-  // Use alpine for trivial one-liners (7MB vs 50MB — much faster pull)
   const cmdStr = command ? command.join(' ') : '';
   const isSmoke = cmdStr.length < 80 && /print\s*\(|['"]hello/i.test(cmdStr);
   const inferredImage = isSmoke ? 'python:3.11-alpine' : 'python:3.11-slim';
   const image = flags.image || (command ? inferredImage : undefined);
   const detach       = flags.detach || false;
-  const fallbackMode = flags.noFallback ? 'none' : (flags.fallback || 'closest');
   const maxRuntimeMs = flags.maxRuntime ? flags.maxRuntime * 60 * 1000 : null;
   const maxCost      = flags.maxCost ?? null;
-
+  const envObj       = parseEnvFlag(flags.env);
   const effectiveTier = normalizeTier(flags.tier);
 
-  // ── Auto GPU selection (no --gpu specified) ────────────────────────────────
   let gpu = flags.gpu ? flags.gpu.toUpperCase().replace('-', '_') : 'auto';
 
   if (gpu === 'auto') {
     console.log(chalk.bold('\n⚡ Running GPU job\n'));
     if (command) console.log(`  ${chalk.bold('Command:')}  ${command.join(' ')}`);
     if (image)   console.log(`  ${chalk.bold('Image:')}    ${image}`);
-    // Show workload estimate so the user knows what Badgr inferred.
     if (command) {
       const profKey = inferProfileFromCommand(cmdStr);
       const prof = _PROFILES[profKey];
@@ -349,7 +385,6 @@ export async function runCommand(config, args, chalk) {
     }
     console.log();
   } else {
-    // Specific GPU requested — show header
     console.log(chalk.bold('\n⚡ Running GPU job\n'));
     if (command) console.log(`  ${chalk.bold('Command:')}  ${command.join(' ')}`);
     if (image)   console.log(`  ${chalk.bold('Image:')}    ${image}`);
@@ -358,6 +393,7 @@ export async function runCommand(config, args, chalk) {
     if (maxCost)          console.log(`  ${chalk.bold('Max cost:')}    $${maxCost.toFixed(2)}`);
     if (flags.maxPrice)   console.log(`  ${chalk.bold('Max price:')}   $${flags.maxPrice.toFixed(2)}/hr`);
     if (detach)           console.log(`  ${chalk.dim('(detached — returns immediately)')}`);
+    if (flags.env?.length) console.log(`  ${chalk.bold('Env:')}      ${flags.env.join(', ')}`);
     console.log();
 
     if (!detach && !flags.maxRuntime && !maxCost) {
@@ -381,6 +417,7 @@ export async function runCommand(config, args, chalk) {
       max_price_per_hour: flags.maxPrice,
       name: flags.name,
       tier: tierOverride || effectiveTier,
+      ...(Object.keys(envObj).length > 0 ? { env: envObj } : {}),
     };
   }
 
@@ -393,15 +430,20 @@ export async function runCommand(config, args, chalk) {
       effectiveTier,
       chalk,
       { thing: 'job', cmd: 'badgr run' },
+      { allowTier2Fallback: !flags.noFallback },
     );
   } catch (err) {
     if (err.isPaymentRequired) {
       console.error(chalk.yellow(err.message));
       const rerun = ['badgr run', ...args].join(' ');
       console.error(chalk.dim(`After payment, rerun:\n  ${rerun}\n`));
-      process.exit(1);
+      process.exitCode = 1;
+      return;
     }
-    throw err;
+    // CapacityError message is pre-formatted with chalk
+    console.error(err.message);
+    process.exitCode = 1;
+    return;
   }
 
   const rcptId = dep.receipt_id || generateReceiptId();
@@ -438,26 +480,16 @@ export async function runCommand(config, args, chalk) {
     return;
   }
 
-  // Wait through startup phases (queued → provisioning → starting → running).
-  // The max-runtime clock does NOT start until this returns.
-  const STARTUP_STATES = new Set(['queued', 'provisioning', 'starting']);
-  if (STARTUP_STATES.has(dep.status)) {
-    console.log();
-    dep = await waitForRunning(config, dep.deployment_id, chalk);
-  }
-
-  if (dep.status === 'failed') {
-    console.error(chalk.red('\n  ✗ Container failed to start (infrastructure error).\n'));
-    console.error(chalk.dim('  The backend retried automatically. All attempts failed.'));
-    console.error(chalk.dim(`  Contact support with receipt ID: ${rcptId}`));
-    console.log();
-    process.exit(1);
-  }
-
+  // ── Teardown helper ────────────────────────────────────────────────────────
+  // Does NOT call process.exit() — only terminates the deployment and updates receipt.
   const ratePerHour = dep.cost_per_hour || 0;
-  console.log(chalk.dim('\n  ── Running command (Ctrl+C to stop) ────────────────────────────\n'));
+  let teardownCalled = false;
+  let attachStart = Date.now();
 
   async function teardown(reason) {
+    if (teardownCalled) return;
+    teardownCalled = true;
+
     const labels = {
       'max-runtime':    chalk.yellow('\n  ⏱ Max runtime reached — stopping job...'),
       'max-cost':       chalk.yellow('\n  💰 Spend cap reached — stopping job...'),
@@ -465,28 +497,77 @@ export async function runCommand(config, args, chalk) {
       'interrupted':    chalk.yellow('\n  Stopping job...'),
     };
     console.log(labels[reason] ?? chalk.yellow('\n  Stopping job...'));
+
     try {
       await terminateDeployment(config, dep.deployment_id);
     } catch {
-      // best-effort
+      // terminateDeployment retries 3×; best-effort if all fail
     }
+
     const runtimeMs = Date.now() - attachStart;
     const finalCost = ratePerHour * (runtimeMs / 3_600_000);
-    updateReceipt(rcptId, { status: reason, teardownStatus: 'terminated', runtimeSeconds: Math.round(runtimeMs / 1000), finalCost });
+    updateReceipt(rcptId, {
+      status: reason,
+      teardownStatus: 'terminated',
+      runtimeSeconds: Math.round(runtimeMs / 1000),
+      finalCost,
+    });
     console.log(chalk.dim(`  Runtime: ${fmtRuntime(runtimeMs)}  •  Est. cost: $${finalCost.toFixed(4)}`));
     console.log(chalk.dim('  Job stopped. Billing ended.\n'));
+  }
+
+  // ── SIGINT handler — installed immediately after we have a deployment ID ───
+  // Covers: queued, provisioning, starting, running phases.
+  // Does NOT call process.exit() inside teardown; exits here after await.
+  let shuttingDown = false;
+
+  async function handleShutdown(reason) {
+    if (shuttingDown) return;
+    shuttingDown = true;
+    try {
+      await teardown(reason);
+    } catch (err) {
+      console.error(chalk.red(`  Warning: teardown may have failed: ${err.message}`));
+    }
     process.exit(reason === 'interrupted' ? 0 : 1);
   }
 
-  const attachStart = Date.now();
+  process.once('SIGINT', () => { void handleShutdown('interrupted'); });
+
+  // ── Wait through startup phases ────────────────────────────────────────────
+  if (STARTUP_STATES.has(dep.status)) {
+    console.log();
+    dep = await waitForRunning(config, dep.deployment_id, chalk);
+  }
+
+  if (dep.status === 'failed') {
+    process.removeListener('SIGINT', handleShutdown);
+    console.error(chalk.red('\n  ✗ Container failed to start (infrastructure error).\n'));
+    console.error(chalk.dim('  The backend retried automatically. All attempts failed.'));
+    console.error(chalk.dim(`  Contact support with receipt ID: ${rcptId}`));
+    console.log();
+    process.exitCode = 1;
+    return;
+  }
+
+  console.log(chalk.dim('\n  ── Running command (Ctrl+C to stop) ────────────────────────────\n'));
+
+  attachStart = Date.now();
   const { status: finalStatus, exitCode, runtimeMs, failureType } = await attachToJob(config, dep.deployment_id, {
     chalk,
     maxRuntimeMs,
     maxCost,
     ratePerHour,
     onTeardown: teardown,
+    isShuttingDown: () => shuttingDown,
   });
 
+  // Remove SIGINT handler — job is done (or SIGINT was handled)
+  process.removeListener('SIGINT', handleShutdown);
+
+  // 'interrupted' = SIGINT handler is managing teardown + exit — don't duplicate
+  if (finalStatus === 'interrupted') return;
+
   console.log();
 
   const finalCost = ratePerHour * (runtimeMs / 3_600_000);
@@ -508,6 +589,12 @@ export async function runCommand(config, args, chalk) {
   }
   console.log(`  ${chalk.bold('Receipt:')}  ${chalk.dim(rcptId)}`);
 
+  // 'capped' = max-runtime or max-cost path; teardown message already printed
+  if (finalStatus === 'capped') {
+    process.exitCode = 1;
+    return;
+  }
+
   if (finalStatus === 'failed' || (exitCode !== null && exitCode !== 0)) {
     if (failureType === 'infrastructure') {
       console.error(chalk.red(`\n  ✗ Machine failure — this is not your code.\n`));
@@ -517,8 +604,11 @@ export async function runCommand(config, args, chalk) {
       console.error(chalk.dim(`  Check logs: badgr logs ${dep.deployment_id}`));
     }
     console.log();
-    process.exit(exitCode ?? 1);
-  } else if ((finalStatus === 'completed' || finalStatus === 'succeeded') && (exitCode === 0 || exitCode === null)) {
+    process.exitCode = exitCode ?? 1;
+    return;
+  }
+
+  if ((finalStatus === 'completed' || finalStatus === 'succeeded') && (exitCode === 0 || exitCode === null)) {
     try {
       await terminateDeployment(config, dep.deployment_id);
     } catch { /* already stopped */ }