From fd835b37ed4e5b8400b5d221820df68e308bb56a Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Mon, 3 Aug 2026 15:59:06 -0500 Subject: [PATCH] feat(ci): schedule Slurm jobs with weighted runner leases Derive each benchmark allocation size from explicit aggregate node counts, recipe resources, or topology; expose it as nodes:N queue metadata; and use node count as a secondary priority signal.\n\nIntegrate with dynamic aggregate-capacity leases from the priority controller while keeping the feature disabled by default. Historical matrix rows without node-count fall back to legacy one-node accounting instead of requesting an unsatisfiable nodes: label. --- .../workflows/benchmark-multinode-tmpl.yml | 49 +++-- .github/workflows/benchmark-tmpl.yml | 45 +++-- .github/workflows/e2e-tests.yml | 3 + .github/workflows/run-sweep.yml | 3 + configs/ci-priority.yaml | 3 + utils/ci_priority.py | 10 + utils/matrix_logic/generate_sweep_configs.py | 178 +++++++++++++++++- .../test_generate_sweep_configs.py | 98 +++++++++- utils/matrix_logic/test_validation.py | 45 +++++ utils/matrix_logic/validation.py | 47 ++++- utils/runner_setup/RUNNER_SETUP.md | 47 +++++ utils/test_ci_priority.py | 27 +++ 12 files changed, 516 insertions(+), 39 deletions(-) diff --git a/.github/workflows/benchmark-multinode-tmpl.yml b/.github/workflows/benchmark-multinode-tmpl.yml index 10411ad069..8a46717eb1 100644 --- a/.github/workflows/benchmark-multinode-tmpl.yml +++ b/.github/workflows/benchmark-multinode-tmpl.yml @@ -6,6 +6,10 @@ on: runner: required: true type: string + node-count: + description: "Total Slurm nodes required by this benchmark" + required: true + type: string priority: description: "Higher-is-sooner CI priority score" required: true @@ -254,19 +258,40 @@ jobs: ${{ fromJSON( vars.PRIORITY_SCHEDULER_ENABLED == 'true' && ( - inputs.skip-queue-pr != '' && - format( - '["self-hosted",{0},{1},{2},{3}]', - toJSON(inputs.runner), - toJSON(format('ci-job-{0}-{1}', inputs.priority, inputs.queue-token)), - toJSON(format('ci-attempt-{0}', github.run_attempt)), - toJSON(format('ci-skip-queue-pr-{0}', inputs.skip-queue-pr)) + vars.NODE_SLOT_SCHEDULER_ENABLED == 'true' && inputs.node-count != '' && + ( + inputs.skip-queue-pr != '' && + format( + '["self-hosted",{0},{1},{2},{3},{4}]', + toJSON(inputs.runner), + toJSON(format('nodes:{0}', inputs.node-count)), + toJSON(format('ci-job-{0}-{1}', inputs.priority, inputs.queue-token)), + toJSON(format('ci-attempt-{0}', github.run_attempt)), + toJSON(format('ci-skip-queue-pr-{0}', inputs.skip-queue-pr)) + ) || + format( + '["self-hosted",{0},{1},{2},{3}]', + toJSON(inputs.runner), + toJSON(format('nodes:{0}', inputs.node-count)), + toJSON(format('ci-job-{0}-{1}', inputs.priority, inputs.queue-token)), + toJSON(format('ci-attempt-{0}', github.run_attempt)) + ) ) || - format( - '["self-hosted",{0},{1},{2}]', - toJSON(inputs.runner), - toJSON(format('ci-job-{0}-{1}', inputs.priority, inputs.queue-token)), - toJSON(format('ci-attempt-{0}', github.run_attempt)) + ( + inputs.skip-queue-pr != '' && + format( + '["self-hosted",{0},{1},{2},{3}]', + toJSON(inputs.runner), + toJSON(format('ci-job-{0}-{1}', inputs.priority, inputs.queue-token)), + toJSON(format('ci-attempt-{0}', github.run_attempt)), + toJSON(format('ci-skip-queue-pr-{0}', inputs.skip-queue-pr)) + ) || + format( + '["self-hosted",{0},{1},{2}]', + toJSON(inputs.runner), + toJSON(format('ci-job-{0}-{1}', inputs.priority, inputs.queue-token)), + toJSON(format('ci-attempt-{0}', github.run_attempt)) + ) ) ) || format('[{0}]', toJSON(inputs.runner)) diff --git a/.github/workflows/benchmark-tmpl.yml b/.github/workflows/benchmark-tmpl.yml index 46a80c8743..3b2f9349a9 100644 --- a/.github/workflows/benchmark-tmpl.yml +++ b/.github/workflows/benchmark-tmpl.yml @@ -202,19 +202,40 @@ jobs: ${{ fromJSON( vars.PRIORITY_SCHEDULER_ENABLED == 'true' && ( - inputs.skip-queue-pr != '' && - format( - '["self-hosted",{0},{1},{2},{3}]', - toJSON(inputs.runner), - toJSON(format('ci-job-{0}-{1}', inputs.priority, inputs.queue-token)), - toJSON(format('ci-attempt-{0}', github.run_attempt)), - toJSON(format('ci-skip-queue-pr-{0}', inputs.skip-queue-pr)) + vars.NODE_SLOT_SCHEDULER_ENABLED == 'true' && + ( + inputs.skip-queue-pr != '' && + format( + '["self-hosted",{0},{1},{2},{3},{4}]', + toJSON(inputs.runner), + toJSON('nodes:1'), + toJSON(format('ci-job-{0}-{1}', inputs.priority, inputs.queue-token)), + toJSON(format('ci-attempt-{0}', github.run_attempt)), + toJSON(format('ci-skip-queue-pr-{0}', inputs.skip-queue-pr)) + ) || + format( + '["self-hosted",{0},{1},{2},{3}]', + toJSON(inputs.runner), + toJSON('nodes:1'), + toJSON(format('ci-job-{0}-{1}', inputs.priority, inputs.queue-token)), + toJSON(format('ci-attempt-{0}', github.run_attempt)) + ) ) || - format( - '["self-hosted",{0},{1},{2}]', - toJSON(inputs.runner), - toJSON(format('ci-job-{0}-{1}', inputs.priority, inputs.queue-token)), - toJSON(format('ci-attempt-{0}', github.run_attempt)) + ( + inputs.skip-queue-pr != '' && + format( + '["self-hosted",{0},{1},{2},{3}]', + toJSON(inputs.runner), + toJSON(format('ci-job-{0}-{1}', inputs.priority, inputs.queue-token)), + toJSON(format('ci-attempt-{0}', github.run_attempt)), + toJSON(format('ci-skip-queue-pr-{0}', inputs.skip-queue-pr)) + ) || + format( + '["self-hosted",{0},{1},{2}]', + toJSON(inputs.runner), + toJSON(format('ci-job-{0}-{1}', inputs.priority, inputs.queue-token)), + toJSON(format('ci-attempt-{0}', github.run_attempt)) + ) ) ) || format('[{0}]', toJSON(inputs.runner)) diff --git a/.github/workflows/e2e-tests.yml b/.github/workflows/e2e-tests.yml index 06604b2f77..bed0f9a94d 100644 --- a/.github/workflows/e2e-tests.yml +++ b/.github/workflows/e2e-tests.yml @@ -266,6 +266,7 @@ jobs: osl: ${{ matrix.config.osl }} max-model-len: ${{ matrix.config.max-model-len }} runner: ${{ matrix.config.runner }} + node-count: ${{ matrix.config.node-count }} priority: ${{ matrix.config.priority }} queue-token: ${{ matrix.config['queue-token'] }} image: ${{ matrix.config.image }} @@ -317,6 +318,7 @@ jobs: osl: ${{ matrix.config.osl }} max-model-len: ${{ matrix.config.max-model-len }} runner: ${{ matrix.config.runner }} + node-count: ${{ matrix.config.node-count }} priority: ${{ matrix.config.priority }} queue-token: ${{ matrix.config['queue-token'] }} image: ${{ matrix.config.image }} @@ -456,6 +458,7 @@ jobs: osl: '0' max-model-len: '0' runner: ${{ matrix.config.runner }} + node-count: ${{ matrix.config.node-count }} priority: ${{ matrix.config.priority }} queue-token: ${{ matrix.config['queue-token'] }} image: ${{ matrix.config.image }} diff --git a/.github/workflows/run-sweep.yml b/.github/workflows/run-sweep.yml index b46c4c5901..6b5cfa43b5 100644 --- a/.github/workflows/run-sweep.yml +++ b/.github/workflows/run-sweep.yml @@ -458,6 +458,7 @@ jobs: osl: ${{ matrix.config.osl }} max-model-len: ${{ matrix.config.max-model-len }} runner: ${{ matrix.config.runner }} + node-count: ${{ matrix.config.node-count }} priority: ${{ matrix.config.priority }} queue-token: ${{ matrix.config['queue-token'] }} skip-queue-pr: ${{ matrix.config['skip-queue-pr'] || '' }} @@ -652,6 +653,7 @@ jobs: osl: '0' max-model-len: '0' runner: ${{ matrix.config.runner }} + node-count: ${{ matrix.config.node-count }} priority: ${{ matrix.config.priority }} queue-token: ${{ matrix.config['queue-token'] }} skip-queue-pr: ${{ matrix.config['skip-queue-pr'] || '' }} @@ -818,6 +820,7 @@ jobs: osl: ${{ matrix.config.osl }} max-model-len: ${{ matrix.config.max-model-len }} runner: ${{ matrix.config.runner }} + node-count: ${{ matrix.config.node-count }} priority: ${{ matrix.config.priority }} queue-token: ${{ matrix.config['queue-token'] }} skip-queue-pr: ${{ matrix.config['skip-queue-pr'] || '' }} diff --git a/configs/ci-priority.yaml b/configs/ci-priority.yaml index 34565e2810..83a21a8117 100644 --- a/configs/ci-priority.yaml +++ b/configs/ci-priority.yaml @@ -6,6 +6,9 @@ version: 1 base-score: 1.0 adjustments: + # Preserve business priority as the primary key, then prefer smaller Slurm + # allocations within an otherwise identical class. + additional-node: -0.001 event: push: 2.0 multi-node: 1.25 diff --git a/utils/ci_priority.py b/utils/ci_priority.py index 89b51d3ed5..540d36e759 100755 --- a/utils/ci_priority.py +++ b/utils/ci_priority.py @@ -113,6 +113,7 @@ def _entry_from_criteria( ) else "" ), + "node-count": entry.get("node-count", 1), } @@ -141,6 +142,15 @@ def calculate_priority( score = _decimal(policy["base-score"]) score += _decimal(adjustments.get("event", {}).get(context.event_name, 0)) + node_count = entry.get("node-count", 1) + if ( + not isinstance(node_count, int) + or isinstance(node_count, bool) + or node_count < 1 + ): + raise ValueError(f"node-count must be a positive integer, got {node_count!r}") + score += _decimal(adjustments.get("additional-node", 0)) * (node_count - 1) + if entry.get("prefill") is not None: score += _decimal(adjustments.get("multi-node", 0)) if entry.get("scenario-type") == "agentic-coding": diff --git a/utils/matrix_logic/generate_sweep_configs.py b/utils/matrix_logic/generate_sweep_configs.py index 86f192b1c8..bd32edb7b2 100644 --- a/utils/matrix_logic/generate_sweep_configs.py +++ b/utils/matrix_logic/generate_sweep_configs.py @@ -1,20 +1,24 @@ +import argparse import fnmatch import json -import argparse +import math +import re import sys from decimal import Decimal from pathlib import Path +import yaml + # Ensure sibling modules are importable regardless of how script is invoked sys.path.insert(0, str(Path(__file__).resolve().parent)) from validation import ( - validate_matrix_entry, - validate_agentic_matrix_entry, + DEFAULT_AGENTIC_DURATION_SECONDS, + Fields, load_config_files, load_runner_file, - Fields, - DEFAULT_AGENTIC_DURATION_SECONDS, + validate_agentic_matrix_entry, + validate_matrix_entry, ) seq_len_stoi = { @@ -82,6 +86,150 @@ def runner_gpus_per_node(runner: str, runner_data: dict) -> int: return runner_hardware_int(runner, runner_data, Fields.GPUS_PER_NODE.value) +def _hardware_family(label: str) -> str: + """Return the GPU family encoded in a runner or cluster label.""" + return label.removeprefix("cluster:").split("-", 1)[0] + + +def scheduling_gpus_per_node(label: str, runner_data: dict) -> int: + """Resolve GPUs per node for an abstract runner or worker hardware label. + + Scheduling labels such as ``b200-multinode`` do not currently duplicate + the hardware facts stored under ``cluster:b200-dgxc``. Fall back to the + GPU family when the exact label has no hardware record, while rejecting + ambiguous families that disagree about node shape. + """ + hardware = runner_hardware(runner_data) + exact = hardware.get(label) + if exact is not None: + return exact[Fields.GPUS_PER_NODE.value] + + family = _hardware_family(label) + matches = { + facts[Fields.GPUS_PER_NODE.value] + for hardware_label, facts in hardware.items() + if _hardware_family(hardware_label) == family + } + if len(matches) == 1: + return matches.pop() + if not matches: + raise ValueError( + f"Cannot resolve {Fields.GPUS_PER_NODE.value} for '{label}'" + ) + raise ValueError( + f"Ambiguous {Fields.GPUS_PER_NODE.value} for '{label}': {sorted(matches)}" + ) + + +def _worker_node_override(worker: dict, setting_name: str) -> int | None: + """Read an explicit role node count from additional settings.""" + pattern = re.compile(rf"^{re.escape(setting_name)}=(\d+)$") + values = [] + for setting in worker.get(Fields.ADDITIONAL_SETTINGS.value, []) or []: + match = pattern.match(setting) + if match: + values.append(int(match.group(1))) + if not values: + return None + if len(set(values)) != 1 or values[0] <= 0: + raise ValueError(f"Conflicting or invalid {setting_name} settings: {values}") + return values[0] + + +def recipe_node_count(prefill: dict, decode: dict) -> int | None: + """Read the authoritative node count from a checked-in srt-slurm recipe.""" + config_files = { + setting.split("=", 1)[1] + for worker in (prefill, decode) + for setting in (worker.get(Fields.ADDITIONAL_SETTINGS.value, []) or []) + if setting.startswith("CONFIG_FILE=") + } + if not config_files: + return None + if len(config_files) != 1: + raise ValueError(f"Conflicting CONFIG_FILE settings: {sorted(config_files)}") + + relative_path = config_files.pop().removeprefix("recipes/") + recipe_path = ( + Path(__file__).resolve().parents[2] + / "benchmarks" + / "multi_node" + / "srt-slurm-recipes" + / relative_path + ) + if not recipe_path.exists(): + # Some srt-slurm recipes live only in the runtime image. Their master + # config topology remains the best available scheduling estimate. + return None + + resources = yaml.safe_load(recipe_path.read_text())["resources"] + if "agg_nodes" in resources: + return int(resources["agg_nodes"]) + if "prefill_nodes" in resources and "decode_nodes" in resources: + return int(resources["prefill_nodes"]) + int(resources["decode_nodes"]) + raise ValueError(f"Recipe has no supported node resource fields: {recipe_path}") + + +def worker_node_count( + worker: dict, + role: str, + runner: str, + runner_data: dict, +) -> int: + """Return physical nodes consumed by one prefill or decode role.""" + override = _worker_node_override(worker, f"{role.upper()}_NODES") + if override is not None: + return override + + hardware_label = worker.get(Fields.HARDWARE.value) or runner + gpus_per_node = scheduling_gpus_per_node(hardware_label, runner_data) + total_gpus = ( + worker[Fields.NUM_WORKER.value] + * worker[Fields.TP.value] + * worker.get(Fields.PP.value, 1) + * worker.get(Fields.PCP_SIZE.value, 1) + ) + return math.ceil(total_gpus / gpus_per_node) + + +def multinode_node_count( + prefill: dict, + decode: dict, + runner: str, + runner_data: dict, +) -> int: + """Return the total Slurm node request represented by a matrix row.""" + recipe_count = recipe_node_count(prefill, decode) + if recipe_count is not None: + return recipe_count + return ( + worker_node_count(prefill, "prefill", runner, runner_data) + + worker_node_count(decode, "decode", runner, runner_data) + ) + + +def add_multinode_node_count( + entry: dict, + runner_data: dict, + num_nodes: int | None, +) -> dict: + """Annotate a multi-node row with its scheduling node count.""" + if not entry[Fields.DISAGG.value] and num_nodes is not None: + entry[Fields.NODE_COUNT.value] = num_nodes + elif num_nodes is not None: + raise ValueError( + f"{Fields.NUM_NODES.value} is not valid for disaggregated entries" + ) + elif runner_data: + entry[Fields.NODE_COUNT.value] = multinode_node_count( + entry[Fields.PREFILL.value], + entry[Fields.DECODE.value], + entry[Fields.RUNNER.value], + runner_data, + ) + return entry + + def effective_gpu_count(benchmark: dict) -> int: """Return GPUs used by a single-node TP/PP/PCP topology.""" return ( @@ -569,6 +717,11 @@ def generate_full_sweep(args, all_config_data, runner_data): Fields.RUN_EVAL.value: False, # Default, may be overridden by mark_eval_entries } entry.update(component_metadata(bmk, val)) + add_multinode_node_count( + entry, + runner_data, + bmk.get(Fields.NUM_NODES.value), + ) validate_matrix_entry(entry, is_multinode) matrix_values.append(entry) @@ -788,6 +941,11 @@ def generate_full_sweep(args, all_config_data, runner_data): if kv_offload_backend is not None: entry[Fields.KV_OFFLOAD_BACKEND.value] = kv_offload_backend entry.update(component_metadata(bmk, val)) + add_multinode_node_count( + entry, + runner_data, + bmk.get(Fields.NUM_NODES.value), + ) validate_agentic_matrix_entry(entry) matrix_values.append(entry) else: @@ -940,6 +1098,11 @@ def generate_test_config_sweep(args, all_config_data, runner_data=None): Fields.RUN_EVAL.value: False, } entry.update(component_metadata(bmk, val)) + add_multinode_node_count( + entry, + runner_data, + bmk.get(Fields.NUM_NODES.value), + ) matrix_values.append(validate_matrix_entry(entry, is_multinode=True)) else: # Single-node config @@ -1086,6 +1249,11 @@ def generate_test_config_sweep(args, all_config_data, runner_data=None): if kv_offload_backend is not None: entry[Fields.KV_OFFLOAD_BACKEND.value] = kv_offload_backend entry.update(component_metadata(bmk, val)) + add_multinode_node_count( + entry, + runner_data, + bmk.get(Fields.NUM_NODES.value), + ) matrix_values.append(validate_agentic_matrix_entry(entry)) else: for conc in conc_values: diff --git a/utils/matrix_logic/test_generate_sweep_configs.py b/utils/matrix_logic/test_generate_sweep_configs.py index 7cbfea79af..9e1daa0847 100644 --- a/utils/matrix_logic/test_generate_sweep_configs.py +++ b/utils/matrix_logic/test_generate_sweep_configs.py @@ -1,21 +1,105 @@ """Comprehensive tests for generate_sweep_configs.py""" -import pytest + import argparse import copy + +import pytest + from generate_sweep_configs import ( MIN_EVAL_CONC, - seq_len_stoi, - seq_len_itos, - seq_len_to_str, + add_multinode_node_count, + apply_node_type_defaults, + expand_config_keys, generate_full_sweep, generate_test_config_sweep, - mark_eval_entries, mark_all_eval_entries, - apply_node_type_defaults, - expand_config_keys, + mark_eval_entries, + multinode_node_count, + seq_len_itos, + seq_len_stoi, + seq_len_to_str, ) +def test_aggregated_multinode_node_count_uses_explicit_num_nodes(): + entry = { + "runner": "unknown", + "disagg": False, + "prefill": {}, + "decode": {}, + } + + add_multinode_node_count(entry, {}, num_nodes=3) + + assert entry["node-count"] == 3 + + +def test_disaggregated_multinode_node_count_rejects_num_nodes(): + entry = { + "runner": "unknown", + "disagg": True, + "prefill": {}, + "decode": {}, + } + + with pytest.raises(ValueError, match="num-nodes.*disaggregated"): + add_multinode_node_count(entry, {}, num_nodes=3) + +def test_multinode_node_count_uses_role_gpu_footprints(sample_runner_config): + prefill = {"num-worker": 3, "tp": 2, "pp": 1, "pcp-size": 1} + decode = {"num-worker": 2, "tp": 8, "pp": 1, "pcp-size": 1} + + assert multinode_node_count( + prefill, decode, "cluster:b300-nv", sample_runner_config + ) == 3 + + +def test_multinode_node_count_honors_explicit_role_node_settings(): + prefill = { + "num-worker": 1, + "tp": 8, + "additional-settings": ["PREFILL_NODES=2"], + } + decode = { + "num-worker": 1, + "tp": 8, + "additional-settings": ["DECODE_NODES=1"], + } + + assert multinode_node_count(prefill, decode, "unknown", {}) == 3 + + +def test_multinode_node_count_resolves_heterogeneous_worker_hardware( + sample_runner_config, +): + prefill = {"hardware": "gb200", "num-worker": 5, "tp": 4} + decode = {"hardware": "h100", "num-worker": 1, "tp": 8} + + assert multinode_node_count( + prefill, decode, "gb200", sample_runner_config + ) == 6 + + +def test_multinode_node_count_prefers_checked_in_recipe_resources( + sample_runner_config, +): + prefill = { + "num-worker": 3, + "tp": 1, + "additional-settings": [ + ( + "CONFIG_FILE=recipes/vllm/kimi-k2.5-fp4/8k1k/" + "disagg-gb200-3p1d-dep4-dep16.yaml" + ) + ], + } + decode = {"num-worker": 1, "tp": 1} + + assert multinode_node_count( + prefill, decode, "cluster:gb200-nv", sample_runner_config + ) == 7 + + # ============================================================================= # Test Fixtures # ============================================================================= diff --git a/utils/matrix_logic/test_validation.py b/utils/matrix_logic/test_validation.py index 3d25d03eef..453285232b 100644 --- a/utils/matrix_logic/test_validation.py +++ b/utils/matrix_logic/test_validation.py @@ -1053,6 +1053,51 @@ def test_aggregated_multinode_allows_kv_p2p_transfer( assert config.kv_p2p_transfer == "nixl" + def test_aggregated_multinode_allows_explicit_num_nodes( + self, + valid_multinode_master_config, + ): + """Aggregated search-space entries may set their Slurm node count.""" + valid_multinode_master_config["disagg"] = False + search_entry = valid_multinode_master_config[ + "scenarios" + ]["fixed-seq-len"][0]["search-space"][0] + search_entry["num-nodes"] = 3 + + config = MultiNodeMasterConfigEntry(**valid_multinode_master_config) + + validated_entry = config.scenarios.fixed_seq_len[0].search_space[0] + assert validated_entry.num_nodes == 3 + + def test_disaggregated_multinode_rejects_num_nodes( + self, + valid_multinode_master_config, + ): + """Disaggregated entries derive nodes from prefill and decode.""" + search_entry = valid_multinode_master_config[ + "scenarios" + ]["fixed-seq-len"][0]["search-space"][0] + search_entry["num-nodes"] = 3 + + with pytest.raises(Exception, match="num-nodes is only valid"): + MultiNodeMasterConfigEntry(**valid_multinode_master_config) + + @pytest.mark.parametrize("num_nodes", [0, -1, True]) + def test_aggregated_multinode_rejects_invalid_num_nodes( + self, + valid_multinode_master_config, + num_nodes, + ): + """Explicit aggregate node counts must be strict positive integers.""" + valid_multinode_master_config["disagg"] = False + search_entry = valid_multinode_master_config[ + "scenarios" + ]["fixed-seq-len"][0]["search-space"][0] + search_entry["num-nodes"] = num_nodes + + with pytest.raises(Exception, match="num-nodes"): + MultiNodeMasterConfigEntry(**valid_multinode_master_config) + def test_component_metadata_rejects_image_as_version(self): """Component versions identify the component, not its container.""" with pytest.raises(Exception, match="not an image reference"): diff --git a/utils/matrix_logic/validation.py b/utils/matrix_logic/validation.py index 04d3ae8b03..53ea840997 100644 --- a/utils/matrix_logic/validation.py +++ b/utils/matrix_logic/validation.py @@ -72,6 +72,8 @@ class Fields(Enum): AVAILABLE_CPU_DRAM_MIB = 'available-cpu-dram-mib' DRAM_UTILIZATION = 'dram-utilization' GPUS_PER_NODE = 'gpus-per-node' + NUM_NODES = 'num-nodes' + NODE_COUNT = 'node-count' DURATION = 'duration' # Matrix entry fields @@ -227,6 +229,9 @@ class MultiNodeMatrixEntry(BaseModel): alias=Fields.SPEC_DECODING.value ) runner: str + node_count: Optional[int] = Field( + default=None, alias=Fields.NODE_COUNT.value, gt=0, strict=True + ) isl: int osl: int prefill: WorkerConfig @@ -318,6 +323,9 @@ class MultiNodeAgenticMatrixEntry(BaseModel): alias=Fields.SPEC_DECODING.value ) runner: str + node_count: Optional[int] = Field( + default=None, alias=Fields.NODE_COUNT.value, gt=0, strict=True + ) prefill: WorkerConfig decode: WorkerConfig conc: list[int] @@ -512,6 +520,8 @@ class MultiNodeSearchSpaceEntry(BaseModel): default="none", alias=Fields.SPEC_DECODING.value) prefill: WorkerConfig decode: WorkerConfig + num_nodes: Optional[int] = Field( + default=None, alias=Fields.NUM_NODES.value, gt=0, strict=True) router: Optional[ComponentMetadata] = None kv_p2p_transfer: Optional[str] = Field( default=None, alias=Fields.KV_P2P_TRANSFER.value, min_length=1 @@ -568,6 +578,8 @@ class AgenticCodingSearchSpaceEntry(BaseModel): default="none", alias=Fields.SPEC_DECODING.value) prefill: Optional[WorkerConfig] = None decode: Optional[WorkerConfig] = None + num_nodes: Optional[int] = Field( + default=None, alias=Fields.NUM_NODES.value, gt=0, strict=True) kv_offloading: Optional[Literal["none", "dram"]] = Field( default=None, alias=Fields.KV_OFFLOADING.value ) @@ -684,9 +696,9 @@ def at_least_one_scenario(self): return self -def _validate_component_metadata_scope(self: BaseModel) -> BaseModel: - """Require unambiguous component metadata across a master config.""" - search_space_entries = [ +def _master_search_space_entries(self: BaseModel) -> list[BaseModel]: + """Return every search-space entry in a master config.""" + return [ entry for scenario_configs in ( self.scenarios.fixed_seq_len, @@ -696,6 +708,11 @@ def _validate_component_metadata_scope(self: BaseModel) -> BaseModel: for entry in scenario_config.search_space ] + +def _validate_component_metadata_scope(self: BaseModel) -> BaseModel: + """Require unambiguous component metadata across a master config.""" + search_space_entries = _master_search_space_entries(self) + for field in (Fields.ROUTER, Fields.KV_P2P_TRANSFER): attribute = field.value.replace("-", "_") top_level_value = getattr(self, attribute, None) @@ -733,6 +750,22 @@ def _validate_component_metadata_scope(self: BaseModel) -> BaseModel: return self +def _validate_num_nodes_scope(self: BaseModel) -> BaseModel: + """Allow explicit node counts only for aggregated multinode entries.""" + search_space_entries = _master_search_space_entries(self) + entries_with_num_nodes = [ + entry for entry in search_space_entries + if getattr(entry, "num_nodes", None) is not None + ] + + if (not self.multinode or self.disagg) and entries_with_num_nodes: + raise ValueError( + f"{Fields.NUM_NODES.value} is only valid when " + f"{Fields.MULTINODE.value}=true and {Fields.DISAGG.value}=false" + ) + return self + + class SingleNodeMasterConfigEntry(BaseModel): """Top-level single node master configuration entry.""" model_config = ConfigDict(extra='forbid', populate_by_name=True) @@ -757,6 +790,10 @@ def validate_agentic_runner(self): def validate_component_metadata_scope(self): return _validate_component_metadata_scope(self) + @model_validator(mode='after') + def validate_num_nodes_scope(self): + return _validate_num_nodes_scope(self) + class MultiNodeMasterConfigEntry(BaseModel): """Top-level multinode master configuration entry.""" @@ -785,6 +822,10 @@ def validate_agentic_runner(self): def validate_component_metadata_scope(self): return _validate_component_metadata_scope(self) + @model_validator(mode='after') + def validate_num_nodes_scope(self): + return _validate_num_nodes_scope(self) + def validate_master_config(master_configs: dict) -> List[dict]: """Validate input master configuration structure.""" diff --git a/utils/runner_setup/RUNNER_SETUP.md b/utils/runner_setup/RUNNER_SETUP.md index 5ac57e0e6c..11793645a4 100644 --- a/utils/runner_setup/RUNNER_SETUP.md +++ b/utils/runner_setup/RUNNER_SETUP.md @@ -180,6 +180,53 @@ self-hosted, Linux, X64, slurm, b200, b200-dsv4, cluster:b200-dgxc, b200-dgxc_00 Labels can be edited later on the runners settings page without re-registering. +### Dynamic Slurm node leases + +The optional node-slot scheduler performs weighted admission across every job +size in one physical Slurm cluster. A job that needs three nodes adds +`nodes:3` to its queued `runs-on` labels. `nodes:N` is request metadata for the +trusted priority controller, not a permanent runner capability label. + +The controller groups runners by their permanent `cluster:` label. To +admit a three-node job, it selects three online, unleased runners from one +compatible cluster and temporarily adds the same `ci-lease-*` label to all +three. After every lease write succeeds, it adds the job's unique `ci-job-*`, +`ci-attempt-*`, and `nodes:3` labels to one of those runners as the anchor. +GitHub can then dispatch only that job, while the other two runners remain idle +as capacity tokens for the Slurm allocation. + +For example, after admitting ten- and eight-node jobs on an 18-node cluster, +the controller has leased all 18 runners and will not publish the unique label +for a queued nine-node job. This enforces the aggregate invariant +`sum(admitted node counts) <= online cluster capacity` across mixed job sizes. + +Lease allocation is serialized and two-phase: reserve every capacity token, +verify the reservation, then publish the anchor's dispatch label. Completion +and periodic reconciliation remove orphaned leases. Every managed GPU workflow +must require its unique `ci-job-*` label; a workflow that targets only a generic +hardware label bypasses admission accounting. + +Set `NODE_SLOT_SCHEDULER_ENABLED=true` only after the deployed priority +controller supports `nodes:N` and `ci-lease-*`. `PRIORITY_SCHEDULER_ENABLED` +must also remain enabled. If either variable is disabled, workflows omit +`nodes:N` and retain the existing unweighted behavior. The priority score also +subtracts `0.001` per additional node so otherwise equal work prefers smaller +allocations without overriding the existing business-priority signals. + +Aggregated multi-node search-space entries can declare `num-nodes`; that value +becomes the generated `node-count` directly. `num-nodes` is rejected for +disaggregated entries because their independent prefill and decode allocations +determine the total. During migration, entries without `num-nodes` derive +`node-count` from, in precedence order: + +1. checked-in srt-slurm recipe `resources`; +2. explicit `PREFILL_NODES` and `DECODE_NODES` settings; or +3. worker GPU footprints divided by the relevant hardware's GPUs per node. + +Runner leases account for GitHub-managed work only. Slurm remains the final +capacity authority for external users, reservations, offline compute nodes, +and allocations submitted outside this admission path. + ## Storage layout The login node (where the runners live) and the Slurm compute nodes (where benchmarks diff --git a/utils/test_ci_priority.py b/utils/test_ci_priority.py index 8bb7b11198..8898c4f937 100644 --- a/utils/test_ci_priority.py +++ b/utils/test_ci_priority.py @@ -54,6 +54,33 @@ def test_main_branch_jobs_receive_an_automatic_boost(): ) == Decimal("3.000") +def test_smaller_node_allocations_win_otherwise_equal_priority_ties(): + entry = {"runner": "h100", "framework": "trt"} + + assert calculate_priority({**entry, "node-count": 1}, POLICY) == Decimal("1.000") + assert calculate_priority({**entry, "node-count": 2}, POLICY) == Decimal("0.999") + assert calculate_priority({**entry, "node-count": 3}, POLICY) == Decimal("0.998") + + +def test_node_count_tiebreaker_survives_classifier_projection(): + entry = {"runner": "h100", "framework": "trt", "node-count": 3} + + assert calculate_priority( + entry, + POLICY, + PriorityContext(criteria=frozenset()), + ) == Decimal("0.998") + + +@pytest.mark.parametrize("node_count", [0, -1, 1.5, True]) +def test_node_count_must_be_a_positive_integer(node_count): + with pytest.raises(ValueError, match="positive integer"): + calculate_priority( + {"runner": "h100", "framework": "trt", "node-count": node_count}, + POLICY, + ) + + def test_patchwork_score_uses_half_up_rounding(): policy = deepcopy(POLICY) policy["labels"]["patchwork"]["score"] = 0.7225