From 0328cd788406f76dec120d998b61d6140300d712 Mon Sep 17 00:00:00 2001 From: Ryan Waldheim Date: Tue, 23 Jun 2026 09:35:03 -0700 Subject: [PATCH 1/2] Handle higher memory milan nodes --- ansible-runner/project | 2 +- modules/coact.py | 50 +++++++++++++++++++++++--- tests/test_slurm_node_memory.py | 64 +++++++++++++++++++++++++++++++++ 3 files changed, 111 insertions(+), 5 deletions(-) create mode 100644 tests/test_slurm_node_memory.py diff --git a/ansible-runner/project b/ansible-runner/project index 9497aed..fb1cb6f 160000 --- a/ansible-runner/project +++ b/ansible-runner/project @@ -1 +1 @@ -Subproject commit 9497aed512ae9bd1d42b088fae748ddc2a24bc9a +Subproject commit fb1cb6ff427baba5756099b9da9c9f40a772c63a diff --git a/modules/coact.py b/modules/coact.py index 52f5b40..c7c2aba 100644 --- a/modules/coact.py +++ b/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 diff --git a/tests/test_slurm_node_memory.py b/tests/test_slurm_node_memory.py new file mode 100644 index 0000000..1946f22 --- /dev/null +++ b/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 == [""] From 3819c5a3160974612c2f73ce6f22cf9ea21a8b10 Mon Sep 17 00:00:00 2001 From: Ryan Waldheim Date: Tue, 23 Jun 2026 09:46:44 -0700 Subject: [PATCH 2/2] revert sdf-ansible update --- ansible-runner/project | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ansible-runner/project b/ansible-runner/project index fb1cb6f..9497aed 160000 --- a/ansible-runner/project +++ b/ansible-runner/project @@ -1 +1 @@ -Subproject commit fb1cb6ff427baba5756099b9da9c9f40a772c63a +Subproject commit 9497aed512ae9bd1d42b088fae748ddc2a24bc9a