From eadf7966481ce14f17b1322f78b3c27854022ac6 Mon Sep 17 00:00:00 2001 From: nilsmechtel Date: Mon, 7 Sep 2026 12:24:05 +0200 Subject: [PATCH] fix(cluster): strip the accelerator_type prefix, not its characters _get_accelerator_type used str.lstrip("accelerator_type:"), which strips any leading character in that set rather than the literal prefix. It is correct for A40, T4 and G only because each starts with an out-of-set character; "accelerator_type:tesla" yields "sla". Co-Authored-By: Claude Opus 5 --- bioengine/cluster/proxy_actor.py | 2 +- tests/test_accelerator_type.py | 26 ++++++++++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) create mode 100644 tests/test_accelerator_type.py diff --git a/bioengine/cluster/proxy_actor.py b/bioengine/cluster/proxy_actor.py index 713e275..1c57242 100644 --- a/bioengine/cluster/proxy_actor.py +++ b/bioengine/cluster/proxy_actor.py @@ -342,7 +342,7 @@ def _get_accelerator_type(self, resources: Dict[str, float]) -> Optional[str]: """ for resource_name in resources: if resource_name.startswith("accelerator_type:"): - return resource_name.lstrip("accelerator_type:") + return resource_name.removeprefix("accelerator_type:") def _get_slurm_job_id(self, resources: Dict[str, float]) -> Optional[str]: """Extract the SLURM job ID from a node's resource dictionary. diff --git a/tests/test_accelerator_type.py b/tests/test_accelerator_type.py new file mode 100644 index 0000000..b06bcb1 --- /dev/null +++ b/tests/test_accelerator_type.py @@ -0,0 +1,26 @@ +"""Unit tests for ``BioEngineProxyActor._get_accelerator_type``. + +Exercises the method against plain resource dicts so no Ray cluster is needed. +""" + +from bioengine.cluster.proxy_actor import BioEngineProxyActor + +_get_accelerator_type = ( + BioEngineProxyActor.__ray_metadata__.modified_class._get_accelerator_type +) + + +def test_reads_the_accelerator_type_resource(): + resources = {"CPU": 8.0, "GPU": 1.0, "accelerator_type:A40": 1.0} + assert _get_accelerator_type(None, resources) == "A40" + + +def test_type_starting_with_a_prefix_character_survives(): + # str.lstrip("accelerator_type:") strips the character *set*, so a type + # whose first characters all appear in the prefix loses them. + resources = {"GPU": 1.0, "accelerator_type:tesla": 1.0} + assert _get_accelerator_type(None, resources) == "tesla" + + +def test_returns_none_without_an_accelerator_resource(): + assert _get_accelerator_type(None, {"CPU": 8.0}) is None