badgr-cli 1.0.42 → 1.0.43

package/tests/template.test.js

@@ -327,18 +327,21 @@ describe('parseTemplateOverrides', () => {
 // Sequence for serveCommand:
 //   callApi('/serve', ...) → dep  (via callWithFallback → callApi)
 //   callApi('/deployments/<id>', ...) → { status: 'running' }  (pre-health check)
-//   global.fetch(endpoint + health_path) → { ok: true }
+//   callApi('/deployments/<id>', ...) → { status: 'running', endpoint_ready: true }  (waitForEndpoint poll)
+//
+// Readiness comes entirely from Badgr's own deployment status — the CLI never
+// fetches the pod/RunPod-proxy endpoint directly.
 
-function setupServe(depOverrides = {}) {
+function setupServe(depOverrides = {}, readyOverrides = {}) {
   api.callApi
     .mockResolvedValueOnce(makeServeDep(depOverrides))  // POST /serve
-    .mockResolvedValueOnce({ status: 'running' });      // pre-health dep status
-  global.fetch = vi.fn().mockResolvedValue({ ok: true, status: 200, json: async () => ({}) });
+    .mockResolvedValueOnce({ status: 'running' })        // pre-health dep status
+    .mockResolvedValueOnce({ status: 'running', endpoint_ready: true, ...readyOverrides }); // waitForEndpoint poll
 }
 
 describe('badgr serve template <name>', () => {
-  it('vllm: dispatches with correct image, max_cost_usd, and health check on /v1/models', async () => {
-    setupServe();
+  it('vllm: dispatches with correct image, max_cost_usd, and readiness on /v1/models', async () => {
+    setupServe({}, { health_path: '/v1/models' });
     const p = serveCommand(config, ['template', 'vllm', '--max-cost', '5'], chalk);
     await vi.advanceTimersByTimeAsync(5000);
     await p;
@@ -347,22 +350,17 @@ describe('badgr serve template <name>', () => {
     expect(route).toBe('/serve');
     expect(opts.body.image).toBe('vllm/vllm-openai:latest');
     expect(opts.body.max_cost_usd).toBe(5);
-    // health check fetch should use the /v1/models path from the template
-    const fetchUrl = global.fetch.mock.calls[0]?.[0] ?? '';
-    expect(fetchUrl).toContain('/v1/models');
     expect(process.exitCode).toBeFalsy();
   });
 
-  it('comfyui: dispatches with ComfyUI image and health check on /system_stats', async () => {
-    setupServe();
+  it('comfyui: dispatches with ComfyUI image and readiness on /system_stats', async () => {
+    setupServe({}, { health_path: '/system_stats' });
     const p = serveCommand(config, ['template', 'comfyui', '--max-cost', '3'], chalk);
     await vi.advanceTimersByTimeAsync(5000);
     await p;
 
     const [, opts] = api.callApi.mock.calls[0];
     expect(opts.body.image).toBe('yanwk/comfyui-boot:cu126-megapak');
-    const fetchUrl = global.fetch.mock.calls[0]?.[0] ?? '';
-    expect(fetchUrl).toContain('/system_stats');
     expect(process.exitCode).toBeFalsy();
   });
 

package/tests/workload-templates.test.js

@@ -485,9 +485,44 @@ describe('trainCommand', () => {
     expect(body.max_runtime_seconds).toBe(60 * 60);
   });
 
-  it('uses unsloth image for unsloth config', async () => {
+  it('blocks unsloth config instead of running a mismatched command', async () => {
     fs.existsSync.mockReturnValue(true);
     fs.readFileSync.mockReturnValue('# unsloth training config\nmodel: llama\n');
+
+    await trainCommand(config, ['config.yaml', '--detach'], chalk);
+
+    expect(process.exitCode).toBe(1);
+    expect(fallback.callWithFallback).not.toHaveBeenCalled();
+  });
+
+  it('blocks generic (unrecognized) config instead of guessing a command', async () => {
+    fs.existsSync.mockReturnValue(true);
+    fs.readFileSync.mockReturnValue('some_key: some_value\nother: 123\n');
+
+    await trainCommand(config, ['config.yaml', '--detach'], chalk);
+
+    expect(process.exitCode).toBe(1);
+    expect(fallback.callWithFallback).not.toHaveBeenCalled();
+  });
+
+  it('uses the trl CLI command for a trl-detected config', async () => {
+    fs.existsSync.mockReturnValue(true);
+    fs.readFileSync.mockReturnValue('trainer: SFTTrainer\nmodel: mistral\n');
+    fallback.callWithFallback.mockResolvedValue(makeRunDep());
+
+    const p = trainCommand(config, ['config.yaml', '--detach'], chalk);
+    await vi.advanceTimersByTimeAsync(100);
+    await p;
+
+    const bodyBuilder = fallback.callWithFallback.mock.calls[0][2];
+    const body = bodyBuilder();
+    expect(body.image).toBe('huggingface/trl-source:latest');
+    expect(body.command[2]).toContain('trl sft --config /tmp/config.yaml');
+  });
+
+  it('uses the axolotl CLI command for an axolotl-detected config', async () => {
+    fs.existsSync.mockReturnValue(true);
+    fs.readFileSync.mockReturnValue('base_model: meta-llama/Llama-2-7b-hf\nsequence_len: 2048\n');
     fallback.callWithFallback.mockResolvedValue(makeRunDep());
 
     const p = trainCommand(config, ['config.yaml', '--detach'], chalk);
@@ -496,7 +531,7 @@ describe('trainCommand', () => {
 
     const bodyBuilder = fallback.callWithFallback.mock.calls[0][2];
     const body = bodyBuilder();
-    expect(body.image).toBe('unslothai/unsloth:latest');
+    expect(body.command[2]).toContain('axolotl train /tmp/config.yaml');
   });
 });