badgr-cli 1.0.21 → 1.0.22

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "badgr-cli",
-  "version": "1.0.21",
+  "version": "1.0.22",
   "description": "Badgr, run or serve GPU workloads from one command",
   "type": "module",
   "bin": {

package/src/badgr.js

@@ -11,6 +11,7 @@ import { runCommand }      from './commands/run.js';
 import { serveCommand }    from './commands/serve.js';
 import { modelsCommand }   from './commands/models.js';
 import { capacityCommand } from './commands/capacity.js';
+import { testCommand }     from './commands/test-run.js';
 
 const HELP = `
 ${chalk.bold('badgr')} — run or serve GPU workloads from one command
@@ -23,9 +24,13 @@ ${chalk.bold('COMMANDS')}
   ${chalk.cyan('badgr logs <id>')}              Stream logs for a running job or endpoint
   ${chalk.cyan('badgr down <id>')}              Stop a deployment and end billing
   ${chalk.cyan('badgr receipts')}               Show cost history
+  ${chalk.cyan('badgr test')}                   Run an end-to-end test (provision → run → teardown)
   ${chalk.cyan('badgr capacity')}               Check what GPU capacity is available right now
 
 ${chalk.bold('EXAMPLES')}
+  ${chalk.dim('# Verify the stack works end-to-end:')}
+  badgr test
+
   ${chalk.dim('# Simplest — Badgr picks the GPU (RunPod, reliable):')}
   badgr run python train.py
   badgr serve meta-llama/Llama-3.1-8B-Instruct
@@ -98,6 +103,7 @@ async function main() {
     case 'receipts': return receiptsCommand(config, rest, chalk);
     case 'models':   return modelsCommand(config, chalk);
     case 'capacity': return capacityCommand(config, rest, chalk);
+    case 'test':     return testCommand(config, chalk);
     // legacy aliases kept for compatibility
     case 'up':       return upCommand(config, rest, chalk);
     case 'config': {

package/src/commands/run.js

@@ -270,12 +270,53 @@ async function attachToJob(config, depId, { chalk, maxRuntimeMs = null, maxCost
   }
 }
 
-// Ask the backend for the best available GPU right now.
-// Defaults to tier 1 (RunPod) for reliability; pass tier='2' for budget providers.
-// Returns { gpu, region, price } or null when nothing is available.
-async function findAutoGpu(config, chalk, tier = '1') {
+/**
+ * Infer a workload profile name from the command the user wants to run.
+ * The backend uses this to enforce a VRAM floor when picking a GPU.
+ *
+ * Note: \b word boundaries are intentionally avoided for keyword checks
+ * because keywords commonly appear inside filenames joined by underscores
+ * (e.g. lora_train.py, run_sdxl.py) where `_` is a word character.
+ */
+export function inferWorkload(command) {
+  if (!command || command.length === 0) return 'general';
+
+  const cmd = command.join(' ').toLowerCase();
+
+  // Trivial one-liner / smoke test
+  if (/print\s*\(|['"]hello/.test(cmd) && cmd.length < 80) return 'smoke_test';
+
+  // LoRA / QLoRA / PEFT fine-tuning
+  if (/lora|qlora|finetune|fine[_-]tun|peft/.test(cmd)) return 'lora_finetune';
+
+  // Diffusion / image generation
+  if (/diffusion|stable.?diff|sdxl|sd.?xl|comfyui|a1111|invoke|kohya/.test(cmd)) return 'image_gen';
+
+  // vLLM / TGI / inference server
+  if (/vllm|[^a-z]tgi[^a-z]|^tgi\b|tgi$|text.generation.inference/.test(cmd)) return 'inference_small';
+
+  // Explicit training script
+  if (/\btrain\.py\b/.test(cmd)) return 'lora_finetune';
+
+  return 'general';
+}
+
+const WORKLOAD_LABELS = {
+  smoke_test:       'smoke test',
+  general:          'GPU job',
+  lora_finetune:    'fine-tuning (40GB+ VRAM)',
+  image_gen:        'image generation',
+  inference_small:  'inference (7B–8B model)',
+  inference_medium: 'inference (30B–34B model)',
+  inference_large:  'inference (70B+ model)',
+};
+
+// Ask the backend for the best available GPU for a given workload.
+// Routes by: workload → min VRAM → provider tier → cheapest match.
+// Returns { gpu, region, price, workload, workload_desc } or null when nothing is available.
+async function findAutoGpu(config, chalk, tier = '1', workload = 'general') {
   try {
-    const params = new URLSearchParams({ max_price: '10', tier });
+    const params = new URLSearchParams({ max_price: '10', tier, workload });
     return await callApi(`/capacity/auto?${params}`, {
       apiKey: config.apiKey,
       baseUrl: config.baseUrl,
@@ -319,17 +360,21 @@ export async function runCommand(config, args, chalk) {
   let gpu = flags.gpu ? flags.gpu.toUpperCase().replace('-', '_') : null;
   let autoRegion = null;
 
+  // Infer workload from the command so the backend can apply the correct VRAM floor.
+  const workload = command ? inferWorkload(command) : 'general';
+  const workloadLabel = WORKLOAD_LABELS[workload] || 'GPU job';
+
   if (!gpu) {
-    console.log(chalk.bold('\n⚡ Running GPU job\n'));
+    console.log(chalk.bold(`\n⚡ Running ${workloadLabel}\n`));
     if (command) console.log(`  ${chalk.bold('Command:')}  ${command.join(' ')}`);
     if (image)   console.log(`  ${chalk.bold('Image:')}    ${image}`);
     if (effectiveTier === '2') console.log(`  ${chalk.dim('(budget mode — searching all providers)')}`);
     console.log();
-    process.stdout.write(chalk.dim('  Finding GPU...'));
+    process.stdout.write(chalk.dim('  Finding best GPU...'));
 
     let best;
     try {
-      best = await findAutoGpu(config, chalk, effectiveTier);
+      best = await findAutoGpu(config, chalk, effectiveTier, workload);
     } catch (err) {
       process.stdout.write('\n');
       console.error(chalk.red(`\n  ✗ Could not find GPU capacity: ${err.message}\n`));
@@ -344,17 +389,17 @@ export async function runCommand(config, args, chalk) {
       process.exit(1);
     }
 
-    console.log(`\n  ${chalk.bold('Best match:')} ${chalk.cyan(best.gpu)} in ${best.region} — ${chalk.green('$' + best.price.toFixed(2) + '/hr')} estimated`);
+    const vramNote = best.min_vram_gb ? chalk.dim(` (${best.min_vram_gb}GB VRAM)`) : '';
+    console.log(`\n  ${chalk.bold('Selected:')}   ${chalk.cyan(best.gpu)} in ${best.region} — ${chalk.green('$' + best.price.toFixed(2) + '/hr')}${vramNote}`);
     console.log();
 
-    if (process.stdin.isTTY) {
-      const answer = await askConfirm(`  Press ${chalk.bold('Enter')} to run, or ${chalk.bold('q')} to cancel: `);
+    if (effectiveTier === '2' && process.stdin.isTTY) {
+      // Budget mode: confirm because the user is being routed to a less reliable provider.
+      const answer = await askConfirm(`  Press ${chalk.bold('Enter')} to run on budget provider, or ${chalk.bold('q')} to cancel: `);
       if (answer.toLowerCase() === 'q') {
         console.log(chalk.dim('\n  Cancelled.\n'));
         process.exit(0);
       }
-    } else {
-      console.log(chalk.dim(`  Auto-selecting ${best.gpu} (non-interactive).`));
     }
 
     gpu = best.gpu;

package/src/commands/serve.js

@@ -28,21 +28,26 @@ export function parseServeArgs(args) {
   return { model, flags };
 }
 
+function _serveStageLabel(elapsedSec) {
+  if (elapsedSec < 45)  return 'Starting vLLM…';
+  if (elapsedSec < 150) return 'Downloading model…';
+  return 'Waiting for /v1/models…';
+}
+
 async function waitForEndpoint(endpointUrl, timeoutMs = 5 * 60 * 1000, chalk) {
-  const deadline = Date.now() + timeoutMs;
-  let attempt = 0;
+  const startMs  = Date.now();
+  const deadline = startMs + timeoutMs;
 
   while (Date.now() < deadline) {
-    attempt++;
     try {
       const res = await fetch(`${endpointUrl}/models`, { signal: AbortSignal.timeout(8000) });
-      if (res.ok) return true;
+      if (res.ok) { process.stdout.write('\n'); return true; }
     } catch {
       // still starting
     }
-    const elapsed = Math.round((Date.now() - (deadline - timeoutMs)) / 1000);
+    const elapsed = Math.round((Date.now() - startMs) / 1000);
     process.stdout.write(
-      `\r  ${chalk.dim(`Health-checking /v1/models… ${elapsed}s`)}   `
+      `\r  ${chalk.dim(_serveStageLabel(elapsed) + ` (${elapsed}s)`)}   `
     );
     await new Promise(r => setTimeout(r, 8000));
   }
@@ -212,12 +217,11 @@ export async function serveCommand(config, args, chalk) {
 
   // ── Result ────────────────────────────────────────────────────────────────
   if (endpointReady) {
-    console.log(chalk.green('\n✓ Endpoint ready\n'));
+    console.log(chalk.green('\n  ✓ Endpoint ready\n'));
   } else {
-    console.log(chalk.yellow('\n⏳ Endpoint still starting\n'));
-    console.log(chalk.dim('  The deployment is provisioned but not yet responding.'));
-    console.log(chalk.dim(`  Check: badgr logs ${dep.deployment_id}`));
-    console.log(chalk.dim(`  Stop if it never starts: badgr down ${dep.deployment_id}`));
+    console.log(chalk.yellow('\n  ⏳ Endpoint still starting — model download may still be in progress.\n'));
+    console.log(`  ${chalk.bold('Stop billing now:')}   ${chalk.dim(`badgr down ${dep.deployment_id}`)}`);
+    console.log(`  ${chalk.bold('Continue watching:')}  ${chalk.dim(`badgr logs ${dep.deployment_id}`)}`);
     console.log();
   }
 

package/src/commands/test-run.js

@@ -0,0 +1,147 @@
+import { requireApiKey } from '../config.js';
+import { callApi, terminateDeployment } from '../api.js';
+import { addReceipt, generateReceiptId } from '../store.js';
+
+// max $1.50/hr × 2 min ≈ $0.05 total spend cap
+const TEST_MAX_PRICE  = 1.50;
+const TEST_MAX_RUNTIME_MS = 2 * 60 * 1000;
+const TEST_COMMAND    = ['python', '-c', "print('hello from badgr')"];
+const TEST_IMAGE      = 'python:3.11-slim';
+const EXPECTED_OUTPUT = 'hello from badgr';
+
+function step(chalk, ok, msg, detail = '') {
+  const icon = ok ? chalk.green('✓') : chalk.red('✗');
+  const suffix = detail ? chalk.dim(` — ${detail}`) : '';
+  console.log(`  ${icon} ${msg}${suffix}`);
+}
+
+async function pollStatus(config, depId, targetStatuses, timeoutMs) {
+  const deadline = Date.now() + timeoutMs;
+  while (Date.now() < deadline) {
+    await new Promise(r => setTimeout(r, 3000));
+    try {
+      const dep = await callApi(`/deployments/${depId}`, {
+        apiKey: config.apiKey,
+        baseUrl: config.baseUrl,
+      });
+      if (targetStatuses.has(dep.status)) return dep;
+    } catch { /* retry */ }
+  }
+  return null;
+}
+
+async function pollLogs(config, depId, expected, timeoutMs) {
+  const deadline = Date.now() + timeoutMs;
+  while (Date.now() < deadline) {
+    await new Promise(r => setTimeout(r, 4000));
+    try {
+      const data = await callApi(`/deployments/${depId}/logs`, {
+        apiKey: config.apiKey,
+        baseUrl: config.baseUrl,
+      });
+      const lines = data?.logs ?? [];
+      if (lines.some(l => l.includes(expected))) return true;
+    } catch { /* retry */ }
+  }
+  return false;
+}
+
+export async function testCommand(config, chalk) {
+  requireApiKey(config);
+
+  console.log(chalk.bold('\n⚡ Running end-to-end test\n'));
+  console.log(chalk.dim(`  Command:  ${TEST_COMMAND.join(' ')}`));
+  console.log(chalk.dim(`  Provider: RunPod (Tier 1 — reliable)`));
+  console.log(chalk.dim(`  Budget:   max $${TEST_MAX_PRICE.toFixed(2)}/hr · 2 minute cap (~$0.05 max)`));
+  console.log();
+
+  const rcptId = generateReceiptId();
+  let depId;
+
+  // ── 1. Provision ─────────────────────────────────────────────────────────
+  process.stdout.write(chalk.dim('  Provisioning GPU...'));
+  let dep;
+  try {
+    dep = await callApi('/run', {
+      method: 'POST',
+      apiKey: config.apiKey,
+      baseUrl: config.baseUrl,
+      body: {
+        command: TEST_COMMAND,
+        image:   TEST_IMAGE,
+        tier:    '1',
+        // smoke_test workload → cheapest reliable GPU with ≥4 GB VRAM
+        gpu:     'RTX_3080',
+        max_price_per_hour: TEST_MAX_PRICE,
+      },
+    });
+    depId = dep.deployment_id;
+    process.stdout.write('\n');
+    step(chalk, true, 'Provisioned', `${dep.deployment_id} on ${dep.gpu_type}`);
+  } catch (err) {
+    process.stdout.write('\n');
+    step(chalk, false, 'Provisioned', err.message);
+    console.log();
+    console.error(chalk.red('  Test failed — could not provision GPU.\n'));
+    process.exit(1);
+  }
+
+  // ── 2. Container started ─────────────────────────────────────────────────
+  process.stdout.write(chalk.dim('  Waiting for container to start...'));
+  const started = await pollStatus(
+    config, depId,
+    new Set(['running', 'failed', 'stopped', 'completed']),
+    TEST_MAX_RUNTIME_MS,
+  );
+  process.stdout.write('\n');
+
+  if (!started || started.status === 'failed') {
+    step(chalk, false, 'Container started', started?.status ?? 'timeout');
+    console.log();
+    console.error(chalk.red('  Test failed — container did not start.\n'));
+    try { await terminateDeployment(config, depId); } catch { /* best-effort */ }
+    process.exit(1);
+  }
+  step(chalk, true, 'Container started');
+
+  // ── 3. Command output ────────────────────────────────────────────────────
+  process.stdout.write(chalk.dim('  Checking command output...'));
+  const gotOutput = await pollLogs(config, depId, EXPECTED_OUTPUT, 90_000);
+  process.stdout.write('\n');
+  if (gotOutput) {
+    step(chalk, true, 'Command printed output');
+  } else {
+    step(chalk, false, 'Command printed output', 'not found in logs (logs may be buffered)');
+  }
+
+  // ── 4. Stop billing ──────────────────────────────────────────────────────
+  process.stdout.write(chalk.dim('  Stopping deployment...'));
+  let stopped = false;
+  try {
+    await terminateDeployment(config, depId);
+    stopped = true;
+  } catch { /* best-effort */ }
+  process.stdout.write('\n');
+  step(chalk, stopped, 'Billing stopped');
+
+  // ── 5. Receipt ───────────────────────────────────────────────────────────
+  addReceipt({
+    receiptId:    rcptId,
+    action:       'badgr test',
+    deploymentId: depId,
+    gpu:          dep.gpu_type,
+    status:       'test_complete',
+    createdAt:    new Date().toISOString(),
+  });
+  step(chalk, true, 'Receipt created', rcptId);
+
+  // ── Summary ──────────────────────────────────────────────────────────────
+  console.log();
+  const passed = stopped;
+  if (passed) {
+    console.log(chalk.green(chalk.bold('  ✓ Test passed\n')));
+  } else {
+    console.log(chalk.red(chalk.bold('  ✗ Test failed\n')));
+    process.exit(1);
+  }
+}

package/tests/commands.test.js

@@ -1,6 +1,7 @@
 import { describe, it, expect, vi } from 'vitest';
-import { parseRunArgs, classifyFailure } from '../src/commands/run.js';
+import { parseRunArgs, classifyFailure, inferWorkload } from '../src/commands/run.js';
 import { parseServeArgs } from '../src/commands/serve.js';
+import { testCommand } from '../src/commands/test-run.js';
 import { rankAlternatives, diffDescription, promptFallback } from '../src/fallback.js';
 
 describe('parseRunArgs', () => {
@@ -214,6 +215,45 @@ describe('parseServeArgs', () => {
   });
 });
 
+describe('testCommand', () => {
+  it('is a function', () => {
+    expect(typeof testCommand).toBe('function');
+  });
+});
+
+describe('inferWorkload', () => {
+  it('returns general for empty command', () => {
+    expect(inferWorkload([])).toBe('general');
+    expect(inferWorkload(null)).toBe('general');
+  });
+
+  it('returns smoke_test for short print() one-liners', () => {
+    expect(inferWorkload(['python', '-c', "print('hello')"])).toBe('smoke_test');
+    expect(inferWorkload(['python', '-c', "print('hello from badgr')"])).toBe('smoke_test');
+  });
+
+  it('returns general for a typical script', () => {
+    expect(inferWorkload(['python', 'script.py'])).toBe('general');
+    expect(inferWorkload(['python', 'run.py', '--epochs', '10'])).toBe('general');
+  });
+
+  it('returns lora_finetune for LoRA/fine-tuning keywords', () => {
+    expect(inferWorkload(['python', 'lora_train.py'])).toBe('lora_finetune');
+    expect(inferWorkload(['python', '-c', 'import peft; finetune()'])).toBe('lora_finetune');
+    expect(inferWorkload(['python', 'train.py', '--epochs', '3'])).toBe('lora_finetune');
+  });
+
+  it('returns image_gen for diffusion keywords', () => {
+    expect(inferWorkload(['python', 'stable_diff.py'])).toBe('image_gen');
+    expect(inferWorkload(['python', 'run_sdxl.py'])).toBe('image_gen');
+  });
+
+  it('returns inference_small for vllm/tgi', () => {
+    expect(inferWorkload(['python', '-m', 'vllm.entrypoints.openai.api_server'])).toBe('inference_small');
+    expect(inferWorkload(['python', 'serve_tgi.py'])).toBe('inference_small');
+  });
+});
+
 describe('promptFallback output', () => {
   it('does not print a Difference line', async () => {
     const lines = [];