diff --git a/src/cloudai/parser/system_parser/slurm_system_parser.py b/src/cloudai/parser/system_parser/slurm_system_parser.py index acf9c539..158e2cb6 100644 --- a/src/cloudai/parser/system_parser/slurm_system_parser.py +++ b/src/cloudai/parser/system_parser/slurm_system_parser.py @@ -43,6 +43,9 @@ def safe_int(value): except ValueError: return None + def str_to_bool(value: str) -> bool: + return value.lower() in ("true", "1", "yes") + name = data.get("name") if not name: raise ValueError("Missing mandatory field: 'name'") @@ -78,7 +81,7 @@ def safe_int(value): gpus_per_node = safe_int(data.get("gpus_per_node")) ntasks_per_node = safe_int(data.get("ntasks_per_node")) - cache_docker_images_locally = data.get("cache_docker_images_locally", False) + cache_docker_images_locally = str_to_bool(data.get("cache_docker_images_locally", "False")) nodes_dict: Dict[str, SlurmNode] = {} updated_partitions: Dict[str, List[SlurmNode]] = {} diff --git a/tests/test_slurm_system_parser.py b/tests/test_slurm_system_parser.py new file mode 100644 index 00000000..f2ee1cb4 --- /dev/null +++ b/tests/test_slurm_system_parser.py @@ -0,0 +1,84 @@ +from typing import Any, Dict + +import pytest +from cloudai.parser.system_parser.slurm_system_parser import SlurmSystemParser +from cloudai.systems.slurm import SlurmSystem + + +@pytest.fixture +def example_data() -> Dict[str, Any]: + return { + "name": "test_system", + "install_path": "/fake/path", + "output_path": "/fake/output", + "default_partition": "main", + "partitions": { + "main": { + "name": "main", + "nodes": ["node-[033-034]"], + "groups": {"group1": {"name": "group1", "nodes": ["node-033"]}}, + }, + "backup": { + "name": "backup", + "nodes": ["node-[01-02]"], + "groups": {"group2": {"name": "group2", "nodes": ["node-01"]}}, + }, + }, + "cache_docker_images_locally": "True", + } + + +def test_parse_slurm_system_parser(example_data): + parser = SlurmSystemParser() + slurm_system = parser.parse(example_data) + + assert isinstance(slurm_system, SlurmSystem) + assert slurm_system.name == "test_system" + assert slurm_system.install_path == "/fake/path" + assert slurm_system.output_path == "/fake/output" + assert slurm_system.default_partition == "main" + assert slurm_system.cache_docker_images_locally is True + assert "main" in slurm_system.partitions + assert "backup" in slurm_system.partitions + assert "group1" in slurm_system.groups["main"] + assert "group2" in slurm_system.groups["backup"] + + +@pytest.mark.parametrize( + "input_value, expected_result", + [ + ("True", True), + ("False", False), + ("true", True), + ("false", False), + ("1", True), + ("0", False), + ("yes", True), + ("no", False), + ], +) +def test_str_to_bool_conversion(input_value, expected_result): + parser = SlurmSystemParser() + result = parser.parse( + { + "name": "test_system", + "install_path": "/fake/path", + "output_path": "/fake/output", + "default_partition": "main", + "partitions": { + "main": { + "name": "main", + "nodes": ["node-[033-034]"], + "groups": {"group1": {"name": "group1", "nodes": ["node-033"]}}, + }, + "backup": { + "name": "backup", + "nodes": ["node-[01-02]"], + "groups": {"group2": {"name": "group2", "nodes": ["node-01"]}}, + }, + }, + "cache_docker_images_locally": input_value, + } + ) + + assert result.cache_docker_images_locally == expected_result