diff --git a/src/thunder_forge/cli.py b/src/thunder_forge/cli.py index 7b5c1a7..78ecfc6 100644 --- a/src/thunder_forge/cli.py +++ b/src/thunder_forge/cli.py @@ -737,12 +737,28 @@ def _format_disk_gb(total_disk_gb: float) -> str: def _sum_aliases_disk_gb(config: ClusterConfig, aliases: list[str]) -> float: + """Sum disk usage of the listed aliases, counting each runtime model once. + + Aliases that share a runtime_model_id (unified layout) point at the same + on-disk model, so they must not be summed per alias. + """ + disk_by_runtime: dict[str, float] = {} + for model in config.models.values(): + if model.runtime_model_id: + disk_by_runtime.setdefault(model.runtime_model_id, model.disk_gb) + seen_runtimes: set[str] = set() total = 0.0 for alias in aliases: model = config.models.get(alias) if model is None: continue - total += model.disk_gb + if model.runtime_model_id: + if model.runtime_model_id in seen_runtimes: + continue + seen_runtimes.add(model.runtime_model_id) + total += disk_by_runtime.get(model.runtime_model_id, model.disk_gb) + else: + total += model.disk_gb return total @@ -2056,6 +2072,21 @@ def _active_model_users(payload: dict[str, object]) -> dict[tuple[str, str], set return active_users +def _model_status_aliases(status: dict) -> list[str]: + aliases = status.get("aliases") + if isinstance(aliases, list) and aliases and all(isinstance(a, str) and a for a in aliases): + return aliases + return [str(status.get("id", status.get("runtime_id", "model")))] + + +def _model_display_id(status: dict) -> str: + aliases = _model_status_aliases(status) + runtime_id = str(status.get("runtime_id", aliases[0])) + if len(aliases) > 1: + return f"{','.join(aliases)} ({runtime_id})" + return aliases[0] + + def _node_idle_seconds(node: dict[str, object]) -> float | None: active_requests = node.get("active_requests") if isinstance(active_requests, int) and active_requests > 0: @@ -2171,12 +2202,14 @@ def _print_cluster_status_rich(payload: dict[str, object], *, console: Console) model_table.add_column("Idle", justify="right") model_table.add_column("RAM", justify="right") for node_name, status in model_statuses: - model_id = str(status.get("id", status.get("runtime_id", "model"))) - active_key = (node_name, model_id) - state = "active" if active_key in active_users else str(status.get("state", "unknown")) + model_id = _model_display_id(status) + matched_users: set[str] = set() + for active_key in ((node_name, alias) for alias in _model_status_aliases(status)): + matched_users |= active_users.get(active_key, set()) + state = "active" if matched_users else str(status.get("state", "unknown")) tf_user = ( - ", ".join(sorted(active_users[active_key])) - if active_key in active_users + ", ".join(sorted(matched_users)) + if matched_users else str(status.get("tf_user", "-")) ) state_style = {"loaded": "green", "loading": "yellow", "cold": "dim"}.get( @@ -2285,12 +2318,16 @@ def _print_cluster_status_payload( if not isinstance(model_status, dict): continue model_idle = _format_model_idle(model_status.get("idle_seconds")) - model_id = str(model_status.get("id", model_status.get("runtime_id", "model"))) - active_key = (str(node.get("name", "node")), model_id) - state = "active" if active_key in active_users else str(model_status.get("state", "unknown")) + model_id = _model_display_id(model_status) + node_name = str(node.get("name", "node")) + active_keys = [(node_name, alias) for alias in _model_status_aliases(model_status)] + matched_users: set[str] = set() + for active_key in active_keys: + matched_users |= active_users.get(active_key, set()) + state = "active" if matched_users else str(model_status.get("state", "unknown")) tf_user = ( - ", ".join(sorted(active_users[active_key])) - if active_key in active_users + ", ".join(sorted(matched_users)) + if matched_users else str(model_status.get("tf_user", "-")) ) typer.echo( diff --git a/src/thunder_forge/cluster/status.py b/src/thunder_forge/cluster/status.py index 0a95892..d60a6e2 100644 --- a/src/thunder_forge/cluster/status.py +++ b/src/thunder_forge/cluster/status.py @@ -25,6 +25,7 @@ def normalize_model_statuses( is_loading = raw_status.get("is_loading") is True status: dict[str, object] = { "id": aliases[0], + "aliases": aliases, "runtime_id": runtime_model_id, "loaded": loaded, "is_loading": is_loading, diff --git a/tests/test_cli_runtime.py b/tests/test_cli_runtime.py index 3a31821..b5f071e 100644 --- a/tests/test_cli_runtime.py +++ b/tests/test_cli_runtime.py @@ -1222,13 +1222,107 @@ def test_cluster_status_reports_inference_health(tmp_path: Path, monkeypatch) -> assert "infer-03: health=ok models=ok" in result.stdout assert "macos_version: 15.6.1" in result.stdout assert "omlx_version: 0.4.2.dev2" in result.stdout - assert "served_models (13G): memory" in result.stdout - assert "hot_loaded_models (13G): memory" in result.stdout + # Disk totals come from the repo tfconfig model registry; only the alias list + # is stable across registry changes, so assert the shape without the size. + served_lines = [line for line in result.stdout.splitlines() if line.strip().startswith("served_models (")] + hot_lines = [line for line in result.stdout.splitlines() if line.strip().startswith("hot_loaded_models (")] + assert len(served_lines) == 1 and served_lines[0].strip().endswith(": memory") + assert len(hot_lines) == 1 and hot_lines[0].strip().endswith(": memory") assert "model: memory tf_user=shag state=active idle=12s ram=13.0G" in result.stdout assert "latest_omlx_version: v0.4.2" in result.stdout assert "omlx_upgrade_hint: no (versions aligned)" in result.stdout +def test_cluster_status_shows_alias_group_for_shared_runtime_model(tmp_path: Path, monkeypatch) -> None: + """Unified layout: one loaded runtime model behind several aliases must show them all.""" + import thunder_forge.cli as cli_module + + payload = { + "ok": True, + "target": "all", + "gateway": None, + "inference": [ + { + "name": "infer-03", + "host": "infer-03.lan", + "health": "ok", + "models": "ok", + "active_requests": 1, + "omlx_version": "0.7.1rc1", + "macos_version": "27.0.1", + "served_models": ["agent", "coder", "memory"], + "hot_loaded_models": ["agent", "coder", "memory"], + "model_statuses": [ + { + "id": "agent", + "aliases": ["agent", "coder", "memory"], + "runtime_id": "Qwen3.8-Flash-Next-oQ4e-mtp", + "state": "loaded", + "idle_seconds": 12.0, + "actual_size": 106 * 1024**3, + } + ], + "errors": [], + } + ], + "summary": {}, + "active_edge_requests": [{"client_id": "shag", "model": "coder"}], + } + monkeypatch.setattr(cli_module, "_fetch_cluster_status_payload", lambda config, *, target: payload) + + result = runner.invoke(app, ["cluster", "status"]) + + assert result.exit_code == 0 + assert ( + "model: agent,coder,memory (Qwen3.8-Flash-Next-oQ4e-mtp) tf_user=shag state=active idle=12s ram=106.0G" + in result.stdout + ) + # Aliases sharing one runtime model must be counted once, not once per alias. + assert "served_models (106.3G): agent, coder, memory" in result.stdout + assert "hot_loaded_models (106.3G): agent, coder, memory" in result.stdout + + +def test_print_cluster_status_rich_shows_alias_group() -> None: + """Rich table must show the full alias group, runtime id, and group-wide active users.""" + from io import StringIO + + from rich.console import Console + + from thunder_forge.cli import _print_cluster_status_rich + + payload = { + "ok": True, + "gateway": None, + "inference": [ + { + "name": "msm3", + "host": "msm3-wifi.lan", + "health": "ok", + "models": "ok", + "model_statuses": [ + { + "id": "agent", + "aliases": ["agent", "coder", "memory"], + "runtime_id": "Qwen3.8-Flash-Next-oQ4e-mtp", + "state": "loaded", + "idle_seconds": 12.0, + "actual_size": 106 * 1024**3, + } + ], + } + ], + "active_edge_requests": [{"client_id": "shag", "model": "coder", "node": "msm3"}], + "summary": {}, + } + output = StringIO() + _print_cluster_status_rich(payload, console=Console(file=output, width=200, color_system=None)) + value = output.getvalue() + + assert "agent,coder,memory" in value + assert "Qwen3.8-Flash-Next-oQ4e-mtp" in value + assert "shag" in value + + def test_cluster_status_replaces_unknown_upstream_versions(monkeypatch) -> None: import thunder_forge.cli as cli_module diff --git a/tests/test_status.py b/tests/test_status.py index b1d7e94..14317f3 100644 --- a/tests/test_status.py +++ b/tests/test_status.py @@ -20,6 +20,7 @@ def test_normalize_model_statuses_maps_aliases_and_idle_time() -> None: assert result == [ { "id": "memory", + "aliases": ["memory"], "runtime_id": "runtime-model", "loaded": True, "is_loading": False, @@ -31,6 +32,26 @@ def test_normalize_model_statuses_maps_aliases_and_idle_time() -> None: ] +def test_normalize_model_statuses_exposes_full_alias_group_for_shared_runtime_model() -> None: + """Unified layout: several aliases share one runtime model; status must show them all.""" + result = normalize_model_statuses( + {"Qwen3.8-Flash-Next-oQ4e-mtp": {"loaded": True, "is_loading": False}}, + map_aliases=lambda _model_ids: ["agent", "coder", "memory"], + now=100.0, + ) + + assert result == [ + { + "id": "agent", + "aliases": ["agent", "coder", "memory"], + "runtime_id": "Qwen3.8-Flash-Next-oQ4e-mtp", + "loaded": True, + "is_loading": False, + "state": "loaded", + } + ] + + def test_normalize_model_statuses_preserves_missing_optional_fields() -> None: result = normalize_model_statuses( {"runtime-model": {"loaded": False, "is_loading": True}}, @@ -41,6 +62,7 @@ def test_normalize_model_statuses_preserves_missing_optional_fields() -> None: assert result == [ { "id": "memory", + "aliases": ["memory"], "runtime_id": "runtime-model", "loaded": False, "is_loading": True,