Merge pull request #385 from NVIDIA/am/slurm-system

Remove node list definition from slurm partition

Author: amaslenn

conf/common/system/example_slurm_cluster.toml

@@ -27,7 +27,6 @@ ntasks_per_node = 8
 
 [[partitions]]
 name = "partition_1"
-nodes = ["node-[001-100]"]
 
   [[partitions.groups]]
   name = "group_1"
@@ -47,7 +46,6 @@ nodes = ["node-[001-100]"]
 
 [[partitions]]
 name = "partition_2"
-nodes = ["node-[101-200]"]
 
 [global_env_vars]
 # NCCL Specific Configurations

src/cloudai/systems/slurm/slurm_system.py

@@ -76,25 +76,8 @@ class SlurmPartition(BaseModel):
 
     model_config = ConfigDict(extra="forbid")
     name: str
-    nodes: List[str]
     groups: List[SlurmGroup] = []
-
-    _slurm_nodes: List[SlurmNode] = []
-
-    @property
-    def slurm_nodes(self) -> List[SlurmNode]:
-        if self._slurm_nodes:
-            return self._slurm_nodes
-
-        node_names = set()
-        for nodes_list in self.nodes:
-            node_names.update(set(parse_node_list(nodes_list)))
-
-        self._slurm_nodes = [
-            SlurmNode(name=node_name, partition=self.name, state=SlurmNodeState.UNKNOWN_STATE)
-            for node_name in node_names
-        ]
-        return self._slurm_nodes
+    slurm_nodes: list[SlurmNode] = Field(default_factory=list[SlurmNode], exclude=True)
 
 
 class SlurmSystem(BaseModel, System):
@@ -147,7 +130,10 @@ def groups(self) -> Dict[str, Dict[str, List[SlurmNode]]]:
                 node_names = set()
                 for group_nodes in group.nodes:
                     node_names.update(set(parse_node_list(group_nodes)))
-                groups[part.name][group.name] = [node for node in part.slurm_nodes if node.name in node_names]
+                groups[part.name][group.name] = [
+                    SlurmNode(name=node_name, partition=self.name, state=SlurmNodeState.UNKNOWN_STATE)
+                    for node_name in node_names
+                ]
 
         return groups
 
@@ -163,7 +149,10 @@ def update(self) -> None:
         commands, and correlating this information to determine the state of each node and the user running jobs on
         each node.
         """
-        self.update_node_states()
+        squeue_output, _ = self.fetch_command_output("squeue -o '%N|%u' --noheader")
+        sinfo_output, _ = self.fetch_command_output("sinfo")
+        node_user_map = self.parse_squeue_output(squeue_output)
+        self.parse_sinfo_output(sinfo_output, node_user_map)
 
     def is_job_running(self, job: BaseJob, retry_threshold: int = 3) -> bool:
         """
@@ -373,7 +362,7 @@ def get_available_nodes_from_group(
         """
         self.validate_partition_and_group(partition_name, group_name)
 
-        self.update_node_states()
+        self.update()
 
         grouped_nodes = self.group_nodes_by_state(partition_name, group_name)
 
@@ -490,18 +479,6 @@ def allocate_nodes(
 
         return allocated_nodes
 
-    def is_node_in_system(self, node_name: str) -> bool:
-        """
-        Check if a given node is part of the Slurm system.
-
-        Args:
-            node_name (str): The name of the node to check.
-
-        Returns:
-            True if the node is part of the system, otherwise False.
-        """
-        return any(any(node.name == node_name for node in part.slurm_nodes) for part in self.partitions)
-
     def scancel(self, job_id: int) -> None:
         """
         Terminates a specified Slurm job by sending a cancellation command.
@@ -511,39 +488,6 @@ def scancel(self, job_id: int) -> None:
         """
         self.cmd_shell.execute(f"scancel {job_id}")
 
-    def update_node_states(self) -> None:
-        """
-        Update the states of nodes in the Slurm system.
-
-        By querying the current state of each node using the 'sinfo' command, and correlates this with 'squeue' to
-        determine which user is running jobs on each node. This method parses the output of these commands, identifies
-        the state of nodes and the users, and updates the corresponding SlurmNode instances in the system.
-        """
-        squeue_output = self.get_squeue()
-        sinfo_output = self.get_sinfo()
-        node_user_map = self.parse_squeue_output(squeue_output)
-        self.parse_sinfo_output(sinfo_output, node_user_map)
-
-    def get_squeue(self) -> str:
-        """
-        Fetch the output from the 'squeue' command.
-
-        Returns
-            str: The stdout from the 'squeue' command execution.
-        """
-        squeue_output, _ = self.fetch_command_output("squeue -o '%N|%u' --noheader")
-        return squeue_output
-
-    def get_sinfo(self) -> str:
-        """
-        Fetch the output from the 'sinfo' command.
-
-        Returns
-            str: The stdout from the 'sinfo' command execution.
-        """
-        sinfo_output, _ = self.fetch_command_output("sinfo")
-        return sinfo_output
-
     def fetch_command_output(self, command: str) -> Tuple[str, str]:
         """
         Execute a system command and return its output.
@@ -614,12 +558,25 @@ def parse_sinfo_output(self, sinfo_output: str, node_user_map: Dict[str, str]) -
                 for part in self.partitions:
                     if part.name != partition:
                         continue
+
+                    found = False
                     for node in part.slurm_nodes:
                         if node.name == node_name:
+                            found = True
                             node.state = state_enum
                             node.user = node_user_map.get(node_name, "N/A")
                             break
 
+                    if not found:
+                        part.slurm_nodes.append(
+                            SlurmNode(
+                                name=node_name,
+                                partition=partition,
+                                state=state_enum,
+                                user=node_user_map.get(node_name, "N/A"),
+                            )
+                        )
+
     def convert_state_to_enum(self, state_str: str) -> SlurmNodeState:
         """
         Convert a Slurm node state string to its corresponding enum member.
@@ -709,13 +666,30 @@ def parse_nodes(self, nodes: List[str]) -> List[str]:
                 group_nodes = self.get_available_nodes_from_group(partition_name, group_name, num_nodes)
                 parsed_nodes += [node.name for node in group_nodes]
             else:
-                # Handle both individual node names and ranges
-                if self.is_node_in_system(node_spec) or "[" in node_spec:
-                    expanded_nodes = parse_node_list(node_spec)
-                    parsed_nodes += expanded_nodes
-                else:
-                    raise ValueError(f"Node '{node_spec}' not found.")
+                expanded_nodes = parse_node_list(node_spec)
+                parsed_nodes += expanded_nodes
 
         # Remove duplicates while preserving order
         parsed_nodes = list(dict.fromkeys(parsed_nodes))
         return parsed_nodes
+
+    def get_nodes_by_spec(self, num_nodes: int, nodes: list[str]) -> Tuple[int, list[str]]:
+        """
+        Retrieve a list of node names based on specifications.
+
+        When nodes is empty, returns `(num_nodes, [])`, otherwise parses the node specifications and returns the number
+        of nodes and a list of node names.
+
+        Args:
+            num_nodes (int): The number of nodes, can't be `0`.
+            nodes (list[str]): A list of node names specifications, slurm format or `PARTITION:GROUP:NUM_NODES`.
+
+        Returns:
+            Tuple[int, list[str]]: The number of nodes and a list of node names.
+        """
+        num_nodes, node_list = num_nodes, []
+        parsed_nodes = self.parse_nodes(nodes)
+        if parsed_nodes:
+            num_nodes = len(parsed_nodes)
+            node_list = parsed_nodes
+        return num_nodes, node_list

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

@@ -42,14 +42,6 @@ def __init__(self, system: SlurmSystem, cmd_args: Dict[str, Any]) -> None:
         """
         super().__init__(system, cmd_args)
         self.system = system
-        if not self.system.default_partition:
-            raise ValueError(
-                "Default partition not set in the Slurm system object. "
-                "The 'default_partition' attribute should be properly defined in the Slurm system configuration. "
-                "Please ensure that 'default_partition' is set correctly in the corresponding system configuration "
-                "(e.g., system.toml)."
-            )
-
         self.docker_image_url = self.cmd_args.get("docker_image_url", "")
 
     @abstractmethod
@@ -130,15 +122,12 @@ def _parse_slurm_args(
             KeyError: If partition or essential node settings are missing.
         """
         job_name = self.job_name(job_name_prefix)
-
-        parsed_nodes = self.system.parse_nodes(tr.nodes)
-        num_nodes = len(parsed_nodes) if parsed_nodes else tr.num_nodes
-        node_list_str = ",".join(parsed_nodes) if parsed_nodes else ""
+        num_nodes, node_list = self.system.get_nodes_by_spec(tr.num_nodes, tr.nodes)
 
         slurm_args = {
             "job_name": job_name,
             "num_nodes": num_nodes,
-            "node_list_str": node_list_str,
+            "node_list_str": ",".join(node_list),
         }
         if tr.time_limit:
             slurm_args["time_limit"] = tr.time_limit

src/cloudai/workloads/nemo_launcher/slurm_command_gen_strategy.py

@@ -36,7 +36,7 @@ def _container_mounts(self, tr: TestRun) -> list[str]:
     def gen_exec_command(self, tr: TestRun) -> str:
         self._prepare_environment(tr.test.cmd_args, tr.test.extra_env_vars, tr.output_path)
 
-        nodes = self.system.parse_nodes(tr.nodes)
+        _, nodes = self.system.get_nodes_by_spec(tr.num_nodes, tr.nodes)
         self._set_node_config(nodes, tr.num_nodes)
 
         tdef: NeMoLauncherTestDefinition = cast(NeMoLauncherTestDefinition, tr.test.test_definition)

src/cloudai/workloads/nemo_run/slurm_command_gen_strategy.py

@@ -84,7 +84,7 @@ def generate_test_command(
             "-y",
         ]
 
-        num_nodes = len(self.system.parse_nodes(tr.nodes)) if tr.nodes else tr.num_nodes
+        num_nodes, _ = self.system.get_nodes_by_spec(tr.num_nodes, tr.nodes)
 
         if cmd_args_dict["trainer"]["num_nodes"] and cmd_args_dict["trainer"]["num_nodes"] > num_nodes:
             err = (

tests/conftest.py

@@ -39,15 +39,13 @@ def slurm_system(tmp_path: Path) -> SlurmSystem:
         partitions=[
             SlurmPartition(
                 name="main",
-                nodes=["node-[033-064]"],
                 groups=[
                     SlurmGroup(name="group1", nodes=["node-[033-048]"]),
                     SlurmGroup(name="group2", nodes=["node-[049-064]"]),
                 ],
             ),
             SlurmPartition(
                 name="backup",
-                nodes=["node0[1-8]"],
                 groups=[
                     SlurmGroup(name="group1", nodes=["node0[1-4]"]),
                     SlurmGroup(name="group2", nodes=["node0[5-8]"]),

tests/slurm_command_gen_strategy/conftest.py

@@ -75,7 +75,7 @@ def test_filename_generation(strategy_fixture: SlurmCommandGenStrategy, testrun_
     assert "test_job" in file_contents
     assert "node1,node2" in file_contents
     assert "srun" in file_contents
-    assert "--mpi=fake-mpi" in file_contents
+    assert f"--mpi={strategy_fixture.system.mpi}" in file_contents
 
 
 def test_num_nodes_and_nodes(strategy_fixture: SlurmCommandGenStrategy):
@@ -131,18 +131,6 @@ def test_time_limit(time_limit: Optional[str], strategy_fixture: SlurmCommandGen
         assert "time_limit" not in slurm_args
 
 
-def test_raises_if_no_default_partition(slurm_system: SlurmSystem):
-    slurm_system.default_partition = ""
-    with pytest.raises(ValueError) as exc_info:
-        MySlurmCommandGenStrategy(slurm_system, {})
-    assert (
-        "Default partition not set in the Slurm system object. "
-        "The 'default_partition' attribute should be properly defined in the Slurm "
-        "system configuration. Please ensure that 'default_partition' is set correctly "
-        "in the corresponding system configuration (e.g., system.toml)."
-    ) in str(exc_info.value)
-
-
 @pytest.mark.parametrize(
     "pre_test,post_test,expected_script_lines",
     [

tests/slurm_command_gen_strategy/test_common_slurm_command_gen_strategy.py

@@ -59,13 +59,14 @@ def grok_test(self) -> GrokTestDefinition:
     @pytest.mark.parametrize("test_fixture", ["gpt_test", "grok_test"])
     def test_gen_exec_command(
         self,
-        slurm_system,
+        slurm_system: SlurmSystem,
         cmd_gen_strategy: JaxToolboxSlurmCommandGenStrategy,
         tmp_path: Path,
         request,
         test_fixture,
     ) -> None:
         test_def = request.getfixturevalue(test_fixture)
+        slurm_system.output_path.mkdir(parents=True, exist_ok=True)
 
         test = Test(test_definition=test_def, test_template=TestTemplate(slurm_system, "name"))
         test_run = TestRun(

tests/slurm_command_gen_strategy/test_jax_toolbox_slurm_command_gen_strategy.py

@@ -56,7 +56,7 @@ def slurm_system(tmp_path: Path):
         install_path=tmp_path,
         output_path=tmp_path,
         default_partition="main",
-        partitions=[SlurmPartition(name="main", nodes=["nodeA001", "nodeB001"])],
+        partitions=[SlurmPartition(name="main")],
     )
     return system