feat(cluster): report VRAM_MB booking in cluster status

bioengine/cli/cluster.py

@@ -125,10 +125,20 @@ async def _run():
                     info.get("used_gpu_memory", 0), info.get("total_gpu_memory", 0)
                 )
                 role = "HEAD" if info.get("head") else "worker"
+                # On a VRAM_MB node the GPU fraction is a count of 0.01 device
+                # handles, so it reads as idle at full reservation; show what
+                # actually bounds packing next to it.
+                total_vram_mb = info.get("total_vram_mb", 0)
+                booked = (
+                    f"booked: {info.get('used_vram_mb', 0):.0f}/{total_vram_mb:.0f} MB  "
+                    if total_vram_mb
+                    else ""
+                )
                 click.echo(
                     f"  {info.get('node_ip')} [{role}] "
                     f"{info.get('accelerator_type', '?')} "
-                    f"GPU: {info.get('used_gpu', 0):.1f}/{info.get('total_gpu', 0):.0f} "
+                    f"GPU: {info.get('used_gpu', 0):.2f}/{info.get('total_gpu', 0):.0f} "
+                    f"{booked}"
                     f"VRAM: {vram}  "
                     f"CPU: {info.get('used_cpu', 0):.0f}/{info.get('total_cpu', 0):.0f}"
                 )

bioengine/cluster/proxy_actor.py

@@ -506,6 +506,8 @@ def get_cluster_state(self) -> Dict[str, Any]:
                     "used_cpu": float,
                     "total_gpu": float,
                     "used_gpu": float,
+                    "total_vram_mb": float,
+                    "used_vram_mb": float,
                     "pending_resources": {  # if check_pending_resources=True
                         "actors": List[Dict],
                         "jobs": List[Dict],
@@ -521,6 +523,8 @@ def get_cluster_state(self) -> Dict[str, Any]:
                         "used_cpu": float,
                         "total_gpu": float,
                         "used_gpu": float,
+                        "total_vram_mb": float,  # VRAM_MB advertised, 0 if none
+                        "used_vram_mb": float,   # VRAM_MB booked by the scheduler
                         "total_gpu_memory": Union[int, str],  # in bytes or "NA"
                         "used_gpu_memory": Union[int, str],  # in bytes or "NA"
                         "total_memory": float,
@@ -561,6 +565,8 @@ def get_cluster_state(self) -> Dict[str, Any]:
                 "used_cpu": 0,
                 "total_gpu": 0,
                 "used_gpu": 0,
+                "total_vram_mb": 0,
+                "used_vram_mb": 0,
                 "total_memory": 0,
                 "used_memory": 0,
                 "total_gpu_memory": 0,
@@ -586,6 +592,12 @@ def get_cluster_state(self) -> Dict[str, Any]:
             available_cpu = available_resources.get("CPU", 0)
             total_gpu = total_resources.get("GPU", 0)
             available_gpu = available_resources.get("GPU", 0)
+            # Where the node advertises VRAM_MB, that resource — not GPU — is
+            # what bounds packing: the AppBuilder books a 0.01 GPU handle per
+            # replica purely to bind a device, so used_gpu reads as near-idle on
+            # a GPU whose VRAM is fully reserved.
+            total_vram_mb = total_resources.get("VRAM_MB", 0)
+            available_vram_mb = available_resources.get("VRAM_MB", 0)
             accelerator_type = (
                 "NA" if total_gpu == 0 else self._get_accelerator_type(total_resources)
             )
@@ -626,6 +638,8 @@ def get_cluster_state(self) -> Dict[str, Any]:
                 "used_cpu": max(0, total_cpu - available_cpu),
                 "total_gpu": total_gpu,
                 "used_gpu": max(0, total_gpu - available_gpu),
+                "total_vram_mb": total_vram_mb,
+                "used_vram_mb": max(0, total_vram_mb - available_vram_mb),
                 "total_gpu_memory": total_gpu_memory,
                 "used_gpu_memory": used_gpu_memory,
                 "total_memory": total_memory,  # in bytes

tests/end_to_end/test_worker.py

@@ -166,6 +166,8 @@ async def test_get_status(
         "used_cpu",
         "total_gpu",
         "used_gpu",
+        "total_vram_mb",
+        "used_vram_mb",
     ]
 
     for field in expected_cluster_fields:
@@ -196,6 +198,8 @@ async def test_get_status(
             "used_cpu",
             "total_gpu",
             "used_gpu",
+            "total_vram_mb",
+            "used_vram_mb",
             "total_gpu_memory",
             "used_gpu_memory",
             "total_memory",

tests/test_cluster_status_vram.py

@@ -0,0 +1,73 @@
+"""Unit tests for the VRAM_MB reservation reported by ``get_cluster_state``.
+
+Drives the plain class behind the Ray actor decorator with stubbed node
+resources, so the VRAM_MB branch is exercised without a GPU cluster.
+"""
+
+from bioengine.cluster.proxy_actor import BioEngineProxyActor
+
+_PLAIN_CLASS = BioEngineProxyActor.__ray_metadata__.modified_class
+
+#: One Europa-shaped head node: an RTX 3090 advertising its VRAM, with four
+#: federated-unet replicas holding 5120 MB and a 0.01 GPU handle each.
+_HEAD = {
+    "node:172.17.0.3": 1.0,
+    "node:__internal_head__": 0.001,
+    "CPU": 8.0,
+    "GPU": 1.0,
+    "VRAM_MB": 24576.0,
+    "memory": 32212254720.0,
+    "object_store_memory": 10000000000.0,
+}
+_HEAD_AVAILABLE = {
+    "CPU": 0.0,
+    "GPU": 0.96,
+    "VRAM_MB": 4096.0,
+    "memory": 4294967296.0,
+    "object_store_memory": 9999126230.0,
+}
+
+
+def _actor(total, available):
+    actor = object.__new__(_PLAIN_CLASS)
+    actor.exclude_head_node = False
+    actor.check_pending_resources = False
+    actor.node_gpu_memory = {}
+    actor.global_state = type(
+        "GlobalState",
+        (),
+        {
+            "total_resources_per_node": staticmethod(lambda: {"n1": total}),
+            "available_resources_per_node": staticmethod(lambda: {"n1": available}),
+        },
+    )()
+    actor._get_per_node_gpu_memory_usage = lambda: ({}, False)
+    return actor
+
+
+def test_vram_booking_is_reported_where_the_gpu_fraction_is_not():
+    status = _actor(_HEAD, _HEAD_AVAILABLE).get_cluster_state()
+
+    node = status["nodes"]["n1"]
+    assert node["total_vram_mb"] == 24576.0
+    assert node["used_vram_mb"] == 20480.0
+    # The reason the field exists: 83% of the GPU is booked and used_gpu says 4%.
+    assert round(node["used_gpu"], 2) == 0.04
+
+    assert status["cluster"]["total_vram_mb"] == 24576.0
+    assert status["cluster"]["used_vram_mb"] == 20480.0
+
+
+def test_no_vram_resource_reports_zero_rather_than_inventing_capacity():
+    total = {k: v for k, v in _HEAD.items() if k != "VRAM_MB"}
+    available = {k: v for k, v in _HEAD_AVAILABLE.items() if k != "VRAM_MB"}
+    available["GPU"] = 0.67
+
+    status = _actor(total, available).get_cluster_state()
+
+    node = status["nodes"]["n1"]
+    assert node["total_vram_mb"] == 0
+    assert node["used_vram_mb"] == 0
+    # Where nothing advertises VRAM_MB the GPU fraction is the real reservation,
+    # so a zero here must not read as an idle GPU.
+    assert round(node["used_gpu"], 2) == 0.33