diff --git a/packages/prime/src/prime_cli/api/training.py b/packages/prime/src/prime_cli/api/training.py index 7b40d65d..a8852fb6 100644 --- a/packages/prime/src/prime_cli/api/training.py +++ b/packages/prime/src/prime_cli/api/training.py @@ -9,8 +9,9 @@ from typing import Any, Dict, List, Optional from pydantic import BaseModel, ConfigDict, Field +from pydantic import ValidationError as PydanticValidationError -from prime_cli.core import APIClient +from prime_cli.core import APIClient, APIError, NotFoundError class HostedTrainingRunResponse(BaseModel): @@ -30,6 +31,42 @@ class AvailableGpuTypesResponse(BaseModel): model_config = ConfigDict(populate_by_name=True) +class FFTModelClusterInfo(BaseModel): + """One cluster the caller can dispatch an FFT run to that already has + the parent model cached. + + Intentionally narrower than the backend `FFTModelClusterInfo` + schema: `cache_synced_at` is dropped because cache freshness is an + implementation detail the user shouldn't reason about — the + dispatch picker already gates on PRESENT-cache clusters, so any + entry surfaced here is warm. + """ + + cluster_id: str = Field(..., alias="clusterId") + cluster_name: str = Field(..., alias="clusterName") + gpu_type: str | None = Field(None, alias="gpuType") + + model_config = ConfigDict(populate_by_name=True) + + +class AvailableFFTModel(BaseModel): + """A model that is cached and ready for FFT dispatch on at least one + eligible PrimeCluster.""" + + name: str = Field(..., description="Model name") + clusters: list[FFTModelClusterInfo] = Field(default_factory=list) + + model_config = ConfigDict(populate_by_name=True) + + +class AvailableFFTModelsResponse(BaseModel): + """Response from GET /v1/training/available-fft-models.""" + + models: list[AvailableFFTModel] = Field(default_factory=list) + + model_config = ConfigDict(populate_by_name=True) + + class HostedTrainingClient: """Client for the hosted full-FT training endpoint.""" @@ -71,6 +108,34 @@ def list_available_gpu_types(self, team_id: Optional[str] = None) -> AvailableGp response = self.client.get("/training/available-gpu-types", params=params) return AvailableGpuTypesResponse.model_validate(response) + def list_available_fft_models(self, team_id: str | None = None) -> list[AvailableFFTModel]: + """GET /v1/training/available-fft-models. Models that are already + cached on at least one PrimeCluster the caller can dispatch a + full-FT run to. + + 404 is swallowed to an empty list so the CLI still renders on + older backends that haven't shipped the endpoint yet. Every + other error (auth failure, forbidden, server errors) propagates + — the caller decides whether to surface or hide it based on + whether the LoRA section already ran. + + A schema-drifted response (pydantic ValidationError) is + re-raised as APIError so the command layer's existing + `except APIError` fallback catches it — otherwise a + non-conforming backend payload would kill the LoRA table too. + """ + params: dict[str, Any] = {} + if team_id: + params["team_id"] = team_id + try: + response = self.client.get("/training/available-fft-models", params=params) + except NotFoundError: + return [] + try: + return AvailableFFTModelsResponse.model_validate(response).models + except PydanticValidationError as exc: + raise APIError(f"Failed to parse available FFT models response: {exc}") from exc + def build_payload_from_toml( cfg: Dict[str, Any], diff --git a/packages/prime/src/prime_cli/commands/rl.py b/packages/prime/src/prime_cli/commands/rl.py index b5f0be35..4315b379 100644 --- a/packages/prime/src/prime_cli/commands/rl.py +++ b/packages/prime/src/prime_cli/commands/rl.py @@ -70,6 +70,7 @@ "effective_training_price_per_mtok?, " "effective_inference_input_price_per_mtok?, " "effective_inference_output_price_per_mtok?, promo_label?}", + ".available_fft_models[]? = {name, clusters[{cluster_id, cluster_name, gpu_type?}]}", ) RL_LIST_JSON_HELP = json_output_help( @@ -1673,69 +1674,168 @@ def _format(list_p: Any, eff_p: Any) -> str: @app.command("models", rich_help_panel="Commands", epilog=RL_MODELS_JSON_HELP) def list_models( output: str = typer.Option("table", "--output", "-o", help="Output format: table or json"), + fft_only: bool = typer.Option( + False, + "--fft-only", + help="Suppress the LoRA/Hosted Training section and only show FFT models.", + ), ) -> None: - """List available models for Hosted Training.""" + """List available models for Hosted Training. + + Renders two sections when both are populated: the shared-cluster LoRA + models (pricing per 1M tokens) and the models pre-cached on FFT + dispatch clusters. The FFT section is silent when the API returns + nothing so users on backends without the endpoint see the pre-4782 + output. + """ + from ..api.training import AvailableFFTModel, HostedTrainingClient + validate_output_format(output, console) try: api_client = APIClient() rl_client = RLClient(api_client) + training_client = HostedTrainingClient(api_client) config = Config() - models = rl_client.list_models(team_id=config.team_id) + models: list = [] + if not fft_only: + models = rl_client.list_models(team_id=config.team_id) + + fft_models: list[AvailableFFTModel] = [] + try: + fft_models = training_client.list_available_fft_models(team_id=config.team_id) + except APIError: + # Never let an FFT fetch failure break the LoRA output — the + # endpoint is younger and may still be rolling out. When the + # caller explicitly asked for FFT-only we bubble up the + # error instead of silently hiding it. + if fft_only: + raise + fft_models = [] if output == "json": - output_data_as_json({"models": [m.model_dump() for m in models]}, console) + if fft_only: + # --fft-only: LoRA wasn't fetched so emitting `models` + # would be misleading. Always include the FFT key + # (empty allowed) so scripts using the flag can rely + # on `.available_fft_models[]` being present. + payload: dict[str, Any] = { + "available_fft_models": [m.model_dump() for m in fft_models] + } + else: + payload = {"models": [m.model_dump() for m in models]} + if fft_models: + payload["available_fft_models"] = [m.model_dump() for m in fft_models] + output_data_as_json(payload, console) return - if not models: - console.print("[yellow]No models available for Hosted Training.[/yellow]") + if not fft_only: + if models: + _render_lora_models_table(models) + elif not fft_models: + # Both sections empty — surface the LoRA fallback so + # the user sees *something*. When FFT has data we + # skip this so the FFT table isn't preceded by a + # misleading "no models available" banner. + _render_empty_lora_message() + + if fft_models: + if not fft_only and models: + # Visual break only when the LoRA table above actually + # rendered — otherwise we'd print a stray blank line. + console.print() + _render_fft_models_table(fft_models) + elif fft_only: + console.print("[yellow]No FFT models available.[/yellow]") console.print( - "[dim]This could mean no healthy Hosted Training clusters are running.[/dim]" + "[dim]No dispatchable clusters have a warm model cache yet, or the " + "endpoint isn't deployed on this backend.[/dim]" ) - return - table = Table( - title="Hosted Training - Models", + except APIError as e: + console.print(f"[red]Error:[/red] {e}") + raise typer.Exit(1) + + +def _render_empty_lora_message() -> None: + """Print the LoRA-empty fallback used when neither section has data.""" + console.print("[yellow]No models available for Hosted Training.[/yellow]") + console.print("[dim]This could mean no healthy Hosted Training clusters are running.[/dim]") + + +def _render_lora_models_table(models: list) -> None: + """Render the classic LoRA Hosted Training model listing. + + Caller is responsible for handling the empty-list case (via + `_render_empty_lora_message`) — this helper assumes at least one + row to render. + """ + table = Table( + title="Hosted Training — LoRA", + title_justify="left", + caption_justify="left", + ) + table.add_column("Model", style="cyan") + table.add_column("Status") + table.add_column("Input", style="green", justify="right") + table.add_column("Output", style="green", justify="right") + table.add_column("Train", style="green", justify="right") + + promo_labels: list[str] = [] + for model in sorted(models, key=lambda model: _model_name_sort_key(model.name)): + if model.at_capacity: + status = "[red]At Capacity[/red]" + else: + status = "[green]Available[/green]" + list_train, eff_train = model.resolve_prices("training") + list_input, eff_input = model.resolve_prices("inference_input") + list_output, eff_output = model.resolve_prices("inference_output") + table.add_row( + model.name, + status, + format_promo_price(list_input, eff_input) or "-", + format_promo_price(list_output, eff_output) or "-", + format_promo_price(list_train, eff_train) or "-", ) - table.add_column("Model", style="cyan") - table.add_column("Status") - table.add_column("Input", style="green", justify="right") - table.add_column("Output", style="green", justify="right") - table.add_column("Train", style="green", justify="right") - - promo_labels: List[str] = [] - for model in sorted(models, key=lambda model: _model_name_sort_key(model.name)): - if model.at_capacity: - status = "[red]At Capacity[/red]" - else: - status = "[green]Available[/green]" - list_train, eff_train = model.resolve_prices("training") - list_input, eff_input = model.resolve_prices("inference_input") - list_output, eff_output = model.resolve_prices("inference_output") - table.add_row( - model.name, - status, - format_promo_price(list_input, eff_input) or "-", - format_promo_price(list_output, eff_output) or "-", - format_promo_price(list_train, eff_train) or "-", - ) - if model.promo_label and model.promo_label not in promo_labels: - promo_labels.append(model.promo_label) + if model.promo_label and model.promo_label not in promo_labels: + promo_labels.append(model.promo_label) + + caption = ( + "[dim]Prices are per 1M tokens. All models support context windows of 64K tokens.[/dim]" + ) + caption_lines = [caption] + if promo_labels: + joined = ", ".join(rich_escape(label) for label in promo_labels) + caption_lines.append(f"[bold yellow]{joined}[/bold yellow]") + table.caption = "\n".join(caption_lines) + console.print(table) - caption = ( - "[dim]Prices are per 1M tokens. All models support context windows of 64K tokens.[/dim]" + +def _render_fft_models_table(fft_models: list) -> None: + """Render the FFT dispatch model listing. + + Each model may be cached on multiple clusters/GPU types; we collapse + those into a single row with cluster+GPU joined so the table stays + scannable when a large model is warm everywhere. + """ + table = Table( + title="Hosted Training — Full Finetuning", + title_justify="left", + caption_justify="left", + ) + table.add_column("Model", style="cyan") + table.add_column("GPU Type(s)", style="magenta") + + for model in sorted(fft_models, key=lambda m: _model_name_sort_key(m.name)): + gpu_types = sorted({c.gpu_type for c in model.clusters if c.gpu_type}) + table.add_row( + model.name, + ", ".join(gpu_types) or "-", ) - caption_lines = [caption] - if promo_labels: - joined = ", ".join(rich_escape(label) for label in promo_labels) - caption_lines.append(f"[bold yellow]{joined}[/bold yellow]") - table.caption = "\n".join(caption_lines) - console.print(table) - except APIError as e: - console.print(f"[red]Error:[/red] {e}") - raise typer.Exit(1) + table.caption = "[dim]Models pre-cached on clusters you can dispatch FFT runs to.[/dim]" + console.print(table) @app.command("gpus", rich_help_panel="Commands") diff --git a/packages/prime/tests/test_rl_models.py b/packages/prime/tests/test_rl_models.py index 1101e769..a0cb11cf 100644 --- a/packages/prime/tests/test_rl_models.py +++ b/packages/prime/tests/test_rl_models.py @@ -5,6 +5,7 @@ import pytest from prime_cli.commands.rl import _model_name_sort_key +from prime_cli.core.client import NotFoundError from prime_cli.main import app from prime_cli.utils.formatters import strip_ansi from typer.testing import CliRunner @@ -37,11 +38,59 @@ def _models_payload() -> Dict[str, Any]: } -def _mock_get_factory(calls: List[str]): - def mock_get(self: Any, endpoint: str, params: Dict[str, Any] | None = None) -> Dict[str, Any]: +def _fft_models_payload() -> dict[str, Any]: + return { + "models": [ + { + "name": "meta-llama/Llama-3.1-8B-Instruct", + "clusters": [ + { + "clusterId": "cluster-a", + "clusterName": "athens", + "gpuType": "H200_141GB", + "cacheSyncedAt": "2026-06-01T10:15:00Z", + }, + { + "clusterId": "cluster-b", + "clusterName": "berlin", + "gpuType": "H100_80GB", + "cacheSyncedAt": "2026-06-02T08:00:00Z", + }, + ], + }, + { + "name": "qwen/qwen3-8b", + "clusters": [ + { + "clusterId": "cluster-a", + "clusterName": "athens", + "gpuType": "H200_141GB", + "cacheSyncedAt": "2026-06-01T10:15:00Z", + } + ], + }, + ] + } + + +def _mock_get_factory( + calls: list[str], + *, + fft_payload: dict[str, Any] | None = None, +): + """Mock APIClient.get for the models command. + + By default returns an empty FFT list so tests that pre-date the FFT + endpoint continue to exercise the LoRA-only rendering path. Pass + ``fft_payload`` to opt into a populated FFT response. + """ + + def mock_get(self: Any, endpoint: str, params: dict[str, Any] | None = None) -> dict[str, Any]: calls.append(endpoint) if endpoint == "/rft/models": return _models_payload() + if endpoint == "/training/available-fft-models": + return fft_payload if fft_payload is not None else {"models": []} raise AssertionError(f"Unexpected endpoint: {endpoint}") return mock_get @@ -60,7 +109,9 @@ def test_models_table_renders_pricing(monkeypatch: pytest.MonkeyPatch) -> None: assert "$3" in result.output # Null pricing renders as a dash. assert "-" in result.output - assert calls == ["/rft/models"] + # LoRA endpoint is always hit; FFT endpoint is polled every time so + # the table shows up transparently when it starts returning data. + assert calls == ["/rft/models", "/training/available-fft-models"] def test_models_json_includes_pricing(monkeypatch: pytest.MonkeyPatch) -> None: @@ -82,7 +133,11 @@ def test_models_handles_backend_without_pricing_fields( """Older backends may not return the pricing fields at all.""" def mock_get(self: Any, endpoint: str, params: Dict[str, Any] | None = None) -> Dict[str, Any]: - return {"models": [{"name": "qwen/qwen3-8b", "atCapacity": False}]} + if endpoint == "/rft/models": + return {"models": [{"name": "qwen/qwen3-8b", "atCapacity": False}]} + if endpoint == "/training/available-fft-models": + return {"models": []} + raise AssertionError(f"Unexpected endpoint: {endpoint}") monkeypatch.setattr("prime_cli.core.APIClient.get", mock_get) @@ -114,6 +169,8 @@ def test_models_table_renders_promo_arrow_and_caption( monkeypatch: pytest.MonkeyPatch, ) -> None: def mock_get(self: Any, endpoint: str, params: Dict[str, Any] | None = None) -> Dict[str, Any]: + if endpoint == "/training/available-fft-models": + return {"models": []} return _promo_payload() monkeypatch.setattr("prime_cli.core.APIClient.get", mock_get) @@ -135,6 +192,21 @@ def mock_get(self: Any, endpoint: str, params: Dict[str, Any] | None = None) -> assert plain.count("Free RFT week") == 1 +def _lora_only_mock(payload: dict[str, Any]): + """Return a mock_get that serves ``payload`` for /rft/models and an + empty FFT list for /training/available-fft-models — the shape most + LoRA-focused tests want.""" + + def mock_get(self: Any, endpoint: str, params: dict[str, Any] | None = None) -> dict[str, Any]: + if endpoint == "/rft/models": + return payload + if endpoint == "/training/available-fft-models": + return {"models": []} + raise AssertionError(f"Unexpected endpoint: {endpoint}") + + return mock_get + + def test_models_table_no_promo_when_effective_equals_original( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -154,10 +226,7 @@ def test_models_table_no_promo_when_effective_equals_original( ] } - def mock_get(self: Any, endpoint: str, params: Dict[str, Any] | None = None) -> Dict[str, Any]: - return payload - - monkeypatch.setattr("prime_cli.core.APIClient.get", mock_get) + monkeypatch.setattr("prime_cli.core.APIClient.get", _lora_only_mock(payload)) result = CliRunner().invoke(app, ["rl", "models"], env={"COLUMNS": "200"}) @@ -186,10 +255,7 @@ def test_models_zero_original_with_promo_does_not_render_free( ] } - def mock_get(self: Any, endpoint: str, params: Dict[str, Any] | None = None) -> Dict[str, Any]: - return payload - - monkeypatch.setattr("prime_cli.core.APIClient.get", mock_get) + monkeypatch.setattr("prime_cli.core.APIClient.get", _lora_only_mock(payload)) result = CliRunner().invoke(app, ["rl", "models"], env={"COLUMNS": "200"}) @@ -228,10 +294,7 @@ def test_models_promo_label_deduplicated_across_models( ] } - def mock_get(self: Any, endpoint: str, params: Dict[str, Any] | None = None) -> Dict[str, Any]: - return payload - - monkeypatch.setattr("prime_cli.core.APIClient.get", mock_get) + monkeypatch.setattr("prime_cli.core.APIClient.get", _lora_only_mock(payload)) result = CliRunner().invoke(app, ["rl", "models"], env={"COLUMNS": "200"}) @@ -263,10 +326,7 @@ def test_models_table_renders_promo_with_list_fields( ] } - def mock_get(self: Any, endpoint: str, params: Dict[str, Any] | None = None) -> Dict[str, Any]: - return payload - - monkeypatch.setattr("prime_cli.core.APIClient.get", mock_get) + monkeypatch.setattr("prime_cli.core.APIClient.get", _lora_only_mock(payload)) result = CliRunner().invoke(app, ["rl", "models"], env={"COLUMNS": "200"}) @@ -282,6 +342,8 @@ def mock_get(self: Any, endpoint: str, params: Dict[str, Any] | None = None) -> def test_models_json_includes_effective_fields(monkeypatch: pytest.MonkeyPatch) -> None: def mock_get(self: Any, endpoint: str, params: Dict[str, Any] | None = None) -> Dict[str, Any]: + if endpoint == "/training/available-fft-models": + return {"models": []} return _promo_payload() monkeypatch.setattr("prime_cli.core.APIClient.get", mock_get) @@ -348,3 +410,345 @@ def test_model_name_sort_key_handles_active_params_case_insensitively() -> None: "org/model-30B-A10b", "org/model-30B", ] + + +def test_models_command_renders_fft_section_when_available( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Both LoRA and FFT tables render side by side when the FFT endpoint + returns any results.""" + calls: list[str] = [] + monkeypatch.setattr( + "prime_cli.core.APIClient.get", + _mock_get_factory(calls, fft_payload=_fft_models_payload()), + ) + + result = CliRunner().invoke(app, ["train", "models"], env={"COLUMNS": "200"}) + + assert result.exit_code == 0, result.output + plain = strip_ansi(result.output) + # LoRA table survives. + assert "LoRA" in plain + assert "qwen/qwen3-8b" in plain + # FFT table shows up too. + assert "Full Finetuning" in plain + assert "meta-llama/Llama-3.1-8B-Instruct" in plain + # Model cached on two clusters → two GPU types collapse into the row. + assert "H100_80GB" in plain + assert "H200_141GB" in plain + # Cluster names are intentionally not rendered in the table; users + # dispatch by gpu_type, not by cluster. + assert "athens" not in plain + assert "berlin" not in plain + # Both endpoints were hit. + assert "/rft/models" in calls + assert "/training/available-fft-models" in calls + + +def test_models_json_output_includes_available_fft_models( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + "prime_cli.core.APIClient.get", + _mock_get_factory([], fft_payload=_fft_models_payload()), + ) + + result = CliRunner().invoke(app, ["train", "models", "--output", "json"]) + + assert result.exit_code == 0, result.output + data = json.loads(result.output) + assert [m["name"] for m in data["models"]] == [ + "qwen/qwen3-8b", + "openai/gpt-oss-20b", + ] + fft = data["available_fft_models"] + assert [m["name"] for m in fft] == [ + "meta-llama/Llama-3.1-8B-Instruct", + "qwen/qwen3-8b", + ] + first = fft[0] + assert [c["cluster_name"] for c in first["clusters"]] == ["athens", "berlin"] + assert first["clusters"][0]["gpu_type"] == "H200_141GB" + # cache_synced_at is intentionally omitted from CLI output. + assert "cache_synced_at" not in first["clusters"][0] + + +def test_models_json_omits_fft_key_when_empty(monkeypatch: pytest.MonkeyPatch) -> None: + """JSON output stays backwards compatible: no available_fft_models + key when the FFT endpoint returns an empty list.""" + monkeypatch.setattr("prime_cli.core.APIClient.get", _mock_get_factory([])) + + result = CliRunner().invoke(app, ["train", "models", "--output", "json"]) + + assert result.exit_code == 0, result.output + data = json.loads(result.output) + assert "models" in data + assert "available_fft_models" not in data + + +def test_models_json_fft_only_always_includes_fft_key( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """With --fft-only, scripts specifically consume + `.available_fft_models[]`. The key must be present even when the + endpoint returned zero models, and the misleading `models: []` key + must be absent since LoRA wasn't fetched.""" + + def mock_get(self: Any, endpoint: str, params: dict[str, Any] | None = None) -> dict[str, Any]: + if endpoint == "/training/available-fft-models": + return {"models": []} + raise AssertionError(f"Unexpected endpoint: {endpoint}") + + monkeypatch.setattr("prime_cli.core.APIClient.get", mock_get) + + result = CliRunner().invoke(app, ["train", "models", "--fft-only", "--output", "json"]) + + assert result.exit_code == 0, result.output + data = json.loads(result.output) + assert data == {"available_fft_models": []} + + +def test_models_json_fft_only_with_data(monkeypatch: pytest.MonkeyPatch) -> None: + """--fft-only --output json with populated data: still no `models` + key; `available_fft_models` carries the payload.""" + + def mock_get(self: Any, endpoint: str, params: dict[str, Any] | None = None) -> dict[str, Any]: + if endpoint == "/training/available-fft-models": + return _fft_models_payload() + raise AssertionError(f"Unexpected endpoint: {endpoint}") + + monkeypatch.setattr("prime_cli.core.APIClient.get", mock_get) + + result = CliRunner().invoke(app, ["train", "models", "--fft-only", "--output", "json"]) + + assert result.exit_code == 0, result.output + data = json.loads(result.output) + assert "models" not in data + assert [m["name"] for m in data["available_fft_models"]] == [ + "meta-llama/Llama-3.1-8B-Instruct", + "qwen/qwen3-8b", + ] + + +def test_models_command_survives_fft_endpoint_404( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Older backends that haven't shipped the FFT endpoint yet should + still get the LoRA listing rendered — the CLI must not crash.""" + + def mock_get(self: Any, endpoint: str, params: dict[str, Any] | None = None) -> dict[str, Any]: + if endpoint == "/rft/models": + return _models_payload() + if endpoint == "/training/available-fft-models": + raise NotFoundError("HTTP 404: available-fft-models not deployed on this backend") + raise AssertionError(f"Unexpected endpoint: {endpoint}") + + monkeypatch.setattr("prime_cli.core.APIClient.get", mock_get) + + result = CliRunner().invoke(app, ["train", "models"], env={"COLUMNS": "200"}) + + assert result.exit_code == 0, result.output + plain = strip_ansi(result.output) + assert "qwen/qwen3-8b" in plain + assert "Full Finetuning" not in plain + + +def test_models_fft_only_suppresses_lora_section( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[str] = [] + + def mock_get(self: Any, endpoint: str, params: dict[str, Any] | None = None) -> dict[str, Any]: + calls.append(endpoint) + if endpoint == "/training/available-fft-models": + return _fft_models_payload() + raise AssertionError(f"Unexpected endpoint: {endpoint}") + + monkeypatch.setattr("prime_cli.core.APIClient.get", mock_get) + + result = CliRunner().invoke(app, ["train", "models", "--fft-only"], env={"COLUMNS": "200"}) + + assert result.exit_code == 0, result.output + plain = strip_ansi(result.output) + assert "Full Finetuning" in plain + assert "meta-llama/Llama-3.1-8B-Instruct" in plain + # LoRA table title should not appear when --fft-only is set. + assert "LoRA" not in plain + # Only the FFT endpoint was fetched. + assert calls == ["/training/available-fft-models"] + + +def test_list_available_fft_models_returns_empty_on_404( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """API client method swallows 404 so `prime train models` can silently + fall back to LoRA-only rendering on backends that haven't shipped the + endpoint yet.""" + from prime_cli.api.training import HostedTrainingClient + from prime_cli.core import APIClient + + def mock_get(self: Any, endpoint: str, params: dict[str, Any] | None = None) -> dict[str, Any]: + raise NotFoundError("HTTP 404: not found") + + monkeypatch.setattr("prime_cli.core.APIClient.get", mock_get) + client = HostedTrainingClient(APIClient()) + assert client.list_available_fft_models(team_id=None) == [] + + +def test_list_available_fft_models_propagates_auth_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Auth failures must not be silently converted to an empty list — + otherwise `prime train models --fft-only` on an expired key would + exit 0 with a misleading 'No FFT models available' message.""" + from prime_cli.api.training import HostedTrainingClient + from prime_cli.core import APIClient + from prime_cli.core.client import UnauthorizedError + + def mock_get(self: Any, endpoint: str, params: dict[str, Any] | None = None) -> dict[str, Any]: + raise UnauthorizedError("API key unauthorized.") + + monkeypatch.setattr("prime_cli.core.APIClient.get", mock_get) + client = HostedTrainingClient(APIClient()) + with pytest.raises(UnauthorizedError): + client.list_available_fft_models(team_id=None) + + +def test_models_fft_only_surfaces_auth_error(monkeypatch: pytest.MonkeyPatch) -> None: + """--fft-only skips the LoRA call, so an auth failure on the FFT + endpoint is the only signal the caller has that their token is + bad. It must exit non-zero with the auth error, not print the + generic 'no models' fallback.""" + from prime_cli.core.client import UnauthorizedError + + def mock_get(self: Any, endpoint: str, params: dict[str, Any] | None = None) -> dict[str, Any]: + if endpoint == "/training/available-fft-models": + raise UnauthorizedError("API key unauthorized.") + raise AssertionError(f"Unexpected endpoint: {endpoint}") + + monkeypatch.setattr("prime_cli.core.APIClient.get", mock_get) + + result = CliRunner().invoke(app, ["train", "models", "--fft-only"], env={"COLUMNS": "200"}) + + assert result.exit_code != 0 + assert "unauthorized" in result.output.lower() + assert "No FFT models available" not in result.output + + +def test_models_default_hides_fft_auth_error_after_lora_succeeds( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """In default mode, if LoRA succeeds but the FFT call fails with + anything other than 404, the LoRA table must still render — the FFT + section is best-effort and shouldn't cascade the primary output.""" + from prime_cli.core.client import UnauthorizedError + + def mock_get(self: Any, endpoint: str, params: dict[str, Any] | None = None) -> dict[str, Any]: + if endpoint == "/rft/models": + return _models_payload() + if endpoint == "/training/available-fft-models": + raise UnauthorizedError("API key unauthorized.") + raise AssertionError(f"Unexpected endpoint: {endpoint}") + + monkeypatch.setattr("prime_cli.core.APIClient.get", mock_get) + + result = CliRunner().invoke(app, ["train", "models"], env={"COLUMNS": "200"}) + + assert result.exit_code == 0, result.output + plain = strip_ansi(result.output) + assert "qwen/qwen3-8b" in plain + assert "Full Finetuning" not in plain + + +def test_list_available_fft_models_converts_pydantic_error_to_apierror( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A schema-drifted backend payload must surface as APIError, not + raw pydantic ValidationError — otherwise the command's `except + APIError` fallback would miss it and the LoRA table would fail to + render alongside the FFT section (Bugbot finding on df0e6269).""" + from prime_cli.api.training import HostedTrainingClient + from prime_cli.core import APIClient, APIError + + def mock_get(self: Any, endpoint: str, params: dict[str, Any] | None = None) -> dict[str, Any]: + # Missing required "clusters" list, wrong shape on "name". + return {"models": [{"name": 12345, "clusters": "not-a-list"}]} + + monkeypatch.setattr("prime_cli.core.APIClient.get", mock_get) + client = HostedTrainingClient(APIClient()) + with pytest.raises(APIError): + client.list_available_fft_models(team_id=None) + + +def test_models_command_survives_fft_schema_drift( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """End-to-end: even if the FFT endpoint returns an unparseable + payload, `prime train models` should still emit the LoRA table + rather than exiting with an unhandled traceback.""" + + def mock_get(self: Any, endpoint: str, params: dict[str, Any] | None = None) -> dict[str, Any]: + if endpoint == "/rft/models": + return _models_payload() + if endpoint == "/training/available-fft-models": + return {"models": [{"name": 12345, "clusters": "not-a-list"}]} + raise AssertionError(f"Unexpected endpoint: {endpoint}") + + monkeypatch.setattr("prime_cli.core.APIClient.get", mock_get) + + result = CliRunner().invoke(app, ["train", "models"], env={"COLUMNS": "200"}) + + assert result.exit_code == 0, result.output + plain = strip_ansi(result.output) + assert "qwen/qwen3-8b" in plain + assert "Full Finetuning" not in plain + + +def test_models_command_suppresses_lora_empty_banner_when_fft_populated( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """LoRA-empty + FFT-populated: the 'No models available for Hosted + Training' banner would mislead readers into thinking the whole + command failed. Only render the FFT section in that case.""" + + def mock_get(self: Any, endpoint: str, params: dict[str, Any] | None = None) -> dict[str, Any]: + if endpoint == "/rft/models": + return {"models": []} + if endpoint == "/training/available-fft-models": + return _fft_models_payload() + raise AssertionError(f"Unexpected endpoint: {endpoint}") + + monkeypatch.setattr("prime_cli.core.APIClient.get", mock_get) + + result = CliRunner().invoke(app, ["train", "models"], env={"COLUMNS": "200"}) + + assert result.exit_code == 0, result.output + plain = strip_ansi(result.output) + # The misleading empty-LoRA banner must NOT appear. + assert "No models available for Hosted Training" not in plain + # FFT section still renders. + assert "Full Finetuning" in plain + assert "meta-llama/Llama-3.1-8B-Instruct" in plain + + +def test_models_command_shows_lora_empty_banner_when_both_empty( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Both sections empty: the LoRA fallback banner still surfaces so + the user isn't left with a completely blank command output.""" + + def mock_get(self: Any, endpoint: str, params: dict[str, Any] | None = None) -> dict[str, Any]: + if endpoint == "/rft/models": + return {"models": []} + if endpoint == "/training/available-fft-models": + return {"models": []} + raise AssertionError(f"Unexpected endpoint: {endpoint}") + + monkeypatch.setattr("prime_cli.core.APIClient.get", mock_get) + + result = CliRunner().invoke(app, ["train", "models"], env={"COLUMNS": "200"}) + + assert result.exit_code == 0, result.output + plain = strip_ansi(result.output) + assert "No models available for Hosted Training" in plain + assert "Full Finetuning" not in plain diff --git a/uv.lock b/uv.lock index f1c79e93..c99fcc3e 100644 --- a/uv.lock +++ b/uv.lock @@ -12,12 +12,10 @@ exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for exclude-newer-span = "P7D" [options.exclude-newer-package] -verifiers = false -tasksets = false -prime-pydantic-config = "2099-12-31T23:00:00Z" prime-tunnel = false prime-sandboxes = false -renderers = "2099-12-31T23:00:00Z" +verifiers = false +tasksets = false harnesses = false prime-evals = false prime = false