From d694d70d6a227608d50af4fcf5358050a8f2fb70 Mon Sep 17 00:00:00 2001 From: nilsmechtel Date: Mon, 7 Sep 2026 12:01:58 +0200 Subject: [PATCH 1/2] feat(cluster): report VRAM_MB booking alongside the GPU fraction On a cluster advertising the VRAM_MB custom resource 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: four federated-unet replicas on Europa report used_gpu 0.04 while holding 20480 of 24576 MB. get_cluster_state now carries total_vram_mb/used_vram_mb per node and rolled up, and the CLI prints the booking next to the fraction on nodes that advertise it. Nodes without VRAM_MB report 0 and their output is unchanged. Co-Authored-By: Claude Opus 5 --- bioengine/cli/cluster.py | 12 ++++- bioengine/cluster/proxy_actor.py | 14 ++++++ tests/end_to_end/test_worker.py | 4 ++ tests/test_cluster_status_vram.py | 73 +++++++++++++++++++++++++++++++ 4 files changed, 102 insertions(+), 1 deletion(-) create mode 100644 tests/test_cluster_status_vram.py diff --git a/bioengine/cli/cluster.py b/bioengine/cli/cluster.py index 20a290a..0da0850 100644 --- a/bioengine/cli/cluster.py +++ b/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}" ) diff --git a/bioengine/cluster/proxy_actor.py b/bioengine/cluster/proxy_actor.py index 713e275..ee8c54f 100644 --- a/bioengine/cluster/proxy_actor.py +++ b/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 diff --git a/tests/end_to_end/test_worker.py b/tests/end_to_end/test_worker.py index 62192e0..6ab0cf2 100644 --- a/tests/end_to_end/test_worker.py +++ b/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", diff --git a/tests/test_cluster_status_vram.py b/tests/test_cluster_status_vram.py new file mode 100644 index 0000000..7818752 --- /dev/null +++ b/tests/test_cluster_status_vram.py @@ -0,0 +1,73 @@ +"""Unit tests for the VRAM_MB reservation reported by ``get_cluster_status``. + +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 From 099075dac9444445d8c60ffe4db4c8b7d00067e9 Mon Sep 17 00:00:00 2001 From: nilsmechtel Date: Mon, 7 Sep 2026 12:07:46 +0200 Subject: [PATCH 2/2] docs(test): name the method the test actually calls --- tests/test_cluster_status_vram.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_cluster_status_vram.py b/tests/test_cluster_status_vram.py index 7818752..ac5703d 100644 --- a/tests/test_cluster_status_vram.py +++ b/tests/test_cluster_status_vram.py @@ -1,4 +1,4 @@ -"""Unit tests for the VRAM_MB reservation reported by ``get_cluster_status``. +"""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.