badgr-cli 1.0.45 → 1.0.47

package/tests/train-lora-dataset.test.js

@@ -0,0 +1,176 @@
+/**
+ * badgr train lora — dataset validation, upload dedup, and --resume.
+ *
+ * A malformed local .jsonl dataset must be rejected before it's uploaded or
+ * a GPU is provisioned. A dataset whose content hasn't changed since a prior
+ * run must not be re-uploaded (upload cache keyed by content hash). --resume
+ * must thread resume_from_checkpoint_url into the job input.
+ */
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { trainLoraCommand, parseTrainLoraArgs } from '../src/commands/train.js';
+
+vi.mock('../src/api.js', async (importOriginal) => {
+  const actual = await importOriginal();
+  return { ...actual, callApi: vi.fn(), uploadBlob: actual.uploadBlob };
+});
+
+vi.mock('../src/store.js', () => ({
+  addReceipt:          vi.fn(),
+  generateReceiptId:   vi.fn(() => 'rcpt-ds-001'),
+  getCachedUploadId:   vi.fn(() => null),
+  setCachedUploadId:   vi.fn(),
+}));
+
+vi.mock('fs', async (importOriginal) => {
+  const actual = await importOriginal();
+  return {
+    ...actual,
+    readFileSync: vi.fn(),
+    existsSync:   vi.fn(() => true),
+  };
+});
+
+import * as api   from '../src/api.js';
+import * as store from '../src/store.js';
+import * as fs    from 'fs';
+
+const chalk = {
+  bold: s => s, dim: s => s, red: s => s, yellow: s => s, green: s => s, cyan: s => s,
+};
+
+const config = { apiKey: 'sk-test', baseUrl: 'https://api.test/v1' };
+
+beforeEach(() => {
+  process.exitCode = undefined;
+  vi.spyOn(console, 'log').mockImplementation(() => {});
+  vi.spyOn(console, 'error').mockImplementation(() => {});
+  vi.spyOn(process.stdout, 'write').mockImplementation(() => true);
+  vi.clearAllMocks();
+  fs.existsSync.mockReturnValue(true);
+  store.getCachedUploadId.mockReturnValue(null);
+  // POST /jobs returns the created job; the subsequent GET /jobs/{id} poll
+  // (pollJobUntilTerminal, 15s real-world interval) must resolve terminal
+  // immediately so these tests don't wait on real timers.
+  api.callApi.mockImplementation((path) => {
+    if (path === '/jobs') return Promise.resolve({ job_id: 'job-1', status: 'queued' });
+    return Promise.resolve({ job_id: 'job-1', status: 'completed', output: {} });
+  });
+  vi.useFakeTimers();
+});
+
+afterEach(() => {
+  vi.useRealTimers();
+  vi.restoreAllMocks();
+  process.exitCode = undefined;
+});
+
+async function runTrainLora(args) {
+  const promise = trainLoraCommand(config, args, chalk);
+  await vi.advanceTimersByTimeAsync(20_000);
+  return promise;
+}
+
+describe('parseTrainLoraArgs --resume', () => {
+  it('parses --resume', () => {
+    const flags = parseTrainLoraArgs(['--base-model', 'x', '--resume', 'https://cdn/ckpt.tar.gz']);
+    expect(flags.resume).toBe('https://cdn/ckpt.tar.gz');
+  });
+});
+
+describe('trainLoraCommand dataset validation', () => {
+  it('rejects a .jsonl dataset with an invalid line before uploading', async () => {
+    fs.readFileSync.mockReturnValue(Buffer.from('{"a":1}\nnot json\n{"b":2}\n'));
+    const fetchSpy = vi.spyOn(globalThis, 'fetch');
+
+    await trainLoraCommand(config, [
+      '--base-model', 'mistralai/Mistral-7B-v0.1',
+      '--dataset', './bad.jsonl',
+      '--max-cost', '10',
+    ], chalk);
+
+    expect(process.exitCode).toBe(1);
+    expect(fetchSpy).not.toHaveBeenCalled();
+    const logged = console.error.mock.calls.flat().join('\n');
+    expect(logged).toContain('not valid JSONL');
+  });
+
+  it('rejects an empty .jsonl dataset', async () => {
+    fs.readFileSync.mockReturnValue(Buffer.from('   \n  \n'));
+
+    await trainLoraCommand(config, [
+      '--base-model', 'mistralai/Mistral-7B-v0.1',
+      '--dataset', './empty.jsonl',
+      '--max-cost', '10',
+    ], chalk);
+
+    expect(process.exitCode).toBe(1);
+    const logged = console.error.mock.calls.flat().join('\n');
+    expect(logged).toContain('empty');
+  });
+
+  it('uploads a valid .jsonl dataset and caches the resulting upload id', async () => {
+    fs.readFileSync.mockReturnValue(Buffer.from('{"a":1}\n{"b":2}\n'));
+    const fetchSpy = vi.spyOn(globalThis, 'fetch').mockResolvedValue({
+      ok: true, json: async () => ({ upload_id: 'up_new_1' }),
+    });
+
+    await runTrainLora([
+      '--base-model', 'mistralai/Mistral-7B-v0.1',
+      '--dataset', './good.jsonl',
+      '--max-cost', '10',
+    ]);
+
+    expect(fetchSpy).toHaveBeenCalledTimes(1);
+    expect(store.setCachedUploadId).toHaveBeenCalledWith(expect.any(String), 'up_new_1', expect.any(Object));
+    expect(api.callApi).toHaveBeenCalledWith('/jobs', expect.objectContaining({
+      body: expect.objectContaining({ input: expect.objectContaining({ dataset_file_id: 'up_new_1' }) }),
+    }));
+  });
+
+  it('skips re-upload when the dataset content matches a cached hash', async () => {
+    fs.readFileSync.mockReturnValue(Buffer.from('{"a":1}\n{"b":2}\n'));
+    store.getCachedUploadId.mockReturnValue('up_cached_1');
+    const fetchSpy = vi.spyOn(globalThis, 'fetch');
+
+    await runTrainLora([
+      '--base-model', 'mistralai/Mistral-7B-v0.1',
+      '--dataset', './good.jsonl',
+      '--max-cost', '10',
+    ]);
+
+    expect(fetchSpy).not.toHaveBeenCalled();
+    expect(api.callApi).toHaveBeenCalledWith('/jobs', expect.objectContaining({
+      body: expect.objectContaining({ input: expect.objectContaining({ dataset_file_id: 'up_cached_1' }) }),
+    }));
+    const logged = console.log.mock.calls.flat().join('\n');
+    expect(logged).toContain('reusing upload');
+  });
+});
+
+describe('trainLoraCommand --resume', () => {
+  it('threads resume_from_checkpoint_url into the job input', async () => {
+    await runTrainLora([
+      '--base-model', 'mistralai/Mistral-7B-v0.1',
+      '--dataset', 'https://example.com/data.jsonl',
+      '--max-cost', '10',
+      '--resume', 'https://cdn/ckpt.tar.gz',
+    ]);
+
+    expect(api.callApi).toHaveBeenCalledWith('/jobs', expect.objectContaining({
+      body: expect.objectContaining({
+        input: expect.objectContaining({ resume_from_checkpoint_url: 'https://cdn/ckpt.tar.gz' }),
+      }),
+    }));
+  });
+
+  it('omits resume_from_checkpoint_url when --resume is not given', async () => {
+    await runTrainLora([
+      '--base-model', 'mistralai/Mistral-7B-v0.1',
+      '--dataset', 'https://example.com/data.jsonl',
+      '--max-cost', '10',
+    ]);
+
+    const call = api.callApi.mock.calls.find(c => c[0] === '/jobs');
+    expect(call[1].body.input.resume_from_checkpoint_url).toBeUndefined();
+  });
+});

package/tests/workload-spec.test.js

@@ -0,0 +1,180 @@
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
+import { mkdtempSync, writeFileSync, rmSync } from 'fs';
+import { tmpdir } from 'os';
+import { join } from 'path';
+import {
+  parseWorkloadYaml,
+  parseInputEntry,
+  validateWorkloadSpec,
+  WorkloadSpecError,
+} from '../src/workloadSpec.js';
+
+let dir;
+
+beforeEach(() => {
+  dir = mkdtempSync(join(tmpdir(), 'badgr-batch-test-'));
+});
+
+afterEach(() => {
+  rmSync(dir, { recursive: true, force: true });
+});
+
+function writeYaml(name, content) {
+  const path = join(dir, name);
+  writeFileSync(path, content);
+  return path;
+}
+
+describe('parseInputEntry', () => {
+  it('splits local:container on the last colon', () => {
+    const { localPath, containerPath } = parseInputEntry('./policy.py:/inputs/policy.py');
+    expect(localPath).toBe('./policy.py');
+    expect(containerPath).toBe('/inputs/policy.py');
+  });
+
+  it('rejects a non-absolute container path', () => {
+    expect(() => parseInputEntry('./policy.py:inputs/policy.py')).toThrow(WorkloadSpecError);
+  });
+
+  it('rejects an entry with no colon', () => {
+    expect(() => parseInputEntry('./policy.py')).toThrow(WorkloadSpecError);
+  });
+});
+
+describe('validateWorkloadSpec', () => {
+  it('requires name, image, command, max_cost, max_runtime_minutes', () => {
+    const errors = validateWorkloadSpec({});
+    expect(errors).toContain('name is required');
+    expect(errors).toContain('image is required');
+    expect(errors).toContain('command is required and must be a non-empty list');
+    expect(errors).toContain('max_cost is required and must be a positive number');
+    expect(errors).toContain('max_runtime_minutes is required and must be a positive number');
+  });
+
+  it('passes on a minimal valid spec', () => {
+    const errors = validateWorkloadSpec({
+      name: 'toy', image: 'busybox', command: ['echo', 'hi'],
+      max_cost: 5, max_runtime_minutes: 10,
+    });
+    expect(errors).toEqual([]);
+  });
+
+  it('flags a malformed inputs entry', () => {
+    const errors = validateWorkloadSpec({
+      name: 'toy', image: 'busybox', command: ['echo'],
+      max_cost: 5, max_runtime_minutes: 10,
+      inputs: ['./policy.py'],
+    });
+    expect(errors.some(e => e.includes('inputs entry'))).toBe(true);
+  });
+
+  it('flags outputs that are not absolute paths', () => {
+    const errors = validateWorkloadSpec({
+      name: 'toy', image: 'busybox', command: ['echo'],
+      max_cost: 5, max_runtime_minutes: 10,
+      outputs: ['outputs/metrics.json'],
+    });
+    expect(errors).toContain('outputs must be a list of absolute container paths');
+  });
+
+  it('requires success_metric.file and key when success_metric is set', () => {
+    const errors = validateWorkloadSpec({
+      name: 'toy', image: 'busybox', command: ['echo'],
+      max_cost: 5, max_runtime_minutes: 10,
+      success_metric: {},
+    });
+    expect(errors.some(e => e.includes('success_metric.file'))).toBe(true);
+    expect(errors.some(e => e.includes('success_metric.key'))).toBe(true);
+  });
+});
+
+describe('parseWorkloadYaml', () => {
+  it('parses the minimal toy-robot-eval example from the spec', () => {
+    const path = writeYaml('workload.yml', `
+name: toy-robot-eval
+image: badgr/physical-ai-eval-lite:latest
+
+command:
+  - python
+  - run_eval.py
+  - --policy
+  - /inputs/policy.py
+  - --scenarios
+  - /inputs/scenarios
+
+inputs:
+  - ./policy.py:/inputs/policy.py
+  - ./scenarios:/inputs/scenarios
+
+outputs:
+  - /outputs/metrics.json
+  - /outputs/videos
+  - /outputs/logs
+
+max_cost: 20
+max_runtime_minutes: 60
+
+success_metric:
+  file: /outputs/metrics.json
+  key: pass_rate
+  higher_is_better: true
+`);
+    const spec = parseWorkloadYaml(path);
+    expect(spec.name).toBe('toy-robot-eval');
+    expect(spec.image).toBe('badgr/physical-ai-eval-lite:latest');
+    expect(spec.command).toEqual(['python', 'run_eval.py', '--policy', '/inputs/policy.py', '--scenarios', '/inputs/scenarios']);
+    expect(spec.inputs).toHaveLength(2);
+    expect(spec.inputs[0].containerPath).toBe('/inputs/policy.py');
+    expect(spec.inputs[0].localPath).toBe(join(dir, 'policy.py'));
+    expect(spec.outputs).toEqual(['/outputs/metrics.json', '/outputs/videos', '/outputs/logs']);
+    expect(spec.maxCost).toBe(20);
+    expect(spec.maxRuntimeMinutes).toBe(60);
+    expect(spec.successMetric).toEqual({ file: '/outputs/metrics.json', key: 'pass_rate', higherIsBetter: true });
+  });
+
+  it('defaults higherIsBetter to true when omitted', () => {
+    const path = writeYaml('workload.yml', `
+name: toy
+image: busybox
+command: [echo, hi]
+max_cost: 1
+max_runtime_minutes: 1
+success_metric:
+  file: /outputs/metrics.json
+  key: score
+`);
+    const spec = parseWorkloadYaml(path);
+    expect(spec.successMetric.higherIsBetter).toBe(true);
+  });
+
+  it('defaults inputs/outputs/env to empty', () => {
+    const path = writeYaml('workload.yml', `
+name: toy
+image: busybox
+command: [echo, hi]
+max_cost: 1
+max_runtime_minutes: 1
+`);
+    const spec = parseWorkloadYaml(path);
+    expect(spec.inputs).toEqual([]);
+    expect(spec.outputs).toEqual([]);
+    expect(spec.env).toEqual({});
+    expect(spec.successMetric).toBeNull();
+  });
+
+  it('throws WorkloadSpecError with all problems on an invalid file', () => {
+    const path = writeYaml('workload.yml', `
+image: busybox
+`);
+    expect(() => parseWorkloadYaml(path)).toThrow(WorkloadSpecError);
+  });
+
+  it('throws WorkloadSpecError on malformed YAML', () => {
+    const path = writeYaml('workload.yml', '{ not: valid: yaml: [');
+    expect(() => parseWorkloadYaml(path)).toThrow(WorkloadSpecError);
+  });
+
+  it('throws WorkloadSpecError when the file does not exist', () => {
+    expect(() => parseWorkloadYaml(join(dir, 'nope.yml'))).toThrow(WorkloadSpecError);
+  });
+});