From 2f4b3e4002b704adfcf542de14368fe00cacbfa3 Mon Sep 17 00:00:00 2001 From: zhiweiz Date: Wed, 1 Apr 2026 17:44:42 +0000 Subject: [PATCH 1/5] Add cross-job checkpoint promotion script When the original RLOR trainer job has been deleted, `forge rft promote` fails with HTTP 404 because it tries to list checkpoints on a job that no longer exists. This script works around the problem by: 1. Reading checkpoints.jsonl to find the DCP state_path for the target step 2. Spinning up a temporary service-mode RLOR trainer with the same base model 3. Loading the cross-job DCP checkpoint into the temporary trainer 4. Exporting a sampler (HF-format) checkpoint from the loaded weights 5. Promoting the sampler checkpoint to a deployable Fireworks model 6. Cleaning up (deleting the temporary trainer) Supports: - Auto-inference of training shape from base model name - Direct promotion fallback when sampler_path exists - Dry-run mode to preview the plan without creating resources - Automatic cleanup of temporary trainer on success or failure Co-authored-by: zhiweiz --- cross_job_promote.py | 487 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 487 insertions(+) create mode 100644 cross_job_promote.py diff --git a/cross_job_promote.py b/cross_job_promote.py new file mode 100644 index 0000000..f84ef6e --- /dev/null +++ b/cross_job_promote.py @@ -0,0 +1,487 @@ +#!/usr/bin/env python3 +"""Cross-job checkpoint promotion for Fireworks RFT. + +When the original RLOR trainer job has been deleted, `forge rft promote` +fails with HTTP 404 because it tries to list checkpoints on a job that +no longer exists. + +This script works around the problem by: + 1. Reading checkpoints.jsonl to find the DCP state_path for the target step. + 2. Spinning up a *temporary* service-mode RLOR trainer with the same base model. + 3. Loading the cross-job DCP checkpoint into the temporary trainer. + 4. Exporting a sampler (HF-format) checkpoint from the loaded weights. + 5. Promoting the sampler checkpoint to a deployable Fireworks model. + 6. Cleaning up (deleting the temporary trainer). + +Prerequisites: + - The Fireworks Training SDK: pip install fireworks-ai[training] + (provides fireworks.training.sdk with TrainerJobManager, FiretitanServiceClient, etc.) + - FIREWORKS_API_KEY environment variable set + +Usage: + python cross_job_promote.py \\ + --checkpoints-jsonl rl_logs/rft-clinical-reference-v2-0-1/checkpoints.jsonl \\ + --step 8 \\ + --output-model-id rft-clinical-reference-v2-0-1 + + # With explicit training shape: + python cross_job_promote.py \\ + --checkpoints-jsonl rl_logs/rft-clinical-reference-v2-0-1/checkpoints.jsonl \\ + --step 8 \\ + --output-model-id rft-clinical-reference-v2-0-1 \\ + --training-shape-id accounts/fireworks/trainingShapes/qwen3-8b-128k +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from pathlib import Path +from typing import Any + + +def _check_sdk(): + """Check that the Training SDK is installed and provide instructions if not.""" + try: + from fireworks.training.sdk import ( # noqa: F401 + FiretitanServiceClient, + FireworksClient, + TrainerJobConfig, + TrainerJobManager, + ) + except ImportError: + print( + "Error: The Fireworks Training SDK is required but not installed.\n" + "\n" + "Install it with:\n" + " pip install fireworks-ai[training]\n" + "\n" + "The Training SDK provides TrainerJobManager, FiretitanServiceClient,\n" + "and other components needed for cross-job checkpoint operations.\n" + "\n" + "If the [training] extra is not available in your version, check:\n" + " https://docs.fireworks.ai/api-reference/training-sdk/overview" + ) + sys.exit(1) + + +# --------------------------------------------------------------------------- +# checkpoints.jsonl helpers +# --------------------------------------------------------------------------- + +def load_checkpoints(jsonl_path: str) -> list[dict[str, Any]]: + """Load all checkpoint entries from a checkpoints.jsonl file.""" + entries = [] + with open(jsonl_path) as f: + for line in f: + line = line.strip() + if line: + entries.append(json.loads(line)) + return entries + + +def find_checkpoint(entries: list[dict], step: int | None) -> dict: + """Find a checkpoint entry by step number. + + If step is None, returns the latest entry that has a state_path + (for cross-job load) or sampler_path (for direct promotion). + """ + if step is not None: + for entry in entries: + if entry.get("step") == step: + return entry + available = sorted(set(e.get("step") for e in entries if e.get("step") is not None)) + raise ValueError( + f"No checkpoint found at step {step}. " + f"Available steps: {available}" + ) + + for entry in reversed(entries): + if entry.get("sampler_path") or entry.get("state_path"): + return entry + + raise ValueError("No checkpoint with state_path or sampler_path found") + + +# --------------------------------------------------------------------------- +# Training shape auto-inference +# --------------------------------------------------------------------------- + +SHAPE_MAP = { + "qwen3-4b": "accounts/fireworks/trainingShapes/qwen3-4b-minimum", + "qwen3-8b": "accounts/fireworks/trainingShapes/qwen3-8b-128k", + "qwen3-32b": "accounts/fireworks/trainingShapes/qwen3-32b-65k", + "qwen3-30b-a3b": "accounts/fireworks/trainingShapes/qwen3-30b-a3b-instruct-2507-128k", + "qwen3-235b": "accounts/fireworks/trainingShapes/qwen3-235b-2507-instruct-128k", + "qwen3-vl-8b": "accounts/fireworks/trainingShapes/qwen3-vl-8b-65k", + "llama-v3p3-70b": "accounts/fireworks/trainingShapes/llama70b-policy", +} + + +def infer_training_shape(base_model: str) -> str: + """Infer a training shape ID from the base model name.""" + model_lower = base_model.lower() + for key, shape in SHAPE_MAP.items(): + if key in model_lower: + return shape + raise ValueError( + f"Cannot infer training shape for base model '{base_model}'.\n" + f"Please specify --training-shape-id explicitly.\n" + f"Known patterns: {list(SHAPE_MAP.keys())}\n" + f"See: https://docs.fireworks.ai/fine-tuning/training-sdk/training-shapes" + ) + + +# --------------------------------------------------------------------------- +# Direct promotion (when sampler_path exists) +# --------------------------------------------------------------------------- + +def try_direct_promote( + api_key: str, + source_job_id: str, + sampler_path: str, + output_model_id: str, + base_url: str, +) -> dict | None: + """Attempt to promote directly using an existing sampler_path. + + Returns the model dict on success, or None if the source job is gone. + """ + from fireworks.training.sdk import FireworksClient + + client = FireworksClient(api_key=api_key, base_url=base_url) + try: + return client.promote_checkpoint( + job_id=source_job_id, + checkpoint_id=sampler_path, + output_model_id=output_model_id, + ) + except Exception as e: + err_str = str(e) + if "404" in err_str or "not found" in err_str.lower(): + return None + raise + + +# --------------------------------------------------------------------------- +# Cross-job promotion via temporary trainer +# --------------------------------------------------------------------------- + +def cross_job_promote( + api_key: str, + checkpoint_entry: dict, + output_model_id: str, + training_shape_id: str, + base_url: str = "https://api.fireworks.ai", +) -> dict: + """Promote a cross-job DCP checkpoint by spinning up a temporary trainer. + + Steps: + 1. Resolve training shape → pinned version + 2. Create temporary service-mode RLOR trainer + 3. Connect FiretitanServiceClient + training client + 4. resolve_checkpoint_path() + load_state() from the dead job's DCP + 5. save_weights_for_sampler_ext() → base sampler checkpoint + 6. promote_checkpoint() → deployable Fireworks model + 7. Delete the temporary trainer + """ + from fireworks.training.sdk import ( + FiretitanServiceClient, + TrainerJobConfig, + TrainerJobManager, + ) + + base_model = checkpoint_entry["base_model"] + source_job_id = checkpoint_entry["source_job_id"] + checkpoint_name = checkpoint_entry["name"] + lora_rank = checkpoint_entry.get("lora_rank", 0) + + print(f"\n{'='*60}") + print("Cross-job checkpoint promotion") + print(f"{'='*60}") + print(f" Base model: {base_model}") + print(f" Source job: {source_job_id}") + print(f" Checkpoint: {checkpoint_name}") + print(f" Step: {checkpoint_entry.get('step')}") + print(f" Output model: {output_model_id}") + print(f" Training shape: {training_shape_id}") + + mgr = TrainerJobManager(api_key=api_key, base_url=base_url) + + # 1. Resolve training shape + print(f"\n[1/7] Resolving training shape...") + profile = mgr.resolve_training_profile(training_shape_id) + print(f" → {profile.training_shape_version}") + + # 2. Create temporary trainer + config = TrainerJobConfig( + base_model=base_model, + training_shape_ref=profile.training_shape_version, + lora_rank=lora_rank, + learning_rate=1e-5, + display_name="cross-job-promote-temp", + ) + + print("[2/7] Creating temporary service-mode trainer (this may take a few minutes)...") + endpoint = mgr.create_and_wait(config, timeout_s=900) + temp_job_id = endpoint.job_id + print(f" Trainer ready: {temp_job_id}") + print(f" Endpoint: {endpoint.base_url}") + + try: + # 3. Connect training client + print("[3/7] Connecting training client...") + service = FiretitanServiceClient( + base_url=endpoint.base_url, + api_key=api_key, + ) + training_client = service.create_training_client( + base_model=base_model, + lora_rank=lora_rank, + ) + print(" Connected") + + # 4. Load cross-job checkpoint + print(f"[4/7] Loading cross-job checkpoint '{checkpoint_name}' from job {source_job_id}...") + checkpoint_ref = training_client.resolve_checkpoint_path( + checkpoint_name, + source_job_id=source_job_id, + ) + training_client.load_state(checkpoint_ref).result() + print(" Weights loaded (optimizer state not restored — not needed for promotion)") + + # 5. Export sampler checkpoint + sampler_name = f"promote-{checkpoint_name}" + print(f"[5/7] Exporting sampler checkpoint '{sampler_name}'...") + sampler_result = training_client.save_weights_for_sampler_ext( + sampler_name, + checkpoint_type="base", + ) + sampler_snapshot = sampler_result.snapshot_name + print(f" Snapshot: {sampler_snapshot}") + + # 6. Promote to model + print(f"[6/7] Promoting to model '{output_model_id}'...") + model = mgr.promote_checkpoint( + job_id=temp_job_id, + checkpoint_id=sampler_snapshot, + output_model_id=output_model_id, + ) + print(f" Promoted! State: {model.get('state', 'unknown')}") + return model + + finally: + # 7. Cleanup + print(f"\n[7/7] Deleting temporary trainer {temp_job_id}...") + try: + mgr.delete(job_id=temp_job_id) + print(" Deleted") + except Exception as e: + print(f" Warning: cleanup failed: {e}") + print(f" Manually delete with: firectl delete rlorTrainerJob {temp_job_id}") + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +def main(): + parser = argparse.ArgumentParser( + description="Cross-job checkpoint promotion for Fireworks RFT", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + + # Promote step 8 from a deleted training run: + python cross_job_promote.py \\ + --checkpoints-jsonl rl_logs/rft-clinical-reference-v2-0-1/checkpoints.jsonl \\ + --step 8 \\ + --output-model-id rft-clinical-reference-v2-0-1 + + # Explicit training shape: + python cross_job_promote.py \\ + --checkpoints-jsonl rl_logs/my-run/checkpoints.jsonl \\ + --step 8 \\ + --output-model-id my-fine-tuned-model \\ + --training-shape-id accounts/fireworks/trainingShapes/qwen3-8b-128k + + # Dry run — see what would happen without creating any resources: + python cross_job_promote.py \\ + --checkpoints-jsonl rl_logs/my-run/checkpoints.jsonl \\ + --step 8 \\ + --output-model-id my-model \\ + --dry-run +""", + ) + + parser.add_argument( + "--checkpoints-jsonl", + required=True, + help="Path to checkpoints.jsonl from the training run", + ) + parser.add_argument( + "--step", + type=int, + default=None, + help="Step number to promote (default: latest checkpoint with state_path or sampler_path)", + ) + parser.add_argument( + "--output-model-id", + required=True, + help="Model ID for the promoted model (1-63 chars, lowercase a-z, 0-9, hyphen)", + ) + parser.add_argument( + "--training-shape-id", + default=None, + help="Training shape ID for the temporary trainer " + "(auto-inferred from base model if omitted). " + "See https://docs.fireworks.ai/fine-tuning/training-sdk/training-shapes", + ) + parser.add_argument( + "--api-key", + default=None, + help="Fireworks API key (default: $FIREWORKS_API_KEY)", + ) + parser.add_argument( + "--base-url", + default="https://api.fireworks.ai", + help="Fireworks API base URL (default: https://api.fireworks.ai)", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Show what would be done without creating any resources", + ) + + args = parser.parse_args() + + # Validate API key + api_key = args.api_key or os.environ.get("FIREWORKS_API_KEY") + if not api_key: + print( + "Error: FIREWORKS_API_KEY not set.\n" + "Set it via --api-key or export FIREWORKS_API_KEY=..." + ) + sys.exit(1) + + # Validate checkpoints file + jsonl_path = args.checkpoints_jsonl + if not Path(jsonl_path).exists(): + print(f"Error: {jsonl_path} not found") + sys.exit(1) + + # Load and find checkpoint + print(f"Loading checkpoints from {jsonl_path}...") + entries = load_checkpoints(jsonl_path) + print(f" Found {len(entries)} checkpoint entries") + + try: + checkpoint = find_checkpoint(entries, args.step) + except ValueError as e: + print(f"Error: {e}") + sys.exit(1) + step = checkpoint.get("step", "?") + print(f" Selected checkpoint at step {step}") + + base_model = checkpoint.get("base_model") + if not base_model: + print("Error: Checkpoint entry missing 'base_model' field") + sys.exit(1) + + source_job_id = checkpoint.get("source_job_id") + if not source_job_id: + print("Error: Checkpoint entry missing 'source_job_id' field") + sys.exit(1) + + # ------------------------------------------------------------------- + # Path A: checkpoint already has sampler_path → try direct promotion + # ------------------------------------------------------------------- + if checkpoint.get("sampler_path"): + sampler_path = checkpoint["sampler_path"] + print(f"\n Checkpoint has sampler_path: {sampler_path}") + + if args.dry_run: + print("\n[dry-run] Would attempt direct promotion:") + print(f" job_id: {source_job_id}") + print(f" checkpoint_id: {sampler_path}") + print(f" output_model_id: {args.output_model_id}") + print("\n If the source job is gone (404), would fall through to") + print(" cross-job promotion via temporary trainer.") + # Fall through to show the cross-job plan too + else: + _check_sdk() + print(" Attempting direct promotion...") + model = try_direct_promote( + api_key=api_key, + source_job_id=source_job_id, + sampler_path=sampler_path, + output_model_id=args.output_model_id, + base_url=args.base_url, + ) + if model is not None: + print(f"\n Direct promotion succeeded!") + print(f" Model state: {model.get('state', 'unknown')}") + return + print(" Direct promotion failed (source job not found)") + print(" Falling back to cross-job promotion via temporary trainer...") + + # ------------------------------------------------------------------- + # Path B: need cross-job promotion via temporary trainer + # ------------------------------------------------------------------- + if not checkpoint.get("state_path"): + print( + f"\nError: Checkpoint at step {step} has no state_path (DCP reference).\n" + "Cross-job promotion requires a DCP checkpoint.\n" + "Only checkpoints saved with kind=STATE or kind=BOTH have a state_path." + ) + sys.exit(1) + + # Resolve training shape + training_shape_id = args.training_shape_id + if not training_shape_id: + try: + training_shape_id = infer_training_shape(base_model) + print(f" Auto-inferred training shape: {training_shape_id}") + except ValueError as e: + print(f"Error: {e}") + sys.exit(1) + + if args.dry_run: + print(f"\n[dry-run] Would perform cross-job promotion:") + print(f" base_model: {base_model}") + print(f" source_job_id: {source_job_id}") + print(f" checkpoint_name: {checkpoint['name']}") + print(f" state_path: {checkpoint['state_path']}") + print(f" training_shape_id: {training_shape_id}") + print(f" output_model_id: {args.output_model_id}") + print(f"\n Steps:") + print(f" 1. Resolve training shape → pinned version") + print(f" 2. Create temp service-mode trainer ({base_model})") + print(f" 3. Wait for trainer to become RUNNING") + print(f" 4. Connect training client (FiretitanServiceClient)") + print(f" 5. resolve_checkpoint_path('{checkpoint['name']}', source_job_id='{source_job_id}')") + print(f" 6. load_state(checkpoint_ref) — weights only, no optimizer") + print(f" 7. save_weights_for_sampler_ext('promote-{checkpoint['name']}', checkpoint_type='base')") + print(f" 8. promote_checkpoint(temp_job_id, snapshot_name, '{args.output_model_id}')") + print(f" 9. Delete temporary trainer") + return + + _check_sdk() + + model = cross_job_promote( + api_key=api_key, + checkpoint_entry=checkpoint, + output_model_id=args.output_model_id, + training_shape_id=training_shape_id, + base_url=args.base_url, + ) + + print(f"\n{'='*60}") + print(f"SUCCESS: Model '{args.output_model_id}' promoted!") + print(f"{'='*60}") + + +if __name__ == "__main__": + main() From b88591066447edfc6fff962d173ee3dffcbc5cee Mon Sep 17 00:00:00 2001 From: zhiweiz Date: Thu, 16 Apr 2026 17:46:37 +0000 Subject: [PATCH 2/5] Add E2E test for Qwen3 LoRA training shapes; update shape map Root cause analysis: - All Qwen3 (legacy) LoRA training shape versions have public=False - This prevents external accounts (like AILabs) from resolving them via resolve_training_profile() - The shapes exist and have latestValidated versions, but those versions aren't marked public - Qwen 3.5 (qwen3p5-*) LoRA shapes are properly configured with public=True and work correctly Failing shapes (public=False on latestValidated): - qwen3-4b-256k-h200-lora - qwen3-8b-256k-h200-lora - qwen3-4b-minimum-h200-lora - qwen3-4b-minimum-h200-forward-lora - qwen3-235b-2507-instruct-128k-b200-lora - qwen3-235b-2507-instruct-128k-b200-forward-only-lora Working shapes (public=True): - qwen3p5-9b-256k-lora - qwen3p5-27b-256k-lora - qwen3p5-35b-a3b-256k-lora - qwen3p5-397b-a17b-256k-lora - qwen3-vl-8b-256k-h200-lora Fix: Mark the latestValidated version as public=True for each failing legacy shape, or direct users to the new qwen3p5 shapes. Also updated cross_job_promote.py shape map to include Qwen 3.5 shapes. Co-authored-by: zhiweiz --- cross_job_promote.py | 7 + test_qwen3_lora_shapes.py | 414 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 421 insertions(+) create mode 100644 test_qwen3_lora_shapes.py diff --git a/cross_job_promote.py b/cross_job_promote.py index f84ef6e..7431054 100644 --- a/cross_job_promote.py +++ b/cross_job_promote.py @@ -110,12 +110,19 @@ def find_checkpoint(entries: list[dict], step: int | None) -> dict: # --------------------------------------------------------------------------- SHAPE_MAP = { + # Qwen 3.5 (current generation) + "qwen3p5-9b": "accounts/fireworks/trainingShapes/qwen3p5-9b-256k", + "qwen3p5-27b": "accounts/fireworks/trainingShapes/qwen3p5-27b-256k", + "qwen3p5-35b-a3b": "accounts/fireworks/trainingShapes/qwen3p5-35b-a3b-256k", + "qwen3p5-397b-a17b": "accounts/fireworks/trainingShapes/qwen3p5-397b-a17b-256k", + # Qwen 3 (legacy, some shapes may not be public) "qwen3-4b": "accounts/fireworks/trainingShapes/qwen3-4b-minimum", "qwen3-8b": "accounts/fireworks/trainingShapes/qwen3-8b-128k", "qwen3-32b": "accounts/fireworks/trainingShapes/qwen3-32b-65k", "qwen3-30b-a3b": "accounts/fireworks/trainingShapes/qwen3-30b-a3b-instruct-2507-128k", "qwen3-235b": "accounts/fireworks/trainingShapes/qwen3-235b-2507-instruct-128k", "qwen3-vl-8b": "accounts/fireworks/trainingShapes/qwen3-vl-8b-65k", + # Llama "llama-v3p3-70b": "accounts/fireworks/trainingShapes/llama70b-policy", } diff --git a/test_qwen3_lora_shapes.py b/test_qwen3_lora_shapes.py new file mode 100644 index 0000000..3271871 --- /dev/null +++ b/test_qwen3_lora_shapes.py @@ -0,0 +1,414 @@ +#!/usr/bin/env python3 +"""E2E test for Qwen3 LoRA training shapes on Fireworks AI. + +Tests that all Qwen3 (and Qwen 3.5) LoRA training shapes are properly +configured for external account usage: + + 1. Shape exists and is readable + 2. Shape has at least one version + 3. Shape has a latestValidated version + 4. The latestValidated version is marked public (required for external accounts) + 5. The shape references a valid deployment shape + 6. The shape has correct trainerMode (LORA_TRAINER or FORWARD_ONLY) + +Usage: + FIREWORKS_API_KEY=... python test_qwen3_lora_shapes.py + FIREWORKS_API_KEY=... python test_qwen3_lora_shapes.py --verbose + FIREWORKS_API_KEY=... python test_qwen3_lora_shapes.py --account ailabs-account-id +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from dataclasses import dataclass, field +from typing import Any + +import httpx + +# All Qwen3 LoRA training shapes (current and legacy). +# These are the shapes that should be available to external customers. +QWEN3_LORA_SHAPES = { + # Qwen 3.5 (current generation, documented) + "qwen3p5-9b-256k-lora": { + "expected_model": "accounts/fireworks/models/qwen3p5-9b", + "expected_mode": "LORA_TRAINER", + "documented": True, + }, + "qwen3p5-27b-256k-lora": { + "expected_model": "accounts/fireworks/models/qwen3p5-27b", + "expected_mode": "LORA_TRAINER", + "documented": True, + }, + "qwen3p5-35b-a3b-256k-lora": { + "expected_model": "accounts/fireworks/models/qwen3p5-35b-a3b", + "expected_mode": "LORA_TRAINER", + "documented": True, + }, + "qwen3p5-397b-a17b-256k-lora": { + "expected_model": "accounts/fireworks/models/qwen3p5-397b-a17b", + "expected_mode": "LORA_TRAINER", + "documented": True, + }, + # Qwen3 legacy LoRA shapes (still referenced internally) + "qwen3-4b-256k-h200-lora": { + "expected_model": "accounts/fireworks/models/qwen3-4b", + "expected_mode": "LORA_TRAINER", + "documented": False, + }, + "qwen3-8b-256k-h200-lora": { + "expected_model": "accounts/fireworks/models/qwen3-8b", + "expected_mode": "LORA_TRAINER", + "documented": False, + }, + "qwen3-4b-minimum-h200-lora": { + "expected_model": "accounts/fireworks/models/qwen3-4b", + "expected_mode": "LORA_TRAINER", + "documented": False, + }, + "qwen3-4b-minimum-h200-forward-lora": { + "expected_model": "accounts/fireworks/models/qwen3-4b", + "expected_mode": "FORWARD_ONLY", + "documented": False, + }, + "qwen3-235b-2507-instruct-128k-b200-lora": { + "expected_model": "accounts/fireworks/models/qwen3-235b-a22b-instruct-2507", + "expected_mode": "LORA_TRAINER", + "documented": False, + }, + "qwen3-235b-2507-instruct-128k-b200-forward-only-lora": { + "expected_model": "accounts/fireworks/models/qwen3-235b-a22b-instruct-2507", + "expected_mode": "FORWARD_ONLY", + "documented": False, + }, + "qwen3-vl-8b-256k-h200-lora": { + "expected_model": "accounts/fireworks/models/qwen3-vl-8b-instruct", + "expected_mode": "LORA_TRAINER", + "documented": False, + }, +} + + +@dataclass +class TestResult: + shape_name: str + passed: bool + checks: list[tuple[str, bool, str]] = field(default_factory=list) + + def add_check(self, name: str, passed: bool, detail: str = ""): + self.checks.append((name, passed, detail)) + if not passed: + self.passed = False + + +def get_headers(api_key: str) -> dict: + return { + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + } + + +def resolve_account_id(api_key: str, base_url: str) -> str: + resp = httpx.get( + f"{base_url}/v1/accounts", + headers=get_headers(api_key), + timeout=30, + ) + resp.raise_for_status() + accounts = resp.json().get("accounts", []) + if not accounts: + raise RuntimeError("No accounts found for this API key") + return accounts[0]["name"].split("/")[-1] + + +def test_shape( + api_key: str, + shape_short_name: str, + expected: dict, + base_url: str, + verbose: bool = False, +) -> TestResult: + """Run all checks on a single training shape.""" + headers = get_headers(api_key) + shape_id = f"accounts/fireworks/trainingShapes/{shape_short_name}" + result = TestResult(shape_name=shape_short_name, passed=True) + + # 1. Shape exists + resp = httpx.get(f"{base_url}/v1/{shape_id}", headers=headers, timeout=30) + result.add_check( + "shape_exists", + resp.status_code == 200, + f"HTTP {resp.status_code}" if resp.status_code != 200 else "OK", + ) + if resp.status_code != 200: + return result + + shape_data = resp.json() + + # 2. Correct trainer mode + actual_mode = shape_data.get("trainerMode", "UNKNOWN") + expected_mode = expected["expected_mode"] + result.add_check( + "trainer_mode", + actual_mode == expected_mode, + f"expected={expected_mode} actual={actual_mode}", + ) + + # 3. Correct base model + actual_model = shape_data.get("baseModel", "") + expected_model = expected["expected_model"] + result.add_check( + "base_model", + actual_model == expected_model, + f"expected={expected_model} actual={actual_model}", + ) + + # 4. Has deployment shape version + deploy_shape = shape_data.get("deploymentShapeVersion", "") + result.add_check( + "deployment_shape", + bool(deploy_shape), + deploy_shape or "MISSING", + ) + + # 5. Has trainer image tag + image_tag = shape_data.get("trainerImageTag", "") + result.add_check( + "trainer_image_tag", + bool(image_tag), + image_tag or "MISSING", + ) + + # 6. Has versions + resp2 = httpx.get( + f"{base_url}/v1/{shape_id}/versions", headers=headers, timeout=30 + ) + result.add_check( + "versions_endpoint", + resp2.status_code == 200, + f"HTTP {resp2.status_code}" if resp2.status_code != 200 else "OK", + ) + if resp2.status_code != 200: + return result + + versions = resp2.json().get("trainingShapeVersions", []) + result.add_check( + "has_versions", + len(versions) > 0, + f"{len(versions)} versions", + ) + if not versions: + return result + + # 7. Has validated versions + validated = [v for v in versions if v.get("validated", False)] + result.add_check( + "has_validated", + len(validated) > 0, + f"{len(validated)} validated out of {len(versions)} total", + ) + + # 8. Has latestValidated version + latest_validated = [v for v in versions if v.get("latestValidated", False)] + result.add_check( + "has_latest_validated", + len(latest_validated) == 1, + f"{len(latest_validated)} latestValidated versions", + ) + + # 9. latestValidated version is PUBLIC (critical for external accounts!) + if latest_validated: + lv = latest_validated[0] + is_public = lv.get("public", False) + result.add_check( + "latest_validated_is_public", + is_public, + f"public={is_public} version={lv['name'].split('/')[-1]}", + ) + + if verbose: + lv_snapshot = lv.get("snapshot", {}) + result.add_check( + "version_has_snapshot", + bool(lv_snapshot), + "snapshot present" if lv_snapshot else "MISSING snapshot", + ) + else: + result.add_check( + "latest_validated_is_public", + False, + "no latestValidated version to check", + ) + + return result + + +def test_shape_resolution_as_account( + api_key: str, + account_id: str, + shape_short_name: str, + base_url: str, +) -> tuple[bool, str]: + """Simulate what the SDK's resolve_training_profile does for an account. + + The SDK finds the latestValidated version with public=True. + If no such version exists, external accounts will fail. + """ + headers = get_headers(api_key) + shape_id = f"accounts/fireworks/trainingShapes/{shape_short_name}" + + resp = httpx.get( + f"{base_url}/v1/{shape_id}/versions", headers=headers, timeout=30 + ) + if resp.status_code != 200: + return False, f"Cannot list versions: HTTP {resp.status_code}" + + versions = resp.json().get("trainingShapeVersions", []) + + # The SDK looks for a version that is both latestValidated AND public + public_validated = [ + v + for v in versions + if v.get("latestValidated", False) and v.get("public", False) + ] + + if public_validated: + v = public_validated[0] + return True, v["name"] + + # Check if there's a latestValidated but it's not public + latest_val = [v for v in versions if v.get("latestValidated", False)] + if latest_val: + v = latest_val[0] + return ( + False, + f"latestValidated exists ({v['name'].split('/')[-1]}) " + f"but public={v.get('public', False)} — " + f"external accounts cannot resolve this shape", + ) + + return False, "No latestValidated version exists" + + +def main(): + parser = argparse.ArgumentParser( + description="E2E test for Qwen3 LoRA training shapes" + ) + parser.add_argument("--api-key", default=None, help="Fireworks API key") + parser.add_argument( + "--base-url", + default="https://api.fireworks.ai", + help="API base URL", + ) + parser.add_argument( + "--verbose", "-v", action="store_true", help="Show all check details" + ) + parser.add_argument( + "--account", + default=None, + help="Account ID to test resolution against (auto-resolved if not set)", + ) + parser.add_argument( + "--only-documented", + action="store_true", + help="Only test shapes that are in the current docs", + ) + args = parser.parse_args() + + api_key = args.api_key or os.environ.get("FIREWORKS_API_KEY") + if not api_key: + print("Error: FIREWORKS_API_KEY not set") + sys.exit(1) + + account_id = args.account + if not account_id: + account_id = resolve_account_id(api_key, args.base_url) + print(f"Testing as account: {account_id}") + + shapes_to_test = { + k: v + for k, v in QWEN3_LORA_SHAPES.items() + if not args.only_documented or v.get("documented", False) + } + + print(f"Testing {len(shapes_to_test)} Qwen3 LoRA training shapes\n") + + all_passed = True + results: list[TestResult] = [] + + for shape_name, expected in shapes_to_test.items(): + documented = expected.get("documented", False) + tag = "[DOCS] " if documented else "[LEGACY]" + + result = test_shape( + api_key=api_key, + shape_short_name=shape_name, + expected=expected, + base_url=args.base_url, + verbose=args.verbose, + ) + results.append(result) + + status = "PASS" if result.passed else "FAIL" + print(f" {status} {tag} {shape_name}") + + if args.verbose or not result.passed: + for check_name, check_passed, detail in result.checks: + marker = " ✓" if check_passed else " ✗" + print(f" {marker} {check_name}: {detail}") + + if not result.passed: + all_passed = False + + # Resolution test + print(f"\n{'='*70}") + print("Shape resolution test (simulating external account)") + print(f"{'='*70}\n") + + resolution_failures = [] + for shape_name in shapes_to_test: + ok, detail = test_shape_resolution_as_account( + api_key=api_key, + account_id=account_id, + shape_short_name=shape_name, + base_url=args.base_url, + ) + status = "PASS" if ok else "FAIL" + documented = shapes_to_test[shape_name].get("documented", False) + tag = "[DOCS] " if documented else "[LEGACY]" + print(f" {status} {tag} {shape_name}") + if not ok: + print(f" → {detail}") + resolution_failures.append((shape_name, detail)) + all_passed = False + + # Summary + print(f"\n{'='*70}") + total = len(shapes_to_test) + passed = sum(1 for r in results if r.passed) + print(f"Shape config: {passed}/{total} passed") + + resolution_passed = total - len(resolution_failures) + print(f"Resolution: {resolution_passed}/{total} passed") + + if resolution_failures: + print(f"\n⚠ Resolution failures (shapes external accounts CANNOT use):") + for name, detail in resolution_failures: + print(f" • {name}: {detail}") + print( + f"\nFix: Mark the latestValidated version as public=True " + f"for each failing shape." + ) + + if all_passed: + print(f"\nAll tests passed!") + return 0 + else: + print(f"\nSome tests FAILED — see details above.") + return 1 + + +if __name__ == "__main__": + sys.exit(main()) From e2fcbda6dd5486c900655d6e47c10557b2da1e72 Mon Sep 17 00:00:00 2001 From: zhiweiz Date: Thu, 16 Apr 2026 18:11:38 +0000 Subject: [PATCH 3/5] Rewrite E2E test: focus on qwen3p5-*-lora shapes and permission chain MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Updated test to focus on the actual customer issue: AILabs can't access qwen3p5-*-lora shapes with a permission error. Test now covers three phases: 1. Shape metadata validation (exists, trainerMode, versions, public/validated) 2. Dependency chain access (base model, deployment shape, deployment shape version — all must be accessible and public) 3. Job lifecycle E2E (opt-in: creates and deletes a real RLOR trainer job) All qwen3p5 LoRA shapes pass from the pyroworks account: - Shape versions have latestValidated=True and public=True - Deployment shape versions are public and validated - Base models are READY and accessible - Job creation succeeds The permission issue for AILabs is likely at the account level: - Accelerator quota (B200/B300 access) - Service-mode RLOR permission - Or SDK version mismatch Run with AILabs API key to reproduce: FIREWORKS_API_KEY= python test_qwen3_lora_shapes.py -v Co-authored-by: zhiweiz --- test_qwen3_lora_shapes.py | 586 +++++++++++++++++++------------------- 1 file changed, 293 insertions(+), 293 deletions(-) diff --git a/test_qwen3_lora_shapes.py b/test_qwen3_lora_shapes.py index 3271871..f6bfac0 100644 --- a/test_qwen3_lora_shapes.py +++ b/test_qwen3_lora_shapes.py @@ -1,20 +1,28 @@ #!/usr/bin/env python3 """E2E test for Qwen3 LoRA training shapes on Fireworks AI. -Tests that all Qwen3 (and Qwen 3.5) LoRA training shapes are properly -configured for external account usage: +Tests that all Qwen 3.5 LoRA training shapes are properly configured +and accessible for a given account: - 1. Shape exists and is readable - 2. Shape has at least one version - 3. Shape has a latestValidated version - 4. The latestValidated version is marked public (required for external accounts) - 5. The shape references a valid deployment shape - 6. The shape has correct trainerMode (LORA_TRAINER or FORWARD_ONLY) + Phase 1 — Shape metadata validation: + - Shape exists and is readable + - Has correct trainerMode, baseModel, deploymentShapeVersion + - Has at least one version with latestValidated=True and public=True + + Phase 2 — Dependency chain access: + - Base model is accessible (GET returns 200) + - Deployment shape and its version are accessible + - Deployment shape version is public and validated + + Phase 3 — Job lifecycle E2E (opt-in via --create-job): + - Creates a service-mode RLOR trainer job with LoRA config + - Verifies job reaches RUNNING state + - Deletes the job Usage: FIREWORKS_API_KEY=... python test_qwen3_lora_shapes.py FIREWORKS_API_KEY=... python test_qwen3_lora_shapes.py --verbose - FIREWORKS_API_KEY=... python test_qwen3_lora_shapes.py --account ailabs-account-id + FIREWORKS_API_KEY=... python test_qwen3_lora_shapes.py --create-job --shape qwen3p5-9b-256k-lora """ from __future__ import annotations @@ -23,99 +31,59 @@ import json import os import sys +import time from dataclasses import dataclass, field from typing import Any import httpx -# All Qwen3 LoRA training shapes (current and legacy). -# These are the shapes that should be available to external customers. -QWEN3_LORA_SHAPES = { - # Qwen 3.5 (current generation, documented) +# ───────────────────────────────────────────────────────────────────── +# Shape catalog +# ───────────────────────────────────────────────────────────────────── + +QWEN3P5_LORA_SHAPES: dict[str, dict[str, Any]] = { "qwen3p5-9b-256k-lora": { "expected_model": "accounts/fireworks/models/qwen3p5-9b", "expected_mode": "LORA_TRAINER", - "documented": True, }, "qwen3p5-27b-256k-lora": { "expected_model": "accounts/fireworks/models/qwen3p5-27b", "expected_mode": "LORA_TRAINER", - "documented": True, }, "qwen3p5-35b-a3b-256k-lora": { "expected_model": "accounts/fireworks/models/qwen3p5-35b-a3b", "expected_mode": "LORA_TRAINER", - "documented": True, }, "qwen3p5-397b-a17b-256k-lora": { "expected_model": "accounts/fireworks/models/qwen3p5-397b-a17b", "expected_mode": "LORA_TRAINER", - "documented": True, - }, - # Qwen3 legacy LoRA shapes (still referenced internally) - "qwen3-4b-256k-h200-lora": { - "expected_model": "accounts/fireworks/models/qwen3-4b", - "expected_mode": "LORA_TRAINER", - "documented": False, - }, - "qwen3-8b-256k-h200-lora": { - "expected_model": "accounts/fireworks/models/qwen3-8b", - "expected_mode": "LORA_TRAINER", - "documented": False, - }, - "qwen3-4b-minimum-h200-lora": { - "expected_model": "accounts/fireworks/models/qwen3-4b", - "expected_mode": "LORA_TRAINER", - "documented": False, - }, - "qwen3-4b-minimum-h200-forward-lora": { - "expected_model": "accounts/fireworks/models/qwen3-4b", - "expected_mode": "FORWARD_ONLY", - "documented": False, - }, - "qwen3-235b-2507-instruct-128k-b200-lora": { - "expected_model": "accounts/fireworks/models/qwen3-235b-a22b-instruct-2507", - "expected_mode": "LORA_TRAINER", - "documented": False, - }, - "qwen3-235b-2507-instruct-128k-b200-forward-only-lora": { - "expected_model": "accounts/fireworks/models/qwen3-235b-a22b-instruct-2507", - "expected_mode": "FORWARD_ONLY", - "documented": False, - }, - "qwen3-vl-8b-256k-h200-lora": { - "expected_model": "accounts/fireworks/models/qwen3-vl-8b-instruct", - "expected_mode": "LORA_TRAINER", - "documented": False, }, } +# ───────────────────────────────────────────────────────────────────── +# Helpers +# ───────────────────────────────────────────────────────────────────── + @dataclass class TestResult: - shape_name: str - passed: bool + name: str + passed: bool = True checks: list[tuple[str, bool, str]] = field(default_factory=list) - def add_check(self, name: str, passed: bool, detail: str = ""): - self.checks.append((name, passed, detail)) - if not passed: + def check(self, label: str, ok: bool, detail: str = ""): + self.checks.append((label, ok, detail)) + if not ok: self.passed = False + return ok -def get_headers(api_key: str) -> dict: - return { - "Authorization": f"Bearer {api_key}", - "Content-Type": "application/json", - } +def headers_for(api_key: str) -> dict: + return {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"} def resolve_account_id(api_key: str, base_url: str) -> str: - resp = httpx.get( - f"{base_url}/v1/accounts", - headers=get_headers(api_key), - timeout=30, - ) + resp = httpx.get(f"{base_url}/v1/accounts", headers=headers_for(api_key), timeout=30) resp.raise_for_status() accounts = resp.json().get("accounts", []) if not accounts: @@ -123,197 +91,265 @@ def resolve_account_id(api_key: str, base_url: str) -> str: return accounts[0]["name"].split("/")[-1] -def test_shape( - api_key: str, - shape_short_name: str, - expected: dict, - base_url: str, - verbose: bool = False, +# ───────────────────────────────────────────────────────────────────── +# Phase 1: Shape metadata validation +# ───────────────────────────────────────────────────────────────────── + +def test_shape_metadata( + api_key: str, shape_short: str, expected: dict, base_url: str, verbose: bool ) -> TestResult: - """Run all checks on a single training shape.""" - headers = get_headers(api_key) - shape_id = f"accounts/fireworks/trainingShapes/{shape_short_name}" - result = TestResult(shape_name=shape_short_name, passed=True) + h = headers_for(api_key) + shape_id = f"accounts/fireworks/trainingShapes/{shape_short}" + r = TestResult(name=shape_short) # 1. Shape exists - resp = httpx.get(f"{base_url}/v1/{shape_id}", headers=headers, timeout=30) - result.add_check( - "shape_exists", - resp.status_code == 200, - f"HTTP {resp.status_code}" if resp.status_code != 200 else "OK", - ) - if resp.status_code != 200: - return result + resp = httpx.get(f"{base_url}/v1/{shape_id}", headers=h, timeout=30) + if not r.check("shape_exists", resp.status_code == 200, f"HTTP {resp.status_code}"): + return r + shape = resp.json() - shape_data = resp.json() - - # 2. Correct trainer mode - actual_mode = shape_data.get("trainerMode", "UNKNOWN") - expected_mode = expected["expected_mode"] - result.add_check( + # 2. Trainer mode + r.check( "trainer_mode", - actual_mode == expected_mode, - f"expected={expected_mode} actual={actual_mode}", + shape.get("trainerMode") == expected["expected_mode"], + f"got {shape.get('trainerMode')}", ) - # 3. Correct base model - actual_model = shape_data.get("baseModel", "") - expected_model = expected["expected_model"] - result.add_check( + # 3. Base model + r.check( "base_model", - actual_model == expected_model, - f"expected={expected_model} actual={actual_model}", + shape.get("baseModel") == expected["expected_model"], + f"got {shape.get('baseModel')}", ) # 4. Has deployment shape version - deploy_shape = shape_data.get("deploymentShapeVersion", "") - result.add_check( - "deployment_shape", - bool(deploy_shape), - deploy_shape or "MISSING", - ) + deploy_sv = shape.get("deploymentShapeVersion", "") + r.check("has_deployment_shape_version", bool(deploy_sv), deploy_sv or "MISSING") # 5. Has trainer image tag - image_tag = shape_data.get("trainerImageTag", "") - result.add_check( - "trainer_image_tag", - bool(image_tag), - image_tag or "MISSING", - ) + r.check("has_trainer_image_tag", bool(shape.get("trainerImageTag")), shape.get("trainerImageTag", "MISSING")) - # 6. Has versions - resp2 = httpx.get( - f"{base_url}/v1/{shape_id}/versions", headers=headers, timeout=30 - ) - result.add_check( - "versions_endpoint", - resp2.status_code == 200, - f"HTTP {resp2.status_code}" if resp2.status_code != 200 else "OK", - ) - if resp2.status_code != 200: - return result + # 6. Versions exist + resp2 = httpx.get(f"{base_url}/v1/{shape_id}/versions", headers=h, timeout=30) + if not r.check("versions_accessible", resp2.status_code == 200, f"HTTP {resp2.status_code}"): + return r versions = resp2.json().get("trainingShapeVersions", []) - result.add_check( - "has_versions", - len(versions) > 0, - f"{len(versions)} versions", - ) - if not versions: - return result - - # 7. Has validated versions - validated = [v for v in versions if v.get("validated", False)] - result.add_check( - "has_validated", - len(validated) > 0, - f"{len(validated)} validated out of {len(versions)} total", - ) + r.check("has_versions", len(versions) > 0, f"{len(versions)} total") + + # 7. Has latestValidated + latest = [v for v in versions if v.get("latestValidated")] + r.check("has_latest_validated", len(latest) == 1, f"{len(latest)} found") + + # 8. latestValidated is public + if latest: + lv = latest[0] + ver_id = lv["name"].split("/")[-1] + r.check("latest_validated_public", lv.get("public", False), f"version={ver_id} public={lv.get('public')}") + r.check("latest_validated_validated", lv.get("validated", False), f"validated={lv.get('validated')}") + else: + r.check("latest_validated_public", False, "no latestValidated to check") - # 8. Has latestValidated version - latest_validated = [v for v in versions if v.get("latestValidated", False)] - result.add_check( - "has_latest_validated", - len(latest_validated) == 1, - f"{len(latest_validated)} latestValidated versions", - ) + return r - # 9. latestValidated version is PUBLIC (critical for external accounts!) - if latest_validated: - lv = latest_validated[0] - is_public = lv.get("public", False) - result.add_check( - "latest_validated_is_public", - is_public, - f"public={is_public} version={lv['name'].split('/')[-1]}", - ) - if verbose: - lv_snapshot = lv.get("snapshot", {}) - result.add_check( - "version_has_snapshot", - bool(lv_snapshot), - "snapshot present" if lv_snapshot else "MISSING snapshot", +# ───────────────────────────────────────────────────────────────────── +# Phase 2: Dependency chain access +# ───────────────────────────────────────────────────────────────────── + +def test_dependency_chain( + api_key: str, shape_short: str, base_url: str, verbose: bool +) -> TestResult: + h = headers_for(api_key) + shape_id = f"accounts/fireworks/trainingShapes/{shape_short}" + r = TestResult(name=f"{shape_short} (deps)") + + resp = httpx.get(f"{base_url}/v1/{shape_id}", headers=h, timeout=30) + if resp.status_code != 200: + r.check("shape_read", False, f"HTTP {resp.status_code}") + return r + shape = resp.json() + + # 1. Base model accessible + base_model = shape.get("baseModel", "") + if base_model: + resp2 = httpx.get(f"{base_url}/v1/{base_model}", headers=h, timeout=30) + r.check("base_model_accessible", resp2.status_code == 200, f"{base_model} → HTTP {resp2.status_code}") + if resp2.status_code == 200: + model_data = resp2.json() + r.check( + "base_model_ready", + model_data.get("state") == "READY", + f"state={model_data.get('state')}", ) - else: - result.add_check( - "latest_validated_is_public", - False, - "no latestValidated version to check", - ) - return result + # 2. Deployment shape version accessible + deploy_sv = shape.get("deploymentShapeVersion", "") + if deploy_sv: + resp3 = httpx.get(f"{base_url}/v1/{deploy_sv}", headers=h, timeout=30) + r.check("deploy_shape_version_accessible", resp3.status_code == 200, f"{deploy_sv} → HTTP {resp3.status_code}") + if resp3.status_code == 200: + dsv_data = resp3.json() + r.check("deploy_shape_version_public", dsv_data.get("public", False), f"public={dsv_data.get('public')}") + r.check("deploy_shape_version_validated", dsv_data.get("validated", False), f"validated={dsv_data.get('validated')}") + + # 3. Deployment shape parent accessible + if deploy_sv: + deploy_parent = "/".join(deploy_sv.split("/")[:4]) + resp4 = httpx.get(f"{base_url}/v1/{deploy_parent}", headers=h, timeout=30) + r.check("deploy_shape_parent_accessible", resp4.status_code == 200, f"{deploy_parent} → HTTP {resp4.status_code}") + + # 4. Training shape version snapshot has correct deployment shape ref + resp5 = httpx.get(f"{base_url}/v1/{shape_id}/versions", headers=h, timeout=30) + if resp5.status_code == 200: + versions = resp5.json().get("trainingShapeVersions", []) + latest = [v for v in versions if v.get("latestValidated")] + if latest: + snapshot = latest[0].get("snapshot", {}) + snap_deploy = snapshot.get("deploymentShapeVersion", "") + r.check( + "version_snapshot_deploy_shape", + snap_deploy == deploy_sv, + f"shape={deploy_sv} snapshot={snap_deploy}", + ) + + return r + +# ───────────────────────────────────────────────────────────────────── +# Phase 3: Job lifecycle E2E +# ───────────────────────────────────────────────────────────────────── -def test_shape_resolution_as_account( +def test_job_lifecycle( api_key: str, + shape_short: str, account_id: str, - shape_short_name: str, base_url: str, -) -> tuple[bool, str]: - """Simulate what the SDK's resolve_training_profile does for an account. - - The SDK finds the latestValidated version with public=True. - If no such version exists, external accounts will fail. - """ - headers = get_headers(api_key) - shape_id = f"accounts/fireworks/trainingShapes/{shape_short_name}" + verbose: bool, + wait_timeout: float = 120, +) -> TestResult: + h = headers_for(api_key) + shape_id = f"accounts/fireworks/trainingShapes/{shape_short}" + r = TestResult(name=f"{shape_short} (job)") - resp = httpx.get( - f"{base_url}/v1/{shape_id}/versions", headers=headers, timeout=30 - ) + resp = httpx.get(f"{base_url}/v1/{shape_id}", headers=h, timeout=30) if resp.status_code != 200: - return False, f"Cannot list versions: HTTP {resp.status_code}" - - versions = resp.json().get("trainingShapeVersions", []) - - # The SDK looks for a version that is both latestValidated AND public - public_validated = [ - v - for v in versions - if v.get("latestValidated", False) and v.get("public", False) - ] - - if public_validated: - v = public_validated[0] - return True, v["name"] - - # Check if there's a latestValidated but it's not public - latest_val = [v for v in versions if v.get("latestValidated", False)] - if latest_val: - v = latest_val[0] - return ( - False, - f"latestValidated exists ({v['name'].split('/')[-1]}) " - f"but public={v.get('public', False)} — " - f"external accounts cannot resolve this shape", + r.check("shape_read", False, f"HTTP {resp.status_code}") + return r + shape = resp.json() + + base_model = shape["baseModel"] + accel_type = shape.get("acceleratorType", "") + accel_count = shape.get("acceleratorCount", 0) + node_count = shape.get("nodeCount", 1) + + # Create RLOR job + create_url = f"{base_url}/v1/accounts/{account_id}/rlorTrainerJobs" + body: dict[str, Any] = { + "serviceMode": True, + "keepAlive": True, + "displayName": f"e2e-test-{shape_short}", + "trainingConfig": { + "baseModel": base_model, + "loraRank": 16, + "learningRate": 1e-5, + }, + "nodeCount": node_count, + } + if accel_type: + body["trainingConfig"]["acceleratorType"] = accel_type + if accel_count: + body["trainingConfig"]["acceleratorCount"] = accel_count + + print(f" Creating job for {shape_short}...") + resp2 = httpx.post(create_url, headers=h, json=body, timeout=120) + if not r.check( + "job_create", + resp2.status_code in (200, 201), + f"HTTP {resp2.status_code}: {resp2.text[:200]}", + ): + return r + + job_data = resp2.json() + job_name = job_data.get("name", "") + job_id = job_name.split("/")[-1] + print(f" Job created: {job_id}") + + try: + # Wait for RUNNING + deadline = time.time() + wait_timeout + final_state = job_data.get("state", "UNKNOWN") + while time.time() < deadline: + time.sleep(5) + resp3 = httpx.get(f"{base_url}/v1/{job_name}", headers=h, timeout=30) + if resp3.status_code != 200: + continue + current = resp3.json() + final_state = current.get("state", "UNKNOWN") + status = current.get("status", {}) + + if verbose: + print(f" state={final_state} status={status}") + + if final_state == "JOB_STATE_RUNNING": + break + if final_state in ("JOB_STATE_FAILED", "JOB_STATE_CANCELLED"): + break + + r.check( + "job_reached_running", + final_state == "JOB_STATE_RUNNING", + f"final_state={final_state}", ) - return False, "No latestValidated version exists" + finally: + print(f" Deleting job {job_id}...") + del_resp = httpx.delete(f"{base_url}/v1/{job_name}", headers=h, timeout=60) + r.check( + "job_delete", + del_resp.status_code in (200, 204), + f"HTTP {del_resp.status_code}", + ) + print(f" Deleted: HTTP {del_resp.status_code}") + + return r + + +# ───────────────────────────────────────────────────────────────────── +# Main +# ───────────────────────────────────────────────────────────────────── + +def print_result(r: TestResult, verbose: bool): + status = "PASS" if r.passed else "FAIL" + print(f" {status} {r.name}") + if verbose or not r.passed: + for label, ok, detail in r.checks: + marker = "✓" if ok else "✗" + print(f" {marker} {label}: {detail}") def main(): parser = argparse.ArgumentParser( - description="E2E test for Qwen3 LoRA training shapes" + description="E2E test for Qwen 3.5 LoRA training shapes" ) parser.add_argument("--api-key", default=None, help="Fireworks API key") + parser.add_argument("--base-url", default="https://api.fireworks.ai", help="API base URL") + parser.add_argument("--verbose", "-v", action="store_true", help="Show all check details") parser.add_argument( - "--base-url", - default="https://api.fireworks.ai", - help="API base URL", - ) - parser.add_argument( - "--verbose", "-v", action="store_true", help="Show all check details" - ) - parser.add_argument( - "--account", + "--shape", default=None, - help="Account ID to test resolution against (auto-resolved if not set)", + help="Test only this shape (e.g. qwen3p5-9b-256k-lora)", ) parser.add_argument( - "--only-documented", + "--create-job", action="store_true", - help="Only test shapes that are in the current docs", + help="Phase 3: actually create+delete an RLOR job (costs GPU time!)", + ) + parser.add_argument( + "--wait-timeout", + type=float, + default=120, + help="Seconds to wait for job to reach RUNNING (default: 120)", ) args = parser.parse_args() @@ -322,92 +358,56 @@ def main(): print("Error: FIREWORKS_API_KEY not set") sys.exit(1) - account_id = args.account - if not account_id: - account_id = resolve_account_id(api_key, args.base_url) - print(f"Testing as account: {account_id}") - - shapes_to_test = { - k: v - for k, v in QWEN3_LORA_SHAPES.items() - if not args.only_documented or v.get("documented", False) - } + account_id = resolve_account_id(api_key, args.base_url) + print(f"Account: {account_id}\n") - print(f"Testing {len(shapes_to_test)} Qwen3 LoRA training shapes\n") + shapes = QWEN3P5_LORA_SHAPES + if args.shape: + if args.shape not in shapes: + print(f"Unknown shape: {args.shape}") + print(f"Available: {', '.join(shapes.keys())}") + sys.exit(1) + shapes = {args.shape: shapes[args.shape]} all_passed = True - results: list[TestResult] = [] - - for shape_name, expected in shapes_to_test.items(): - documented = expected.get("documented", False) - tag = "[DOCS] " if documented else "[LEGACY]" - - result = test_shape( - api_key=api_key, - shape_short_name=shape_name, - expected=expected, - base_url=args.base_url, - verbose=args.verbose, - ) - results.append(result) - - status = "PASS" if result.passed else "FAIL" - print(f" {status} {tag} {shape_name}") - - if args.verbose or not result.passed: - for check_name, check_passed, detail in result.checks: - marker = " ✓" if check_passed else " ✗" - print(f" {marker} {check_name}: {detail}") - if not result.passed: + # ── Phase 1 ────────────────────────────────────────────────────── + print(f"Phase 1: Shape metadata ({len(shapes)} shapes)") + print("-" * 60) + for name, expected in shapes.items(): + r = test_shape_metadata(api_key, name, expected, args.base_url, args.verbose) + print_result(r, args.verbose) + if not r.passed: all_passed = False - # Resolution test - print(f"\n{'='*70}") - print("Shape resolution test (simulating external account)") - print(f"{'='*70}\n") - - resolution_failures = [] - for shape_name in shapes_to_test: - ok, detail = test_shape_resolution_as_account( - api_key=api_key, - account_id=account_id, - shape_short_name=shape_name, - base_url=args.base_url, - ) - status = "PASS" if ok else "FAIL" - documented = shapes_to_test[shape_name].get("documented", False) - tag = "[DOCS] " if documented else "[LEGACY]" - print(f" {status} {tag} {shape_name}") - if not ok: - print(f" → {detail}") - resolution_failures.append((shape_name, detail)) + # ── Phase 2 ────────────────────────────────────────────────────── + print(f"\nPhase 2: Dependency chain access") + print("-" * 60) + for name in shapes: + r = test_dependency_chain(api_key, name, args.base_url, args.verbose) + print_result(r, args.verbose) + if not r.passed: all_passed = False - # Summary - print(f"\n{'='*70}") - total = len(shapes_to_test) - passed = sum(1 for r in results if r.passed) - print(f"Shape config: {passed}/{total} passed") - - resolution_passed = total - len(resolution_failures) - print(f"Resolution: {resolution_passed}/{total} passed") - - if resolution_failures: - print(f"\n⚠ Resolution failures (shapes external accounts CANNOT use):") - for name, detail in resolution_failures: - print(f" • {name}: {detail}") - print( - f"\nFix: Mark the latestValidated version as public=True " - f"for each failing shape." - ) + # ── Phase 3 (opt-in) ───────────────────────────────────────────── + if args.create_job: + print(f"\nPhase 3: Job lifecycle E2E") + print("-" * 60) + for name in shapes: + r = test_job_lifecycle( + api_key, name, account_id, args.base_url, args.verbose, args.wait_timeout + ) + print_result(r, args.verbose) + if not r.passed: + all_passed = False + # ── Summary ────────────────────────────────────────────────────── + print(f"\n{'='*60}") if all_passed: - print(f"\nAll tests passed!") - return 0 + print("All tests PASSED") else: - print(f"\nSome tests FAILED — see details above.") - return 1 + print("Some tests FAILED — see details above") + return 0 if all_passed else 1 if __name__ == "__main__": From 253af767b9db3dded819f687d86bf85fcadf71c0 Mon Sep 17 00:00:00 2001 From: zhiweiz Date: Thu, 16 Apr 2026 18:55:21 +0000 Subject: [PATCH 4/5] Rewrite test to use real Training SDK code paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Now uses the actual fireworks.training.sdk to exercise: 1. resolve_training_profile() - the exact API the SDK calls (GET /versions?filter=latest_validated=true&pageSize=1) 2. TrainerJobConfig validation with shape ref 3. Job lifecycle via mgr.create() with ?trainingShape= query param Key findings from reading the SDK source: - resolve_training_profile() filters versions with 'latest_validated=true' server-side - If the server returns no versions, it raises: 'No latest validated training-shape version was returned' - Job creation passes the shape ref as a QUERY PARAMETER (?trainingShape=...), not in the request body - The server does shape-to-infra resolution and validation All 4 qwen3p5-*-lora shapes pass resolution and config validation. The permission issue for AILabs is server-side — either: - The 'latest_validated=true' filter returns empty for their account - Or the job creation with ?trainingShape= fails at server validation Run with any account's API key to diagnose: FIREWORKS_API_KEY= python test_qwen3_lora_shapes.py -v Co-authored-by: zhiweiz --- test_qwen3_lora_shapes.py | 447 ++++++++++++++++++-------------------- 1 file changed, 213 insertions(+), 234 deletions(-) diff --git a/test_qwen3_lora_shapes.py b/test_qwen3_lora_shapes.py index f6bfac0..edbf32d 100644 --- a/test_qwen3_lora_shapes.py +++ b/test_qwen3_lora_shapes.py @@ -1,23 +1,21 @@ #!/usr/bin/env python3 -"""E2E test for Qwen3 LoRA training shapes on Fireworks AI. +"""E2E test for Qwen 3.5 LoRA training shapes on Fireworks AI. Tests that all Qwen 3.5 LoRA training shapes are properly configured -and accessible for a given account: +and accessible using the exact SDK code paths: - Phase 1 — Shape metadata validation: - - Shape exists and is readable - - Has correct trainerMode, baseModel, deploymentShapeVersion - - Has at least one version with latestValidated=True and public=True + Phase 1 — resolve_training_profile(): + Uses the real Training SDK to resolve each shape, which calls: + GET /v1/{shape}/versions?filter=latest_validated=true&pageSize=1 + Verifies the returned profile has valid fields. - Phase 2 — Dependency chain access: - - Base model is accessible (GET returns 200) - - Deployment shape and its version are accessible - - Deployment shape version is public and validated + Phase 2 — TrainerJobConfig validation: + Builds a TrainerJobConfig from the resolved profile and validates + it passes the SDK's pre-flight checks. Phase 3 — Job lifecycle E2E (opt-in via --create-job): - - Creates a service-mode RLOR trainer job with LoRA config - - Verifies job reaches RUNNING state - - Deletes the job + Creates a real service-mode RLOR trainer job, waits for RUNNING, + then deletes it. Usage: FIREWORKS_API_KEY=... python test_qwen3_lora_shapes.py @@ -28,14 +26,22 @@ from __future__ import annotations import argparse -import json import os import sys import time +import logging from dataclasses import dataclass, field from typing import Any -import httpx +from fireworks.training.sdk import ( + FireworksClient, + TrainerJobConfig, + TrainerJobManager, +) +from fireworks.training.sdk.fireworks_client import TrainingShapeProfile + +logger = logging.getLogger(__name__) + # ───────────────────────────────────────────────────────────────────── # Shape catalog @@ -62,7 +68,7 @@ # ───────────────────────────────────────────────────────────────────── -# Helpers +# Test result tracking # ───────────────────────────────────────────────────────────────────── @dataclass @@ -71,147 +77,149 @@ class TestResult: passed: bool = True checks: list[tuple[str, bool, str]] = field(default_factory=list) - def check(self, label: str, ok: bool, detail: str = ""): + def check(self, label: str, ok: bool, detail: str = "") -> bool: self.checks.append((label, ok, detail)) if not ok: self.passed = False return ok -def headers_for(api_key: str) -> dict: - return {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"} - - -def resolve_account_id(api_key: str, base_url: str) -> str: - resp = httpx.get(f"{base_url}/v1/accounts", headers=headers_for(api_key), timeout=30) - resp.raise_for_status() - accounts = resp.json().get("accounts", []) - if not accounts: - raise RuntimeError("No accounts found for this API key") - return accounts[0]["name"].split("/")[-1] +def print_result(r: TestResult, verbose: bool): + status = "PASS" if r.passed else "FAIL" + print(f" {status} {r.name}") + if verbose or not r.passed: + for label, ok, detail in r.checks: + marker = "✓" if ok else "✗" + print(f" {marker} {label}: {detail}") # ───────────────────────────────────────────────────────────────────── -# Phase 1: Shape metadata validation +# Phase 1: resolve_training_profile (SDK code path) # ───────────────────────────────────────────────────────────────────── -def test_shape_metadata( - api_key: str, shape_short: str, expected: dict, base_url: str, verbose: bool -) -> TestResult: - h = headers_for(api_key) +def test_resolve_profile( + client: FireworksClient, + shape_short: str, + expected: dict, + verbose: bool, +) -> tuple[TestResult, TrainingShapeProfile | None]: + """Exercise the SDK's resolve_training_profile() and check the result.""" shape_id = f"accounts/fireworks/trainingShapes/{shape_short}" r = TestResult(name=shape_short) + profile = None - # 1. Shape exists - resp = httpx.get(f"{base_url}/v1/{shape_id}", headers=h, timeout=30) - if not r.check("shape_exists", resp.status_code == 200, f"HTTP {resp.status_code}"): - return r - shape = resp.json() + try: + profile = client.resolve_training_profile(shape_id) + except Exception as e: + r.check("resolve_training_profile", False, str(e)) + return r, None - # 2. Trainer mode + r.check("resolve_training_profile", True, "OK") + + # Verify profile fields + r.check( + "training_shape_version", + bool(profile.training_shape_version), + profile.training_shape_version or "EMPTY", + ) r.check( "trainer_mode", - shape.get("trainerMode") == expected["expected_mode"], - f"got {shape.get('trainerMode')}", + profile.trainer_mode == expected["expected_mode"], + f"got={profile.trainer_mode} expected={expected['expected_mode']}", ) - - # 3. Base model r.check( - "base_model", - shape.get("baseModel") == expected["expected_model"], - f"got {shape.get('baseModel')}", + "trainer_image_tag", + bool(profile.trainer_image_tag), + profile.trainer_image_tag or "EMPTY", + ) + r.check( + "deployment_shape_version", + bool(profile.deployment_shape_version), + profile.deployment_shape_version or "EMPTY", + ) + r.check( + "accelerator_type", + bool(profile.accelerator_type), + profile.accelerator_type or "EMPTY", + ) + r.check( + "supports_lora", + profile.supports_lora, + f"supports_lora={profile.supports_lora}", ) - # 4. Has deployment shape version - deploy_sv = shape.get("deploymentShapeVersion", "") - r.check("has_deployment_shape_version", bool(deploy_sv), deploy_sv or "MISSING") - - # 5. Has trainer image tag - r.check("has_trainer_image_tag", bool(shape.get("trainerImageTag")), shape.get("trainerImageTag", "MISSING")) - - # 6. Versions exist - resp2 = httpx.get(f"{base_url}/v1/{shape_id}/versions", headers=h, timeout=30) - if not r.check("versions_accessible", resp2.status_code == 200, f"HTTP {resp2.status_code}"): - return r - - versions = resp2.json().get("trainingShapeVersions", []) - r.check("has_versions", len(versions) > 0, f"{len(versions)} total") - - # 7. Has latestValidated - latest = [v for v in versions if v.get("latestValidated")] - r.check("has_latest_validated", len(latest) == 1, f"{len(latest)} found") - - # 8. latestValidated is public - if latest: - lv = latest[0] - ver_id = lv["name"].split("/")[-1] - r.check("latest_validated_public", lv.get("public", False), f"version={ver_id} public={lv.get('public')}") - r.check("latest_validated_validated", lv.get("validated", False), f"validated={lv.get('validated')}") - else: - r.check("latest_validated_public", False, "no latestValidated to check") + if verbose: + r.check( + "accelerator_count", + True, + f"{profile.accelerator_count}", + ) + r.check( + "node_count", + True, + f"{profile.node_count}", + ) + r.check( + "max_supported_context_length", + True, + f"{profile.max_supported_context_length}", + ) + r.check( + "pipeline_parallelism", + True, + f"{profile.pipeline_parallelism}", + ) - return r + return r, profile # ───────────────────────────────────────────────────────────────────── -# Phase 2: Dependency chain access +# Phase 2: TrainerJobConfig validation # ───────────────────────────────────────────────────────────────────── -def test_dependency_chain( - api_key: str, shape_short: str, base_url: str, verbose: bool +def test_config_validation( + profile: TrainingShapeProfile, + shape_short: str, + expected: dict, + verbose: bool, ) -> TestResult: - h = headers_for(api_key) - shape_id = f"accounts/fireworks/trainingShapes/{shape_short}" - r = TestResult(name=f"{shape_short} (deps)") + """Build a TrainerJobConfig from the profile and validate it.""" + r = TestResult(name=f"{shape_short} (config)") - resp = httpx.get(f"{base_url}/v1/{shape_id}", headers=h, timeout=30) - if resp.status_code != 200: - r.check("shape_read", False, f"HTTP {resp.status_code}") + try: + config = TrainerJobConfig( + base_model=expected["expected_model"], + training_shape_ref=profile.training_shape_version, + lora_rank=16, + learning_rate=1e-5, + ) + config.validate() + r.check("config_validate", True, "OK") + except Exception as e: + r.check("config_validate", False, str(e)) return r - shape = resp.json() - - # 1. Base model accessible - base_model = shape.get("baseModel", "") - if base_model: - resp2 = httpx.get(f"{base_url}/v1/{base_model}", headers=h, timeout=30) - r.check("base_model_accessible", resp2.status_code == 200, f"{base_model} → HTTP {resp2.status_code}") - if resp2.status_code == 200: - model_data = resp2.json() - r.check( - "base_model_ready", - model_data.get("state") == "READY", - f"state={model_data.get('state')}", - ) - # 2. Deployment shape version accessible - deploy_sv = shape.get("deploymentShapeVersion", "") - if deploy_sv: - resp3 = httpx.get(f"{base_url}/v1/{deploy_sv}", headers=h, timeout=30) - r.check("deploy_shape_version_accessible", resp3.status_code == 200, f"{deploy_sv} → HTTP {resp3.status_code}") - if resp3.status_code == 200: - dsv_data = resp3.json() - r.check("deploy_shape_version_public", dsv_data.get("public", False), f"public={dsv_data.get('public')}") - r.check("deploy_shape_version_validated", dsv_data.get("validated", False), f"validated={dsv_data.get('validated')}") - - # 3. Deployment shape parent accessible - if deploy_sv: - deploy_parent = "/".join(deploy_sv.split("/")[:4]) - resp4 = httpx.get(f"{base_url}/v1/{deploy_parent}", headers=h, timeout=30) - r.check("deploy_shape_parent_accessible", resp4.status_code == 200, f"{deploy_parent} → HTTP {resp4.status_code}") - - # 4. Training shape version snapshot has correct deployment shape ref - resp5 = httpx.get(f"{base_url}/v1/{shape_id}/versions", headers=h, timeout=30) - if resp5.status_code == 200: - versions = resp5.json().get("trainingShapeVersions", []) - latest = [v for v in versions if v.get("latestValidated")] - if latest: - snapshot = latest[0].get("snapshot", {}) - snap_deploy = snapshot.get("deploymentShapeVersion", "") - r.check( - "version_snapshot_deploy_shape", - snap_deploy == deploy_sv, - f"shape={deploy_sv} snapshot={snap_deploy}", - ) + # Verify shape-owned fields are NOT set (they come from the shape) + r.check( + "no_accelerator_type", + config.accelerator_type is None, + f"accelerator_type={config.accelerator_type}", + ) + r.check( + "no_accelerator_count", + config.accelerator_count is None, + f"accelerator_count={config.accelerator_count}", + ) + r.check( + "no_node_count", + config.node_count is None, + f"node_count={config.node_count}", + ) + r.check( + "training_shape_ref_set", + config.training_shape_ref == profile.training_shape_version, + config.training_shape_ref or "EMPTY", + ) return r @@ -221,80 +229,51 @@ def test_dependency_chain( # ───────────────────────────────────────────────────────────────────── def test_job_lifecycle( - api_key: str, + mgr: TrainerJobManager, + profile: TrainingShapeProfile, shape_short: str, - account_id: str, - base_url: str, + expected: dict, verbose: bool, wait_timeout: float = 120, ) -> TestResult: - h = headers_for(api_key) - shape_id = f"accounts/fireworks/trainingShapes/{shape_short}" + """Create a real RLOR trainer job using the SDK, wait for RUNNING, delete.""" r = TestResult(name=f"{shape_short} (job)") - resp = httpx.get(f"{base_url}/v1/{shape_id}", headers=h, timeout=30) - if resp.status_code != 200: - r.check("shape_read", False, f"HTTP {resp.status_code}") - return r - shape = resp.json() - - base_model = shape["baseModel"] - accel_type = shape.get("acceleratorType", "") - accel_count = shape.get("acceleratorCount", 0) - node_count = shape.get("nodeCount", 1) - - # Create RLOR job - create_url = f"{base_url}/v1/accounts/{account_id}/rlorTrainerJobs" - body: dict[str, Any] = { - "serviceMode": True, - "keepAlive": True, - "displayName": f"e2e-test-{shape_short}", - "trainingConfig": { - "baseModel": base_model, - "loraRank": 16, - "learningRate": 1e-5, - }, - "nodeCount": node_count, - } - if accel_type: - body["trainingConfig"]["acceleratorType"] = accel_type - if accel_count: - body["trainingConfig"]["acceleratorCount"] = accel_count - - print(f" Creating job for {shape_short}...") - resp2 = httpx.post(create_url, headers=h, json=body, timeout=120) - if not r.check( - "job_create", - resp2.status_code in (200, 201), - f"HTTP {resp2.status_code}: {resp2.text[:200]}", - ): - return r - - job_data = resp2.json() - job_name = job_data.get("name", "") - job_id = job_name.split("/")[-1] - print(f" Job created: {job_id}") + config = TrainerJobConfig( + base_model=expected["expected_model"], + training_shape_ref=profile.training_shape_version, + lora_rank=16, + learning_rate=1e-5, + display_name=f"e2e-test-{shape_short}", + ) + job_id = None try: - # Wait for RUNNING + print(f" Creating job for {shape_short}...") + created = mgr.create(config) + job_id = created.job_id + r.check("job_create", True, f"job_id={job_id}") + print(f" Job created: {job_id}") + + # Poll for RUNNING deadline = time.time() + wait_timeout - final_state = job_data.get("state", "UNKNOWN") + final_state = "UNKNOWN" while time.time() < deadline: time.sleep(5) - resp3 = httpx.get(f"{base_url}/v1/{job_name}", headers=h, timeout=30) - if resp3.status_code != 200: - continue - current = resp3.json() - final_state = current.get("state", "UNKNOWN") - status = current.get("status", {}) - - if verbose: - print(f" state={final_state} status={status}") - - if final_state == "JOB_STATE_RUNNING": - break - if final_state in ("JOB_STATE_FAILED", "JOB_STATE_CANCELLED"): - break + try: + job = mgr.get(job_id) + final_state = job.get("state", "UNKNOWN") + if verbose: + print(f" state={final_state}") + if final_state == "JOB_STATE_RUNNING": + break + if final_state in ("JOB_STATE_FAILED", "JOB_STATE_CANCELLED"): + msg = job.get("status", {}).get("message", "") + r.check("job_reached_running", False, f"state={final_state} msg={msg}") + return r + except Exception as e: + if verbose: + print(f" poll error: {e}") r.check( "job_reached_running", @@ -302,15 +281,17 @@ def test_job_lifecycle( f"final_state={final_state}", ) + except Exception as e: + r.check("job_create", False, str(e)[:200]) finally: - print(f" Deleting job {job_id}...") - del_resp = httpx.delete(f"{base_url}/v1/{job_name}", headers=h, timeout=60) - r.check( - "job_delete", - del_resp.status_code in (200, 204), - f"HTTP {del_resp.status_code}", - ) - print(f" Deleted: HTTP {del_resp.status_code}") + if job_id: + print(f" Deleting job {job_id}...") + try: + mgr.delete(job_id) + r.check("job_delete", True, "OK") + print(f" Deleted") + except Exception as e: + r.check("job_delete", False, str(e)[:200]) return r @@ -319,47 +300,35 @@ def test_job_lifecycle( # Main # ───────────────────────────────────────────────────────────────────── -def print_result(r: TestResult, verbose: bool): - status = "PASS" if r.passed else "FAIL" - print(f" {status} {r.name}") - if verbose or not r.passed: - for label, ok, detail in r.checks: - marker = "✓" if ok else "✗" - print(f" {marker} {label}: {detail}") - - def main(): parser = argparse.ArgumentParser( - description="E2E test for Qwen 3.5 LoRA training shapes" + description="E2E test for Qwen 3.5 LoRA training shapes", ) parser.add_argument("--api-key", default=None, help="Fireworks API key") - parser.add_argument("--base-url", default="https://api.fireworks.ai", help="API base URL") - parser.add_argument("--verbose", "-v", action="store_true", help="Show all check details") - parser.add_argument( - "--shape", - default=None, - help="Test only this shape (e.g. qwen3p5-9b-256k-lora)", - ) + parser.add_argument("--base-url", default="https://api.fireworks.ai") + parser.add_argument("--verbose", "-v", action="store_true") + parser.add_argument("--shape", default=None, help="Test only this shape") parser.add_argument( "--create-job", action="store_true", - help="Phase 3: actually create+delete an RLOR job (costs GPU time!)", - ) - parser.add_argument( - "--wait-timeout", - type=float, - default=120, - help="Seconds to wait for job to reach RUNNING (default: 120)", + help="Phase 3: create+delete a real RLOR job (costs GPU time!)", ) + parser.add_argument("--wait-timeout", type=float, default=120) + parser.add_argument("--debug", action="store_true", help="Enable SDK debug logging") args = parser.parse_args() + if args.debug: + logging.basicConfig(level=logging.DEBUG) + elif args.verbose: + logging.basicConfig(level=logging.INFO) + api_key = args.api_key or os.environ.get("FIREWORKS_API_KEY") if not api_key: print("Error: FIREWORKS_API_KEY not set") sys.exit(1) - account_id = resolve_account_id(api_key, args.base_url) - print(f"Account: {account_id}\n") + mgr = TrainerJobManager(api_key=api_key, base_url=args.base_url) + print(f"Account: {mgr.account_id}\n") shapes = QWEN3P5_LORA_SHAPES if args.shape: @@ -370,38 +339,48 @@ def main(): shapes = {args.shape: shapes[args.shape]} all_passed = True + profiles: dict[str, TrainingShapeProfile] = {} - # ── Phase 1 ────────────────────────────────────────────────────── - print(f"Phase 1: Shape metadata ({len(shapes)} shapes)") + # ── Phase 1: resolve_training_profile ───────────────────────────── + print(f"Phase 1: resolve_training_profile ({len(shapes)} shapes)") print("-" * 60) for name, expected in shapes.items(): - r = test_shape_metadata(api_key, name, expected, args.base_url, args.verbose) + r, profile = test_resolve_profile(mgr, name, expected, args.verbose) print_result(r, args.verbose) if not r.passed: all_passed = False + if profile: + profiles[name] = profile - # ── Phase 2 ────────────────────────────────────────────────────── - print(f"\nPhase 2: Dependency chain access") + # ── Phase 2: TrainerJobConfig validation ────────────────────────── + print(f"\nPhase 2: TrainerJobConfig validation") print("-" * 60) - for name in shapes: - r = test_dependency_chain(api_key, name, args.base_url, args.verbose) + for name, expected in shapes.items(): + if name not in profiles: + print(f" SKIP {name} (config) — no profile from phase 1") + continue + r = test_config_validation(profiles[name], name, expected, args.verbose) print_result(r, args.verbose) if not r.passed: all_passed = False - # ── Phase 3 (opt-in) ───────────────────────────────────────────── + # ── Phase 3: Job lifecycle (opt-in) ─────────────────────────────── if args.create_job: print(f"\nPhase 3: Job lifecycle E2E") print("-" * 60) - for name in shapes: + for name, expected in shapes.items(): + if name not in profiles: + print(f" SKIP {name} (job) — no profile from phase 1") + continue r = test_job_lifecycle( - api_key, name, account_id, args.base_url, args.verbose, args.wait_timeout + mgr, profiles[name], name, expected, + args.verbose, args.wait_timeout, ) print_result(r, args.verbose) if not r.passed: all_passed = False - # ── Summary ────────────────────────────────────────────────────── + # ── Summary ─────────────────────────────────────────────────────── print(f"\n{'='*60}") if all_passed: print("All tests PASSED") From 7ac4464db48cad84c4bf308df48e1ce5fcd242c7 Mon Sep 17 00:00:00 2001 From: zhiweiz Date: Thu, 16 Apr 2026 22:31:07 +0000 Subject: [PATCH 5/5] Add cookbook-test to gitignore Co-authored-by: zhiweiz --- .gitignore | 1 + cookbook-test | 1 + 2 files changed, 2 insertions(+) create mode 160000 cookbook-test diff --git a/.gitignore b/.gitignore index 6f747cc..7cd7846 100644 --- a/.gitignore +++ b/.gitignore @@ -127,3 +127,4 @@ dist .yarn/build-state.yml .yarn/install-state.gz .pnp.* +cookbook-test/ diff --git a/cookbook-test b/cookbook-test new file mode 160000 index 0000000..4129995 --- /dev/null +++ b/cookbook-test @@ -0,0 +1 @@ +Subproject commit 4129995e3e50b6bd601a1041f9b33a147a46d822