diff --git a/README.md b/README.md index 84ec176b..84705888 100644 --- a/README.md +++ b/README.md @@ -16,12 +16,13 @@ Slurm is a highly configurable open source workload manager. See the [Slurm proj 8. [Slurm Job Accounting](#slurm-job-accounting) 1. [Cost Reporting](#cost-reporting) 9. [Topology](#topology) - 10. [GB200/GB300 IMEX Support](#gb200gb300-imex-support) - 11. [Setting KeepAlive in CycleCloud](#setting-keepalive) - 12. [Custom Scheduler Names With High-Availability](#custom-scheduler-names-with-high-availability) - 13. [Slurmrestd](#slurmrestd) - 14. [Node Health Checks](#node-health-checks) - 15. [Monitoring](#monitoring) + 10. [Scaling with scale_m1 (GB200/GB300)](#scaling-with-scale_m1-gb200gb300) + 11. [GB200/GB300 IMEX Support](#gb200gb300-imex-support) + 12. [Setting KeepAlive in CycleCloud](#setting-keepalive) + 13. [Custom Scheduler Names With High-Availability](#custom-scheduler-names-with-high-availability) + 14. [Slurmrestd](#slurmrestd) + 15. [Node Health Checks](#node-health-checks) + 16. [Monitoring](#monitoring) 1. [AzSlurm Exporter](#azslurm-exporter) 1. [Exported Metrics](#exported-metrics) 2. [Configure Exporter Port](#configure-exporter-port) @@ -364,6 +365,108 @@ BlockSizes=5 This either prints out the topology in slurm topology format or creates an output file with the topology. +### Scaling with scale_m1 (GB200/GB300) +`scale_m1` is the InterconnectGroups Milestone 1 (M1) scaling tool for GB200/GB300 (gbx00) clusters. It scales a Slurm partition to an exact size while pruning the generated topology so that only fully populated racks are kept. It is essentially `azslurm scale` specialized for rack-aligned gbx00 hardware. + +It is installed on the scheduler node automatically by the azure-slurm installer (alongside `azslurm`) and runs under the same virtual environment at `/opt/azurehpc/slurm/venv`. It is not installed on execute nodes. Logs are written to `/opt/azurehpc/slurm/logs/scale_m1.log`. + +#### Required slurm.conf settings +To use GB200/GB300, the block topology plugin must be enabled and the gbx00 partition must be excluded from auto-suspend: + +``` +TopologyParam=BlockAsNodeRank +TopologyPlugin=topology/block +# exclude any partition that uses gbx00 +SuspendExcParts=gpu +``` + +**Do not** add these lines to the template's **Slurm Configuration** parameter (`additional_slurm_config` / `AdditionalSlurmConfig`). `TopologyPlugin=topology/block` requires a populated `topology.conf` to exist before `slurmctld` starts; applying it at cluster creation, before any nodes are scaled and the topology is generated, prevents `slurmctld` from starting. + +Instead, apply them manually on the scheduler **after** the nodes are scaled and the block topology file is in place. Add the lines to `/etc/slurm/site_specific.conf` (which is included by `slurm.conf`) and run `scontrol reconfigure`. Replace `gpu` in `SuspendExcParts` with the name of the partition that uses gbx00 nodes if it differs. The full ordering is shown in the workflow below. + +#### Workflow +With the cluster running and an SSH session open on the scheduler node, the end-to-end process is: + +1. **Scale the partition** to the target size — reserve the nodes, then power up to the target. +2. **Generate the block topology and put it in place** at `/sched//topology.conf`. +3. **Enable the block topology plugin** by adding the required settings to `/etc/slurm/site_specific.conf`, then run `scontrol reconfigure`. +4. **Release the ready nodes** from the reservation so jobs can be scheduled on them. + +When releasing nodes in the final step, remove only the nodes that are powered up and ready to run jobs. Leave all powered-off or unhealthy nodes in the reservation, and do **not** delete the reservation — deleting it would also release the not-yet-ready nodes back to the cluster. + +```bash +# 1. Ensure that Slurm will not suspend these nodes by adding the following line the following lines to /etc/slurm/site_specific.conf (included by slurm.conf): +# SuspendExcParts=gpu +$EDITOR /etc/slurm/site_specific.conf +#Apply the configuration +scontrol reconfigure + +# 2. Reserve the partition's nodes under the scale_m1 reservation +scale_m1 create_reservation -p gpu + +# 3. Power up 28 racks (504 nodes) to the target size +scale_m1 power_up --racks 28 + +# 4. Generate the block topology for the partition (--block_size 18 = one rack per block) +azslurm topology -n --block --block_size 18 -p gpu > topology.conf + +# 5. Review topology.conf, then copy it into place. +# For a cluster named "ccw": +cp topology.conf /sched/ccw/topology.conf + +# 6. Enable the block topology plugin adding the the following lines to /etc/slurm/site_specific.conf (the same file as step 1) +# +# TopologyParam=BlockAsNodeRank +# TopologyPlugin=topology/block +$EDITOR /etc/slurm/site_specific.conf +#Apply the configuration +scontrol reconfigure + +# 7. Remove only the powered-up, ready nodes from the reservation so workloads can run on them. +# Leave all powered-off or unhealthy nodes in the reservation; do not delete the reservation. +# Update the reservation so it contains only the nodes that are NOT yet ready. For example, +# if ccw-1-3-gpu-[1-3] are still powering up or unhealthy, keep just those reserved: +scontrol update ReservationName=scale_m1 Nodes=ccw-1-3-gpu-[1-3] +``` + +#### Overprovisioning to absorb unhealthy hardware (advanced) +In rare cases you may want to power up extra nodes so that unhealthy hardware can be discarded while still reaching the target size. `power_up` accepts an overprovision buffer, and `prune` / `prune_now` then trim back to the target, keeping only fully populated racks: + +```bash +# Power up 28 racks plus a 6-rack (108-node) buffer to absorb unhealthy hardware +scale_m1 power_up --racks 28 --overprovision-racks 6 + +# Prune back to 28 racks. `prune` prints the nodes to terminate; `prune_now` also terminates them +scale_m1 prune --racks 28 > to_terminate.txt +scale_m1 prune_now --racks 28 +``` + +After pruning to the target size, continue with the topology generation, configuration, and reservation steps from the workflow above. + +#### Subcommands +| Subcommand | Purpose | +| --- | --- | +| `create_reservation` | Create the `scale_m1` reservation over reservable nodes in a partition. | +| `power_up` | Power up nodes to reach the target size (plus an optional overprovision buffer). | +| `prune` | Print the list of nodes to terminate to reach the target size. | +| `prune_now` | Prune to the target size and terminate the excess nodes. | + +#### Sizing by racks or nodes +`power_up`, `prune`, and `prune_now` accept the target size as either nodes or racks. One rack is 18 nodes, so `--racks N` is equivalent to `--target-count (N * 18)`: + +- `-n` / `--target-count ` — target an exact node count. +- `--racks ` — target a whole number of racks (1 rack = 18 nodes). + +`--target-count` and `--racks` are mutually exclusive; exactly one must be supplied, and the count must be positive. + +`power_up` additionally accepts an overprovision buffer, expressed in nodes or racks. The buffer defaults to 0 when omitted: + +- `-b` / `--overprovision ` — overprovision an exact number of extra nodes. +- `--overprovision-racks ` — overprovision a whole number of extra racks. + +`--overprovision` and `--overprovision-racks` are mutually exclusive and must be non-negative. + + ### GB200/GB300 IMEX Support Cyclecloud Slurm clusters now include prolog and epilog scripts to enable and cleanup IMEX service on a per-job basis. The prolog script will attempt to kill an existing IMEX service before configuring a new instance that will be specific to the new, submitted job. The epilog script terminates the IMEX service. By default, these scripts will run for GB200/GB300 nodes and not run for non-GB200/GB300 nodes. A configurable parameter `slurm.imex.enabled` has been added to the slurm cluster configuration template to allow non-GB200/GB300 nodes to enable IMEX support for their jobs or allow GB200/GB300 nodes to disable IMEX support for their jobs. ``` diff --git a/azure-slurm/install.sh b/azure-slurm/install.sh index ae7c8f78..6b6059ad 100755 --- a/azure-slurm/install.sh +++ b/azure-slurm/install.sh @@ -116,6 +116,23 @@ setup_install_dir() { chmod +x $INSTALL_DIR/*.sh } +setup_scale_m1() { + # scale_m1 is the InterconnectGroups (gbx00) scaling tool. It is installed on + # the scheduler exactly like azslurm: a venv-resident executable wired to the + # venv python (so it can import slurmcc) plus a ~/bin symlink. It is + # intentionally NOT installed on execute nodes. Re-running install.sh + # re-installs it cleanly (install -m and ln -sf are idempotent). + if [ ! -f scale_m1/scale_to_n_nodes.py ]; then + echo "scale_m1 entrypoint not found in package; skipping scale_m1 install." >&2 + return 0 + fi + install -m 0755 scale_m1/scale_to_n_nodes.py $VENV/bin/scale_m1 + if [ ! -e ~/bin ]; then + mkdir ~/bin + fi + ln -sf $VENV/bin/scale_m1 ~/bin/ +} + init_azslurm_config() { which jetpack || (echo "Jetpack is not installed. Please run this from a CycleCloud node, or pass in --no-jetpack if you intend to install this outside of CycleCloud provisioned nodes." && exit 1) @@ -198,6 +215,8 @@ main() { setup_venv # setup the install dir - logs and logging.conf, some permissions. setup_install_dir + # install the scale_m1 (gbx00) scaling tool on the scheduler, like azslurm. + setup_scale_m1 # setup the azslurmd but do not start it. setup_azslurmd # If there is no jetpack, we have to stop here. diff --git a/azure-slurm/package.py b/azure-slurm/package.py index e68cfd82..731e7616 100644 --- a/azure-slurm/package.py +++ b/azure-slurm/package.py @@ -166,7 +166,11 @@ def _add(name: str, path: Optional[str] = None, mode: Optional[int] = None) -> N _add("sbin/suspend_program.sh", "sbin/suspend_program.sh") _add("sbin/get_acct_info.sh", "sbin/get_acct_info.sh") _add("logging.conf", "conf/logging.conf") - + # scale_m1 is the InterconnectGroups (gbx00) scaling tool. It is a sibling + # project of azure-slurm and ships only its runtime entrypoint (mock.py is + # test-only and intentionally excluded). It runs under the azure-slurm venv. + _add("scale_m1/scale_to_n_nodes.py", "../scale_m1/scale_to_n_nodes.py", mode=0o100755) + if __name__ == "__main__": diff --git a/scale_m1/cli_args_test.py b/scale_m1/cli_args_test.py new file mode 100644 index 00000000..c0657ad8 --- /dev/null +++ b/scale_m1/cli_args_test.py @@ -0,0 +1,150 @@ +"""Tests for scale_m1 CLI argument handling: racks-vs-nodes sizing, +overprovision resolution, and the shebang/constant invariants.""" +import os + +import pytest + +import scale_to_n_nodes +from scale_to_n_nodes import ( + NODES_PER_RACK, + build_parser, + mock_available, + resolve_overprovision, + resolve_target_count, +) + + +def _parse(argv): + return build_parser().parse_args(argv) + + +# --- target count resolution (--racks vs -n/--target-count) --- + + +def test_racks_flag_sets_target_to_racks_times_18(): + args = _parse(["power_up", "--racks", "28"]) + assert resolve_target_count(args) == 28 * NODES_PER_RACK + + +def test_target_count_still_supported(): + args = _parse(["power_up", "--target-count", "504"]) + assert resolve_target_count(args) == 504 + + +def test_racks_and_target_count_together_errors(): + # argparse mutually exclusive group exits with code 2. + with pytest.raises(SystemExit): + _parse(["power_up", "--target-count", "504", "--racks", "28"]) + + +def test_neither_racks_nor_target_errors(): + # the group is required=True, so omitting both exits. + with pytest.raises(SystemExit): + _parse(["power_up"]) + + +def test_racks_must_be_positive(): + args = _parse(["prune", "--racks", "0"]) + with pytest.raises(SystemExit): + resolve_target_count(args) + + +def test_negative_target_count_errors(): + args = _parse(["prune", "--target-count", "-5"]) + with pytest.raises(SystemExit): + resolve_target_count(args) + + +# --- overprovision resolution --- + + +def test_overprovision_racks_sets_buffer(): + args = _parse(["power_up", "--racks", "28", "--overprovision-racks", "6"]) + assert resolve_overprovision(args) == 6 * NODES_PER_RACK + + +def test_overprovision_nodes_supported(): + args = _parse(["power_up", "--target-count", "504", "--overprovision", "108"]) + assert resolve_overprovision(args) == 108 + + +def test_overprovision_racks_zero_accepted(): + args = _parse(["power_up", "--racks", "28", "--overprovision-racks", "0"]) + assert resolve_overprovision(args) == 0 + + +def test_overprovision_racks_negative_errors(): + args = _parse(["power_up", "--racks", "28", "--overprovision-racks", "-1"]) + with pytest.raises(SystemExit): + resolve_overprovision(args) + + +def test_overprovision_omitted_defaults_to_zero(): + args = _parse(["power_up", "--racks", "28"]) + assert resolve_overprovision(args) == 0 + + +def test_overprovision_negative_nodes_errors(): + args = _parse(["power_up", "--target-count", "504", "--overprovision", "-1"]) + with pytest.raises(SystemExit): + resolve_overprovision(args) + + +def test_overprovision_racks_and_overprovision_together_errors(): + with pytest.raises(SystemExit): + _parse( + ["power_up", "--racks", "28", "--overprovision", "1", "--overprovision-racks", "1"] + ) + + +# --- racks supported on prune / prune_now too --- + + +@pytest.mark.parametrize("cmd", ["prune", "prune_now", "power_up"]) +def test_racks_supported_on_all_scaling_subcommands(cmd): + args = _parse([cmd, "--racks", "10"]) + assert resolve_target_count(args) == 10 * NODES_PER_RACK + + +# --- conditional --mock flag on prune --- + + +def test_prune_mock_flag_available_when_mock_importable(): + # In the source tree mock.py is importable, so --mock is registered. + assert mock_available() is True + args = _parse(["prune", "--racks", "10", "--mock"]) + assert args.mock_topology is True + + +def test_prune_mock_defaults_false_when_omitted(): + args = _parse(["prune", "--racks", "10"]) + assert args.mock_topology is False + + +def test_prune_mock_hidden_when_mock_unavailable(monkeypatch): + monkeypatch.setattr(scale_to_n_nodes, "mock_available", lambda: False) + with pytest.raises(SystemExit): + _parse(["prune", "--racks", "10", "--mock"]) + + +# --- shebang / constant invariants --- + + +def test_nodes_per_rack_constant_value(): + assert NODES_PER_RACK == 18 + + +def test_round_up_to_rack_uses_constant(): + scaler = scale_to_n_nodes.NodeScaler.__new__(scale_to_n_nodes.NodeScaler) + assert scaler.round_up_to_rack(1) == NODES_PER_RACK + assert scaler.round_up_to_rack(NODES_PER_RACK) == NODES_PER_RACK + assert scaler.round_up_to_rack(NODES_PER_RACK + 1) == 2 * NODES_PER_RACK + + +def test_shebang_uses_unversioned_python3(): + source = os.path.join(os.path.dirname(__file__), "scale_to_n_nodes.py") + with open(source) as fh: + first_line = fh.readline().strip() + assert first_line.startswith("#!") + assert first_line.endswith("/python3"), first_line + assert "python3.11" not in first_line diff --git a/scale_m1/mock.py b/scale_m1/mock.py new file mode 100644 index 00000000..c31958e9 --- /dev/null +++ b/scale_m1/mock.py @@ -0,0 +1,360 @@ +import logging +import shlex +import subprocess +import time + +from scale_to_n_nodes import SlurmCommands, AzslurmTopology, get_healthy_nodes, Clock, setup_logging + + +setup_logging() +log = logging.getLogger(__name__) + + +def parse_command(x: str) -> dict: + """Parses sinfo/scontrol commands into a dictionary. Only really relevant for + our specific use cases, DO NOT REUSE + "command a=b -c 'd e'" become {'command': 'command', 'a': 'b', '-c': 'd e'} + """ + ret = {} + toks = shlex.split(x) + i = 0 + while i < len(toks): + tok = toks[i] + if "=" in tok: + key, value = tok.split("=", 1) + elif tok.startswith("-") and i + 1 < len(toks): + key = tok + value = toks[i + 1] + i = i + 1 + else: + key = value = tok + ret[key] = value + i = i + 1 + return ret + + +class MockClock(Clock): + def __init__(self, start_time: float = None): + super().__init__() + self.current_time = start_time if start_time else time.time() + + def sleep(self, seconds: float): + log.info(f"[TEST MODE] Mock sleep for {seconds} seconds") + self.current_time += seconds + + def time(self) -> float: + return self.current_time + + +class MockSlurmCommands(SlurmCommands): + """Mock version of Slurm. Typical use case + mock_slurm = MockSlurmCommands("/tmp/topology.conf") + mock_slurm.create_nodes([f"hpc-{x+1}" for x in range(18)]) + mock_slurm.simulate_failed_converge(["hpc-3"]) + mock_slurm.run_command("scontrol update nodename=hpc-3,hpc-4 state=power_up) + # node will not actually enter POWERING_UP until this is called, to simulate + # the actual delay in slurm + mock_slurm.update_states() + # sinfo works - so you can also just query states + assert mock_slurm.run_command("sinfo -t powering_up -N -o '%N').stdout.strip().split() == ["hpc-3", "hpc-4"] + # call update states again and one node is now powered_up, the other is powering_down + mock_slurm.update_states() + assert mock_slurm.node_dicts["hpc-3"]["power_state"] == "POWERED_UP" + assert mock_slurm.node_dicts["hpc-4"]["power_state"] == "POWERING_DOWN" + assert mock_slurm.run_command("sinfo -t powering_down -N -o '%N').stdout.strip() == "hpc-4" + """ + def __init__(self, topology_file: str = "/etc/slurm/topology.conf"): + self.topology_file = topology_file + self.last_topology = "" + log.info("Running in TEST MODE - all commands will be mocked") + self.nodes_dict: dict[str, dict] = {} + self.reservation_name: str = "" + self.reservation_nodes: list[str] = [] + self.reservation_partition: str = "" + + def read_topology(self) -> str: + """Read the current SLURM topology file.""" + with open(self.topology_file, 'r', encoding='utf-8') as f: + topology = f.read() + log.info(f"[TEST MODE] Read mock topology file: {self.topology_file}") + return topology + + + def create_nodes(self, partition: str, count: int): + """Create mock nodes for testing.""" + log.info(f"[TEST MODE] Creating {count} mock nodes in partition '{partition}'") + assert count >= 18, "Count must be greater than 18 for mock nodes" + for i in range(count): + node_name = f"{partition}-{i + 1}" + self.nodes_dict[node_name] = { + 'partition': partition, + 'state': 'IDLE', + 'power_state': 'POWERED_DOWN', + 'reservation_name': '' + } + log.info(f"[TEST MODE] Created {len(self.nodes_dict)} total mock nodes") + + def show_hostlist(self, nodes: list[str]) -> str: + assert sorted(nodes, key=lambda x: int(x.split("-")[-1])) == nodes + return ",".join(nodes) + + def show_hostnames(self, node_str: str) -> list[str]: + return node_str.split(",") + + + def alloc_nodes(self, nodes: list[str]) -> None: + for node in nodes: + assert self.nodes_dict[node]["power_state"] == "POWERED_UP", f"invalid state {self.nodes_dict[node]}" + self.nodes_dict[node]["state"] = "allocated" + + def drain_nodes(self, nodes: list[str]) -> None: + for node in nodes: + assert self.nodes_dict[node]["power_state"] == "POWERED_UP", f"invalid state {self.nodes_dict[node]}" + self.nodes_dict[node]["state"] = "draining" + + def down_nodes(self, nodes: list[str]) -> None: + for node in nodes: + assert self.nodes_dict[node]["power_state"] == "POWERED_UP", f"invalid state {self.nodes_dict[node]}" + self.nodes_dict[node]["state"] = "down" + + def reserve_nodes(self, nodes: list[str], reservation_name: str = "scale_m1") -> None: + for node in nodes: + self.nodes_dict[node]["reservation_name"] = reservation_name + + def simulate_failed_converge(self, failed_nodes: list): + """Simulate some nodes as unhealthy for testing.""" + log.info(f"[TEST MODE] Simulating failed converge for nodes: {failed_nodes}") + for node in failed_nodes: + if node in self.nodes_dict: + self.nodes_dict[node]['simulate_failure'] = True + log.info(f"[TEST MODE] Updated states for failed nodes: {failed_nodes}") + + def update_states(self, count=9): + """Update mock node states based on simulated failures.""" + log.info(f"[TEST MODE] Updating {count} mock node states") + for node_name, node_data in self.nodes_dict.items(): + if node_data['power_state'] == 'POWERING_UP': + if node_data.get('simulate_failure'): + # # TODO will mess with reservation + # and yes it did. + # node_data['state'] = 'DOWN' + node_data['power_state'] = 'POWERING_DOWN' + else: + node_data['power_state'] = 'POWERED_UP' + count -= 1 + elif node_data['power_state'] == 'POWERING_DOWN': + node_data['power_state'] = 'POWERED_DOWN' + count -= 1 + # Update one node and return + elif node_data['power_state'] == 'PRE_POWERING_UP': + node_data['power_state'] = 'POWERING_UP' + count -= 1 + elif node_data['power_state'] == 'PRE_POWERING_DOWN': + node_data['power_state'] = 'POWERING_DOWN' + count -= 1 + # Update one node and return + if count <= 0: + log.info("[TEST MODE] Mock node states updated") + return + log.info(f"[TEST MODE] Mock {count} node states remaining") + + def sinfo(self, cmd: str, cmd_parsed: dict) -> subprocess.CompletedProcess: + if '-t' in cmd_parsed: + states = cmd_parsed['-t'].upper().split(',') + ret = [] + for node_name, node_dict in self.nodes_dict.items(): + if (node_dict['power_state'].upper() in states or + node_dict['state'].upper() in states): + ret.append(node_name) + elif "RESERVED" in states and node_dict.get("reservation_name"): + # handle reserved by looking at res name NOT state - drained etc can be reserved. + ret.append(node_name) + nodes_str = "\n".join(ret) + "\n" + elif cmd_parsed['-o'] == '%N': + nodes_str = "\n".join(self.nodes_dict.keys()) + "\n" + elif cmd_parsed['-o'] == '%N %T': + ret = [] + target_nodes = cmd_parsed['-n'].split(',') + for node_name, node_dict in self.nodes_dict.items(): + if node_name not in target_nodes: + continue + state = node_dict['state'].upper() + assert state in ['IDLE', 'DOWN', 'DRAINING', 'DRAINED', 'RESERVED', 'ALLOCATED'], f"Unknown state: {state}" + if node_dict['power_state']=="POWERED_UP": + pass + elif node_dict['power_state']=="POWERED_DOWN": + state += '~' + elif node_dict['power_state']=="POWERING_UP": + state += '#' + elif node_dict['power_state']=="POWERING_DOWN": + state += '%' + elif node_dict['power_state'].startswith("PRE_"): + state += '!' + ret.append(f"{node_name} {state}") + nodes_str = "\n".join(ret) + "\n" + else: + raise ValueError(f"Unsupported sinfo output format: {cmd_parsed['-o']}") + + return subprocess.CompletedProcess(cmd, 0, nodes_str, "") + + def scontrol_create_reservation(self, cmd: str, cmd_parsed: dict) ->subprocess.CompletedProcess: + assert not self.reservation_nodes + self.reservation_name = cmd_parsed["ReservationName"] + + self.reservation_nodes = cmd_parsed["Nodes"].split(",") + self.reservation_partition = cmd_parsed["PartitionName"] + + for node_name in self.reservation_nodes: + node_dict = self.nodes_dict[node_name] + assert node_dict['partition'] == cmd_parsed["PartitionName"] + assert node_dict["state"].upper() not in ["MIXED", "ALLOCATED"] + node_dict['reservation_name'] = self.reservation_name + + log.info(f"[TEST MODE] Mock reservation created: {self.reservation_name}") + return subprocess.CompletedProcess(cmd, 0, "Mock reservation created", "") + + def scontrol_update_reservation(self, cmd: str, cmd_parsed: dict) -> subprocess.CompletedProcess: + assert cmd_parsed["ReservationName"] == self.reservation_name, f"{cmd_parsed} res_name={self.reservation_name}" + self.reservation_nodes = cmd_parsed["Nodes"].split(",") + for node in self.reservation_nodes: + node_dict = self.nodes_dict[node] + if node in self.reservation_nodes: + node_dict["state"] = "RESERVED" + elif node_dict["state"] == "RESERVED": + node_dict["state"] = "IDLE" + return subprocess.CompletedProcess(cmd, 0, "Mock reservation updated", "") + + def scontrol_update_nodename(self, cmd: str, cmd_parsed: dict) -> subprocess.CompletedProcess: + states = {'power_up':'PRE_POWERING_UP', 'power_down':'PRE_POWERING_DOWN'} + node_list = cmd_parsed.get("NodeName", "").split(',') + new_state = cmd_parsed.get("State", "").lower() + for node_name in node_list: + if new_state == "power_up" and self.nodes_dict[node_name]['power_state'] == "POWERED_UP": + continue + if new_state == "power_up": + self.nodes_dict[node_name]["created_a_vm"] = True + self.nodes_dict[node_name]['power_state'] = states[new_state] + + log.info(f"[TEST MODE] Mock node {node_name} power state updated to {new_state}") + + # mimic M1 assertion that all VMSS requests end up having a multiple of 18 + if new_state == "power_up" and any([n.get("created_a_vm") for n in self.nodes_dict.values()]): + has_a_vm = 0 + for node_name, node_dict in self.nodes_dict.items(): + if node_dict["power_state"] in ["PRE_POWERING_UP", "POWERING_UP", "POWERED_UP"]: + has_a_vm += 1 + assert has_a_vm % 18 == 0, f"This would result in an improper m1 allocation - {has_a_vm} total vms" + + return subprocess.CompletedProcess(cmd, 0, "Mock nodes updated", "") + + def _power_up(self, node_names: list[str]) -> None: + for node_name in node_names: + self.nodes_dict[node_name]['power_state'] = "PRE_POWERING_UP" + + def scontrol_show_reservation(self, cmd: str, cmd_parsed: dict) -> subprocess.CompletedProcess: + if self.reservation_name == cmd.split()[-1]: + res_count = len(self.reservation_nodes) + output = f"""ReservationName={self.reservation_name} + Nodes={','.join(self.reservation_nodes)} + NodeCnt={res_count} + PartitionName={self.reservation_partition}""" + return subprocess.CompletedProcess(cmd, 0, output, "") + else: + return subprocess.CompletedProcess(cmd, 1, "Reservation not found", "") + + def run_command(self, cmd: str, check: bool = True) -> subprocess.CompletedProcess: + """Mock SLURM and system commands for test mode.""" + ret = self._run_command(cmd) + if ret.returncode != 0 and check: + raise subprocess.CalledProcessError(ret.returncode, cmd, output=ret.output, error=ret.stderr) + return ret + + def _run_command(self, cmd: str) -> subprocess.CompletedProcess: + log.info(f"[TEST MODE] Would run: {cmd}") + cmd_parsed = parse_command(cmd) + + if "azslurm" in cmd_parsed: + return subprocess.CompletedProcess(cmd, 0, "", "") + + if "scontrol" in cmd_parsed: + if "create" in cmd_parsed: + if "reservation" in cmd_parsed: + return self.scontrol_create_reservation(cmd, cmd_parsed) + elif "update" in cmd_parsed: + if "NodeName" in cmd_parsed and "State" in cmd_parsed: + return self.scontrol_update_nodename(cmd, cmd_parsed) + elif "ReservationName" in cmd_parsed: + return self.scontrol_update_reservation(cmd, cmd_parsed) + elif "partitionname" in cmd_parsed: + return subprocess.CompletedProcess(cmd, 0, "", "") + elif "reconfigure" in cmd: + self.last_topology = self.read_topology() + return subprocess.CompletedProcess(cmd, 0, "Mock reconfigure executed", "") + elif "show" in cmd_parsed: + if "reservation" in cmd_parsed: + return self.scontrol_show_reservation(cmd, cmd_parsed) + elif "hostnames" in cmd_parsed: + assert len(cmd.strip().split()) > 3 + nodes = cmd.split()[-1].split(',') + return subprocess.CompletedProcess(cmd, 0, "\n".join(nodes) + "\n", "") + elif "topology" in cmd_parsed: + return subprocess.CompletedProcess(cmd, 0, self.last_topology, "") + elif "sinfo" in cmd_parsed: + return self.sinfo(cmd, cmd_parsed) + + raise RuntimeError(f"Unknown cmd {cmd}") + + def node_state_counts(self) -> dict[str, int]: + counts = {"idle": 0, "reserved": 0, "powered_up": 0, "powered_down": 0, "draining": 0, + "down": 0, "allocated": 0, "created_a_vm": 0, "total": len(self.nodes_dict)} + for node_dict in self.nodes_dict.values(): + power_state = node_dict["power_state"].lower() + if power_state not in counts: + counts[power_state] = 0 + counts[power_state] = counts[power_state] + 1 + + state = node_dict["state"].lower() + if state not in counts: + counts[state] = 0 + counts[state] = counts[state] + 1 + if node_dict.get("reservation_name"): + counts["reserved"] += 1 + if node_dict.get("created_a_vm"): + counts["created_a_vm"] += 1 + return counts + + +class MockAzslurmTopology(AzslurmTopology): + def __init__(self, slurm_commands: SlurmCommands): + super().__init__() + self.slurm_commands = slurm_commands + + def generate_topology(self, partition: str, topology_file: str) -> str: + """Create a mock topology file for testing.""" + mock_topology_lines = [] + nodes = get_healthy_nodes(partition, self.slurm_commands) + total_nodes = len(nodes) + log.info(f"[TEST MODE] Creating mock topology with {total_nodes} nodes") + blocks = {} + mock_topology_lines.append(f"# Mock topology for testing") + for node in nodes: + node_index = int(node.split("-")[-1]) + block_index = (node_index - 1) // 18 + if block_index not in blocks: + blocks[block_index] = [] + blocks[block_index].append(node) + for block, node_list in blocks.items(): + block_name = f"block_{block + 1:03d}" + node_list_str = ",".join(node_list) + + mock_topology_lines.append(f"BlockName={block_name} Nodes={node_list_str}") + mock_topology_lines.append("BlockSizes=1") + + mock_topology = "\n".join(mock_topology_lines) + "\n" + + with open(topology_file, 'w', encoding='utf-8') as f: + f.write(mock_topology) + log.info(f"[TEST MODE] Created mock topology file: {topology_file}") + return mock_topology + + diff --git a/scale_m1/mock_test.py b/scale_m1/mock_test.py new file mode 100644 index 00000000..befa7844 --- /dev/null +++ b/scale_m1/mock_test.py @@ -0,0 +1,63 @@ +from mock import MockAzslurmTopology, MockSlurmCommands + +import scale_to_n_nodes + + +def test_scontrol_power_down(): + mock_slurm_commands = MockSlurmCommands(topology_file="/tmp/topology.txt") + mock_slurm_commands.create_nodes(partition="gpu", count=18) + mock_slurm_commands.run_command("scontrol update NodeName=gpu-6 State=POWER_DOWN") + result = mock_slurm_commands.run_command("sinfo -p gpu -t powering_down -o '%N'") + assert result.stdout.strip() == "" + mock_slurm_commands.update_states(50) + result = mock_slurm_commands.run_command("sinfo -p gpu -t powering_down -o '%N'") + assert result.stdout.strip() == "gpu-6" + + +def test_reservation(): + mock_slurm_commands = MockSlurmCommands(topology_file="/tmp/topology.txt") + mock_slurm_commands.create_nodes(partition="gpu", count=18) + mock_slurm_commands.run_command("scontrol create reservation PartitionName=gpu ReservationName=test_resv Nodes=gpu-1,gpu-2 Duration=60") + result = mock_slurm_commands.run_command("scontrol show reservation test_resv") + assert "ReservationName=test_resv" in result.stdout + assert "Nodes=gpu-1,gpu-2" in result.stdout + + +def test_generate_single_block(): + mock_slurm_commands = MockSlurmCommands(topology_file="/tmp/topology.txt") + mock_slurm_commands.create_nodes(partition="gpu", count=18) + mock_topology = MockAzslurmTopology(mock_slurm_commands) + mock_slurm_commands.run_command("scontrol update NodeName=gpu-1,gpu-2,gpu-3,gpu-4,gpu-5,gpu-6,gpu-7,gpu-8,gpu-9,gpu-10,gpu-11,gpu-12,gpu-13,gpu-14,gpu-15,gpu-16,gpu-17,gpu-18 State=POWER_UP") + mock_slurm_commands.update_states(50) + mock_slurm_commands.update_states(50) + result = mock_slurm_commands.run_command('sinfo -p gpu -t powered_down,powering_up,powering_down,power_down,drain,drained,draining,unknown,down,no_respond,fail,reboot -o "%N" -h -N') + assert result.stdout.strip() == '' + mock_topology.generate_topology("gpu", "/tmp/topology.txt") + with open("/tmp/topology.txt", "r") as f: + content = f.read() + assert "# Mock topology for testing\nBlockName=block_001 Nodes=gpu-1,gpu-2,gpu-3,gpu-4,gpu-5,gpu-6,gpu-7,gpu-8,gpu-9,gpu-10,gpu-11,gpu-12,gpu-13,gpu-14,gpu-15,gpu-16,gpu-17,gpu-18\nBlockSizes=1" in content + + +def test_sinfo(): + sc = MockSlurmCommands(topology_file="/tmp/topology.txt") + sc.create_nodes(partition="gpu", count=18) + sc._power_up(["gpu-1", "gpu-2", "gpu-3", "gpu-4"]) + sc.update_states(4) + sc.update_states(4) + powered_down = [f"gpu-{n+1}" for n in range(4, 18)] + assert scale_to_n_nodes.get_healthy_nodes("gpu", sc) == ["gpu-1", "gpu-2", "gpu-3", "gpu-4"] + assert scale_to_n_nodes.get_powered_up_nodes("gpu", sc) == ["gpu-1", "gpu-2", "gpu-3", "gpu-4"] + assert scale_to_n_nodes.get_powered_down("gpu", sc) == powered_down + assert scale_to_n_nodes.get_reservable_nodes("gpu", sc) == ["gpu-1", "gpu-2", "gpu-3", "gpu-4"] + powered_down + + sc.drain_nodes(["gpu-1"]) + sc.reserve_nodes(["gpu-2"]) + sc.alloc_nodes(["gpu-3"]) + sc.down_nodes(["gpu-4"]) + + assert scale_to_n_nodes.get_healthy_nodes("gpu", sc) == ["gpu-2", "gpu-3"] + assert scale_to_n_nodes.get_powered_up_nodes("gpu", sc) == ["gpu-1", "gpu-2", "gpu-3", "gpu-4"] + assert scale_to_n_nodes.get_powered_down("gpu", sc) == powered_down + # remove gpu-2 (reserved) and gpu-3 (allocated) + assert scale_to_n_nodes.get_reservable_nodes("gpu", sc) == ["gpu-1", "gpu-4"] + powered_down + diff --git a/scale_m1/scale_to_n_nodes.py b/scale_m1/scale_to_n_nodes.py new file mode 100644 index 00000000..a96fcd32 --- /dev/null +++ b/scale_m1/scale_to_n_nodes.py @@ -0,0 +1,782 @@ +#!/opt/azurehpc/slurm/venv/bin/python3 +""" +Scale SLURM cluster to exactly N nodes with topology optimization. +Usage: ./scale_m1 -p gpu -n/--target-count 504 -b/--overprovision 108 +""" + +import argparse +import json +import logging +import logging.handlers +import math +import os +import shutil +import subprocess +import sys +import time as timelib + +from abc import ABC, abstractmethod +from pathlib import Path +from typing import List, Dict, Tuple, Union + + +# Add the parent directory to Python path to import slurmcc modules +sys.path.insert(0, str(Path(__file__).parent)) + +from slurmcc.topology import output_block_nodelist +from slurmcc import util as slutil + + +# Number of nodes in a single rack. Racks are the unit of physical provisioning +# for gb200/gb300 (gbx00) clusters; users may size requests in racks instead of +# raw node counts. +NODES_PER_RACK = 18 + + +def setup_logging(verbose: bool = False) -> None: + """Configure logging with both console and rotating file handlers.""" + logger = logging.getLogger() + logger.setLevel(logging.DEBUG) + logger.handlers.clear() + formatter = logging.Formatter( + '%(asctime)s - %(levelname)s - %(message)s', + datefmt='%Y-%m-%d %H:%M:%S' + ) + + # Console handler + console_handler = logging.StreamHandler() + console_handler.setLevel(logging.DEBUG if verbose else logging.INFO) + console_handler.setFormatter(formatter) + logger.addHandler(console_handler) + + # Rotating file handler + if os.getenv('SCALE_M1_LOG_FILE'): + log_file = os.getenv('SCALE_M1_LOG_FILE') + else: + log_file = "/opt/azurehpc/slurm/logs/scale_m1.log" + # Rotating file handler + file_handler = logging.handlers.RotatingFileHandler( + log_file, + maxBytes=1024 * 1024 * 5, # 5MB + backupCount=5, + encoding='utf-8' + ) + file_handler.setLevel(logging.DEBUG) + file_handler.setFormatter(formatter) + logger.addHandler(file_handler) + + logger.info(f"Logging initialized. Console level: {'DEBUG' if verbose else 'INFO'}, File: {log_file}") + + +log = logging.getLogger(__name__) + + +__LAST_MSG = "" +def cache_log(msg: str) -> None: + "Don't repeat the last message" + global __LAST_MSG + if __LAST_MSG != msg: + log.info(msg) + __LAST_MSG = msg + + +class Clock(ABC): + @abstractmethod + def sleep(self, seconds: int): + """Sleep for a given number of seconds.""" + ... + + @abstractmethod + def time(self) -> float: + """Get the current time in seconds since the epoch.""" + ... + + +class RealClock(Clock): + def sleep(self, seconds: int): + """Sleep for a given number of seconds.""" + timelib.sleep(seconds) + + def time(self) -> float: + + """Get the current time in seconds since the epoch.""" + return timelib.time() + + +CLOCK: Clock = RealClock() + +class SlurmM1Error(Exception): + """Base class for SLURM M1 errors.""" + pass + + +class NodeInvalidStateError(SlurmM1Error): + """Raised when nodes are in an invalid state for scaling.""" + def __init__(self, message: str): + super().__init__(message) + self.message = message + + def __str__(self): + return f"NodeInvalidStateError: {self.message}" + + +class SlurmCommandError(SlurmM1Error): + """Raised when a SLURM command fails.""" + def __init__(self, command: str, returncode: int, output: str): + super().__init__(f"Command '{command}' failed with return code {returncode}") + self.command = command + self.returncode = returncode + self.output = output + + def __str__(self): + return f"SlurmCommandError: Command '{self.command}' failed with return code {self.returncode}. Output: {self.output}" + + +class NoReservationError(SlurmM1Error): + pass + + +class AzslurmTopology(ABC): + @abstractmethod + def generate_topology(self, partition: str, topology_file: str) -> str: + """Generate the SLURM topology configuration.""" + ... + + +class AzslurmTopologyImpl(AzslurmTopology): + def generate_topology(self, partition: str, topology_file: str) -> str: + """Generate the SLURM topology configuration.""" + cmd = f"azslurm topology -n -p {partition} -o {topology_file}" + try: + result = slutil.run(cmd, shell=True, timeout=300) + log.info(f"Topology generated: {topology_file}") + return result.stdout + except subprocess.CalledProcessError as e: + log.error(f"Failed to generate topology: {e}") + raise + + +class SlurmCommands(ABC): + @abstractmethod + def run_command(self, cmd: str, check: bool = True) -> subprocess.CompletedProcess: + ... + + def update_states(self, count: int = 9): + pass # overridden for testing + + @abstractmethod + def show_hostlist(self, nodes: list[str]) -> str: + ... + + @abstractmethod + def show_hostnames(self, node_str: str) -> list[str]: + ... + + +class SlurmCommandsImpl(SlurmCommands): + def run_command(self, cmd: str, check: bool = True) -> subprocess.CompletedProcess: + """Run a SLURM command and return the result.""" + try: + return slutil.run(cmd, shell=True, timeout=300, check=check) + except subprocess.CalledProcessError as e: + log.error(f"Command failed: {cmd}\nError: {e}") + raise + except subprocess.TimeoutExpired as e: + log.error(f"Command timed out: {cmd}\nError: {e}") + raise + + def show_hostlist(self, nodes: list[str]) -> str: + return slutil.to_hostlist(nodes) + + def show_hostnames(self, node_str: str) -> list[str]: + return slutil.from_hostlist(node_str) + + +class NodeScaler: + def __init__(self, target_count: int, overprovision: int, + slurm_commands: SlurmCommands, azslurm_topology: AzslurmTopology, + topology_file: str = os.path.realpath("/etc/slurm/topology.conf"), + reservation_name: str = f"scale_m1"): + self.target_count = target_count + self.overprovision = overprovision + self.topology_file = topology_file + self.slurm_commands = slurm_commands + self.azslurm_topology = azslurm_topology + self.reservation_name = reservation_name + self.partition, self.reserved_nodes = parse_reservation(reservation_name, slurm_commands) + + def round_up_to_rack(self, number: int) -> int: + """Round up to the nearest whole rack (multiple of NODES_PER_RACK).""" + return math.ceil(number / NODES_PER_RACK) * NODES_PER_RACK + + def validate_nodes(self) -> None: + cmd = f"sinfo -p {self.partition} -t powering_down,powering_up -h -o '%N'" + result = self.run_command(cmd) + powering_nodes = result.stdout.strip() + if powering_nodes: + raise NodeInvalidStateError("Some nodes are in POWERING_DOWN or POWERING_UP state, cannot proceed with scaling.") + + def run_command(self, cmd: str, check: bool = True) -> subprocess.CompletedProcess: + """Run a command and return the result. In test mode, return mocked responses.""" + return self.slurm_commands.run_command(cmd, check=check) + + def power_up_nodes(self, node_count: int) -> str: + """ + Power up at least node_count healthy nodes, while guaranteeing + we are within the rack boundaries. + """ + # get our list of actually reserved nodess + reserved_nodes = set(self.reserved_nodes) + + powered_down = get_powered_down(self.partition, self.slurm_commands) + currently_powered_up_nodes = len(get_powered_up_nodes(self.partition, self.slurm_commands)) + current_healthy_set = set(get_healthy_nodes(self.partition, self.slurm_commands)) + + if len(current_healthy_set) >= self.target_count: + log.info(f"Target count already reached. {len(current_healthy_set)} >= {self.target_count}") + return "" + + new_healthy_delta = node_count - len(current_healthy_set) + if new_healthy_delta <= 0: + log.warning(f"No need to create any nodes {node_count} {len(current_healthy_set)}") + return "" + + # we need a whole number of racks powered up (after considering powered and unhealthy) + new_powered_up_target = self.round_up_to_rack(currently_powered_up_nodes + new_healthy_delta) + power_up_delta = new_powered_up_target - currently_powered_up_nodes + + # we can only power up reserved and powered down nodes + reserved_and_powered_down = set(powered_down).intersection(reserved_nodes) + sorted_powered_down_nodes = sorted_nodes(reserved_and_powered_down) + + if len(sorted_powered_down_nodes) < power_up_delta: + raise SlurmM1Error(f"There are not enough nodes in a powered down state to fulfill this request - required {power_up_delta} < {len(sorted_powered_down_nodes)} {locals()}") + + assert power_up_delta > 0 + power_up_delta_nodes = sorted_powered_down_nodes[:power_up_delta] + + power_up_delta_nodes_str = self.slurm_commands.show_hostlist(power_up_delta_nodes) + cmd = f"scontrol update NodeName={power_up_delta_nodes_str} State=power_up" + + if os.getenv("DRYRUN"): + d = dict(locals()) + for k, v in d.items(): + print(f" {k} = {v if not isinstance(v, (set, list)) else len(v)}") + print(f"Would start {len(power_up_delta_nodes)} nodes via {cmd}") + sys.exit(0) + try: + log.info(cmd) + self.run_command(cmd) + log.info("Nodes powering up initiated successfully") + self.run_command(f"azslurm resume --node-list {power_up_delta_nodes_str}") + return power_up_delta_nodes_str + except subprocess.CalledProcessError as e: + log.error("Failed to power up nodes") + raise SlurmCommandError(cmd, e.returncode, e.output) + + def wait_for_nodes_to_start(self, node_str: str, timeout: int = 7200) -> None: + """Wait for reserved nodes to start. In test mode, simulate the wait.""" + self.wait_for_nodes(timeout, "start", node_str) + + def generate_topology(self, suffix: str = "") -> None: + """Generate topology configuration file using azslurm CLI.""" + log.info("Generating topology configuration using azslurm CLI...") + self.azslurm_topology.generate_topology(self.partition, self.topology_file + suffix) + + def get_ordered_blocks(self, suffix: str = "") -> List[Dict]: + """Get blocks ordered by number of healthy nodes. In test mode, return mock data.""" + try: + json_output = output_block_nodelist(self.topology_file + suffix, table=False) + blocks = json.loads(json_output) + log.info(f"Found {len(blocks)} blocks in topology in {self.topology_file + suffix}") + return blocks + except Exception as e: + log.error(f"Failed to parse topology blocks: {e}") + return [] + + def calculate_nodes_to_terminate(self, blocks: List[Dict]) -> Tuple[List[str], int]: + """Calculate which nodes need to be terminated to reach target count.""" + # Count total healthy nodes + total_nodes = sum(block['size'] for block in blocks) + nodes_to_remove = total_nodes - self.target_count + + if nodes_to_remove <= 0: + log.info(f"Already at or below target count ({total_nodes} <= {self.target_count})") + return [], total_nodes + + log.info(f"Need to remove {nodes_to_remove} nodes from {total_nodes} total") + + nodes_to_terminate = [] + remaining_to_remove = nodes_to_remove + + # Start with smallest blocks (already sorted by size) + for block in blocks: + if remaining_to_remove <= 0: + break + + block_size = block['size'] + reserved_within_block = [x for x in block['nodelist'] if x in self.reserved_nodes][ :remaining_to_remove] + + # Terminate entire block, as long as it is in the reservation + nodes_to_terminate.extend(reserved_within_block) + remaining_to_remove -= len(reserved_within_block) + assert remaining_to_remove >= 0 + if len(reserved_within_block) == block_size: + log.info(f"Terminating entire block {block['blockname']} ({block_size} nodes)") + else: + log.info(f"Terminating partial block {block['blockname']} ({len(reserved_within_block)} nodes)") + + + final_count = total_nodes - len(nodes_to_terminate) + log.info(f"Will terminate {len(nodes_to_terminate)} nodes, resulting in {final_count} nodes") + + return nodes_to_terminate, final_count + + def terminate_nodes(self, nodes_str: str) -> None: + """Terminate specified nodes using scontrol.""" + # Convert list to comma-separated string + cmd = f"scontrol update NodeName={nodes_str} State=POWER_DOWN" + + try: + self.run_command(cmd) + log.info("Successfully initiated node termination") + except subprocess.CalledProcessError as e: + raise SlurmCommandError(cmd, e.returncode, e.output) + + def wait_for_nodes_to_terminate(self, node_str: str, timeout: int = 1200) -> None: + """Wait for reserved nodes to terminate. In test mode, simulate the wait.""" + self.wait_for_nodes(timeout, "terminate", node_str) + + def wait_for_nodes(self, timeout: int, action: str, nodes: str) -> None: + """Wait for reserved nodes to start. In test mode, simulate the wait.""" + assert nodes + log.info(f"Waiting for nodes to {action} {nodes}...") + start_time = CLOCK.time() + filter_state = 'powering_up' if action == 'start' else 'powering_down' + terminal_states = ['POWERED_UP'] if action == 'start' else ['POWERED_DOWN', 'POWERING_DOWN'] + secondary_terminal_states = ['POWERED_DOWN', 'POWERING_DOWN'] if action == 'start' else [] + not_responsive_nodes = set() + terminal_state_nodes_for_logging = set() + secondary_state_nodes_for_logging = set() + powering_up_nodes = set() + while CLOCK.time() - start_time < timeout: + nodes_not_in_terminal_state = set() + state_counts: dict[str, int] = {} + try: + cmd = f"sinfo -n {nodes} -h -N -o '%N %T'" + result = self.run_command(cmd) + except subprocess.CalledProcessError as e: + log.warning(f"Error checking node status: {e}. Will retry") + continue + + lines = [x for x in result.stdout.splitlines() if x.strip()] + assert len(lines) >= 1 + for line in lines: + if not line.strip(): + continue + + node, state = line.split() + + power_state = "POWERED_UP" + if not state[-1].isalpha(): + if state[-1] == "~": + power_state = "POWERED_DOWN" + elif state[-1] == "%": + power_state = "POWERING_DOWN" + elif state[-1] == "#": + power_state = "POWERING_UP" + if action == "start": + powering_up_nodes.add(node) + elif state[-1] == "*": + if node not in not_responsive_nodes: + not_responsive_nodes.add(node) + log.warning(f"Node {node} is not responsive") + elif state[-1] == "!": + log.debug("Intermediate state found for node %s", node) + power_state = "INTERMEDIATE" + pass + else: + log.warning(f"Unknown power state symbol '{state[-1]}' for node {node}") + power_state = "UNKNOWN" + + state_counts[power_state] = state_counts.get(power_state, 0) + 1 + + if power_state in terminal_states: + terminal_state_nodes_for_logging.add(node) + elif power_state in secondary_terminal_states and node in powering_up_nodes: + if node not in secondary_state_nodes_for_logging: + secondary_state_nodes_for_logging.add(node) + log.warning(f"Node {node} has moved from state {terminal_states} to state {power_state}") + # This is considered a successful state + # Specifically a node that had powered up is now powering down/powered down + else: + nodes_not_in_terminal_state.add(node) + + if not nodes_not_in_terminal_state: + log.info(f"All nodes have finished {filter_state}") + return + + unexpected_states = terminal_state_nodes_for_logging.intersection(nodes_not_in_terminal_state) + if unexpected_states: + log.error(f"The following nodes {unexpected_states} previously reached their terminal state {terminal_states} but are no longer in that/those state(s).") + + cache_log(f"Current power states: {state_counts}") + CLOCK.sleep(5) + self.slurm_commands.update_states() # Update states in test mode + + log.error(f"Timeout waiting for nodes to {action} after {timeout} seconds") + raise SlurmM1Error(f"Timeout waiting for nodes to {action}") + + def reconfigure_slurm(self) -> None: + """Reconfigure SLURM to apply topology changes.""" + try: + log.info("Reconfiguring SLURM...") + self.run_command("scontrol reconfigure") + log.info("SLURM reconfiguration completed") + except subprocess.CalledProcessError as e: + raise SlurmCommandError("scontrol reconfigure", e.returncode, e.output) + + + def power_up(self) -> None: + """Execute the complete scaling workflow.""" + log.info(f"Starting cluster scaling:") + log.info(f" Partition: {self.partition}") + log.info(f" Target nodes: {self.target_count}") + log.info(f" Overprovision: {self.overprovision}") + # Step 0: Validate nodes + self.validate_nodes() + + # Step 1: Create reservation + target_plus_buffer = self.target_count + self.overprovision + + # Step 2: Wait for nodes to start + power_up_list = self.power_up_nodes(target_plus_buffer) + + if power_up_list: + self.wait_for_nodes_to_start(power_up_list) + + # Step 3: Check healthy nodes + healthy_nodes = get_healthy_nodes(self.partition, self.slurm_commands) + + if len(healthy_nodes) < self.target_count: + deficit = self.target_count - len(healthy_nodes) + new_overprovision = self.overprovision + deficit + log.error(f"Insufficient healthy nodes: {len(healthy_nodes)} < {self.target_count}") + log.error(f"Try again with overprovision: {new_overprovision}") + raise NodeInvalidStateError(f"Insufficient healthy nodes: {len(healthy_nodes)} < {self.target_count}") + + log.info(f"Sufficient healthy nodes: {len(healthy_nodes)} >= {self.target_count}") + + def prune(self) -> str: + """Returns a suggested hostlist of nodes to terminate""" + # Step 4: Generate initial topology + self.generate_topology(".pre-pruning") + + # Step 5: Get ordered blocks + blocks = self.get_ordered_blocks(".pre-pruning") + if not blocks: + log.error("No blocks found in topology") + raise SlurmM1Error("No blocks found in topology") + + # Step 6: Calculate and terminate excess nodes + nodes_to_terminate, final_count = self.calculate_nodes_to_terminate(blocks) + if not nodes_to_terminate: + return "" + log.info(f"Suggesting powering down {len(nodes_to_terminate)} nodes") + short_list = self.slurm_commands.show_hostlist(nodes_to_terminate) + + return short_list + + def prune_now(self) -> None: + """automates the pruning (termination), topology regeneration and reconfiguration of slurm""" + to_prune = self.prune() + if to_prune: + self.terminate_nodes(to_prune) + + # Wait for nodes to power down + log.info("Waiting for nodes to power down...") + self.wait_for_nodes_to_terminate(to_prune) + + # Step 7: Re-generate topology + self.generate_topology() + else: + log.info(f"Copying {self.topology_file + '.pre-pruning'} to {self.topology_file}") + shutil.copyfile(self.topology_file + ".pre-pruning", self.topology_file) + log.info("No nodes need to be terminated") + + # Step 8: Reconfigure SLURM + self.reconfigure_slurm() + + log.info(f"Successfully scaled cluster to {self.target_count} nodes!") + log.info(f"Run 'scontrol delete reservation {self.reservation_name}' to release the nodes to the cluster.") + + + +def get_powered_up_nodes(partition, slurm_commands) -> List[str]: + """Get list of healthy and idle nodes in the partition.""" + return get_nodes_excluding_states(partition, slurm_commands, + "powered_down,powering_up,powering_down,power_down") + + +def get_powered_down(partition, slurm_commands) -> List[str]: + sinfo_cmd = f'sinfo -p {partition} -o "%N" -h -N -t powered_down' + return sorted_nodes(slurm_commands.run_command(sinfo_cmd).stdout.strip().split('\n')) + + +def get_healthy_nodes(partition: str, slurm_commands: SlurmCommands) -> List[str]: + return get_nodes_excluding_states(partition, slurm_commands, + "powered_down,powering_up,powering_down,power_down,drain,drained,draining,unknown,down,no_respond,fail,reboot") + + +def get_nodes_excluding_states(partition: str, slurm_commands: SlurmCommands, exclude_states: str) -> List[str]: + """Get list of healthy and idle nodes in the partition.""" + partition_cmd = f'-p {partition} ' + sinfo_cmd = f'sinfo -p {partition} -o "%N" -h -N' + hosts = set(slurm_commands.run_command(sinfo_cmd).stdout.strip().split('\n')) + + sinfo_cmd = f'sinfo {partition_cmd}-t {exclude_states} -o "%N" -h -N' + excluded_host = set(slurm_commands.run_command(sinfo_cmd).stdout.strip().split('\n')) + log.debug(f"Excluding {len(excluded_host)} nodes with states {exclude_states}") + + nodes = sorted_nodes(list(hosts - excluded_host)) + log.debug(f"Found {len(nodes)} nodes after excluding states {exclude_states}") + return nodes + + +def get_reservable_nodes(partition: str, slurm_commands: SlurmCommands, filter_states: str = "reserved,maint,allocated,mixed"): + return get_nodes_excluding_states(partition, slurm_commands, filter_states) + + +def sorted_nodes(nodes: Union[list[str], set[str]]) -> list[str]: + nodes = list(nodes) + if len(nodes) == 1 and nodes[0].strip() == "": + return [] + return sorted(nodes, key=lambda x: int(x.split("-")[-1])) + + +def create_reservation(reservation_name: str, partition: str, slurm_commands: SlurmCommands, filter_states: str = "reserved,maint,allocated,mixed") -> bool: + """Create a SLURM reservation for the specified number of nodes.""" + slurm_commands.run_command(f"scontrol update partitionname={partition} state=DOWN") + log.info("Set partition %s DOWN and sleeping 15 seconds to ensure valid sinfo information for proper reservation creation", partition) + CLOCK.sleep(15) + reservable_nodes = get_reservable_nodes(partition, slurm_commands, filter_states) + if not reservable_nodes: + raise SlurmM1Error("There are no reservable nodes.") + reservation_nodes_short = slurm_commands.show_hostlist(reservable_nodes) + cmd = (f"scontrol create reservation ReservationName={reservation_name} " + f"PartitionName={partition} Nodes={reservation_nodes_short} " + f"StartTime=now Duration=infinite User=root") + try: + slurm_commands.run_command(cmd) + log.info(f"Successfully created reservation: {reservation_name}") + return True + except subprocess.CalledProcessError as e: + log.error(f"Failed to update reservation: {e}") + raise SlurmCommandError(cmd, e.returncode, e.output) + finally: + log.info("Setting partition %s UP", partition) + slurm_commands.run_command(f"scontrol update partitionname={partition} state=UP") + + +def parse_reservation(reservation_name: str, slurm_commands: SlurmCommands) -> Tuple[str, list[str]]: + cmd_nodes = f"scontrol show reservation {reservation_name}" + + result = slurm_commands.run_command(cmd_nodes, check=False) + if result.returncode != 0: + raise NoReservationError(reservation_name) + log.info(f"Reservation nodes: {result.stdout.strip()}") + + # Parse something like: "Nodes=ccw-gpu-[1-612]" + reserved_nodes_str = None + partition = None + assert "PartitionName" in result.stdout + for line in result.stdout.split(): + if line.startswith("Nodes="): + reserved_nodes_str = line.split("=", 1)[1] + elif line.startswith("PartitionName="): + partition = line.split("=", 1)[1] + if not reserved_nodes_str : + log.error("failed to get nodes from reservation") + raise SlurmCommandError(cmd_nodes, 1, f"No nodes found in reservation - {result.stdout}") + if not partition: + log.error("failed to get partition from reservation") + raise SlurmCommandError(cmd_nodes, 1, f"No partition found in reservation - {result.stdout}") + return partition, sorted_nodes(slurm_commands.show_hostnames(reserved_nodes_str)) + + +def add_target_args(sub: argparse.ArgumentParser) -> None: + """Add the mutually exclusive node-count vs rack-count target args. + + Exactly one of -n/--target-count or --racks must be supplied. + """ + group = sub.add_mutually_exclusive_group(required=True) + group.add_argument('-n', '--target-count', type=int, default=None, + help='Target number of nodes') + group.add_argument('--racks', type=int, default=None, + help=f'Target number of racks (1 rack = {NODES_PER_RACK} nodes)') + + +def resolve_target_count(args: argparse.Namespace) -> int: + """Resolve the effective target node count from --target-count or --racks. + + Exits with a non-zero status if the resolved count is not a positive int. + """ + if getattr(args, 'racks', None) is not None: + if args.racks <= 0: + log.error("Racks must be positive") + sys.exit(1) + return args.racks * NODES_PER_RACK + if args.target_count is None or args.target_count <= 0: + log.error("Target count must be positive") + sys.exit(1) + return args.target_count + + +def resolve_overprovision(args: argparse.Namespace) -> int: + """Resolve the effective overprovision node count for power_up. + + Supports --overprovision (nodes) or --overprovision-racks (racks). When + neither is supplied the overprovision buffer defaults to 0. + """ + if getattr(args, 'overprovision_racks', None) is not None: + if args.overprovision_racks < 0: + log.error("Overprovision racks must be non-negative") + sys.exit(1) + return args.overprovision_racks * NODES_PER_RACK + overprovision = getattr(args, 'overprovision', None) + if overprovision is None: + return 0 + if overprovision < 0: + log.error("Overprovision must be non-negative") + sys.exit(1) + return overprovision + + +def mock_available() -> bool: + """Return True if the test-only `mock` module can be imported. + + `mock.py` is a test fixture that ships only in the source tree, not in the + installed package, so the --mock option is hidden in production. + """ + import importlib.util + return importlib.util.find_spec("mock") is not None + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description='Scale SLURM cluster to exactly N nodes with topology optimization.', + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=''' +Examples: + ./scale_m1 create_reservation -p gpu + ./scale_m1 power_up -n 504 -b 108 + ./scale_m1 prune -n 504 > to_terminate.txt + ... + ./scale_m1 prune_now -n 504 + ''' + ) + + subparser = parser.add_subparsers() + create_res = subparser.add_parser("create_reservation") + create_res.add_argument('-p', '--partition', required=True, + help='SLURM partition name') + create_res.add_argument('--reservation', required=False, help="Optional: use an existing reservation", + default="scale_m1") + create_res.add_argument('--filter-states', required=False, + help="Comma-separated list of node states to exclude when finding reservable nodes (default: 'reserved,maint,allocated,mixed')", + default="reserved,maint,allocated,mixed") + create_res.add_argument('--verbose', '-v', action='store_true', + help='Enable verbose logging') + create_res.set_defaults(cmd="create_reservation") + + prune = subparser.add_parser("prune") + prune.set_defaults(cmd="prune", overprovision=0, mock_topology=False, timeout=1800) + add_target_args(prune) + prune.add_argument('--verbose', '-v', action='store_true', + help='Enable verbose logging') + prune.add_argument('--reservation', required=False, help="Optional: use an existing reservation", + default="scale_m1") + if mock_available(): + prune.add_argument('--mock', dest='mock_topology', action='store_true', + help='Use mock topology instead of querying the cluster (test only)') + + prune_now = subparser.add_parser("prune_now") + prune_now.set_defaults(cmd="prune_now", overprovision=0, mock_topology=False, timeout=1800) + add_target_args(prune_now) + prune_now.add_argument('--verbose', '-v', action='store_true', + help='Enable verbose logging') + prune_now.add_argument('--reservation', required=False, help="Optional: use an existing reservation", + default="scale_m1") + + power_up = subparser.add_parser("power_up") + power_up.set_defaults(cmd="power_up") + add_target_args(power_up) + overprovision_group = power_up.add_mutually_exclusive_group(required=False) + overprovision_group.add_argument('-b', '--overprovision', type=int, default=None, + help='Number of additional nodes to overprovision (default: 0)') + overprovision_group.add_argument('--overprovision-racks', type=int, default=None, + dest='overprovision_racks', + help=f'Number of additional racks to overprovision (1 rack = {NODES_PER_RACK} nodes)') + power_up.add_argument('--reservation', required=False, help="Optional: use an existing reservation", + default="scale_m1") + power_up.add_argument('--verbose', '-v', action='store_true', + help='Enable verbose logging') + power_up.add_argument('--mock-topology', '-m', action='store_true', + help='Enable mock topology for testing') + + return parser + + +def main() -> None: + parser = build_parser() + + args = parser.parse_args() + + setup_logging(args.verbose) + + slurm_commands = SlurmCommandsImpl() + + # ensure any down~ nodes are back to idle~ + slurm_commands.run_command("azslurm return_to_idle") + + if args.cmd == "create_reservation": + create_reservation(args.reservation, args.partition, slurm_commands, args.filter_states) + return + + # Resolve racks/nodes inputs into effective node counts. + args.target_count = resolve_target_count(args) + args.overprovision = resolve_overprovision(args) + + azslurm_topology: AzslurmTopology + if args.mock_topology: + from mock import MockAzslurmTopology + azslurm_topology = MockAzslurmTopology(slurm_commands) + else: + azslurm_topology = AzslurmTopologyImpl() + + try: + scaler = NodeScaler(args.target_count, args.overprovision, + slurm_commands, azslurm_topology, reservation_name=args.reservation) + if args.cmd == "power_up": + scaler.power_up() + elif args.cmd == "prune": + print(scaler.prune()) + elif args.cmd == "prune_now": + scaler.prune_now() + + except NoReservationError as e: + log.error(f"Please create the reservation {args.reservation} first by running `scale_m1 create_reservation --partition PARTITION` first") + sys.exit(1) + except SlurmM1Error as e: + log.error(f"Scaling failed: {e}") + sys.exit(1) + except KeyboardInterrupt: + log.info("Scaling interrupted by user") + sys.exit(130) + except Exception as e: + log.exception(f"Unexpected error during scaling: {e}") + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/scale_m1/scale_to_nodes_test.py b/scale_m1/scale_to_nodes_test.py new file mode 100644 index 00000000..622e68c7 --- /dev/null +++ b/scale_m1/scale_to_nodes_test.py @@ -0,0 +1,493 @@ +from scale_to_n_nodes import NodeScaler, NodeInvalidStateError +from mock import MockSlurmCommands, MockAzslurmTopology, MockClock +import scale_to_n_nodes +import os +import pytest + + +scale_to_n_nodes.CLOCK = MockClock() + + +@pytest.fixture(autouse=True) +def run_around_tests(): + # just create a topology that already existed + + with open("/tmp/topology.conf", "w") as fw: + fw.write("preexisting") + with open("/tmp/topology.conf.pre-pruning", "w") as fw: + fw.write("preexisting") + yield + + +def test_exit_early_if_powering_down(): + mock_slurm_commands = MockSlurmCommands("/tmp/topology.conf") + mock_slurm_commands.create_nodes(partition="gpu", count=30) + mock_slurm_commands.run_command("scontrol update NodeName=gpu-6 State=power_down") + mock_slurm_commands.update_states(50) + + scale_to_n_nodes.create_reservation("scale_m1", "gpu", mock_slurm_commands) + nodescaler = NodeScaler(target_count=10, overprovision=0, slurm_commands=mock_slurm_commands,azslurm_topology=None) + try: + nodescaler.power_up() + nodescaler.prune_now() + assert False, "Expected NodeInvalidStateError" + except NodeInvalidStateError as e: + assert e.message == "Some nodes are in POWERING_DOWN or POWERING_UP state, cannot proceed with scaling." + _post_test(mock_slurm_commands, + powered_up=0, + created_a_vm=0, + topology="preexisting") + + +def test_exit_early_if_powering_up(): + mock_slurm_commands = MockSlurmCommands("/tmp/topology.conf") + mock_slurm_commands.create_nodes(partition="gpu", count=30) + mock_slurm_commands._power_up(["gpu-6"]) + mock_slurm_commands.update_states() + + scale_to_n_nodes.create_reservation("scale_m1", "gpu", mock_slurm_commands) + nodescaler = NodeScaler(target_count=10, overprovision=0, slurm_commands=mock_slurm_commands,azslurm_topology=None) + try: + nodescaler.power_up() + nodescaler.prune_now() + assert False, "Expected NodeInvalidStateError" + except NodeInvalidStateError as e: + assert e.message == "Some nodes are in POWERING_DOWN or POWERING_UP state, cannot proceed with scaling." + _post_test(mock_slurm_commands, + powered_up=1, + created_a_vm=0, + topology="preexisting") + + +def _post_test(mock_slurm_commands: MockSlurmCommands, + powered_up: int, + created_a_vm: int, + topology: str, + topology_pre_pruning: str = "", + reserved: int = -1, + draining: int = 0, + allocated: int = 0, + down: int = 0) -> None: + mock_slurm_commands.update_states(1000) + topology_pre_pruning = topology_pre_pruning or topology + state_counts = mock_slurm_commands.node_state_counts() + total = state_counts.pop("total") + powered_down = total - powered_up + if reserved < 0: + reserved = total - allocated + orig_state_counts = str(state_counts) + idle = total - draining - allocated - down + assert state_counts.pop("draining") == draining, orig_state_counts + assert state_counts.pop("down") == down, orig_state_counts + assert state_counts.pop("allocated") == allocated, orig_state_counts + assert state_counts.pop("reserved") == reserved, orig_state_counts + assert state_counts.pop("idle") == idle, orig_state_counts + assert state_counts.pop("powered_up") == powered_up, orig_state_counts + assert state_counts.pop("powered_down") == powered_down, orig_state_counts + assert state_counts.pop("created_a_vm") == created_a_vm, orig_state_counts + + assert not state_counts, f"unchecked states! {state_counts}" + with open("/tmp/topology.conf", "r") as f: + actual_topology = f.read() + + if topology == "preexisting": + assert topology == actual_topology + else: + if not topology.startswith("# Mock topology for testing"): + topology = "# Mock topology for testing\n" + topology + if not topology.endswith("BlockSizes=1"): + topology += "\nBlockSizes=1" + assert actual_topology.strip() == topology.strip() + + if topology_pre_pruning: + with open("/tmp/topology.conf.pre-pruning", "r") as f: + actual_topology = f.read() + if topology_pre_pruning != "preexisting": + if not topology_pre_pruning.startswith("# Mock topology for testing"): + topology_pre_pruning = "# Mock topology for testing\n" + topology_pre_pruning + if not topology_pre_pruning.endswith("BlockSizes=1"): + topology_pre_pruning += "\nBlockSizes=1" + assert actual_topology.strip() == topology_pre_pruning.strip() + + +def test_basic_scaling_og(): + mock_slurm_commands = MockSlurmCommands("/tmp/topology.conf") + mock_slurm_commands.create_nodes(partition="gpu", count=18) + mock_azslurm_topology = MockAzslurmTopology(mock_slurm_commands) + + scale_to_n_nodes.create_reservation("scale_m1", "gpu", mock_slurm_commands) + nodescaler = NodeScaler(target_count=18, overprovision=0, topology_file="/tmp/topology.conf", slurm_commands=mock_slurm_commands, azslurm_topology=mock_azslurm_topology) + + # Run the scaling process + nodescaler.power_up() + nodescaler.prune_now() + + _post_test(mock_slurm_commands, + powered_up=18, + created_a_vm=18, + topology="""BlockName=block_001 Nodes=gpu-1,gpu-2,gpu-3,gpu-4,gpu-5,gpu-6,gpu-7,gpu-8,gpu-9,gpu-10,gpu-11,gpu-12,gpu-13,gpu-14,gpu-15,gpu-16,gpu-17,gpu-18""") + +def test_nothing_to_do_og(): + mock_slurm_commands = MockSlurmCommands("/tmp/topology.conf") + mock_slurm_commands.create_nodes(partition="gpu", count=18) + mock_azslurm_topology = MockAzslurmTopology(mock_slurm_commands) + + scale_to_n_nodes.create_reservation("scale_m1", "gpu", mock_slurm_commands) + nodescaler = NodeScaler(target_count=18, overprovision=0, topology_file="/tmp/topology.conf", slurm_commands=mock_slurm_commands, azslurm_topology=mock_azslurm_topology) + for n in range(18): + mock_slurm_commands._power_up([f"gpu-{n+1}"]) + mock_slurm_commands.update_states(100) + mock_slurm_commands.update_states(100) + assert mock_slurm_commands.node_state_counts()["powered_up"] == 18 + # Run the scaling process + nodescaler.power_up() + nodescaler.prune_now() + + # there is no need to even create a reservation: all idle + _post_test(mock_slurm_commands, + powered_up=18, + created_a_vm=0, + topology="""# Mock topology for testing +BlockName=block_001 Nodes=gpu-1,gpu-2,gpu-3,gpu-4,gpu-5,gpu-6,gpu-7,gpu-8,gpu-9,gpu-10,gpu-11,gpu-12,gpu-13,gpu-14,gpu-15,gpu-16,gpu-17,gpu-18 +BlockSizes=1""") + + +def test_basic_scaling_some_existing(): + mock_slurm_commands = MockSlurmCommands("/tmp/topology.conf") + mock_slurm_commands.create_nodes(partition="gpu", count=18) + mock_azslurm_topology = MockAzslurmTopology(mock_slurm_commands) + + scale_to_n_nodes.create_reservation("scale_m1", "gpu", mock_slurm_commands) + nodescaler = NodeScaler(target_count=18, overprovision=0, topology_file="/tmp/topology.conf", slurm_commands=mock_slurm_commands, azslurm_topology=mock_azslurm_topology) + mock_slurm_commands._power_up(["gpu-1"]) + mock_slurm_commands.update_states() + mock_slurm_commands.update_states() + assert mock_slurm_commands.nodes_dict["gpu-1"]["power_state"] == "POWERED_UP" + assert scale_to_n_nodes.get_healthy_nodes("gpu", mock_slurm_commands) == ["gpu-1"] + # Run the scaling process + nodescaler.power_up() + nodescaler.prune_now() + _post_test(mock_slurm_commands, + powered_up=18, + created_a_vm=17, + topology="BlockName=block_001 Nodes=gpu-1,gpu-2,gpu-3,gpu-4,gpu-5,gpu-6,gpu-7,gpu-8,gpu-9,gpu-10,gpu-11,gpu-12,gpu-13,gpu-14,gpu-15,gpu-16,gpu-17,gpu-18") + + +def test_basic_scaling_running_jobs(): + mock_slurm_commands = MockSlurmCommands("/tmp/topology.conf") + mock_slurm_commands.create_nodes(partition="gpu", count=18) + mock_azslurm_topology = MockAzslurmTopology(mock_slurm_commands) + mock_slurm_commands._power_up(["gpu-1"]) + mock_slurm_commands.update_states() + mock_slurm_commands.update_states() + mock_slurm_commands.alloc_nodes(["gpu-1"]) + + scale_to_n_nodes.create_reservation("scale_m1", "gpu", mock_slurm_commands) + nodescaler = NodeScaler(target_count=18, overprovision=0, topology_file="/tmp/topology.conf", slurm_commands=mock_slurm_commands, azslurm_topology=mock_azslurm_topology) + + # Run the scaling process + nodescaler.power_up() + nodescaler.prune_now() + _post_test(mock_slurm_commands, + powered_up=18, + created_a_vm=17, + allocated=1, + topology="BlockName=block_001 Nodes=gpu-1,gpu-2,gpu-3,gpu-4,gpu-5,gpu-6,gpu-7,gpu-8,gpu-9,gpu-10,gpu-11,gpu-12,gpu-13,gpu-14,gpu-15,gpu-16,gpu-17,gpu-18") + + +def test_basic_scaling_drained_fail(): + """ + power up a node manually, and leave it in a draining state. + This forces the allocator to try and allocate 36 nodes instead of 18, as one that is powered up is "unhealthy" + i.e. not in idle,allocated,or mixed. + However - I only created 18 nodes, so we should just fail with a "not enough nodes to allocate" message. + """ + mock_slurm_commands = MockSlurmCommands("/tmp/topology.conf") + mock_slurm_commands.create_nodes(partition="gpu", count=18) + mock_azslurm_topology = MockAzslurmTopology(mock_slurm_commands) + mock_slurm_commands._power_up(["gpu-1"]) + mock_slurm_commands.update_states() + mock_slurm_commands.update_states() + mock_slurm_commands.drain_nodes(["gpu-1"]) + + scale_to_n_nodes.create_reservation("scale_m1", "gpu", mock_slurm_commands) + nodescaler = NodeScaler(target_count=18, overprovision=0, topology_file="/tmp/topology.conf", slurm_commands=mock_slurm_commands, azslurm_topology=mock_azslurm_topology) + + # mock_slurm_commands.run_command("") + # Run the scaling process + try: + nodescaler.power_up() + nodescaler.prune_now() + assert False + except scale_to_n_nodes.SlurmM1Error as e: + assert "There are not enough nodes in a powered down state" in str(e) + _post_test(mock_slurm_commands, + powered_up=1, + draining=1, + created_a_vm=0, + topology="preexisting") + + +def test_basic_scaling_drained_success(): + """ + Identical as above except we pre-create 36 nodes so it can succeed. + """ + mock_slurm_commands = MockSlurmCommands("/tmp/topology.conf") + mock_slurm_commands.create_nodes(partition="gpu", count=36) + mock_azslurm_topology = MockAzslurmTopology(mock_slurm_commands) + mock_slurm_commands._power_up(["gpu-1"]) + mock_slurm_commands.update_states() + mock_slurm_commands.update_states() + mock_slurm_commands.drain_nodes(["gpu-1"]) + + scale_to_n_nodes.create_reservation("scale_m1", "gpu", mock_slurm_commands) + nodescaler = NodeScaler(target_count=18, overprovision=0, topology_file="/tmp/topology.conf", slurm_commands=mock_slurm_commands, azslurm_topology=mock_azslurm_topology) + + # Run the scaling process + + nodescaler.power_up() + nodescaler.prune_now() + + _post_test(mock_slurm_commands, + powered_up=19, + created_a_vm=35, # need 19, which means we need 36 - but less 1 drained + draining=1, + topology="""BlockName=block_002 Nodes=gpu-19,gpu-20,gpu-21,gpu-22,gpu-23,gpu-24,gpu-25,gpu-26,gpu-27,gpu-28,gpu-29,gpu-30,gpu-31,gpu-32,gpu-33,gpu-34,gpu-35,gpu-36 +BlockSizes=1""", + topology_pre_pruning="""# Mock topology for testing +BlockName=block_001 Nodes=gpu-2,gpu-3,gpu-4,gpu-5,gpu-6,gpu-7,gpu-8,gpu-9,gpu-10,gpu-11,gpu-12,gpu-13,gpu-14,gpu-15,gpu-16,gpu-17,gpu-18 +BlockName=block_002 Nodes=gpu-19,gpu-20,gpu-21,gpu-22,gpu-23,gpu-24,gpu-25,gpu-26,gpu-27,gpu-28,gpu-29,gpu-30,gpu-31,gpu-32,gpu-33,gpu-34,gpu-35,gpu-36 +BlockSizes=1""") + + +def test_reproduce_scale_to_19(): + """ + Identical as above except we pre-create 36 nodes so it can succeed. + 10 drained + 4 allocated + 14 idle + goal -> 36 in powered up. + i.e. only starting 8 + """ + mock_slurm_commands = MockSlurmCommands("/tmp/topology.conf") + mock_slurm_commands.create_nodes(partition="gpu", count=50) + mock_azslurm_topology = MockAzslurmTopology(mock_slurm_commands) + for i in range(28): + mock_slurm_commands._power_up([f"gpu-{i+1}"]) + mock_slurm_commands.update_states(50) + mock_slurm_commands.update_states(50) + + for i in range(10): + mock_slurm_commands.drain_nodes([f"gpu-{i+1}"]) + + mock_slurm_commands.alloc_nodes([f"gpu-11", "gpu-12", "gpu-13", "gpu-14"]) + + assert 18 == len(scale_to_n_nodes.get_healthy_nodes("gpu", mock_slurm_commands)) + assert mock_slurm_commands.nodes_dict["gpu-20"]["state"] == "IDLE" + assert mock_slurm_commands.nodes_dict["gpu-20"]["power_state"] == "POWERED_UP" + + scale_to_n_nodes.create_reservation("scale_m1", "gpu", mock_slurm_commands) + nodescaler = NodeScaler(target_count=19, overprovision=0, topology_file="/tmp/topology.conf", slurm_commands=mock_slurm_commands, azslurm_topology=mock_azslurm_topology) + + # Run the scaling process + + nodescaler.power_up() + nodescaler.prune_now() + + _post_test(mock_slurm_commands, + powered_up=19 + 10, # targe=19 + 10 draining + created_a_vm=8, + draining=10, + allocated=4, + topology="""# Mock topology for testing +BlockName=block_001 Nodes=gpu-11,gpu-12,gpu-13,gpu-14 +BlockName=block_002 Nodes=gpu-22,gpu-23,gpu-24,gpu-25,gpu-26,gpu-27,gpu-28,gpu-29,gpu-30,gpu-31,gpu-32,gpu-33,gpu-34,gpu-35,gpu-36 +BlockSizes=1""", + topology_pre_pruning="""# Mock topology for testing +BlockName=block_001 Nodes=gpu-11,gpu-12,gpu-13,gpu-14,gpu-15,gpu-16,gpu-17,gpu-18 +BlockName=block_002 Nodes=gpu-19,gpu-20,gpu-21,gpu-22,gpu-23,gpu-24,gpu-25,gpu-26,gpu-27,gpu-28,gpu-29,gpu-30,gpu-31,gpu-32,gpu-33,gpu-34,gpu-35,gpu-36 +BlockSizes=1""") + + +def test_scale_not_enough_healthy_nodes(): + mock_slurm_commands = MockSlurmCommands("/tmp/topology.conf") + mock_slurm_commands.create_nodes(partition="gpu", count=18) + mock_slurm_commands.simulate_failed_converge(["gpu-1", "gpu-2", "gpu-3"]) # Simulate some nodes as unhealthy + + scale_to_n_nodes.create_reservation("scale_m1", "gpu", mock_slurm_commands) + nodescaler = NodeScaler(target_count=17, overprovision=0, slurm_commands=mock_slurm_commands, azslurm_topology="preexisting") + + try: + nodescaler.power_up() + nodescaler.prune_now() + assert False, "Expected NodeInvalidStateError" + except NodeInvalidStateError as e: + assert e.message == "Insufficient healthy nodes: 15 < 17" + _post_test(mock_slurm_commands, + powered_up=15, + created_a_vm=18, + topology="preexisting") + + +def test_scale_not_enough_healthy_nodes_fail_all(): + mock_slurm_commands = MockSlurmCommands("/tmp/topology.conf") + mock_slurm_commands.create_nodes(partition="gpu", count=18) + node_list = [f"gpu-{i}" for i in range(1, 19)] + mock_slurm_commands.simulate_failed_converge(node_list) # Simulate some nodes as unhealthy + + scale_to_n_nodes.create_reservation("scale_m1", "gpu", mock_slurm_commands) + nodescaler = NodeScaler(target_count=17, overprovision=0, slurm_commands=mock_slurm_commands, azslurm_topology="preexisting") + + try: + nodescaler.power_up() + nodescaler.prune_now() + assert False, "Expected NodeInvalidStateError" + except NodeInvalidStateError as e: + assert e.message == "Insufficient healthy nodes: 0 < 17" + + _post_test(mock_slurm_commands, + powered_up=0, + created_a_vm=18, + topology="preexisting") + + +def test_successful_over_alloc(): + mock_slurm_commands = MockSlurmCommands("/tmp/topology.conf") + mock_slurm_commands.create_nodes(partition="gpu", count=40) + node_list = [f"gpu-{i}" for i in range(1, 6)] + [f"gpu-{i}" for i in range(20, 26)] + mock_slurm_commands.simulate_failed_converge(node_list) # Simulate some nodes as unhealthy + mock_azslurm_topology = MockAzslurmTopology(mock_slurm_commands) + + scale_to_n_nodes.create_reservation("scale_m1", "gpu", mock_slurm_commands) + nodescaler = NodeScaler(target_count=17, overprovision=18, topology_file="/tmp/topology.conf", slurm_commands=mock_slurm_commands, azslurm_topology=mock_azslurm_topology) + + nodescaler.power_up() + nodescaler.prune_now() + + _post_test(mock_slurm_commands, + powered_up=17, + created_a_vm=36, + topology="""# Mock topology for testing +BlockName=block_001 Nodes=gpu-6,gpu-7,gpu-8,gpu-9,gpu-10,gpu-11,gpu-12,gpu-13,gpu-14,gpu-15,gpu-16,gpu-17,gpu-18 +BlockName=block_002 Nodes=gpu-33,gpu-34,gpu-35,gpu-36 +BlockSizes=1""", + topology_pre_pruning="""# Mock topology for testing +BlockName=block_001 Nodes=gpu-6,gpu-7,gpu-8,gpu-9,gpu-10,gpu-11,gpu-12,gpu-13,gpu-14,gpu-15,gpu-16,gpu-17,gpu-18 +BlockName=block_002 Nodes=gpu-19,gpu-26,gpu-27,gpu-28,gpu-29,gpu-30,gpu-31,gpu-32,gpu-33,gpu-34,gpu-35,gpu-36 +BlockSizes=1""") + + +def test_termination_with_running_jobs(): + mock_slurm_commands = MockSlurmCommands("/tmp/topology.conf") + mock_slurm_commands.create_nodes(partition="gpu", count=19) + for n in range(18): + mock_slurm_commands._power_up([f"gpu-{n + 1}"]) + mock_slurm_commands.update_states(50) + mock_slurm_commands.update_states(50) + + for n in range(18): + mock_slurm_commands.alloc_nodes([f"gpu-{n + 1}"]) + + mock_azslurm_topology = MockAzslurmTopology(mock_slurm_commands) + + scale_to_n_nodes.create_reservation("scale_m1", "gpu", mock_slurm_commands) + nodescaler = NodeScaler(target_count=1, overprovision=0, topology_file="/tmp/topology.conf", slurm_commands=mock_slurm_commands, azslurm_topology=mock_azslurm_topology) + nodescaler.power_up() + nodescaler.prune_now() + + _post_test(mock_slurm_commands, + powered_up=18, + created_a_vm=0, + allocated=18, + reserved=1, + topology="""# Mock topology for testing +BlockName=block_001 Nodes=gpu-1,gpu-2,gpu-3,gpu-4,gpu-5,gpu-6,gpu-7,gpu-8,gpu-9,gpu-10,gpu-11,gpu-12,gpu-13,gpu-14,gpu-15,gpu-16,gpu-17,gpu-18 +BlockSizes=1""") + + +def test_basic_scaling_large_delete(): + mock_slurm_commands = MockSlurmCommands("/tmp/topology.conf") + mock_slurm_commands.create_nodes(partition="gpu", count=54) + mock_azslurm_topology = MockAzslurmTopology(mock_slurm_commands) + + scale_to_n_nodes.create_reservation("scale_m1", "gpu", mock_slurm_commands) + nodescaler = NodeScaler(target_count=18, overprovision=36, topology_file="/tmp/topology.conf", slurm_commands=mock_slurm_commands, azslurm_topology=mock_azslurm_topology) + + # Run the scaling process + nodescaler.power_up() + nodescaler.prune_now() + + _post_test(mock_slurm_commands, + powered_up=18, + created_a_vm=54, + topology="""# Mock topology for testing +BlockName=block_003 Nodes=gpu-37,gpu-38,gpu-39,gpu-40,gpu-41,gpu-42,gpu-43,gpu-44,gpu-45,gpu-46,gpu-47,gpu-48,gpu-49,gpu-50,gpu-51,gpu-52,gpu-53,gpu-54 +BlockSizes=1""", + topology_pre_pruning="""# Mock topology for testing +BlockName=block_001 Nodes=gpu-1,gpu-2,gpu-3,gpu-4,gpu-5,gpu-6,gpu-7,gpu-8,gpu-9,gpu-10,gpu-11,gpu-12,gpu-13,gpu-14,gpu-15,gpu-16,gpu-17,gpu-18 +BlockName=block_002 Nodes=gpu-19,gpu-20,gpu-21,gpu-22,gpu-23,gpu-24,gpu-25,gpu-26,gpu-27,gpu-28,gpu-29,gpu-30,gpu-31,gpu-32,gpu-33,gpu-34,gpu-35,gpu-36 +BlockName=block_003 Nodes=gpu-37,gpu-38,gpu-39,gpu-40,gpu-41,gpu-42,gpu-43,gpu-44,gpu-45,gpu-46,gpu-47,gpu-48,gpu-49,gpu-50,gpu-51,gpu-52,gpu-53,gpu-54 +BlockSizes=1""") + + +def test_basic_scaling_with_reserved_noop(): + """ + leave gpu-[19-36] as still reserved - it should be a no-op + """ + mock_slurm_commands = MockSlurmCommands("/tmp/topology.conf") + mock_slurm_commands.create_nodes(partition="gpu", count=54) + mock_azslurm_topology = MockAzslurmTopology(mock_slurm_commands) + for n in range(18, 36): + mock_slurm_commands._power_up([f"gpu-{n+1}"]) + mock_slurm_commands.update_states(30) + mock_slurm_commands.update_states(30) + for n in range(18, 36): + mock_slurm_commands.reserve_nodes([f"gpu-{n+1}"]) + + scale_to_n_nodes.create_reservation("scale_m1", "gpu", mock_slurm_commands) + nodescaler = NodeScaler(target_count=18, overprovision=36, topology_file="/tmp/topology.conf", slurm_commands=mock_slurm_commands, azslurm_topology=mock_azslurm_topology) + + # Run the scaling process + nodescaler.power_up() + nodescaler.prune_now() + + _post_test(mock_slurm_commands, + powered_up=18, + created_a_vm=0, # 18 were already started, it is a noop + topology="""# Mock topology for testing +BlockName=block_002 Nodes=gpu-19,gpu-20,gpu-21,gpu-22,gpu-23,gpu-24,gpu-25,gpu-26,gpu-27,gpu-28,gpu-29,gpu-30,gpu-31,gpu-32,gpu-33,gpu-34,gpu-35,gpu-36 +BlockSizes=1""") + + +def test_basic_scaling_with_reserved_og(): + """ + leave gpu-[19-30] as still reserved - it should include 19-30 and others + """ + mock_slurm_commands = MockSlurmCommands("/tmp/topology.conf") + mock_slurm_commands.create_nodes(partition="gpu", count=54) + mock_azslurm_topology = MockAzslurmTopology(mock_slurm_commands) + for n in range(18, 30): + mock_slurm_commands._power_up([f"gpu-{n+1}"]) + mock_slurm_commands.update_states(30) + mock_slurm_commands.update_states(30) + for n in range(18, 30): + mock_slurm_commands.reserve_nodes([f"gpu-{n+1}"]) + + scale_to_n_nodes.create_reservation("scale_m1", "gpu", mock_slurm_commands) + nodescaler = NodeScaler(target_count=18, overprovision=36, topology_file="/tmp/topology.conf", slurm_commands=mock_slurm_commands, azslurm_topology=mock_azslurm_topology) + + # Run the scaling process + nodescaler.power_up() + nodescaler.prune_now() + _post_test(mock_slurm_commands, + powered_up=18, + created_a_vm=42, # 12 were already started + topology="""# Mock topology for testing +BlockName=block_002 Nodes=gpu-19,gpu-20,gpu-21,gpu-22,gpu-23,gpu-24,gpu-25,gpu-26,gpu-27,gpu-28,gpu-29,gpu-30 +BlockName=block_003 Nodes=gpu-49,gpu-50,gpu-51,gpu-52,gpu-53,gpu-54 +BlockSizes=1""", + topology_pre_pruning="""# Mock topology for testing +BlockName=block_001 Nodes=gpu-1,gpu-2,gpu-3,gpu-4,gpu-5,gpu-6,gpu-7,gpu-8,gpu-9,gpu-10,gpu-11,gpu-12,gpu-13,gpu-14,gpu-15,gpu-16,gpu-17,gpu-18 +BlockName=block_002 Nodes=gpu-19,gpu-20,gpu-21,gpu-22,gpu-23,gpu-24,gpu-25,gpu-26,gpu-27,gpu-28,gpu-29,gpu-30,gpu-31,gpu-32,gpu-33,gpu-34,gpu-35,gpu-36 +BlockName=block_003 Nodes=gpu-37,gpu-38,gpu-39,gpu-40,gpu-41,gpu-42,gpu-43,gpu-44,gpu-45,gpu-46,gpu-47,gpu-48,gpu-49,gpu-50,gpu-51,gpu-52,gpu-53,gpu-54 +BlockSizes=1""") \ No newline at end of file