diff --git a/truss/cli/loops_commands.py b/truss/cli/loops_commands.py index bff71cd68..6c25a01f3 100644 --- a/truss/cli/loops_commands.py +++ b/truss/cli/loops_commands.py @@ -277,6 +277,201 @@ def _render_loops_runs(runs: List[Dict[str, Any]]) -> None: console.print(table) +@loops.command(name="metrics") +@click.option( + "--deployment-id", + type=str, + required=False, + help="Loops deployment ID to fetch metrics for.", +) +@click.option( + "--base-model", + type=str, + required=False, + help=( + "Base model name; resolves to the caller's active Loops deployment " + "for that model. Errors if zero or more than one active deployment " + "matches." + ), +) +@click.option("--remote", type=str, required=False, help="Remote to use.") +@common.common_options() +def loops_metrics( + deployment_id: Optional[str], base_model: Optional[str], remote: Optional[str] +) -> None: + """Show utilization + request metrics for a Loops deployment. + + Pass exactly one of ``--deployment-id`` or ``--base-model``. Prints a + single-snapshot view of GPU/CPU/memory utilization per pod plus + service-level request metrics from the Knative queue-proxy. + """ + if deployment_id and base_model: + raise click.UsageError("Pass either --deployment-id or --base-model, not both.") + if not deployment_id and not base_model: + raise click.UsageError( + "Pass --deployment-id or --base-model to identify a Loops deployment." + ) + + if not remote: + remote = remote_cli.inquire_remote_name() + remote_provider: BasetenRemote = cast( + BasetenRemote, RemoteFactory.create(remote=remote) + ) + + resolved_deployment_id = deployment_id or _resolve_deployment_id_by_base_model( + remote_provider, cast(str, base_model) + ) + payload = remote_provider.api.get_loops_deployment_metrics( + deployment_id=resolved_deployment_id + ) + _render_loops_metrics(payload) + + +def _resolve_deployment_id_by_base_model( + remote_provider: BasetenRemote, base_model: str +) -> str: + deployments = remote_provider.api.list_loops_deployments() + matches = [d for d in deployments if d.get("base_model") == base_model] + if not matches: + raise click.UsageError( + f"No active Loops deployment found for base model {base_model!r}." + ) + if len(matches) > 1: + ids = ", ".join(d.get("id", "?") for d in matches) + raise click.UsageError( + f"Multiple active Loops deployments match {base_model!r}: {ids}. " + "Pass --deployment-id to disambiguate." + ) + return matches[0]["id"] + + +def _render_loops_metrics(payload: Dict[str, Any]) -> None: + metrics = payload.get("metrics") or {} + deployment_id = payload.get("deployment_id", "") + + console.print( + f"[bold]Loops deployment metrics[/bold] for [cyan]{deployment_id}[/cyan]" + ) + + service_table = _build_service_metrics_table(metrics) + console.print(service_table) + + per_node_table = _build_per_node_table(metrics.get("per_node_metrics") or []) + if per_node_table is not None: + console.print(per_node_table) + else: + console.print("No per-node compute metrics in this window.", style="yellow") + + +def _build_service_metrics_table(metrics: Dict[str, Any]) -> rich.table.Table: + table = rich.table.Table( + show_header=True, + header_style="bold magenta", + title="Service-level (latest)", + box=rich.table.box.ROUNDED, + border_style="blue", + ) + table.add_column("Metric", style="cyan") + table.add_column("Value", justify="right") + + rate = _latest_value(metrics.get("inference_volume")) + concurrent = _latest_value(metrics.get("concurrent_requests")) + latencies = (metrics.get("response_time_stats") or [None])[-1] or {} + + table.add_row("Request rate (RPS)", _fmt_float(rate)) + table.add_row("Concurrent requests", _fmt_float(concurrent)) + table.add_row("Latency p50 (ms)", _fmt_float(latencies.get("p50"))) + table.add_row("Latency p95 (ms)", _fmt_float(latencies.get("p95"))) + table.add_row("Latency p99 (ms)", _fmt_float(latencies.get("p99"))) + + by_status = (metrics.get("inference_volume_by_status") or [None])[-1] or {} + table.add_row("2xx rate", _fmt_float(by_status.get("status_2xx"))) + table.add_row("4xx rate", _fmt_float(by_status.get("status_4xx"))) + table.add_row("5xx rate", _fmt_float(by_status.get("status_5xx"))) + + return table + + +def _build_per_node_table(nodes: List[Dict[str, Any]]) -> Optional[rich.table.Table]: + if not nodes: + return None + table = rich.table.Table( + show_header=True, + header_style="bold magenta", + title="Per-node compute (latest)", + box=rich.table.box.ROUNDED, + border_style="blue", + ) + table.add_column("Node", style="cyan") + table.add_column("CPU (cores)", justify="right") + table.add_column("CPU mem", justify="right") + table.add_column("GPU util", justify="right") + table.add_column("GPU mem", justify="right") + table.add_column("Eph. storage util", justify="right") + + for node in nodes: + cpu = _fmt_float(_latest_value(node.get("cpu_usage"))) + cpu_mem = _fmt_bytes(_latest_value(node.get("cpu_memory_usage_bytes"))) + gpu_util = _fmt_gpu_map(node.get("gpu_utilization") or {}, as_percent=True) + gpu_mem = _fmt_gpu_map(node.get("gpu_memory_usage_bytes") or {}, as_bytes=True) + eph_storage = node.get("ephemeral_storage") or {} + eph_util = _fmt_percent(_latest_value(eph_storage.get("utilization"))) + table.add_row( + str(node.get("node_id", "")), cpu, cpu_mem, gpu_util, gpu_mem, eph_util + ) + return table + + +def _latest_value(series: Optional[List[Dict[str, Any]]]) -> Optional[float]: + if not series: + return None + last = series[-1] + val = last.get("value") + return float(val) if val is not None else None + + +def _fmt_float(value: Optional[float]) -> str: + if value is None: + return "—" + return f"{value:.2f}" + + +def _fmt_percent(value: Optional[float]) -> str: + if value is None: + return "—" + return f"{value * 100:.1f}%" + + +def _fmt_bytes(value: Optional[float]) -> str: + if value is None: + return "—" + for unit, divisor in (("TB", 2**40), ("GB", 2**30), ("MB", 2**20), ("KB", 2**10)): + if value >= divisor: + return f"{value / divisor:.2f} {unit}" + return f"{value:.0f} B" + + +def _fmt_gpu_map( + per_gpu: Dict[str, List[Dict[str, Any]]], + *, + as_percent: bool = False, + as_bytes: bool = False, +) -> str: + """Latest value per GPU rank, joined as ``"0:val 1:val"``.""" + if not per_gpu: + return "—" + parts: List[str] = [] + for rank in sorted(per_gpu.keys(), key=lambda r: int(r) if r.isdigit() else r): + latest = _latest_value(per_gpu[rank]) + if as_percent: + parts.append(f"{rank}:{_fmt_percent(latest)}") + elif as_bytes: + parts.append(f"{rank}:{_fmt_bytes(latest)}") + else: + parts.append(f"{rank}:{_fmt_float(latest)}") + return " ".join(parts) + + def _render_loops_samplers(samplers: List[Dict[str, Any]]) -> None: if not samplers: console.print("No Loops samplers found.", style="yellow") diff --git a/truss/remote/baseten/api.py b/truss/remote/baseten/api.py index 0564c5ac1..e2b8f5656 100644 --- a/truss/remote/baseten/api.py +++ b/truss/remote/baseten/api.py @@ -1264,3 +1264,15 @@ def deactivate_loops_deployment(self, deployment_id: str) -> None: self._rest_api_client.post( f"v1/loops/deployments/{deployment_id}/deactivate", body={} ) + + def get_loops_deployment_metrics( + self, + deployment_id: str, + start_epoch_millis: Optional[int] = None, + end_epoch_millis: Optional[int] = None, + ) -> Dict[str, Any]: + """Fetch utilization + request metrics for a Loops trainer deployment.""" + return self._rest_api_client.post( + f"v1/loops/deployments/{deployment_id}/metrics", + body=self._prepare_time_range_query(start_epoch_millis, end_epoch_millis), + ) diff --git a/truss/tests/cli/test_loops_cli.py b/truss/tests/cli/test_loops_cli.py index cb7e5696d..812b11a1e 100644 --- a/truss/tests/cli/test_loops_cli.py +++ b/truss/tests/cli/test_loops_cli.py @@ -714,3 +714,177 @@ def test_checkpoints_deploy_rejects_checkpoint_ids_with_config(mock_remote, tmp_ ) assert result.exit_code != 0 mock_create.assert_not_called() + + +# ── loops metrics ──────────────────────────────────────────────────────────── + + +_FAKE_METRICS_PAYLOAD = { + "deployment_id": "trn_abc123", + "metrics": { + "inference_volume": [{"timestamp": "2026-05-13T12:00:00Z", "value": 12.5}], + "concurrent_requests": [{"timestamp": "2026-05-13T12:00:00Z", "value": 3.0}], + "response_time_stats": [ + {"timestamp": "2026-05-13T12:00:00Z", "p50": 0.04, "p95": 0.18, "p99": 0.31} + ], + "inference_volume_by_status": [ + { + "timestamp": "2026-05-13T12:00:00Z", + "status_2xx": 11.0, + "status_4xx": 0.5, + "status_5xx": 1.0, + } + ], + "gpu_memory_usage_bytes": { + "0": [{"timestamp": "2026-05-13T12:00:00Z", "value": 5e9}] + }, + "gpu_utilization": { + "0": [{"timestamp": "2026-05-13T12:00:00Z", "value": 0.75}] + }, + "cpu_usage": [{"timestamp": "2026-05-13T12:00:00Z", "value": 1.4}], + "cpu_memory_usage_bytes": [ + {"timestamp": "2026-05-13T12:00:00Z", "value": 2 * 2**30} + ], + "ephemeral_storage": { + "usage_bytes": [{"timestamp": "2026-05-13T12:00:00Z", "value": 3 * 2**30}], + "utilization": [{"timestamp": "2026-05-13T12:00:00Z", "value": 0.42}], + }, + "per_node_metrics": [ + { + "node_id": "g00r0", + "gpu_memory_usage_bytes": { + "0": [{"timestamp": "2026-05-13T12:00:00Z", "value": 5e9}] + }, + "gpu_utilization": { + "0": [{"timestamp": "2026-05-13T12:00:00Z", "value": 0.75}] + }, + "cpu_usage": [{"timestamp": "2026-05-13T12:00:00Z", "value": 1.4}], + "cpu_memory_usage_bytes": [ + {"timestamp": "2026-05-13T12:00:00Z", "value": 2 * 2**30} + ], + "ephemeral_storage": { + "usage_bytes": [ + {"timestamp": "2026-05-13T12:00:00Z", "value": 3 * 2**30} + ], + "utilization": [ + {"timestamp": "2026-05-13T12:00:00Z", "value": 0.42} + ], + }, + } + ], + }, +} + + +def test_metrics_with_deployment_id(mock_remote): + mock_remote.api.get_loops_deployment_metrics.return_value = _FAKE_METRICS_PAYLOAD + + result = _invoke( + [ + "loops", + "metrics", + "--deployment-id", + "trn_abc123", + "--remote", + "test_remote", + ], + mock_remote, + ) + + assert result.exit_code == 0, result.output + mock_remote.api.get_loops_deployment_metrics.assert_called_once_with( + deployment_id="trn_abc123" + ) + # Latest snapshot values are present in the rendered output. + assert "12.5" in result.output or "12.50" in result.output + assert "g00r0" in result.output + + +def test_metrics_with_base_model_resolves_active_deployment(mock_remote): + mock_remote.api.list_loops_deployments.return_value = [ + {"id": "trn_abc123", "base_model": "Qwen/Qwen3-8B"} + ] + mock_remote.api.get_loops_deployment_metrics.return_value = _FAKE_METRICS_PAYLOAD + + result = _invoke( + [ + "loops", + "metrics", + "--base-model", + "Qwen/Qwen3-8B", + "--remote", + "test_remote", + ], + mock_remote, + ) + + assert result.exit_code == 0, result.output + mock_remote.api.get_loops_deployment_metrics.assert_called_once_with( + deployment_id="trn_abc123" + ) + + +def test_metrics_with_base_model_no_match_fails(mock_remote): + mock_remote.api.list_loops_deployments.return_value = [] + + result = _invoke( + [ + "loops", + "metrics", + "--base-model", + "Qwen/Qwen3-8B", + "--remote", + "test_remote", + ], + mock_remote, + ) + + assert result.exit_code != 0 + mock_remote.api.get_loops_deployment_metrics.assert_not_called() + + +def test_metrics_with_base_model_multiple_matches_fails(mock_remote): + mock_remote.api.list_loops_deployments.return_value = [ + {"id": "trn_a", "base_model": "Qwen/Qwen3-8B"}, + {"id": "trn_b", "base_model": "Qwen/Qwen3-8B"}, + ] + + result = _invoke( + [ + "loops", + "metrics", + "--base-model", + "Qwen/Qwen3-8B", + "--remote", + "test_remote", + ], + mock_remote, + ) + + assert result.exit_code != 0 + mock_remote.api.get_loops_deployment_metrics.assert_not_called() + + +def test_metrics_requires_one_selector(mock_remote): + result = _invoke(["loops", "metrics", "--remote", "test_remote"], mock_remote) + + assert result.exit_code != 0 + assert "--deployment-id" in result.output or "--base-model" in result.output + + +def test_metrics_rejects_both_selectors(mock_remote): + result = _invoke( + [ + "loops", + "metrics", + "--deployment-id", + "trn_abc123", + "--base-model", + "Qwen/Qwen3-8B", + "--remote", + "test_remote", + ], + mock_remote, + ) + + assert result.exit_code != 0