ENG-4782: prime train models — show FFT models by default

packages/prime/src/prime_cli/api/training.py

@@ -6,11 +6,12 @@
 token; admin role is gated server-side.
 """
 
+from datetime import datetime
 from typing import Any, Dict, List, Optional
 
 from pydantic import BaseModel, ConfigDict, Field
 
-from prime_cli.core import APIClient
+from prime_cli.core import APIClient, APIError, NotFoundError, UnauthorizedError
 
 
 class HostedTrainingRunResponse(BaseModel):
@@ -30,6 +31,38 @@ 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 — mirrors the backend
+    `FFTModelClusterInfo` in `platform/backend/app/packages/training/schemas.py`.
+    """
+
+    cluster_id: str = Field(..., alias="clusterId")
+    cluster_name: str = Field(..., alias="clusterName")
+    gpu_type: Optional[str] = Field(None, alias="gpuType")
+    cache_synced_at: Optional[datetime] = Field(None, alias="cacheSyncedAt")
+
+    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 +104,35 @@ 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: Optional[str] = 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.
+
+        Returns an empty list when the caller has no eligible clusters
+        (403), when the endpoint is not deployed yet on the target
+        backend (404) or when auth is missing (401) — the CLI is called
+        in a "silent when empty" context alongside the LoRA listing and
+        should not crash the primary output on a still-rolling-out
+        endpoint.
+        """
+        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, UnauthorizedError):
+            return []
+        except APIError as exc:
+            # 403 (not a team member) surfaces as APIError with an
+            # HTTP 403 prefix. Treat like NotFound so a personal caller
+            # querying with an unrelated team_id doesn't fail the whole
+            # command.
+            if "HTTP 403" in str(exc):
+                return []
+            raise
+        return AvailableFFTModelsResponse.model_validate(response).models
+
 
 def build_payload_from_toml(
     cfg: Dict[str, Any],

packages/prime/src/prime_cli/commands/rl.py

@@ -70,6 +70,8 @@
     "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?, cache_synced_at?}]}",
 )
 
 RL_LIST_JSON_HELP = json_output_help(
@@ -1673,69 +1675,148 @@ 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)
+            payload: Dict[str, Any] = {"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:
+            _render_lora_models_table(models)
+
+        if fft_models:
+            _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_lora_models_table(models: List) -> None:
+    """Render the classic LoRA Hosted Training model listing."""
+    if not models:
+        console.print("[yellow]No models available for Hosted Training.[/yellow]")
+        console.print("[dim]This could mean no healthy Hosted Training clusters are running.[/dim]")
+        return
+
+    table = Table(
+        title="Hosted Training - Models",
+    )
+    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")
+    table.add_column("Model", style="cyan")
+    table.add_column("GPU Type(s)", style="magenta")
+    table.add_column("Available On", style="green")
+    table.add_column("Cache Synced", style="dim", justify="right")
+
+    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})
+        cluster_names = sorted({c.cluster_name for c in model.clusters})
+        latest_sync = max(
+            (c.cache_synced_at for c in model.clusters if c.cache_synced_at),
+            default=None,
+        )
+        table.add_row(
+            model.name,
+            ", ".join(gpu_types) or "-",
+            ", ".join(cluster_names) or "-",
+            latest_sync.strftime("%Y-%m-%d %H:%M") if latest_sync else "-",
         )
-        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. "
+        "Cache Synced shows the most recent cluster manifest.[/dim]"
+    )
+    console.print(table)
 
 
 @app.command("gpus", rich_help_panel="Commands")