Handle higher memory nodes within milan partition

modules/coact.py

@@ -513,6 +513,17 @@ def slurm_import(ctx, print_output, debug, username, password_file, batch, data,
 class SlurmImporter(GraphQlMixin):
     """Handles the slurm import logic."""
 
+    # High-memory node specifications (GB)
+    # sdfmilan[269-272] have 4x standard milan memory (480GB * 4 = 1920GB)
+    # These are currently an exception as most nodes have the same resources
+    # as others within their cluster.
+    HIGH_MEMORY_NODES = {
+        "sdfmilan0269": 1920,
+        "sdfmilan0270": 1920,
+        "sdfmilan0271": 1920,
+        "sdfmilan0272": 1920,
+    }
+
     def __init__(self, username: str, password_file: str, verbose: bool = False, exit_on_error: bool = False):
         self.username = username
         self.password_file = password_file
@@ -673,6 +684,25 @@ def output_json(self, jobs: list, indent: int = 2):
         """Output jobs as JSON."""
         click.echo(json.dumps(jobs, indent=indent, default=datetime_converter))
 
+    def parse_slurm_nodelist(self, nodelist: str) -> list[str]:
+        """Parse SLURM compressed node list format into individual node names."""
+        if '[' not in nodelist:
+            return [nodelist]
+        match = re.match(r'([a-z]+)(\[[\d,\-]+\])', nodelist)
+        if not match:
+            return [nodelist]
+        prefix = match.group(1)
+        ranges_str = match.group(2).strip('[]')
+        nodes = []
+        for part in ranges_str.split(','):
+            if '-' in part:
+                start, end = map(int, part.split('-'))
+                for num in range(start, end + 1):
+                    nodes.append(f"{prefix}{num:04d}")
+            else:
+                nodes.append(f"{prefix}{int(part):04d}")
+        return nodes
+
     def convert(self, index: dict, parts: list, default_facility: str = "shared", default_repo: str = "default") -> Optional[dict]:
         """Convert a line of sacct output to a job dictionary."""
 
@@ -697,7 +727,7 @@ def kilos_to_int(s: str) -> int:
             else:
                 raise Exception("Can't parse %s" % s)
 
-        def calc_resource_hours(startTs, endTs, tres: str, cluster: dict, alloc_nodes: Optional[int], ncpus: Optional[int]) -> tuple:
+        def calc_resource_hours(startTs, endTs, tres: str, cluster: dict, alloc_nodes: Optional[int], ncpus: Optional[int], nodelist: Optional[str]) -> tuple:
             elapsed_secs = (endTs - startTs).total_seconds()
             # min time
             if elapsed_secs <= 0:
@@ -712,18 +742,29 @@ def calc_resource_hours(startTs, endTs, tres: str, cluster: dict, alloc_nodes: O
                         k = "gpu"
                     if alloc_nodes > 0:
                         used[k] = kilos_to_int(v) * 1.0 / alloc_nodes
+
+            # Adjust cluster memory for high-memory nodes
+            adjusted_cluster = cluster.copy()
+            if nodelist:
+                parsed_nodes = self.parse_slurm_nodelist(nodelist)
+                high_mem_nodes = [n for n in parsed_nodes if n in self.HIGH_MEMORY_NODES]
+                if high_mem_nodes:
+                    mem_gb = self.HIGH_MEMORY_NODES[high_mem_nodes[0]]
+                    adjusted_cluster["mem"] = mem_gb * 1073741824
+                    logger.info(f"    Adjusted memory for high-mem node {high_mem_nodes[0]}: {mem_gb}GB")
+
             # if node is exclusive
             # max % of cpu, mem or gpu's for servers
             ratios = {}
             max_ratio = 0
             for resource in ("cpu", "gpu", "mem"):
                 if resource in used:
-                    ratios[resource] = used[resource] / cluster[resource]
+                    ratios[resource] = used[resource] / adjusted_cluster[resource]
                     if ratios[resource] > max_ratio:
                         max_ratio = ratios[resource]
-                    logger.debug(f"    {resource}: used {used[resource]} / {cluster[resource]} -> {ratios[resource]:.5}")
+                    logger.debug(f"    {resource}: used {used[resource]} / {adjusted_cluster[resource]} -> {ratios[resource]:.5}")
             compute_time = elapsed_secs * ncpus / 3600.0
-            resource_time = elapsed_secs * alloc_nodes * max_ratio * cluster["cpu"] / 3600.0
+            resource_time = elapsed_secs * alloc_nodes * max_ratio * adjusted_cluster["cpu"] / 3600.0
             if self.verbose:
                 click.echo(f"  calc time: {elapsed_secs}s compute_hours: {resource_time:.5} core_hours: {compute_time:.5}")
             return resource_time, elapsed_secs
@@ -745,6 +786,7 @@ def calc_resource_hours(startTs, endTs, tres: str, cluster: dict, alloc_nodes: O
             resource_hours, elapsed_secs = calc_resource_hours(
                 startTs=startTs, endTs=endTs, tres=d["AllocTRES"],
                 alloc_nodes=alloc_nodes, ncpus=ncpus, cluster=self._clusters[d["Partition"]],
+                nodelist=d.get("NodeList")
             )
         else:
             resource_hours = 0.0

tests/test_slurm_node_memory.py

@@ -0,0 +1,64 @@
+"""
+Unit tests for SLURM node memory adjustment functionality.
+
+Tests the parsing of SLURM compressed node lists and memory adjustment
+for high-memory nodes (sdfmilan[269-272]).
+"""
+
+from modules.coact import SlurmImporter
+
+
+class TestSlurmNodelistParsing:
+    """Test SLURM nodelist parsing functionality."""
+
+    def setup_method(self):
+        """Create a SlurmImporter instance for testing."""
+        self.importer = SlurmImporter(
+            username="test",
+            password_file="test",
+            verbose=False,
+            exit_on_error=False
+        )
+
+    def test_parse_slurm_nodelist_single(self):
+        """Test parsing a single node name."""
+        result = self.importer.parse_slurm_nodelist("sdfmilan0271")
+        assert result == ["sdfmilan0271"]
+
+    def test_parse_slurm_nodelist_range(self):
+        """Test parsing a SLURM node range."""
+        result = self.importer.parse_slurm_nodelist("sdfmilan[269-272]")
+        expected = ["sdfmilan0269", "sdfmilan0270", "sdfmilan0271", "sdfmilan0272"]
+        assert result == expected
+
+    def test_parse_slurm_nodelist_list(self):
+        """Test parsing a comma-separated list of nodes."""
+        result = self.importer.parse_slurm_nodelist("sdfmilan[006,011,027]")
+        expected = ["sdfmilan0006", "sdfmilan0011", "sdfmilan0027"]
+        assert result == expected
+
+    def test_parse_slurm_nodelist_mixed(self):
+        """Test parsing a mixed range and list."""
+        result = self.importer.parse_slurm_nodelist("sdfmilan[001-003,010,020-022]")
+        expected = [
+            "sdfmilan0001", "sdfmilan0002", "sdfmilan0003",
+            "sdfmilan0010",
+            "sdfmilan0020", "sdfmilan0021", "sdfmilan0022"
+        ]
+        assert result == expected
+
+    def test_parse_slurm_nodelist_different_prefix(self):
+        """Test parsing with different node prefix."""
+        result = self.importer.parse_slurm_nodelist("sdfrome[001-003]")
+        expected = ["sdfrome0001", "sdfrome0002", "sdfrome0003"]
+        assert result == expected
+
+    def test_parse_slurm_nodelist_unparseable(self):
+        """Test that unparseable format returns original string."""
+        result = self.importer.parse_slurm_nodelist("invalid[format")
+        assert result == ["invalid[format"]
+
+    def test_parse_slurm_nodelist_empty(self):
+        """Test parsing empty string."""
+        result = self.importer.parse_slurm_nodelist("")
+        assert result == [""]