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