badgr-cli 1.0.20 → 1.0.22

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "badgr-cli",
-  "version": "1.0.20",
+  "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,13 +24,21 @@ ${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('# Simplest — Badgr picks the GPU:')}
+  ${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
 
+  ${chalk.dim('# Opt into cheaper budget providers (Vast.ai etc.):')}
+  badgr run python train.py --cheap
+  badgr serve meta-llama/Llama-3.1-8B-Instruct --cheap
+
   ${chalk.dim('# Pin a specific GPU:')}
   badgr run python train.py --gpu A100
   badgr serve meta-llama/Llama-3.1-8B-Instruct --gpu L40S
@@ -44,10 +53,12 @@ ${chalk.bold('EXAMPLES')}
   badgr receipts dep-abc123
 
 ${chalk.bold('badgr run OPTIONS')}
-  --gpu <type>           GPU type (default: auto — Badgr picks cheapest available)
+  --gpu <type>           GPU type (default: auto — Badgr picks best available on RunPod)
+  --cheap                Search budget providers too (Vast.ai etc.) for lower prices
   --image <image>        Docker image (default: python:3.11-slim)
   --count <n>            Number of GPUs (default: 1)
   --region US|EU|AU      Region preference
+  --tier 1|2             Provider tier: 1 = reliable (default), 2 = budget
   --max-price <$/hr>     Hard spend cap per GPU-hour
   --max-runtime <min>    Auto-stop after N minutes (recommended)
   --max-cost <$>         Auto-stop when spend reaches this amount
@@ -55,8 +66,10 @@ ${chalk.bold('badgr run OPTIONS')}
 
 ${chalk.bold('badgr serve OPTIONS')}
   --gpu <type>           GPU type (default: auto — inferred from model size)
+  --cheap                Search budget providers too (Vast.ai etc.) for lower prices
   --count <n>            Number of GPUs (default: 1)
   --region US|EU|AU      Region preference
+  --tier 1|2             Provider tier: 1 = reliable (default), 2 = budget
   --max-price <$/hr>     Hard spend cap per GPU-hour
   --no-wait              Skip endpoint health check
 
@@ -90,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

@@ -19,6 +19,7 @@ export function parseRunArgs(args) {
     if (args[i] === '--count')          { flags.count = parseInt(args[++i], 10); i++; continue; }
     if (args[i] === '--region')         { flags.region = args[++i]; i++; continue; }
     if (args[i] === '--tier')           { flags.tier = args[++i]; i++; continue; }
+    if (args[i] === '--cheap')          { flags.cheap = true; i++; continue; }
     if (args[i] === '--max-price')      { flags.maxPrice = parseFloat(args[++i]); i++; continue; }
     if (args[i] === '--name')           { flags.name = args[++i]; i++; continue; }
     if (args[i] === '--detach')          { flags.detach = true; i++; continue; }
@@ -269,11 +270,53 @@ async function attachToJob(config, depId, { chalk, maxRuntimeMs = null, maxCost
   }
 }
 
-// Ask the backend for the cheapest available GPU right now.
-// Returns { gpu, region, price } or null when nothing is available.
-async function findAutoGpu(config, chalk) {
+/**
+ * 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' });
+    const params = new URLSearchParams({ max_price: '10', tier, workload });
     return await callApi(`/capacity/auto?${params}`, {
       apiKey: config.apiKey,
       baseUrl: config.baseUrl,
@@ -307,20 +350,31 @@ export async function runCommand(config, args, chalk) {
   const maxRuntimeMs = flags.maxRuntime ? flags.maxRuntime * 60 * 1000 : null;
   const maxCost      = flags.maxCost ?? null;
 
+  // Resolve effective tier: --cheap and --tier 2 opt into budget providers;
+  // everything else defaults to tier 1 (RunPod only) for reliability.
+  const effectiveTier = (flags.cheap || flags.tier === '2' || flags.tier === 'tier2' || flags.tier === 'tier-2')
+    ? '2'
+    : (flags.tier || '1');
+
   // ── Auto GPU selection (no --gpu specified) ────────────────────────────────
   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);
+      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`));
@@ -335,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;
@@ -385,7 +439,7 @@ export async function runCommand(config, args, chalk) {
       ...(effectiveRegion ? { region: effectiveRegion } : {}),
       max_price_per_hour: flags.maxPrice,
       name: flags.name,
-      ...(flags.tier ? { tier: flags.tier } : {}),
+      tier: effectiveTier,
     };
   }
 

package/src/commands/serve.js

@@ -18,6 +18,7 @@ export function parseServeArgs(args) {
     if (args[i] === '--count')      { flags.count = parseInt(args[++i], 10); i++; continue; }
     if (args[i] === '--region')     { flags.region = args[++i]; i++; continue; }
     if (args[i] === '--tier')       { flags.tier = args[++i]; i++; continue; }
+    if (args[i] === '--cheap')      { flags.cheap = true; i++; continue; }
     if (args[i] === '--max-price')  { flags.maxPrice = parseFloat(args[++i]); i++; continue; }
     if (args[i] === '--name')       { flags.name = args[++i]; i++; continue; }
     if (args[i] === '--no-wait')    { flags.noWait = true; i++; continue; }
@@ -27,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));
   }
@@ -64,9 +70,15 @@ export async function serveCommand(config, args, chalk) {
   const gpu = flags.gpu ? flags.gpu.toUpperCase().replace('-', '_') : 'AUTO';
   const gpuLabel = gpu === 'AUTO' ? 'auto' : gpu;
 
+  // Default to tier 1 (RunPod) for reliability; --cheap or --tier 2 opts into budget providers.
+  const effectiveTier = (flags.cheap || flags.tier === '2' || flags.tier === 'tier2' || flags.tier === 'tier-2')
+    ? '2'
+    : (flags.tier || '1');
+
   console.log(chalk.bold('\nServing model\n'));
   console.log(`  ${chalk.bold('Model:')}  ${model}`);
   console.log(`  ${chalk.bold('GPU:')}    ${gpuLabel}`);
+  if (effectiveTier === '2') console.log(`  ${chalk.dim('(budget mode — searching all providers)')}`);
   console.log();
   process.stdout.write(chalk.dim('  Finding GPU capacity...\n'));
 
@@ -80,7 +92,7 @@ export async function serveCommand(config, args, chalk) {
       ...(regionOverride || effectiveRegion ? { region: regionOverride || effectiveRegion } : {}),
       max_price_per_hour: flags.maxPrice,
       name: flags.name,
-      ...(flags.tier ? { tier: flags.tier } : {}),
+      tier: effectiveTier,
     };
   }
 
@@ -205,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', () => {
@@ -71,6 +72,21 @@ describe('parseRunArgs', () => {
     const { flags } = parseRunArgs(['python', 'train.py', '--max-cost', '5.00']);
     expect(flags.maxCost).toBe(5.0);
   });
+
+  it('parses --cheap flag', () => {
+    const { flags } = parseRunArgs(['python', 'train.py', '--cheap']);
+    expect(flags.cheap).toBe(true);
+  });
+
+  it('--cheap defaults to falsy when not passed', () => {
+    const { flags } = parseRunArgs(['python', 'train.py']);
+    expect(flags.cheap).toBeFalsy();
+  });
+
+  it('parses --tier 2 flag', () => {
+    const { flags } = parseRunArgs(['python', 'train.py', '--tier', '2']);
+    expect(flags.tier).toBe('2');
+  });
 });
 
 describe('classifyFailure', () => {
@@ -187,6 +203,55 @@ describe('parseServeArgs', () => {
     const { model } = parseServeArgs(['--gpu', 'RTX_4090']);
     expect(model).toBeNull();
   });
+
+  it('parses --cheap flag', () => {
+    const { flags } = parseServeArgs(['my/model', '--cheap']);
+    expect(flags.cheap).toBe(true);
+  });
+
+  it('--cheap defaults to falsy when not passed', () => {
+    const { flags } = parseServeArgs(['my/model']);
+    expect(flags.cheap).toBeFalsy();
+  });
+});
+
+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', () => {