Merge pull request #535 from NVIDIA/am/upd-before-access

Optimize slurm updates

Author: amaslenn

src/cloudai/parser.py

@@ -50,13 +50,18 @@ def __init__(self, system_config_path: Path) -> None:
         """
         logging.debug(f"Initializing parser with: {system_config_path=}")
         self.system_config_path = system_config_path
+        self._system: Optional[System] = None
 
     @property
     def system(self) -> System:
+        if self._system:
+            return self._system
+
         try:
-            return self.parse_system(self.system_config_path)
+            self._system = self.parse_system(self.system_config_path)
         except SystemConfigParsingError:
             exit(1)  # exit right away to keep error message readable for users
+        return self._system
 
     def parse(
         self,

src/cloudai/systems/slurm/slurm_system.py

@@ -415,10 +415,10 @@ def get_available_nodes_from_group(
             ValueError: If the partition or group is not found, or if the requested number of nodes exceeds the
                 available nodes.
         """
-        self.validate_partition_and_group(partition_name, group_name)
-
         self.update()
 
+        self.validate_partition_and_group(partition_name, group_name)
+
         grouped_nodes = self.group_nodes_by_state(partition_name, group_name)
 
         try:

src/cloudai/systems/slurm/strategy/slurm_command_gen_strategy.py

@@ -14,6 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 from abc import abstractmethod
 from datetime import datetime
 from pathlib import Path
@@ -44,6 +45,8 @@ def __init__(self, system: SlurmSystem, cmd_args: Dict[str, Any]) -> None:
         self.system = system
         self.docker_image_url = self.cmd_args.get("docker_image_url", "")
 
+        self._node_spec_cache: dict[str, tuple[int, list[str]]] = {}
+
     @abstractmethod
     def _container_mounts(self, tr: TestRun) -> list[str]:
         """Return CommandGenStrategy specific container mounts for the test run."""
@@ -132,7 +135,7 @@ def _parse_slurm_args(
             KeyError: If partition or essential node settings are missing.
         """
         job_name = self.job_name(job_name_prefix)
-        num_nodes, node_list = self.system.get_nodes_by_spec(tr.num_nodes, tr.nodes)
+        num_nodes, node_list = self.get_cached_nodes_spec(tr)
 
         slurm_args = {
             "job_name": job_name,
@@ -311,7 +314,7 @@ def _ranks_mapping_cmd(self, slurm_args: dict[str, Any], tr: TestRun) -> str:
 
     def _metadata_cmd(self, slurm_args: dict[str, Any], tr: TestRun) -> str:
         (tr.output_path.absolute() / "metadata").mkdir(parents=True, exist_ok=True)
-        num_nodes, _ = self.system.get_nodes_by_spec(tr.num_nodes, tr.nodes)
+        num_nodes, _ = self.get_cached_nodes_spec(tr)
         metadata_script_path = "/cloudai_install"
         if "image_path" not in slurm_args:
             metadata_script_path = str(self.system.install_path.absolute())
@@ -432,7 +435,7 @@ def _append_sbatch_directives(self, batch_script_content: List[str], args: Dict[
         )
 
     def _append_nodes_related_directives(self, content: List[str], args: Dict[str, Any], tr: TestRun) -> Optional[Path]:
-        num_nodes, node_list = self.system.get_nodes_by_spec(tr.num_nodes, tr.nodes)
+        num_nodes, node_list = self.get_cached_nodes_spec(tr)
 
         if node_list:
             content.append("#SBATCH --distribution=arbitrary")
@@ -480,3 +483,19 @@ def gen_srun_success_check(self, tr: TestRun) -> str:
             str: The generated command to check the success of the test run.
         """
         return ""
+
+    def get_cached_nodes_spec(self, tr: TestRun) -> tuple[int, list[str]]:
+        """
+        Get nodes for a test run, using cache when available.
+
+        It is needed to avoid multiple calls to the system.get_nodes_by_spec method which in turn queries the Slurm API.
+        For a single test run it is not required, we can get actual nodes status only once.
+        """
+        cache_key = f"{tr.current_iteration}:{tr.step}:{tr.num_nodes}:{','.join(tr.nodes)}"
+
+        if cache_key in self._node_spec_cache:
+            logging.debug(f"Using cached node allocation for {cache_key}: {self._node_spec_cache[cache_key]}")
+            return self._node_spec_cache[cache_key]
+
+        self._node_spec_cache[cache_key] = self.system.get_nodes_by_spec(tr.num_nodes, tr.nodes)
+        return self._node_spec_cache[cache_key]

tests/test_slurm_system.py

@@ -21,10 +21,12 @@
 
 import pytest
 
-from cloudai import BaseJob
+from cloudai import BaseJob, Test, TestRun, TestTemplate
 from cloudai.systems import SlurmSystem
 from cloudai.systems.slurm import SlurmNode, SlurmNodeState
 from cloudai.systems.slurm.slurm_system import parse_node_list
+from cloudai.systems.slurm.strategy.slurm_command_gen_strategy import SlurmCommandGenStrategy
+from cloudai.workloads.nccl_test import NCCLCmdArgs, NCCLTestDefinition
 
 
 def test_parse_squeue_output(slurm_system):
@@ -388,3 +390,143 @@ def test_with_commas(self, slurm_system: SlurmSystem):
     def test_colon_invalid_syntax(self, slurm_system: SlurmSystem, spec: str):
         with pytest.raises(ValueError):
             slurm_system.parse_nodes([spec])
+
+
+class TestGetNodesBySpec:
+    def test_empty_nodes_list(self, slurm_system: SlurmSystem):
+        num_nodes, node_list = slurm_system.get_nodes_by_spec(3, [])
+        assert num_nodes == 3
+        assert node_list == []
+
+    @pytest.mark.parametrize(
+        "in_nnodes,in_nodes,exp_nnodes,exp_nodes",
+        [
+            (2, ["node0[1-3]"], 3, ["node01", "node02", "node03"]),
+            (4, ["node01,node02"], 2, ["node01", "node02"]),
+            (1, ["node01,node02"], 2, ["node01", "node02"]),
+        ],
+    )
+    @patch("cloudai.systems.slurm.slurm_system.SlurmSystem.parse_nodes")
+    def test_explicit_node_names(
+        self,
+        mock_parse_nodes: Mock,
+        slurm_system: SlurmSystem,
+        in_nnodes: int,
+        in_nodes: list[str],
+        exp_nnodes: int,
+        exp_nodes: list[str],
+    ):
+        mock_parse_nodes.return_value = exp_nodes
+
+        num_nodes, node_list = slurm_system.get_nodes_by_spec(in_nnodes, in_nodes)
+
+        mock_parse_nodes.assert_called_once_with(in_nodes)
+        assert num_nodes == exp_nnodes
+        assert node_list == exp_nodes
+
+
+class ConcreteSlurmStrategy(SlurmCommandGenStrategy):
+    def _container_mounts(self, tr: TestRun) -> list[str]:
+        return []
+
+    def generate_test_command(self, env_vars, cmd_args, tr):
+        return ["test_command"]
+
+    def job_name(self, job_name_prefix: str) -> str:
+        return "job_name"
+
+
+@pytest.fixture
+def test_run(slurm_system: SlurmSystem) -> TestRun:
+    test_run = TestRun(
+        name="test_run",
+        test=Test(
+            test_definition=NCCLTestDefinition(
+                name="test_run", description="test_run", test_template_name="nccl", cmd_args=NCCLCmdArgs()
+            ),
+            test_template=TestTemplate(slurm_system),
+        ),
+        num_nodes=2,
+        nodes=["main:group1:2"],
+        output_path=slurm_system.output_path,
+    )
+
+    test_run.output_path.mkdir(parents=True, exist_ok=True)
+
+    return test_run
+
+
+class TestSlurmCommandGenStrategyCache:
+    @patch("cloudai.systems.slurm.SlurmSystem.get_nodes_by_spec")
+    def test_strategy_caching(self, mock_get_nodes: Mock, slurm_system: SlurmSystem, test_run: TestRun):
+        mock_get_nodes.return_value = (2, ["node01", "node02"])
+
+        strategy = ConcreteSlurmStrategy(slurm_system, {})
+
+        # First call to get nodes
+        res = strategy.get_cached_nodes_spec(test_run)
+        assert mock_get_nodes.call_count == 1
+        assert res == (2, ["node01", "node02"])
+
+        # Second call with same parameters should use cache
+        res = strategy.get_cached_nodes_spec(test_run)
+        assert mock_get_nodes.call_count == 1
+        assert res == (2, ["node01", "node02"])
+
+        # Different node spec should call get_nodes_by_spec again
+        test_run.num_nodes = 1
+        test_run.nodes = []
+        strategy.get_cached_nodes_spec(test_run)
+        assert mock_get_nodes.call_count == 2
+
+        test_run.num_nodes = 2
+        test_run.nodes = ["node01", "node03"]
+        strategy.get_cached_nodes_spec(test_run)
+        assert mock_get_nodes.call_count == 3
+
+    @patch("cloudai.systems.slurm.SlurmSystem.get_nodes_by_spec")
+    def test_per_test_isolation(self, mock_get_nodes: Mock, slurm_system: SlurmSystem, test_run: TestRun):
+        mock_get_nodes.side_effect = [(2, ["node01", "node02"]), (2, ["node03", "node04"])]
+
+        # Simulate two different test cases
+        strategy1, strategy2 = ConcreteSlurmStrategy(slurm_system, {}), ConcreteSlurmStrategy(slurm_system, {})
+
+        res = strategy1.get_cached_nodes_spec(test_run)
+        assert mock_get_nodes.call_count == 1
+        assert res == (2, ["node01", "node02"])
+
+        res = strategy2.get_cached_nodes_spec(test_run)
+        assert mock_get_nodes.call_count == 2
+        assert res == (2, ["node03", "node04"])
+
+        assert strategy1._node_spec_cache != strategy2._node_spec_cache, "Caches should be different"
+
+    @patch("cloudai.systems.slurm.SlurmSystem.get_nodes_by_spec")
+    def test_per_iteration_isolation(self, mock_get_nodes: Mock, slurm_system: SlurmSystem, test_run: TestRun):
+        mock_get_nodes.side_effect = [(2, ["node01", "node02"]), (2, ["node03", "node04"])]
+
+        strategy = ConcreteSlurmStrategy(slurm_system, {})
+
+        res = strategy.get_cached_nodes_spec(test_run)
+        assert mock_get_nodes.call_count == 1
+        assert res == (2, ["node01", "node02"])
+
+        test_run.current_iteration = 1
+        res = strategy.get_cached_nodes_spec(test_run)
+        assert mock_get_nodes.call_count == 2
+        assert res == (2, ["node03", "node04"])
+
+    @patch("cloudai.systems.slurm.SlurmSystem.get_nodes_by_spec")
+    def test_per_step_isolation(self, mock_get_nodes: Mock, slurm_system: SlurmSystem, test_run: TestRun):
+        mock_get_nodes.side_effect = [(2, ["node01", "node02"]), (2, ["node03", "node04"])]
+
+        strategy = ConcreteSlurmStrategy(slurm_system, {})
+
+        res = strategy.get_cached_nodes_spec(test_run)
+        assert mock_get_nodes.call_count == 1
+        assert res == (2, ["node01", "node02"])
+
+        test_run.step = 1
+        res = strategy.get_cached_nodes_spec(test_run)
+        assert mock_get_nodes.call_count == 2
+        assert res == (2, ["node03", "node04"])