scale_m1: add --racks and update docs

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/<clustername>/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 <nodes>` — target an exact node count.
+- `--racks <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 <nodes>` — overprovision an exact number of extra nodes.
+- `--overprovision-racks <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.
 ```

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.

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__":

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