From f427ee2223d35ae7cdc085abe59fba2e1c9512af Mon Sep 17 00:00:00 2001 From: bogdan-petkovic Date: Tue, 23 Jun 2026 07:13:13 -0500 Subject: [PATCH 1/8] [AIROCMLIR-375] Run tests in parallel across multiple GPUs Signed-off-by: bogdan-petkovic --- mlir/utils/jenkins/Jenkinsfile | 22 ++- mlir/utils/jenkins/run_e2e_multigpu.py | 198 ++++++++++++++++++++++ mlir/utils/performance/attentionSweeps.py | 22 ++- mlir/utils/performance/gpu_topology.py | 96 +++++++++++ mlir/utils/performance/parameterSweeps.py | 55 ++++-- 5 files changed, 372 insertions(+), 21 deletions(-) create mode 100644 mlir/utils/jenkins/run_e2e_multigpu.py create mode 100644 mlir/utils/performance/gpu_topology.py diff --git a/mlir/utils/jenkins/Jenkinsfile b/mlir/utils/jenkins/Jenkinsfile index c8655bdf5ded..d5731ae1a4cc 100644 --- a/mlir/utils/jenkins/Jenkinsfile +++ b/mlir/utils/jenkins/Jenkinsfile @@ -883,6 +883,19 @@ int setLitWorkerCount() { return limit_lit_workers } +// Run the rocMLIR lit suite sharded across all GPUs on the node: one lit +// process per GPU, each isolated via ROCR_VISIBLE_DEVICES. On single-GPU or +// heterogeneous nodes the driver falls back to a single lit run, preserving +// legacy behavior. totalJobs is split evenly across shards so per-GPU +// concurrency stays bounded (avoiding the oversubscription issues #1845/#1841). +void runShardedE2E(int totalJobs) { + dir('build') { + sh "python3 ${env.WORKSPACE}/mlir/utils/jenkins/run_e2e_multigpu.py --build-dir . " + + "--total-jobs ${totalJobs} " + + "\"--lit-args=-v --time-tests --timeout=3600 --max-failures=1\"" + } +} + void build_fixedE2ETests(String codepath) { // Limit the number of lit workers for gfx908, gfx90a to (8, 30) on CI as a workaround for issue #1845 and #1841 int limit_lit_workers = setLitWorkerCount() @@ -900,15 +913,17 @@ void build_fixedE2ETests(String codepath) { void check_randomE2ETests(String codepath) { // Limit the number of lit workers for gfx908, gfx90a to (8, 30) on CI as a workaround for issue #1845 and #1841 int limit_lit_workers = setLitWorkerCount() - buildProject('check-rocmlir', """ + // Build the test dependencies only; the lit suite is run separately via the + // multi-GPU sharding driver below. + buildProject('check-rocmlir-build-only', """ -DROCMLIR_DRIVER_PR_E2E_TEST_ENABLED=0 -DROCMLIR_DRIVER_E2E_TEST_ENABLED=1 -DROCK_E2E_TEST_ENABLED=1 -DROCMLIR_DRIVER_RANDOM_DATA_SEED=1 -DROCMLIR_DRIVER_TEST_GPU_VALIDATION=0 - -DLLVM_LIT_ARGS='-v --time-tests --timeout=3600 --max-failures=1 -j ${limit_lit_workers}' -DCMAKE_EXPORT_COMPILE_COMMANDS=1 """) + runShardedE2E(limit_lit_workers) } void parameterSweep(String CONFIG, String sweepType = "default") { @@ -1307,7 +1322,8 @@ pipeline { build_fixedE2ETests("${CODEPATH}") preMergeCheck("${CODEPATH}") timeout(time: 60, activity: true, unit: 'MINUTES') { - sh 'cd build; ninja check-mlir check-rocmlir' + sh 'cd build; ninja check-mlir' + runShardedE2E(setLitWorkerCount()) } } } diff --git a/mlir/utils/jenkins/run_e2e_multigpu.py b/mlir/utils/jenkins/run_e2e_multigpu.py new file mode 100644 index 000000000000..3a48549b6b59 --- /dev/null +++ b/mlir/utils/jenkins/run_e2e_multigpu.py @@ -0,0 +1,198 @@ +#!/usr/bin/env python3 +# Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions. +# See https://llvm.org/LICENSE.txt for license information. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +"""Run the rocMLIR lit test suite in parallel across all GPUs on a node. + +CI nodes commonly expose several GPUs but the lit suite (e.g. `check-rocmlir`) +historically runs every test on GPU 0. This driver shards the suite with lit's +native `--num-shards`/`--run-shard` mechanism, launching one lit process per GPU +and isolating each to its device via ROCR_VISIBLE_DEVICES. The total set of +tests is partitioned, so per-GPU concurrency stays bounded (avoiding the +oversubscription hangs seen with a single global `-j`) while all GPUs are used. + +Distribution is only enabled on homogeneous nodes (all GPUs share one gfx +architecture); single-GPU and heterogeneous nodes fall back to one lit run, +preserving today's behavior. +""" + +from __future__ import annotations + +import argparse +import os +import subprocess +import sys +import time +from typing import List, Optional + +# Reuse the GPU detection helper from the performance scripts. +sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'performance')) +from gpu_topology import select_gpu_ids # noqa: E402 + + +def default_lit_path(build_dir: str) -> str: + return os.path.join(build_dir, 'external', 'llvm-project', 'llvm', 'bin', 'llvm-lit') + + +def build_shard_command(lit: str, lit_args: List[str], jobs: int, num_shards: int, shard: int, + test_paths: List[str]) -> List[str]: + """Build a single lit invocation for shard `shard` (1-based) of `num_shards`.""" + cmd = [sys.executable, lit, '-j', str(jobs)] + if num_shards > 1: + cmd += ['--num-shards', str(num_shards), '--run-shard', str(shard)] + cmd += lit_args + cmd += test_paths + return cmd + + +def resolve_jobs_per_shard(args: argparse.Namespace, num_shards: int) -> int: + """Pick the lit worker count for each shard. + + `--jobs-per-gpu` wins if given. Otherwise `--total-jobs` is split evenly + across shards, so a single-GPU node keeps the legacy global concurrency while + multi-GPU nodes keep per-GPU pressure bounded (total/num_shards). Falls back + to a conservative default when neither is given. + """ + if args.jobs_per_gpu is not None: + return max(1, args.jobs_per_gpu) + if args.total_jobs is not None: + return max(1, args.total_jobs // num_shards) + return 8 + + +def run(args: argparse.Namespace) -> int: + lit = args.lit or default_lit_path(args.build_dir) + test_paths = args.test_paths or [os.path.join(args.build_dir, 'mlir', 'test')] + + gpu_ids, _gpu_arch, gpu_msg = select_gpu_ids(args.gpus) + print(f"[run_e2e_multigpu] {gpu_msg}", flush=True) + + # `select_gpu_ids` returns [None] for the single-GPU / heterogeneous / unknown + # cases; treat those as one un-pinned lit run (legacy behavior). + single_gpu = gpu_ids == [None] + shard_gpus: List[Optional[int]] = [None] if single_gpu else gpu_ids + num_shards = len(shard_gpus) + jobs_per_shard = resolve_jobs_per_shard(args, num_shards) + print(f"[run_e2e_multigpu] {num_shards} shard(s), {jobs_per_shard} lit workers each", + flush=True) + + procs = [] + log_paths = [] + for idx, gpu_id in enumerate(shard_gpus): + cmd = build_shard_command(lit, args.lit_args, jobs_per_shard, num_shards, idx + 1, + test_paths) + env = os.environ.copy() + if gpu_id is not None: + env['ROCR_VISIBLE_DEVICES'] = str(gpu_id) + env.pop('HIP_VISIBLE_DEVICES', None) + label = f"GPU {gpu_id}" if gpu_id is not None else "single" + print(f"[run_e2e_multigpu] shard {idx + 1}/{num_shards} on {label}: {' '.join(cmd)}", + flush=True) + if args.dry_run: + continue + log_path = os.path.join(args.build_dir, f"e2e-shard-{idx}.log") + log_paths.append((label, log_path)) + log_file = open(log_path, 'wb') + procs.append((label, log_file, + subprocess.Popen(cmd, env=env, stdout=log_file, stderr=subprocess.STDOUT))) + + if args.dry_run: + return 0 + + failures = [] + aborted = [] + pending = list(range(len(procs))) + while pending: + time.sleep(1) + for i in list(pending): + label, log_file, proc = procs[i] + rc = proc.poll() + if rc is None: + continue + log_file.close() + pending.remove(i) + if rc != 0: + failures.append((label, rc)) + + # Fail-fast: once any shard fails, terminate the rest so the run aborts + # promptly (matching the previous single-lit --max-failures=1 behavior). + if args.fail_fast and failures and pending: + for i in pending: + procs[i][2].terminate() + for i in pending: + label, log_file, proc = procs[i] + try: + proc.wait(timeout=15) + except subprocess.TimeoutExpired: + proc.kill() + log_file.close() + aborted.append(label) + pending = [] + + # Surface every shard's output in the CI console. + for label, log_path in log_paths: + print(f"\n===== lit output: {label} ({log_path}) =====", flush=True) + with open(log_path, 'r', errors='replace') as f: + sys.stdout.write(f.read()) + + if failures: + summary = ', '.join(f"{label} (exit {rc})" for label, rc in failures) + print(f"\n[run_e2e_multigpu] FAILED shards: {summary}", flush=True) + if aborted: + print(f"[run_e2e_multigpu] aborted (fail-fast): {', '.join(aborted)}", flush=True) + return 1 + print(f"\n[run_e2e_multigpu] all {num_shards} shard(s) passed", flush=True) + return 0 + + +def main() -> int: + parser = argparse.ArgumentParser( + description='Run the rocMLIR lit suite sharded across all GPUs on the node.') + parser.add_argument('test_paths', + nargs='*', + default=None, + help='lit test path(s) to run (default: /mlir/test)') + parser.add_argument('--build-dir', + default='build', + help='Build directory (default: %(default)s)') + parser.add_argument('--lit', + default=None, + help='Path to llvm-lit (default: /external/llvm-project/' + 'llvm/bin/llvm-lit)') + parser.add_argument('--jobs-per-gpu', + type=int, + default=None, + help='lit workers per GPU shard (overrides --total-jobs)') + parser.add_argument('--total-jobs', + type=int, + default=None, + help='Total lit workers split evenly across shards (default per-shard: 8)') + parser.add_argument('--gpus', + type=int, + nargs='+', + default=None, + help='Physical GPU ids to shard across (default: auto-detect homogeneous ' + 'GPUs, else a single run)') + parser.add_argument('--lit-args', + type=str, + default='-v --time-tests --timeout=3600', + help='Extra arguments passed to each lit invocation (default: %(default)r)') + parser.add_argument('--dry-run', + action='store_true', + help='Print the per-shard lit commands without running them') + parser.add_argument('--fail-fast', + dest='fail_fast', + action='store_true', + default=True, + help='Abort remaining shards once any shard fails (default)') + parser.add_argument('--no-fail-fast', + dest='fail_fast', + action='store_false', + help='Let all shards run to completion even if one fails') + args = parser.parse_args() + args.lit_args = args.lit_args.split() + return run(args) + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/mlir/utils/performance/attentionSweeps.py b/mlir/utils/performance/attentionSweeps.py index 7f31223c8ce8..a5bc07deca60 100755 --- a/mlir/utils/performance/attentionSweeps.py +++ b/mlir/utils/performance/attentionSweeps.py @@ -38,6 +38,7 @@ get_codegen_flags_for_codepath, ) from amd_arch_db import GemmFeatures, has_feature, lookup_arch_info +from gpu_topology import select_gpu_ids # GLOBAL VARIABLES DATA_TYPES_ATTENTION = initialize_dtypes_attn() @@ -329,6 +330,12 @@ def main(): parser.add_argument('--quiet', action='store_true') parser.add_argument('--debug-fails', action='store_true') parser.add_argument('-j', '--jobs', type=int, default=(os.cpu_count() or 1)) + parser.add_argument('--gpus', + type=int, + nargs='+', + default=None, + help="Physical GPU ids to spread work across. Default: auto-detect " + "all GPUs when they share one architecture, otherwise use a single GPU.") parser.add_argument('--mlir-build-dir', type=str, default=find_mlir_build_dir()) parser.add_argument('--samples', type=int, default=1000) parser.add_argument('--codepath', @@ -348,12 +355,16 @@ def main(): if args.mlir_build_dir is None: args.mlir_build_dir = find_mlir_build_dir() - arch = get_arch() + gpu_ids, gpu_arch, gpu_msg = select_gpu_ids(args.gpus) + print(f"[attentionSweeps] GPU distribution: {gpu_msg}") + # When work is pinned to a same-arch GPU group, compile for that group's arch + # (and query its CU/chiplet counts) so kernels match the GPUs they run on. + arch = gpu_arch or get_arch() + rep_device = gpu_ids[0] if gpu_ids and gpu_ids[0] is not None else 0 chip_match = GFX_CHIP_RE.search(arch) if chip_match is None: raise RuntimeError(f"Could not find GFX chip in arch string: {arch}") chip = chip_match.group(0) - num_cu = get_num_cu() paths = create_paths(None, args.mlir_build_dir) options = Options(debug_fails=args.debug_fails, debug=args.debug, @@ -361,10 +372,11 @@ def main(): arch=arch, flags=[], concurrent_tests=args.jobs, - num_cu=num_cu, - num_chiplets=get_num_chiplets(), + num_cu=get_num_cu(rep_device), + num_chiplets=get_num_chiplets(rep_device), log_failures=args.log_failures, - test_timeout_sec=args.test_timeout_sec) + test_timeout_sec=args.test_timeout_sec, + gpu_ids=tuple(gpu_ids)) if not args.quiet: print(f"Sampling {args.samples} configurations from attention space...") diff --git a/mlir/utils/performance/gpu_topology.py b/mlir/utils/performance/gpu_topology.py new file mode 100644 index 000000000000..cfd31c546c42 --- /dev/null +++ b/mlir/utils/performance/gpu_topology.py @@ -0,0 +1,96 @@ +# Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions. +# See https://llvm.org/LICENSE.txt for license information. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +"""Shared helpers for distributing GPU test workloads across multiple devices. + +CI nodes commonly expose several GPUs. These helpers detect the visible GPUs, +confirm they share a single architecture (a prerequisite for safely reusing the +same compiled kernels), and build per-process environments that isolate work to +one device via ROCR_VISIBLE_DEVICES (the same mechanism tuningRunner.py uses). +""" + +from __future__ import annotations + +import os +from typing import List, Optional, Tuple + + +def _hip_check(call_result): + """Unwrap a hip-python call result, raising on a non-success status.""" + from hip import hip + err = call_result[0] + result = call_result[1:] + if len(result) == 1: + result = result[0] + if isinstance(err, hip.hipError_t) and err != hip.hipError_t.hipSuccess: + raise RuntimeError(str(err)) + return result + + +def get_per_device_archs() -> List[str]: + """Return the gfx architecture name of every visible GPU, indexed by id. + + Uses hip-python, matching the rest of the test tooling (the lit configs and + perfRunner already require it); callers treat a failure as "use one GPU". + """ + from hip import hip + archs = [] + device_count = _hip_check(hip.hipGetDeviceCount()) + for device in range(device_count): + props = hip.hipDeviceProp_t() + _hip_check(hip.hipGetDeviceProperties(props, device)) + archs.append(props.gcnArchName.decode('utf-8')) + return archs + + +def select_gpu_ids( + requested: Optional[List[int]] = None) -> Tuple[List[Optional[int]], Optional[str], str]: + """Decide which GPUs to spread work across. + + Returns ``(gpu_ids, arch, message)``. ``gpu_ids == [None]`` means run on a + single, unpinned GPU (and ``arch`` is ``None``). Otherwise ``gpu_ids`` lists + the physical devices to isolate work to, all sharing architecture ``arch``. + """ + # Respect a caller that already pinned visibility (e.g. lit shards, manual + # runs); don't second-guess their device selection. + if os.environ.get('ROCR_VISIBLE_DEVICES') or os.environ.get('HIP_VISIBLE_DEVICES'): + return [None], None, "respecting pre-set *_VISIBLE_DEVICES; using a single GPU" + + try: + archs = get_per_device_archs() + except Exception as e: # noqa: BLE001 - any GPU/runtime issue means fall back + return [None], None, f"GPU enumeration failed ({e}); using the default GPU" + + count = len(archs) + if count <= 1: + return [None], None, "single GPU detected; running on one GPU" + + if requested: + selected_archs = {archs[i] for i in requested if 0 <= i < count} + if len(selected_archs) != 1: + return [None], None, (f"requested GPUs {requested} are not a single arch " + f"({sorted(selected_archs)}); using the default GPU") + arch = next(iter(selected_archs)) + return list(requested), arch, f"using requested GPUs {requested} ({arch})" + + if len(set(archs)) > 1: + return [None], None, (f"mixed GPU architectures ({sorted(set(archs))}); " + "using a single GPU") + + return list(range(count)), archs[0], f"distributing across {count} GPUs ({archs[0]})" + + +def make_isolated_gpu_env(gpu_id: Optional[int]) -> Optional[dict]: + """Build an environment dict that isolates a child process to one GPU. + + Returns ``None`` when ``gpu_id`` is ``None`` so callers can pass it straight + through to ``subprocess``/``asyncio`` APIs to mean "inherit the environment". + ROCR_VISIBLE_DEVICES isolates at the HSA/ROCr level (below HIP); we clear + HIP_VISIBLE_DEVICES to avoid the two layers disagreeing. + """ + if gpu_id is None: + return None + env = os.environ.copy() + env["ROCR_VISIBLE_DEVICES"] = str(gpu_id) + env.pop("HIP_VISIBLE_DEVICES", None) + return env diff --git a/mlir/utils/performance/parameterSweeps.py b/mlir/utils/performance/parameterSweeps.py index 2f265ccd2c55..1cdbf65ef598 100755 --- a/mlir/utils/performance/parameterSweeps.py +++ b/mlir/utils/performance/parameterSweeps.py @@ -29,6 +29,8 @@ from amd_arch_db import GemmFeatures, has_feature, lookup_arch_info +from gpu_topology import select_gpu_ids, make_isolated_gpu_env + @dataclass(frozen=True) class Options: @@ -43,6 +45,9 @@ class Options: num_chiplets: int log_failures: bool = False test_timeout_sec: int = 600 + # Physical GPU ids to round-robin work across. A single ``None`` entry keeps + # the legacy single-GPU behavior (inherit the caller's device visibility). + gpu_ids: Tuple[Optional[int], ...] = (None,) async def _kill_process(proc: asyncio.subprocess.Process): @@ -293,9 +298,17 @@ class TestResult(enum.Enum): FAIL = 3 -async def test_config(config, options: Options, paths: Paths) -> TestResult: +async def test_config(config, + options: Options, + paths: Paths, + gpu_id: Optional[int] = None) -> TestResult: """Runs the given configuration and returns whether it successfully concluded, - failed validation, or was inapplicable.""" + failed validation, or was inapplicable. + + When ``gpu_id`` is set, the GPU-executing stage (mlir-runner) is isolated to + that device via ROCR_VISIBLE_DEVICES so concurrent configs spread across all + GPUs. The generation/applicability/lowering stages are CPU-only and keep the + inherited environment.""" if isinstance(config, MLIROnlyConfig): rocmlir_gen_opts = config.generate_mlir_driver_commandline(options.flags) else: @@ -366,7 +379,8 @@ async def test_config(config, options: Options, paths: Paths) -> TestResult: *mlir_cpu_runner_args, stdin=runner_from_lowering, stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE) + stderr=asyncio.subprocess.PIPE, + env=make_isolated_gpu_env(gpu_id)) os.close(runner_from_lowering) _, lowering_errs = await lowering.communicate(input=high_level) @@ -424,10 +438,10 @@ def grouper(iterable: Iterable[IterType], n: int): yield chunk -async def drop_good_config(config, options: Options, paths: Paths): +async def drop_good_config(config, options: Options, paths: Paths, gpu_id: Optional[int] = None): """Test the given `params`, returning the corresponding `config` on failure and `None` on success or inapplicability""" - result = await test_config(config, options, paths) + result = await test_config(config, options, paths, gpu_id) if not options.quiet: if isinstance(config, MLIROnlyConfig): print(f"{result.name}: {config!r}") @@ -452,9 +466,11 @@ async def sweep_parameters(param_iter: Iterable[IterType], to_config: Callable[[ failing_configs = [] passed = 0 invalid = 0 + gpu_ids = options.gpu_ids or (None,) configs = (c for c in (to_config(p, options) for p in param_iter)) - for configs in grouper((drop_good_config(c, options, paths) for c in configs), - options.concurrent_tests): + tasks = (drop_good_config(c, options, paths, gpu_ids[i % len(gpu_ids)]) + for i, c in enumerate(configs)) + for configs in grouper(tasks, options.concurrent_tests): configs_future = asyncio.gather(*configs) try: configs_results = await configs_future @@ -706,7 +722,14 @@ def main() -> bool: '-j', type=int, default=(len(os.sched_getaffinity(0)) // 2), - help="Number of jobs to run in parallel (default %(default)s)") + help="Total number of jobs to run in parallel across all GPUs " + "(default %(default)s)") + parser.add_argument('--gpus', + type=int, + nargs='+', + default=None, + help="Physical GPU ids to spread work across. Default: auto-detect " + "all GPUs when they share one architecture, otherwise use a single GPU.") parser.add_argument('--test-timeout-sec', type=int, default=600, @@ -723,7 +746,13 @@ def main() -> bool: if args.mlir_build_dir is None: args.mlir_build_dir = perfRunner.find_mlir_build_dir() - arch = get_arch() + gpu_ids, gpu_arch, gpu_msg = select_gpu_ids(args.gpus) + print(f"[parameterSweeps] GPU distribution: {gpu_msg}") + # When work is pinned to a same-arch GPU group, compile for that group's arch + # (and query its CU/chiplet counts) so kernels match the GPUs they run on. + arch = gpu_arch or get_arch() + rep_device = gpu_ids[0] if gpu_ids and gpu_ids[0] is not None else 0 + codepath, rocmlir_gen_flags = infer_codegen_flags_from_arch(arch, args.codepath) if codepath == 'unknown': if args.config == 'perf_config': @@ -735,7 +764,6 @@ def main() -> bool: # For non-perf-config sweeps, let rocmlir-gen infer features from --arch. rocmlir_gen_flags = [] - num_cu = get_num_cu() options = Options(debug=args.debug, quiet=args.quiet, log_failures=args.log_failures, @@ -743,9 +771,10 @@ def main() -> bool: arch=arch, flags=rocmlir_gen_flags, concurrent_tests=args.jobs, - num_cu=num_cu, - num_chiplets=get_num_chiplets(), - test_timeout_sec=args.test_timeout_sec) + num_cu=get_num_cu(rep_device), + num_chiplets=get_num_chiplets(rep_device), + test_timeout_sec=args.test_timeout_sec, + gpu_ids=tuple(gpu_ids)) paths = perfRunner.create_paths(None, args.mlir_build_dir) From fd1bc54aad3a3c642a1441c90446c2f0780cc9b0 Mon Sep 17 00:00:00 2001 From: bogdan-petkovic Date: Thu, 9 Jul 2026 02:55:50 -0500 Subject: [PATCH 2/8] [AIROCMLIR-375] Fix multi-GPU E2E driver hitting Jenkins activity timeout Signed-off-by: bogdan-petkovic --- mlir/utils/jenkins/run_e2e_multigpu.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/mlir/utils/jenkins/run_e2e_multigpu.py b/mlir/utils/jenkins/run_e2e_multigpu.py index 3a48549b6b59..c239764f010e 100644 --- a/mlir/utils/jenkins/run_e2e_multigpu.py +++ b/mlir/utils/jenkins/run_e2e_multigpu.py @@ -76,6 +76,21 @@ def run(args: argparse.Namespace) -> int: print(f"[run_e2e_multigpu] {num_shards} shard(s), {jobs_per_shard} lit workers each", flush=True) + # Single shard: stream lit output straight to the console (live per-test + # progress, keeps CI activity timeouts alive). Multi-shard buffers per-GPU. + if num_shards == 1: + gpu_id = shard_gpus[0] + cmd = build_shard_command(lit, args.lit_args, jobs_per_shard, 1, 1, test_paths) + env = os.environ.copy() + if gpu_id is not None: + env['ROCR_VISIBLE_DEVICES'] = str(gpu_id) + env.pop('HIP_VISIBLE_DEVICES', None) + label = f"GPU {gpu_id}" if gpu_id is not None else "single" + print(f"[run_e2e_multigpu] shard 1/1 on {label}: {' '.join(cmd)}", flush=True) + if args.dry_run: + return 0 + return subprocess.call(cmd, env=env) + procs = [] log_paths = [] for idx, gpu_id in enumerate(shard_gpus): @@ -102,8 +117,19 @@ def run(args: argparse.Namespace) -> int: failures = [] aborted = [] pending = list(range(len(procs))) + # Heartbeat so the console keeps emitting output during the otherwise-silent + # buffered run, preventing Jenkins `timeout(activity: true)` from firing. + start = time.time() + last_beat = start + heartbeat_secs = 30 while pending: time.sleep(1) + now = time.time() + if now - last_beat >= heartbeat_secs: + last_beat = now + print(f"[run_e2e_multigpu] still running: {len(pending)}/{len(procs)} " + f"shard(s) active, {int(now - start)}s elapsed", + flush=True) for i in list(pending): label, log_file, proc = procs[i] rc = proc.poll() From c3ac08b5ec162e4481147cb310588b9ae8f5ceb8 Mon Sep 17 00:00:00 2001 From: bogdan-petkovic Date: Thu, 9 Jul 2026 04:39:14 -0500 Subject: [PATCH 3/8] fix yapf format checks Signed-off-by: bogdan-petkovic --- mlir/utils/jenkins/run_e2e_multigpu.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/mlir/utils/jenkins/run_e2e_multigpu.py b/mlir/utils/jenkins/run_e2e_multigpu.py index c239764f010e..78ac502898e1 100644 --- a/mlir/utils/jenkins/run_e2e_multigpu.py +++ b/mlir/utils/jenkins/run_e2e_multigpu.py @@ -127,9 +127,10 @@ def run(args: argparse.Namespace) -> int: now = time.time() if now - last_beat >= heartbeat_secs: last_beat = now - print(f"[run_e2e_multigpu] still running: {len(pending)}/{len(procs)} " - f"shard(s) active, {int(now - start)}s elapsed", - flush=True) + print( + f"[run_e2e_multigpu] still running: {len(pending)}/{len(procs)} " + f"shard(s) active, {int(now - start)}s elapsed", + flush=True) for i in list(pending): label, log_file, proc = procs[i] rc = proc.poll() From f4cdd787ca8c62ca0b777a9a3d9116e8643f6715 Mon Sep 17 00:00:00 2001 From: bogdan-petkovic Date: Mon, 13 Jul 2026 03:55:16 -0500 Subject: [PATCH 4/8] [AIROCMLIR-375] Pass per-GPU lit worker count to sharded E2E driver Signed-off-by: bogdan-petkovic --- mlir/utils/jenkins/Jenkinsfile | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/mlir/utils/jenkins/Jenkinsfile b/mlir/utils/jenkins/Jenkinsfile index d5731ae1a4cc..4ebe609cfe41 100644 --- a/mlir/utils/jenkins/Jenkinsfile +++ b/mlir/utils/jenkins/Jenkinsfile @@ -886,12 +886,14 @@ int setLitWorkerCount() { // Run the rocMLIR lit suite sharded across all GPUs on the node: one lit // process per GPU, each isolated via ROCR_VISIBLE_DEVICES. On single-GPU or // heterogeneous nodes the driver falls back to a single lit run, preserving -// legacy behavior. totalJobs is split evenly across shards so per-GPU -// concurrency stays bounded (avoiding the oversubscription issues #1845/#1841). -void runShardedE2E(int totalJobs) { +// legacy behavior. jobsPerGpu is the lit worker count applied to each GPU shard +// (the historical per-GPU cap from setLitWorkerCount()), so per-GPU concurrency +// stays at the proven level instead of being starved by dividing a machine-wide +// total across shards (avoiding the oversubscription issues #1845/#1841). +void runShardedE2E(int jobsPerGpu) { dir('build') { sh "python3 ${env.WORKSPACE}/mlir/utils/jenkins/run_e2e_multigpu.py --build-dir . " + - "--total-jobs ${totalJobs} " + + "--jobs-per-gpu ${jobsPerGpu} " + "\"--lit-args=-v --time-tests --timeout=3600 --max-failures=1\"" } } From 331bee44716db9534e42cd551409e3b44efbf3ff Mon Sep 17 00:00:00 2001 From: bogdan-petkovic Date: Mon, 20 Jul 2026 05:00:34 -0500 Subject: [PATCH 5/8] [AIROCMLIR-375] Address review: install gpu_topology, validate GPU ids, add tests Signed-off-by: bogdan-petkovic --- mlir/utils/jenkins/run_e2e_multigpu.py | 2 +- mlir/utils/performance/CMakeLists.txt | 1 + mlir/utils/performance/gpu_topology.py | 12 +- .../performance/tests/test_gpu_topology.py | 140 ++++++++++++++++++ 4 files changed, 151 insertions(+), 4 deletions(-) create mode 100644 mlir/utils/performance/tests/test_gpu_topology.py diff --git a/mlir/utils/jenkins/run_e2e_multigpu.py b/mlir/utils/jenkins/run_e2e_multigpu.py index 78ac502898e1..8c48c9f29170 100644 --- a/mlir/utils/jenkins/run_e2e_multigpu.py +++ b/mlir/utils/jenkins/run_e2e_multigpu.py @@ -202,7 +202,7 @@ def main() -> int: 'GPUs, else a single run)') parser.add_argument('--lit-args', type=str, - default='-v --time-tests --timeout=3600', + default='-v --time-tests --timeout=3600 --max-failures=1', help='Extra arguments passed to each lit invocation (default: %(default)r)') parser.add_argument('--dry-run', action='store_true', diff --git a/mlir/utils/performance/CMakeLists.txt b/mlir/utils/performance/CMakeLists.txt index c28085f74d58..e6d2c6ac5e89 100644 --- a/mlir/utils/performance/CMakeLists.txt +++ b/mlir/utils/performance/CMakeLists.txt @@ -10,6 +10,7 @@ set(PERFORMANCE_SCRIPTS convertRocBlasToPerfRunner.py createFusionPerformanceReports.py createPerformanceReports.py + gpu_topology.py perfCommonUtils.py perfRunner.py parameterSweeps.py diff --git a/mlir/utils/performance/gpu_topology.py b/mlir/utils/performance/gpu_topology.py index cfd31c546c42..b42bdd7da486 100644 --- a/mlir/utils/performance/gpu_topology.py +++ b/mlir/utils/performance/gpu_topology.py @@ -66,12 +66,18 @@ def select_gpu_ids( return [None], None, "single GPU detected; running on one GPU" if requested: - selected_archs = {archs[i] for i in requested if 0 <= i < count} + # Preserve order but drop duplicates so we never shard the same device twice. + unique_requested = list(dict.fromkeys(requested)) + invalid = [i for i in unique_requested if not 0 <= i < count] + if invalid: + return [None], None, (f"requested GPU ids {invalid} are out of range " + f"(node has {count} GPU(s)); using the default GPU") + selected_archs = {archs[i] for i in unique_requested} if len(selected_archs) != 1: - return [None], None, (f"requested GPUs {requested} are not a single arch " + return [None], None, (f"requested GPUs {unique_requested} are not a single arch " f"({sorted(selected_archs)}); using the default GPU") arch = next(iter(selected_archs)) - return list(requested), arch, f"using requested GPUs {requested} ({arch})" + return unique_requested, arch, f"using requested GPUs {unique_requested} ({arch})" if len(set(archs)) > 1: return [None], None, (f"mixed GPU architectures ({sorted(set(archs))}); " diff --git a/mlir/utils/performance/tests/test_gpu_topology.py b/mlir/utils/performance/tests/test_gpu_topology.py new file mode 100644 index 000000000000..4e68d164c838 --- /dev/null +++ b/mlir/utils/performance/tests/test_gpu_topology.py @@ -0,0 +1,140 @@ +# Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions. +# See https://llvm.org/LICENSE.txt for license information. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +""" +Tests for gpu_topology.py. + +These cover GPU selection (single/multi/heterogeneous, explicit requests, and +error fallbacks) and per-process GPU isolation. They run in CI without a real +GPU: hip is mocked and get_per_device_archs is monkeypatched per scenario. +""" +import sys +from pathlib import Path + +# Ensure we can import gpu_topology (lives in mlir/utils/performance). +_test_dir = Path(__file__).resolve().parent +_sys_path_parent = str(_test_dir.parent) +if _sys_path_parent not in sys.path: + sys.path.insert(0, _sys_path_parent) +# Mock hip so gpu_topology can be imported/exercised without ROCm (CI has no GPU). +exec( + open(_test_dir / "mock_hip.py").read(), { + "__file__": str(_test_dir / "mock_hip.py"), + "sys": sys + }) + +import gpu_topology # noqa: E402 - must run after mock_hip + + +def _clear_visible_devices(monkeypatch): + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising=False) + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising=False) + + +def _fake_archs(monkeypatch, archs): + monkeypatch.setattr(gpu_topology, "get_per_device_archs", lambda: list(archs)) + + +class TestSelectGpuIds: + """Tests for select_gpu_ids.""" + + def test_respects_preset_rocr_visible_devices(self, monkeypatch): + _clear_visible_devices(monkeypatch) + monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "1") + gpu_ids, arch, msg = gpu_topology.select_gpu_ids() + assert gpu_ids == [None] + assert arch is None + assert "VISIBLE_DEVICES" in msg + + def test_respects_preset_hip_visible_devices(self, monkeypatch): + _clear_visible_devices(monkeypatch) + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "0") + gpu_ids, arch, _ = gpu_topology.select_gpu_ids() + assert gpu_ids == [None] + assert arch is None + + def test_enumeration_failure_falls_back(self, monkeypatch): + _clear_visible_devices(monkeypatch) + + def _boom(): + raise RuntimeError("no driver") + + monkeypatch.setattr(gpu_topology, "get_per_device_archs", _boom) + gpu_ids, arch, msg = gpu_topology.select_gpu_ids() + assert gpu_ids == [None] + assert arch is None + assert "failed" in msg + + def test_single_gpu(self, monkeypatch): + _clear_visible_devices(monkeypatch) + _fake_archs(monkeypatch, ["gfx942"]) + gpu_ids, arch, _ = gpu_topology.select_gpu_ids() + assert gpu_ids == [None] + assert arch is None + + def test_homogeneous_multi_gpu(self, monkeypatch): + _clear_visible_devices(monkeypatch) + _fake_archs(monkeypatch, ["gfx942"] * 8) + gpu_ids, arch, _ = gpu_topology.select_gpu_ids() + assert gpu_ids == [0, 1, 2, 3, 4, 5, 6, 7] + assert arch == "gfx942" + + def test_mixed_archs_no_request_falls_back(self, monkeypatch): + _clear_visible_devices(monkeypatch) + _fake_archs(monkeypatch, ["gfx942", "gfx1100"]) + gpu_ids, arch, msg = gpu_topology.select_gpu_ids() + assert gpu_ids == [None] + assert arch is None + assert "mixed" in msg + + def test_requested_valid_homogeneous_subset(self, monkeypatch): + _clear_visible_devices(monkeypatch) + _fake_archs(monkeypatch, ["gfx942"] * 4) + gpu_ids, arch, _ = gpu_topology.select_gpu_ids(requested=[1, 3]) + assert gpu_ids == [1, 3] + assert arch == "gfx942" + + def test_requested_out_of_range_falls_back(self, monkeypatch): + _clear_visible_devices(monkeypatch) + _fake_archs(monkeypatch, ["gfx942", "gfx942"]) + gpu_ids, arch, msg = gpu_topology.select_gpu_ids(requested=[0, 99]) + assert gpu_ids == [None] + assert arch is None + assert "out of range" in msg + + def test_requested_deduplicated(self, monkeypatch): + _clear_visible_devices(monkeypatch) + _fake_archs(monkeypatch, ["gfx942"] * 4) + gpu_ids, arch, _ = gpu_topology.select_gpu_ids(requested=[2, 2, 0]) + assert gpu_ids == [2, 0] + assert arch == "gfx942" + + def test_requested_mixed_archs_falls_back(self, monkeypatch): + _clear_visible_devices(monkeypatch) + _fake_archs(monkeypatch, ["gfx942", "gfx1100"]) + gpu_ids, arch, msg = gpu_topology.select_gpu_ids(requested=[0, 1]) + assert gpu_ids == [None] + assert arch is None + assert "single arch" in msg + + +class TestMakeIsolatedGpuEnv: + """Tests for make_isolated_gpu_env.""" + + def test_none_returns_none(self): + assert gpu_topology.make_isolated_gpu_env(None) is None + + def test_sets_rocr_and_clears_hip(self, monkeypatch): + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "5") + env = gpu_topology.make_isolated_gpu_env(2) + assert env is not None + assert env["ROCR_VISIBLE_DEVICES"] == "2" + assert "HIP_VISIBLE_DEVICES" not in env + + +class TestGetPerDeviceArchs: + """Tests for get_per_device_archs (exercises the mocked hip path).""" + + def test_returns_arch_per_device(self): + archs = gpu_topology.get_per_device_archs() + assert archs == ["gfx900"] From f244dd04e5e55efcdef320f6b93cc06d97140a50 Mon Sep 17 00:00:00 2001 From: bogdan-petkovic Date: Wed, 29 Jul 2026 04:17:31 -0500 Subject: [PATCH 6/8] [AIROCMLIR-375] Move multi-GPU E2E driver into the performance scripts Signed-off-by: bogdan-petkovic --- mlir/utils/jenkins/Jenkinsfile | 29 +++++------ mlir/utils/performance/CMakeLists.txt | 1 + mlir/utils/performance/gpu_topology.py | 2 +- .../run_e2e_multigpu.py | 52 ++++++------------- 4 files changed, 31 insertions(+), 53 deletions(-) rename mlir/utils/{jenkins => performance}/run_e2e_multigpu.py (76%) diff --git a/mlir/utils/jenkins/Jenkinsfile b/mlir/utils/jenkins/Jenkinsfile index ac8c0f7c325d..25c197c5627e 100644 --- a/mlir/utils/jenkins/Jenkinsfile +++ b/mlir/utils/jenkins/Jenkinsfile @@ -1092,18 +1092,15 @@ int setLitWorkerCount() { return limit_lit_workers } -// Run the rocMLIR lit suite sharded across all GPUs on the node: one lit -// process per GPU, each isolated via ROCR_VISIBLE_DEVICES. On single-GPU or -// heterogeneous nodes the driver falls back to a single lit run, preserving -// legacy behavior. jobsPerGpu is the lit worker count applied to each GPU shard -// (the historical per-GPU cap from setLitWorkerCount()), so per-GPU concurrency -// stays at the proven level instead of being starved by dividing a machine-wide -// total across shards (avoiding the oversubscription issues #1845/#1841). Run via -// shStrict so the driver's stdout is mirrored to the per-row log, letting -// withHealthyNode classify GPU hangs and retry only this node. +// Run the lit suite sharded across the node's GPUs: one lit process per GPU, +// each pinned via ROCR_VISIBLE_DEVICES; single-GPU and heterogeneous nodes fall +// back to one run. jobsPerGpu keeps the per-GPU cap from setLitWorkerCount() +// instead of splitting it across shards (see #1845/#1841). shStrict mirrors the +// output into the per-row log so withHealthyNode can retry just this node. +// ci-performance-scripts copies the driver and gpu_topology.py into ./bin. void runShardedE2E(int jobsPerGpu) { dir('build') { - shStrict "python3 ${env.WORKSPACE}/mlir/utils/jenkins/run_e2e_multigpu.py --build-dir . " + + shStrict "python3 ./bin/run_e2e_multigpu.py --build-dir . " + "--jobs-per-gpu ${jobsPerGpu} " + "\"--lit-args=-v --time-tests --timeout=3600 --max-failures=1\"" } @@ -1112,7 +1109,7 @@ void runShardedE2E(int jobsPerGpu) { void build_fixedE2ETests(String codepath) { // Limit the number of lit workers for gfx908, gfx90a to (8, 30) on CI as a workaround for issue #1845 and #1841 int limit_lit_workers = setLitWorkerCount() - buildProject("check-mlir-build-only check-rocmlir-build-only${params.nightly ? ' hipblaslt-benchmark-driver' : ''}", """ + buildProject("check-mlir-build-only check-rocmlir-build-only ci-performance-scripts${params.nightly ? ' hipblaslt-benchmark-driver' : ''}", """ -DROCMLIR_DRIVER_PR_E2E_TEST_ENABLED=${params.nightly ? '0' : '1'} -DROCMLIR_DRIVER_E2E_TEST_ENABLED=${params.nightly ? '1' : '0'} -DROCK_E2E_TEST_ENABLED=${params.nightly ? '1' : '0'} @@ -1126,12 +1123,10 @@ void build_fixedE2ETests(String codepath) { void check_randomE2ETests(String codepath) { // Limit the number of lit workers for gfx908, gfx90a to (8, 30) on CI as a workaround for issue #1845 and #1841 int limit_lit_workers = setLitWorkerCount() - // Configure and build the E2E deps without running the tests, then run the lit suite via the - // multi-GPU sharding driver. The driver uses shStrict so its stdout is mirrored to the per-row - // log (withHealthyNode classifies GPU hangs there and retries only this node). Running - // check-rocmlir directly through cmakeBuild would bypass shStrict and force a whole-job - // re-kick instead. - buildProject('check-rocmlir-build-only', """ + // Build the E2E deps without running the tests, then run the suite through the sharding + // driver: going through cmakeBuild would bypass shStrict and force a whole-job re-kick + // instead of a single-node retry. + buildProject('check-rocmlir-build-only ci-performance-scripts', """ -DROCMLIR_DRIVER_PR_E2E_TEST_ENABLED=0 -DROCMLIR_DRIVER_E2E_TEST_ENABLED=1 -DROCK_E2E_TEST_ENABLED=1 diff --git a/mlir/utils/performance/CMakeLists.txt b/mlir/utils/performance/CMakeLists.txt index e6d2c6ac5e89..c0834b09e2fb 100644 --- a/mlir/utils/performance/CMakeLists.txt +++ b/mlir/utils/performance/CMakeLists.txt @@ -17,6 +17,7 @@ set(PERFORMANCE_SCRIPTS attentionSweeps.py perfRegressionReport.py reportUtils.py + run_e2e_multigpu.py tuningRunner.py rocmlir_metrics.txt handleNewConfigs.py diff --git a/mlir/utils/performance/gpu_topology.py b/mlir/utils/performance/gpu_topology.py index b42bdd7da486..1dbc469f0b3f 100644 --- a/mlir/utils/performance/gpu_topology.py +++ b/mlir/utils/performance/gpu_topology.py @@ -66,7 +66,7 @@ def select_gpu_ids( return [None], None, "single GPU detected; running on one GPU" if requested: - # Preserve order but drop duplicates so we never shard the same device twice. + # Never shard the same device twice. unique_requested = list(dict.fromkeys(requested)) invalid = [i for i in unique_requested if not 0 <= i < count] if invalid: diff --git a/mlir/utils/jenkins/run_e2e_multigpu.py b/mlir/utils/performance/run_e2e_multigpu.py similarity index 76% rename from mlir/utils/jenkins/run_e2e_multigpu.py rename to mlir/utils/performance/run_e2e_multigpu.py index 8c48c9f29170..dc63397effd0 100644 --- a/mlir/utils/jenkins/run_e2e_multigpu.py +++ b/mlir/utils/performance/run_e2e_multigpu.py @@ -4,16 +4,11 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception """Run the rocMLIR lit test suite in parallel across all GPUs on a node. -CI nodes commonly expose several GPUs but the lit suite (e.g. `check-rocmlir`) -historically runs every test on GPU 0. This driver shards the suite with lit's -native `--num-shards`/`--run-shard` mechanism, launching one lit process per GPU -and isolating each to its device via ROCR_VISIBLE_DEVICES. The total set of -tests is partitioned, so per-GPU concurrency stays bounded (avoiding the -oversubscription hangs seen with a single global `-j`) while all GPUs are used. - -Distribution is only enabled on homogeneous nodes (all GPUs share one gfx -architecture); single-GPU and heterogeneous nodes fall back to one lit run, -preserving today's behavior. +CI nodes often expose several GPUs, but the lit suite historically runs every +test on GPU 0. This driver partitions the suite with lit's `--num-shards` / +`--run-shard`, running one lit process per GPU, each pinned to its device via +ROCR_VISIBLE_DEVICES. Sharding is only used on homogeneous nodes; single-GPU and +mixed-architecture nodes fall back to a single lit run. """ from __future__ import annotations @@ -25,9 +20,7 @@ import time from typing import List, Optional -# Reuse the GPU detection helper from the performance scripts. -sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'performance')) -from gpu_topology import select_gpu_ids # noqa: E402 +from gpu_topology import make_isolated_gpu_env, select_gpu_ids def default_lit_path(build_dir: str) -> str: @@ -48,10 +41,8 @@ def build_shard_command(lit: str, lit_args: List[str], jobs: int, num_shards: in def resolve_jobs_per_shard(args: argparse.Namespace, num_shards: int) -> int: """Pick the lit worker count for each shard. - `--jobs-per-gpu` wins if given. Otherwise `--total-jobs` is split evenly - across shards, so a single-GPU node keeps the legacy global concurrency while - multi-GPU nodes keep per-GPU pressure bounded (total/num_shards). Falls back - to a conservative default when neither is given. + `--jobs-per-gpu` caps concurrency per GPU; `--total-jobs` instead splits a + machine-wide budget across the shards. """ if args.jobs_per_gpu is not None: return max(1, args.jobs_per_gpu) @@ -67,24 +58,19 @@ def run(args: argparse.Namespace) -> int: gpu_ids, _gpu_arch, gpu_msg = select_gpu_ids(args.gpus) print(f"[run_e2e_multigpu] {gpu_msg}", flush=True) - # `select_gpu_ids` returns [None] for the single-GPU / heterogeneous / unknown - # cases; treat those as one un-pinned lit run (legacy behavior). - single_gpu = gpu_ids == [None] - shard_gpus: List[Optional[int]] = [None] if single_gpu else gpu_ids + # [None] means one un-pinned lit run (single-GPU / heterogeneous nodes). + shard_gpus: List[Optional[int]] = gpu_ids num_shards = len(shard_gpus) jobs_per_shard = resolve_jobs_per_shard(args, num_shards) print(f"[run_e2e_multigpu] {num_shards} shard(s), {jobs_per_shard} lit workers each", flush=True) - # Single shard: stream lit output straight to the console (live per-test - # progress, keeps CI activity timeouts alive). Multi-shard buffers per-GPU. + # A single shard streams straight to the console so CI sees live progress; + # multiple shards are buffered per-GPU and dumped once they finish. if num_shards == 1: gpu_id = shard_gpus[0] cmd = build_shard_command(lit, args.lit_args, jobs_per_shard, 1, 1, test_paths) - env = os.environ.copy() - if gpu_id is not None: - env['ROCR_VISIBLE_DEVICES'] = str(gpu_id) - env.pop('HIP_VISIBLE_DEVICES', None) + env = make_isolated_gpu_env(gpu_id) label = f"GPU {gpu_id}" if gpu_id is not None else "single" print(f"[run_e2e_multigpu] shard 1/1 on {label}: {' '.join(cmd)}", flush=True) if args.dry_run: @@ -96,10 +82,7 @@ def run(args: argparse.Namespace) -> int: for idx, gpu_id in enumerate(shard_gpus): cmd = build_shard_command(lit, args.lit_args, jobs_per_shard, num_shards, idx + 1, test_paths) - env = os.environ.copy() - if gpu_id is not None: - env['ROCR_VISIBLE_DEVICES'] = str(gpu_id) - env.pop('HIP_VISIBLE_DEVICES', None) + env = make_isolated_gpu_env(gpu_id) label = f"GPU {gpu_id}" if gpu_id is not None else "single" print(f"[run_e2e_multigpu] shard {idx + 1}/{num_shards} on {label}: {' '.join(cmd)}", flush=True) @@ -117,8 +100,8 @@ def run(args: argparse.Namespace) -> int: failures = [] aborted = [] pending = list(range(len(procs))) - # Heartbeat so the console keeps emitting output during the otherwise-silent - # buffered run, preventing Jenkins `timeout(activity: true)` from firing. + # Heartbeat: keep the console alive during the buffered run so Jenkins' + # timeout(activity: true) does not fire. start = time.time() last_beat = start heartbeat_secs = 30 @@ -141,8 +124,7 @@ def run(args: argparse.Namespace) -> int: if rc != 0: failures.append((label, rc)) - # Fail-fast: once any shard fails, terminate the rest so the run aborts - # promptly (matching the previous single-lit --max-failures=1 behavior). + # Once a shard fails, stop the rest so the run aborts promptly. if args.fail_fast and failures and pending: for i in pending: procs[i][2].terminate() From c9e14cdde65bca1a7d8dc8bf0148d25968bccc16 Mon Sep 17 00:00:00 2001 From: bogdan-petkovic Date: Mon, 3 Aug 2026 06:25:53 -0500 Subject: [PATCH 7/8] [AIROCMLIR-375] Make GPU pinning reach lit tests and cap host concurrency Signed-off-by: bogdan-petkovic --- mlir/test/e2e/lit.cfg.py | 4 +- mlir/test/fusion/e2e/lit.cfg.py | 4 +- mlir/test/lit.cfg.py | 4 +- mlir/utils/performance/gpu_topology.py | 189 ++++++++++++++++-- mlir/utils/performance/run_e2e_multigpu.py | 30 ++- .../performance/tests/test_gpu_topology.py | 81 ++++++++ .../tests/test_run_e2e_multigpu.py | 80 ++++++++ .../performance/tests/test_tuningRunner.py | 6 +- mlir/utils/performance/tuningRunner.py | 151 +------------- 9 files changed, 387 insertions(+), 162 deletions(-) create mode 100644 mlir/utils/performance/tests/test_run_e2e_multigpu.py diff --git a/mlir/test/e2e/lit.cfg.py b/mlir/test/e2e/lit.cfg.py index 43e5c9718e49..756e9223a151 100644 --- a/mlir/test/e2e/lit.cfg.py +++ b/mlir/test/e2e/lit.cfg.py @@ -38,8 +38,10 @@ config.substitutions.append(('%arch', config.arch)) config.substitutions.append(('%pv', config.populate_validation)) +# lit builds a clean environment for tests, so ROCR_VISIBLE_DEVICES has to be listed +# explicitly or the per-GPU pinning from run_e2e_multigpu.py never reaches them. llvm_config.with_system_environment( - ['HOME', 'INCLUDE', 'LIB', 'TMP', 'TEMP', 'HIP_VISIBLE_DEVICES']) + ['HOME', 'INCLUDE', 'LIB', 'TMP', 'TEMP', 'HIP_VISIBLE_DEVICES', 'ROCR_VISIBLE_DEVICES']) # When multiple GPUs are present, limit HIP to device 0 to ensure # compiled binaries match the execution device diff --git a/mlir/test/fusion/e2e/lit.cfg.py b/mlir/test/fusion/e2e/lit.cfg.py index cebdae4eced2..033bcd019943 100644 --- a/mlir/test/fusion/e2e/lit.cfg.py +++ b/mlir/test/fusion/e2e/lit.cfg.py @@ -36,8 +36,10 @@ config.substitutions.append(('%arch', config.arch)) config.substitutions.append(('%pv', config.populate_validation)) +# lit builds a clean environment for tests, so ROCR_VISIBLE_DEVICES has to be listed +# explicitly or the per-GPU pinning from run_e2e_multigpu.py never reaches them. llvm_config.with_system_environment( - ['HOME', 'INCLUDE', 'LIB', 'TMP', 'TEMP', 'HIP_VISIBLE_DEVICES']) + ['HOME', 'INCLUDE', 'LIB', 'TMP', 'TEMP', 'HIP_VISIBLE_DEVICES', 'ROCR_VISIBLE_DEVICES']) # When multiple GPUs are present, limit HIP to device 0 to ensure # compiled binaries match the execution device diff --git a/mlir/test/lit.cfg.py b/mlir/test/lit.cfg.py index ecf63326d695..f4943c374f29 100644 --- a/mlir/test/lit.cfg.py +++ b/mlir/test/lit.cfg.py @@ -39,8 +39,10 @@ config.substitutions.append(('%arch', config.arch)) config.substitutions.append(('%pv', config.populate_validation)) +# lit builds a clean environment for tests, so ROCR_VISIBLE_DEVICES has to be listed +# explicitly or the per-GPU pinning from run_e2e_multigpu.py never reaches them. llvm_config.with_system_environment( - ['HOME', 'INCLUDE', 'LIB', 'TMP', 'TEMP', 'HIP_VISIBLE_DEVICES']) + ['HOME', 'INCLUDE', 'LIB', 'TMP', 'TEMP', 'HIP_VISIBLE_DEVICES', 'ROCR_VISIBLE_DEVICES']) # When multiple GPUs are present, limit HIP to device 0 to ensure # compiled binaries match the execution device diff --git a/mlir/utils/performance/gpu_topology.py b/mlir/utils/performance/gpu_topology.py index 1dbc469f0b3f..4f75750eb590 100644 --- a/mlir/utils/performance/gpu_topology.py +++ b/mlir/utils/performance/gpu_topology.py @@ -1,18 +1,25 @@ # Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions. # See https://llvm.org/LICENSE.txt for license information. # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -"""Shared helpers for distributing GPU test workloads across multiple devices. +"""Shared GPU and NUMA topology helpers for distributing work across devices. -CI nodes commonly expose several GPUs. These helpers detect the visible GPUs, -confirm they share a single architecture (a prerequisite for safely reusing the -same compiled kernels), and build per-process environments that isolate work to -one device via ROCR_VISIBLE_DEVICES (the same mechanism tuningRunner.py uses). +Collects the primitives used by the tuning runner, the parameter sweeps and the +sharded E2E driver: discovering the GPUs on a node, deciding which of them to +spread work over, isolating a child process to one device, and splitting the +host's CPUs between them. + +This module deliberately keeps its imports light (stdlib only at module scope, +hip-python loaded on demand) so that callers which merely want to fall back to a +single GPU do not need a working ROCm stack to start up. """ from __future__ import annotations +import json import os -from typing import List, Optional, Tuple +import subprocess +from dataclasses import dataclass +from typing import Dict, List, Optional, Set, Tuple def _hip_check(call_result): @@ -30,8 +37,8 @@ def _hip_check(call_result): def get_per_device_archs() -> List[str]: """Return the gfx architecture name of every visible GPU, indexed by id. - Uses hip-python, matching the rest of the test tooling (the lit configs and - perfRunner already require it); callers treat a failure as "use one GPU". + Uses hip-python rather than rocm-smi because only HIP reports the gfx name; + rocm-smi reports the card SKU, which is not what kernels are compiled for. """ from hip import hip archs = [] @@ -43,6 +50,109 @@ def get_per_device_archs() -> List[str]: return archs +@dataclass(frozen=True) +class Gpu: + """Information about a GPU.""" + gpu_id: int + sku: str + numa_node: int + + +@dataclass(frozen=True) +class GpuTopology: + """System GPU topology with NUMA mappings.""" + gpus: Dict[int, Gpu] # GPU ID -> Gpu + + def get_numa_node(self, gpu_id: int) -> int: + """Get NUMA node for a GPU.""" + return self.gpus[gpu_id].numa_node + + def validate_homogeneity(self, gpu_ids: List[int]) -> bool: + """Validate that all selected GPUs are of the same model.""" + if len(gpu_ids) <= 1: + return True + + skus = {self.gpus[gpu_id].sku for gpu_id in gpu_ids} + return len(skus) == 1 + + @staticmethod + def discover() -> 'GpuTopology': + """Query GPU topology using rocm-smi. + + rocm-smi reports physical device IDs regardless of environment variables (e.g., ROCR_VISIBLE_DEVICES and HIP_VISIBLE_DEVICES). + """ + # rocm-smi can take ~20s to enumerate large multi-GPU systems, so allow + # a generous timeout to avoid spurious TimeoutExpired failures. + output = subprocess.check_output( + ["rocm-smi", "--showproductname", "--showtoponuma", "--json"], + text=True, + stderr=subprocess.DEVNULL, + timeout=60) + data = json.loads(output) + + gpus = {} + for key, value in data.items(): + if key.startswith("card"): + gpu_id = int(key.replace("card", "")) + + sku = value["Card SKU"] + + numa_node_str = value.get("(Topology) Numa Node") + numa_node = int(numa_node_str) if numa_node_str is not None else 0 + + gpus[gpu_id] = Gpu(gpu_id=gpu_id, sku=sku, numa_node=numa_node) + + if not gpus: + raise RuntimeError("rocm-smi returned no GPU cards") + + return GpuTopology(gpus=gpus) + + +@dataclass(frozen=True) +class NumaTopology: + """System NUMA topology with CPU mappings.""" + numa_to_cpus: Dict[int, List[int]] # NUMA node -> list of CPU IDs + + def get_cpus_for_numa_node(self, numa_node: int) -> List[int]: + """Get CPUs belonging to a NUMA node.""" + return self.numa_to_cpus[numa_node] + + @staticmethod + def discover() -> 'NumaTopology': + """Discover NUMA topology for CPUs. + + Returns a topology where all CPUs are on node 0 if discovery fails or system is non-NUMA. + """ + numa_to_cpus: Dict[int, List[int]] = {} + numa_base = "/sys/devices/system/node" + + if os.path.exists(numa_base): + for entry in os.listdir(numa_base): + if entry.startswith("node") and entry[4:].isdigit(): + node_id = int(entry[4:]) + cpulist_path = os.path.join(numa_base, entry, "cpulist") + with open(cpulist_path, 'r') as f: + numa_to_cpus[node_id] = NumaTopology._parse_cpu_list(f.read()) + + # Fallback: single node with all CPUs + if not numa_to_cpus: + numa_to_cpus[0] = list(range(os.cpu_count() or 1)) + + return NumaTopology(numa_to_cpus=numa_to_cpus) + + @staticmethod + def _parse_cpu_list(cpu_list_str: str) -> List[int]: + """Parse CPU list string like '0-55,112-167' into list of CPU IDs.""" + cpus = [] + for part in cpu_list_str.strip().split(','): + if '-' in part: + start, end = part.split('-', 1) + cpus.extend(range(int(start), int(end) + 1)) + else: + cpus.append(int(part)) + return cpus + + def select_gpu_ids( requested: Optional[List[int]] = None) -> Tuple[List[Optional[int]], Optional[str], str]: """Decide which GPUs to spread work across. @@ -86,17 +196,68 @@ def select_gpu_ids( return list(range(count)), archs[0], f"distributing across {count} GPUs ({archs[0]})" -def make_isolated_gpu_env(gpu_id: Optional[int]) -> Optional[dict]: - """Build an environment dict that isolates a child process to one GPU. +def set_isolated_gpu_env(env: Dict[str, str], gpu_id: int) -> None: + """Modify environment to isolate subprocess to one physical GPU. + + Sets ROCR_VISIBLE_DEVICES at the HSA/ROCr level, providing complete isolation for all higher layers including HIP. + """ + env["ROCR_VISIBLE_DEVICES"] = str(gpu_id) + env.pop("HIP_VISIBLE_DEVICES", None) # Remove HIP_VISIBLE_DEVICES to avoid conflicts + + +def make_isolated_gpu_env(gpu_id: Optional[int]) -> Optional[Dict[str, str]]: + """Create environment that isolates subprocess to one physical GPU. Returns ``None`` when ``gpu_id`` is ``None`` so callers can pass it straight through to ``subprocess``/``asyncio`` APIs to mean "inherit the environment". - ROCR_VISIBLE_DEVICES isolates at the HSA/ROCr level (below HIP); we clear - HIP_VISIBLE_DEVICES to avoid the two layers disagreeing. """ if gpu_id is None: return None env = os.environ.copy() - env["ROCR_VISIBLE_DEVICES"] = str(gpu_id) - env.pop("HIP_VISIBLE_DEVICES", None) + set_isolated_gpu_env(env, gpu_id) return env + + +def _usable_cpus() -> Set[int]: + """CPUs this process is allowed to run on (respects cpuset restrictions).""" + try: + return set(os.sched_getaffinity(0)) + except AttributeError: # sched_getaffinity is Linux-only + return set(range(os.cpu_count() or 1)) + + +def usable_cpu_count() -> int: + """Number of CPUs this process is allowed to run on.""" + return len(_usable_cpus()) + + +def allocate_cpus_per_gpu(gpu_ids: List[int], gpu_topology: GpuTopology, + numa_topology: NumaTopology) -> Dict[int, int]: + """Split the host's CPUs across ``gpu_ids``, keeping each GPU on its own NUMA node. + + CPUs outside this process' affinity mask are left out, so a container that + was given a slice of the machine does not hand out threads it cannot use. + """ + usable = _usable_cpus() + gpus_by_node: Dict[int, List[int]] = {} + for gpu_id in gpu_ids: + gpus_by_node.setdefault(gpu_topology.get_numa_node(gpu_id), []).append(gpu_id) + + allocation: Dict[int, int] = {} + for node, gpus_on_node in gpus_by_node.items(): + cpus_on_node = len(usable.intersection(numa_topology.get_cpus_for_numa_node(node))) + threads_each = max(1, cpus_on_node // len(gpus_on_node)) + for gpu_id in gpus_on_node: + allocation[gpu_id] = threads_each + + return allocation + + +def scale_cpu_allocation(allocation: Dict[int, int], limit: int) -> Dict[int, int]: + """Scale ``allocation`` down proportionally so its total stays within ``limit``.""" + total = sum(allocation.values()) + if total <= limit: + return dict(allocation) + + scale = limit / total + return {gpu_id: max(1, int(count * scale)) for gpu_id, count in allocation.items()} diff --git a/mlir/utils/performance/run_e2e_multigpu.py b/mlir/utils/performance/run_e2e_multigpu.py index dc63397effd0..aac6455ce4a0 100644 --- a/mlir/utils/performance/run_e2e_multigpu.py +++ b/mlir/utils/performance/run_e2e_multigpu.py @@ -20,7 +20,7 @@ import time from typing import List, Optional -from gpu_topology import make_isolated_gpu_env, select_gpu_ids +from gpu_topology import make_isolated_gpu_env, select_gpu_ids, usable_cpu_count def default_lit_path(build_dir: str) -> str: @@ -51,6 +51,20 @@ def resolve_jobs_per_shard(args: argparse.Namespace, num_shards: int) -> int: return 8 +def cap_jobs_to_host_cpus(jobs_per_shard: int, + num_shards: int, + budget: Optional[int] = None) -> int: + """Clamp the per-shard worker count so all shards together fit the host. + + The per-GPU caps exist to protect the GPU, so multiplying one by the number of + GPUs can ask for far more parallelism than the machine has cores. `budget` + defaults to the CPUs this process may run on. + """ + if budget is None: + budget = usable_cpu_count() + return max(1, min(jobs_per_shard, budget // num_shards)) + + def run(args: argparse.Namespace) -> int: lit = args.lit or default_lit_path(args.build_dir) test_paths = args.test_paths or [os.path.join(args.build_dir, 'mlir', 'test')] @@ -61,7 +75,14 @@ def run(args: argparse.Namespace) -> int: # [None] means one un-pinned lit run (single-GPU / heterogeneous nodes). shard_gpus: List[Optional[int]] = gpu_ids num_shards = len(shard_gpus) - jobs_per_shard = resolve_jobs_per_shard(args, num_shards) + requested_jobs = resolve_jobs_per_shard(args, num_shards) + jobs_per_shard = cap_jobs_to_host_cpus(requested_jobs, num_shards, args.max_total_jobs) + if jobs_per_shard < requested_jobs: + budget = args.max_total_jobs if args.max_total_jobs is not None else usable_cpu_count() + print( + f"[run_e2e_multigpu] capping {requested_jobs} to {jobs_per_shard} workers per shard " + f"to stay within {budget} host CPU(s)", + flush=True) print(f"[run_e2e_multigpu] {num_shards} shard(s), {jobs_per_shard} lit workers each", flush=True) @@ -176,6 +197,11 @@ def main() -> int: type=int, default=None, help='Total lit workers split evenly across shards (default per-shard: 8)') + parser.add_argument('--max-total-jobs', + type=int, + default=None, + help='Upper bound on lit workers across all shards ' + '(default: the number of CPUs this process may use)') parser.add_argument('--gpus', type=int, nargs='+', diff --git a/mlir/utils/performance/tests/test_gpu_topology.py b/mlir/utils/performance/tests/test_gpu_topology.py index 4e68d164c838..e217ad2b5efe 100644 --- a/mlir/utils/performance/tests/test_gpu_topology.py +++ b/mlir/utils/performance/tests/test_gpu_topology.py @@ -138,3 +138,84 @@ class TestGetPerDeviceArchs: def test_returns_arch_per_device(self): archs = gpu_topology.get_per_device_archs() assert archs == ["gfx900"] + + +def _topologies(gpu_to_node, node_to_cpus): + """Build a (GpuTopology, NumaTopology) pair from plain mappings.""" + gpus = { + gpu_id: gpu_topology.Gpu(gpu_id=gpu_id, sku="mock-sku", numa_node=node) + for gpu_id, node in gpu_to_node.items() + } + return (gpu_topology.GpuTopology(gpus=gpus), + gpu_topology.NumaTopology(numa_to_cpus=node_to_cpus)) + + +class TestGpuTopology: + """Tests for GpuTopology (moved here from tuningRunner).""" + + def test_get_numa_node(self): + gpus, _ = _topologies({0: 0, 1: 1}, {0: [0], 1: [1]}) + assert gpus.get_numa_node(1) == 1 + + def test_homogeneity_single_gpu_is_always_homogeneous(self): + gpus, _ = _topologies({0: 0}, {0: [0]}) + assert gpus.validate_homogeneity([0]) + + def test_homogeneity_mixed_skus(self): + gpus = gpu_topology.GpuTopology( + gpus={ + 0: gpu_topology.Gpu(gpu_id=0, sku="a", numa_node=0), + 1: gpu_topology.Gpu(gpu_id=1, sku="b", numa_node=0), + }) + assert not gpus.validate_homogeneity([0, 1]) + + +class TestAllocateCpusPerGpu: + """Tests for allocate_cpus_per_gpu.""" + + def test_splits_node_cpus_between_gpus_on_that_node(self, monkeypatch): + monkeypatch.setattr(gpu_topology, "_usable_cpus", lambda: set(range(16))) + gpus, numa = _topologies({0: 0, 1: 0}, {0: list(range(16))}) + assert gpu_topology.allocate_cpus_per_gpu([0, 1], gpus, numa) == {0: 8, 1: 8} + + def test_gpus_on_separate_nodes_get_their_own_cpus(self, monkeypatch): + monkeypatch.setattr(gpu_topology, "_usable_cpus", lambda: set(range(16))) + gpus, numa = _topologies({0: 0, 1: 1}, {0: list(range(8)), 1: list(range(8, 16))}) + assert gpu_topology.allocate_cpus_per_gpu([0, 1], gpus, numa) == {0: 8, 1: 8} + + def test_respects_affinity_mask(self, monkeypatch): + # Only 4 of the node's 16 CPUs are usable (e.g. a container CPU limit). + monkeypatch.setattr(gpu_topology, "_usable_cpus", lambda: set(range(4))) + gpus, numa = _topologies({0: 0, 1: 0}, {0: list(range(16))}) + assert gpu_topology.allocate_cpus_per_gpu([0, 1], gpus, numa) == {0: 2, 1: 2} + + def test_never_allocates_zero(self, monkeypatch): + monkeypatch.setattr(gpu_topology, "_usable_cpus", lambda: {0}) + gpus, numa = _topologies({0: 0, 1: 0, 2: 0, 3: 0}, {0: [0]}) + assert gpu_topology.allocate_cpus_per_gpu([0, 1, 2, 3], gpus, numa) == { + 0: 1, + 1: 1, + 2: 1, + 3: 1 + } + + +class TestScaleCpuAllocation: + """Tests for scale_cpu_allocation.""" + + def test_limit_above_total_is_unchanged(self): + assert gpu_topology.scale_cpu_allocation({0: 8, 1: 8}, 32) == {0: 8, 1: 8} + + def test_scales_down_proportionally(self): + assert gpu_topology.scale_cpu_allocation({0: 8, 1: 8}, 8) == {0: 4, 1: 4} + + def test_never_scales_below_one(self): + assert gpu_topology.scale_cpu_allocation({0: 8, 1: 8}, 1) == {0: 1, 1: 1} + + +class TestUsableCpuCount: + """Tests for usable_cpu_count.""" + + def test_matches_affinity_mask(self, monkeypatch): + monkeypatch.setattr(gpu_topology, "_usable_cpus", lambda: {0, 1, 2}) + assert gpu_topology.usable_cpu_count() == 3 diff --git a/mlir/utils/performance/tests/test_run_e2e_multigpu.py b/mlir/utils/performance/tests/test_run_e2e_multigpu.py new file mode 100644 index 000000000000..c850e66a257b --- /dev/null +++ b/mlir/utils/performance/tests/test_run_e2e_multigpu.py @@ -0,0 +1,80 @@ +# Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions. +# See https://llvm.org/LICENSE.txt for license information. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +""" +Tests for run_e2e_multigpu.py. + +These cover the pure decisions the driver makes before launching lit: how many +workers each shard gets, how that is clamped to the host's CPUs, and what the +per-shard lit command line looks like. No GPU or build tree is required. +""" +import argparse +import sys +from pathlib import Path + +# Ensure we can import run_e2e_multigpu (lives in mlir/utils/performance). +_test_dir = Path(__file__).resolve().parent +_sys_path_parent = str(_test_dir.parent) +if _sys_path_parent not in sys.path: + sys.path.insert(0, _sys_path_parent) + +import run_e2e_multigpu # noqa: E402 + + +def _args(jobs_per_gpu=None, total_jobs=None): + return argparse.Namespace(jobs_per_gpu=jobs_per_gpu, total_jobs=total_jobs) + + +class TestResolveJobsPerShard: + """Tests for resolve_jobs_per_shard.""" + + def test_jobs_per_gpu_wins(self): + assert run_e2e_multigpu.resolve_jobs_per_shard(_args(jobs_per_gpu=20, total_jobs=8), + 4) == 20 + + def test_total_jobs_is_split_across_shards(self): + assert run_e2e_multigpu.resolve_jobs_per_shard(_args(total_jobs=32), 4) == 8 + + def test_default_when_neither_given(self): + assert run_e2e_multigpu.resolve_jobs_per_shard(_args(), 4) == 8 + + def test_never_returns_zero(self): + assert run_e2e_multigpu.resolve_jobs_per_shard(_args(total_jobs=2), 8) == 1 + assert run_e2e_multigpu.resolve_jobs_per_shard(_args(jobs_per_gpu=0), 1) == 1 + + +class TestCapJobsToHostCpus: + """Tests for cap_jobs_to_host_cpus.""" + + def test_no_cap_when_host_has_room(self): + assert run_e2e_multigpu.cap_jobs_to_host_cpus(20, 2, budget=256) == 20 + + def test_caps_when_shards_would_oversubscribe_the_host(self): + # 8 shards x 64 workers = 512 requested, but only 128 CPUs are available. + assert run_e2e_multigpu.cap_jobs_to_host_cpus(64, 8, budget=128) == 16 + + def test_never_caps_below_one(self): + assert run_e2e_multigpu.cap_jobs_to_host_cpus(64, 8, budget=4) == 1 + + def test_defaults_to_process_cpu_budget(self, monkeypatch): + monkeypatch.setattr(run_e2e_multigpu, "usable_cpu_count", lambda: 16) + assert run_e2e_multigpu.cap_jobs_to_host_cpus(64, 4) == 4 + + +class TestBuildShardCommand: + """Tests for build_shard_command.""" + + def test_single_shard_has_no_sharding_flags(self): + cmd = run_e2e_multigpu.build_shard_command('llvm-lit', ['-v'], 8, 1, 1, ['tests']) + assert cmd == [sys.executable, 'llvm-lit', '-j', '8', '-v', 'tests'] + + def test_multi_shard_passes_shard_flags(self): + cmd = run_e2e_multigpu.build_shard_command('llvm-lit', ['-v'], 20, 4, 3, ['tests']) + assert cmd[cmd.index('--num-shards') + 1] == '4' + assert cmd[cmd.index('--run-shard') + 1] == '3' + assert cmd[cmd.index('-j') + 1] == '20' + + def test_lit_args_and_paths_come_last(self): + cmd = run_e2e_multigpu.build_shard_command('llvm-lit', ['-v', '--time-tests'], 8, 1, 1, + ['a', 'b']) + assert cmd[-4:] == ['-v', '--time-tests', 'a', 'b'] diff --git a/mlir/utils/performance/tests/test_tuningRunner.py b/mlir/utils/performance/tests/test_tuningRunner.py index ba2b4f76e6ce..7d1aaed493cd 100644 --- a/mlir/utils/performance/tests/test_tuningRunner.py +++ b/mlir/utils/performance/tests/test_tuningRunner.py @@ -28,11 +28,11 @@ }) import tuningRunner # noqa: E402 - must run after mock_hip +from gpu_topology import NumaTopology # noqa: E402 from tuningRunner import ( # noqa: E402 ConfigState, TuningState, TuningStateFile, TunedConfigsCache, Options, get_state_filepath, - verify_mode_flags, format_error, get_config_class, get_git_commit_hash, NumaTopology, Operation, - NumaNodeLock, resolve_verify_mode, canonicalize_test_vector, DebugFileWriter, TuningResult, - tune_config) + verify_mode_flags, format_error, get_config_class, get_git_commit_hash, Operation, NumaNodeLock, + resolve_verify_mode, canonicalize_test_vector, DebugFileWriter, TuningResult, tune_config) from perfRunner import ( # noqa: E402 GemmConfiguration, ConvConfiguration, AttentionConfiguration, ConvGemmConfiguration, GemmGemmConfiguration, PerfConfiguration, canonicalize_config) diff --git a/mlir/utils/performance/tuningRunner.py b/mlir/utils/performance/tuningRunner.py index 15eced61f0f8..89c28745025c 100755 --- a/mlir/utils/performance/tuningRunner.py +++ b/mlir/utils/performance/tuningRunner.py @@ -50,6 +50,14 @@ from tqdm import tqdm import perfRunner +from gpu_topology import ( + GpuTopology, + NumaTopology, + allocate_cpus_per_gpu, + make_isolated_gpu_env, + scale_cpu_allocation, + set_isolated_gpu_env, +) from perfCommonUtils import CORRECT_RESULT_RE, Operation from perfRunner import ( AttentionConfiguration, @@ -224,114 +232,6 @@ class TuningError(Exception): pass -# ============================================================================= -# System Topology Discovery -# ============================================================================= - - -@dataclass(frozen=True) -class Gpu: - """Information about a GPU.""" - gpu_id: int - sku: str - numa_node: int - - -@dataclass(frozen=True) -class GpuTopology: - """System GPU topology with NUMA mappings.""" - gpus: Dict[int, Gpu] # GPU ID -> Gpu - - def get_numa_node(self, gpu_id: int) -> int: - """Get NUMA node for a GPU.""" - return self.gpus[gpu_id].numa_node - - def validate_homogeneity(self, gpu_ids: List[int]) -> bool: - """Validate that all selected GPUs are of the same model.""" - if len(gpu_ids) <= 1: - return True - - skus = {self.gpus[gpu_id].sku for gpu_id in gpu_ids} - return len(skus) == 1 - - @staticmethod - def discover() -> 'GpuTopology': - """Query GPU topology using rocm-smi. - - rocm-smi reports physical device IDs regardless of environment variables (e.g., ROCR_VISIBLE_DEVICES and HIP_VISIBLE_DEVICES). - """ - # rocm-smi can take ~20s to enumerate large multi-GPU systems, so allow - # a generous timeout to avoid spurious TimeoutExpired failures. - output = subprocess.check_output( - ["rocm-smi", "--showproductname", "--showtoponuma", "--json"], - text=True, - stderr=subprocess.DEVNULL, - timeout=60) - data = json.loads(output) - - gpus = {} - for key, value in data.items(): - if key.startswith("card"): - gpu_id = int(key.replace("card", "")) - - sku = value["Card SKU"] - - numa_node_str = value.get("(Topology) Numa Node") - numa_node = int(numa_node_str) if numa_node_str is not None else 0 - - gpus[gpu_id] = Gpu(gpu_id=gpu_id, sku=sku, numa_node=numa_node) - - if not gpus: - raise RuntimeError("rocm-smi returned no GPU cards") - - return GpuTopology(gpus=gpus) - - -@dataclass(frozen=True) -class NumaTopology: - """System NUMA topology with CPU mappings.""" - numa_to_cpus: Dict[int, List[int]] # NUMA node -> list of CPU IDs - - def get_cpus_for_numa_node(self, numa_node: int) -> List[int]: - """Get CPUs belonging to a NUMA node.""" - return self.numa_to_cpus[numa_node] - - @staticmethod - def discover() -> 'NumaTopology': - """Discover NUMA topology for CPUs. - - Returns a topology where all CPUs are on node 0 if discovery fails or system is non-NUMA. - """ - numa_to_cpus: Dict[int, List[int]] = {} - numa_base = "/sys/devices/system/node" - - if os.path.exists(numa_base): - for entry in os.listdir(numa_base): - if entry.startswith("node") and entry[4:].isdigit(): - node_id = int(entry[4:]) - cpulist_path = os.path.join(numa_base, entry, "cpulist") - with open(cpulist_path, 'r') as f: - numa_to_cpus[node_id] = NumaTopology._parse_cpu_list(f.read()) - - # Fallback: single node with all CPUs - if not numa_to_cpus: - numa_to_cpus[0] = list(range(os.cpu_count() or 1)) - - return NumaTopology(numa_to_cpus=numa_to_cpus) - - @staticmethod - def _parse_cpu_list(cpu_list_str: str) -> List[int]: - """Parse CPU list string like '0-55,112-167' into list of CPU IDs.""" - cpus = [] - for part in cpu_list_str.strip().split(','): - if '-' in part: - start, end = part.split('-', 1) - cpus.extend(range(int(start), int(end) + 1)) - else: - cpus.append(int(part)) - return cpus - - # ============================================================================= # State Management # ============================================================================= @@ -841,27 +741,14 @@ def __post_init__(self): def _compute_thread_allocation(self) -> Dict[int, int]: """Determine how many compile threads each GPU should use based on NUMA topology.""" - # Group GPUs by their NUMA node - gpus_by_node: Dict[int, List[int]] = {} - for gpu_id in self.options.gpu_ids: - node = self.gpu_topology.get_numa_node(gpu_id) - gpus_by_node.setdefault(node, []).append(gpu_id) - - # Allocate CPUs from each node proportionally to GPUs on that node - allocation: Dict[int, int] = {} - for node, gpus_on_node in gpus_by_node.items(): - cpus_on_node = len(self.numa_topology.get_cpus_for_numa_node(node)) - threads_each = max(1, cpus_on_node // len(gpus_on_node)) - for gpu_id in gpus_on_node: - allocation[gpu_id] = threads_each + allocation = allocate_cpus_per_gpu(self.options.gpu_ids, self.gpu_topology, + self.numa_topology) # Apply user-specified CPU limit if provided if self.options.num_cpus is not None: total_allocated = sum(allocation.values()) if self.options.num_cpus < total_allocated: - scale_factor = self.options.num_cpus / total_allocated - for gpu_id in allocation: - allocation[gpu_id] = max(1, int(allocation[gpu_id] * scale_factor)) + allocation = scale_cpu_allocation(allocation, self.options.num_cpus) else: logger.info( f"--num-cpus={self.options.num_cpus} exceeds optimal {total_allocated}, using optimal allocation" @@ -1167,22 +1054,6 @@ def get_git_commit_hash() -> str: return "unknown" -def set_isolated_gpu_env(env: Dict[str, str], gpu_id: int) -> None: - """Modify environment to isolate subprocess to one physical GPU. - - Sets ROCR_VISIBLE_DEVICES at the HSA/ROCr level, providing complete isolation for all higher layers including HIP. - """ - env["ROCR_VISIBLE_DEVICES"] = str(gpu_id) - env.pop("HIP_VISIBLE_DEVICES", None) # Remove HIP_VISIBLE_DEVICES to avoid conflicts - - -def make_isolated_gpu_env(gpu_id: int) -> Dict[str, str]: - """Create environment that isolates subprocess to one physical GPU.""" - env = os.environ.copy() - set_isolated_gpu_env(env, gpu_id) - return env - - def resolve_verify_mode(verify_mode: str, config: PerfConfiguration) -> str: """Resolve the effective verify mode.""" if verify_mode == "gpu" and not isinstance(config, GPU_VALIDATION_CONFIGS): From f1ffcb18721fa9078f84723255f29b5bc5b3faa1 Mon Sep 17 00:00:00 2001 From: bogdan-petkovic Date: Mon, 10 Aug 2026 07:06:30 -0500 Subject: [PATCH 8/8] [AIROCMLIR-375] Stream shard output and isolate the driver from GPU faults Signed-off-by: bogdan-petkovic --- mlir/utils/jenkins/Jenkinsfile | 8 +- mlir/utils/performance/gpu_topology.py | 32 +++- mlir/utils/performance/run_e2e_multigpu.py | 143 ++++++++++++------ .../performance/tests/test_gpu_topology.py | 8 +- .../performance/tests/test_parameterSweeps.py | 82 ++++++++++ 5 files changed, 218 insertions(+), 55 deletions(-) create mode 100644 mlir/utils/performance/tests/test_parameterSweeps.py diff --git a/mlir/utils/jenkins/Jenkinsfile b/mlir/utils/jenkins/Jenkinsfile index b0beaae4ea96..465bef306612 100644 --- a/mlir/utils/jenkins/Jenkinsfile +++ b/mlir/utils/jenkins/Jenkinsfile @@ -1097,9 +1097,11 @@ int setLitWorkerCount() { // Run the lit suite sharded across the node's GPUs: one lit process per GPU, // each pinned via ROCR_VISIBLE_DEVICES; single-GPU and heterogeneous nodes fall -// back to one run. jobsPerGpu keeps the per-GPU cap from setLitWorkerCount() -// instead of splitting it across shards (see #1845/#1841). shStrict mirrors the -// output into the per-row log so withHealthyNode can retry just this node. +// back to one run. jobsPerGpu is the per-GPU upper bound (the setLitWorkerCount() +// cap that keeps the GPU itself out of trouble); the driver lowers it further +// when the shards together would exceed the host's CPU count, so raising +// setLitWorkerCount() has no effect once the host is the limit. shStrict mirrors +// the output into the per-row log so withHealthyNode can retry just this node. // ci-performance-scripts copies the driver and gpu_topology.py into ./bin. void runShardedE2E(int jobsPerGpu) { dir('build') { diff --git a/mlir/utils/performance/gpu_topology.py b/mlir/utils/performance/gpu_topology.py index 4f75750eb590..071c3a59392c 100644 --- a/mlir/utils/performance/gpu_topology.py +++ b/mlir/utils/performance/gpu_topology.py @@ -190,8 +190,11 @@ def select_gpu_ids( return unique_requested, arch, f"using requested GPUs {unique_requested} ({arch})" if len(set(archs)) > 1: - return [None], None, (f"mixed GPU architectures ({sorted(set(archs))}); " - "using a single GPU") + # Pin device 0 rather than leaving the choice open: callers compile for the + # architecture reported here, and an unpinned run would execute on device 0 + # regardless, so anything else risks compiling for the wrong target. + return [0], archs[0], (f"mixed GPU architectures ({sorted(set(archs))}); " + f"using GPU 0 ({archs[0]})") return list(range(count)), archs[0], f"distributing across {count} GPUs ({archs[0]})" @@ -261,3 +264,28 @@ def scale_cpu_allocation(allocation: Dict[int, int], limit: int) -> Dict[int, in scale = limit / total return {gpu_id: max(1, int(count * scale)) for gpu_id, count in allocation.items()} + + +def main() -> int: + """Print ``select_gpu_ids`` as JSON. + + Detection loads the ROCm runtime, which callers may not want to keep alive + for the rest of their run, so this entry point lets them ask for the answer + from a process that exits right away. + """ + import argparse + + parser = argparse.ArgumentParser(description='Report which GPUs to spread work across.') + parser.add_argument('--gpus', + type=int, + nargs='+', + default=None, + help='Physical GPU ids to use (default: auto-detect)') + args = parser.parse_args() + print(json.dumps(select_gpu_ids(args.gpus))) + return 0 + + +if __name__ == '__main__': + import sys + sys.exit(main()) diff --git a/mlir/utils/performance/run_e2e_multigpu.py b/mlir/utils/performance/run_e2e_multigpu.py index aac6455ce4a0..0ae628787849 100644 --- a/mlir/utils/performance/run_e2e_multigpu.py +++ b/mlir/utils/performance/run_e2e_multigpu.py @@ -9,24 +9,56 @@ `--run-shard`, running one lit process per GPU, each pinned to its device via ROCR_VISIBLE_DEVICES. Sharding is only used on homogeneous nodes; single-GPU and mixed-architecture nodes fall back to a single lit run. + +Shard output is forwarded line by line rather than collected at the end: CI +watchdogs key off console activity, and a run that dies mid-way must still leave +behind the output that explains why. Note that all shards share one lit exec +root, so they race on `.lit_test_times.txt`; that only perturbs the ordering +heuristic of a later run. """ from __future__ import annotations import argparse +import json import os +import signal import subprocess import sys +import threading import time -from typing import List, Optional +from typing import List, Optional, Tuple -from gpu_topology import make_isolated_gpu_env, select_gpu_ids, usable_cpu_count +from gpu_topology import make_isolated_gpu_env, usable_cpu_count def default_lit_path(build_dir: str) -> str: return os.path.join(build_dir, 'external', 'llvm-project', 'llvm', 'bin', 'llvm-lit') +def select_gpu_ids_out_of_process( + requested: Optional[List[int]]) -> Tuple[List[Optional[int]], Optional[str], str]: + """Ask a child process which GPUs to use, and what to report about them. + + Enumeration goes through HIP, which leaves the ROCm runtime loaded and + attached to the KFD for the lifetime of the process. This driver has to + outlive a wedged shard so it can report and clean up, and ROCr aborts every + process holding a context when a GPU faults, so the runtime is kept in a + child that exits immediately. A child that fails means "use one GPU", the + same fallback taken when enumeration itself fails. + """ + helper = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'gpu_topology.py') + cmd = [sys.executable, helper] + if requested: + cmd += ['--gpus'] + [str(g) for g in requested] + try: + out = subprocess.run(cmd, capture_output=True, text=True, timeout=180, check=True) + gpu_ids, arch, message = json.loads(out.stdout) + except Exception as e: # noqa: BLE001 - any failure means fall back to one GPU + return [None], None, f"GPU detection failed ({e}); using the default GPU" + return gpu_ids, arch, message + + def build_shard_command(lit: str, lit_args: List[str], jobs: int, num_shards: int, shard: int, test_paths: List[str]) -> List[str]: """Build a single lit invocation for shard `shard` (1-based) of `num_shards`.""" @@ -65,11 +97,31 @@ def cap_jobs_to_host_cpus(jobs_per_shard: int, return max(1, min(jobs_per_shard, budget // num_shards)) +def _forward_output(label: str, proc: subprocess.Popen, lock: threading.Lock) -> None: + """Tag and forward one shard's output so interleaved shards stay tellable apart.""" + for line in proc.stdout: + with lock: + sys.stdout.write(f"[{label}] {line}") + sys.stdout.flush() + + +def _terminate_tree(proc: subprocess.Popen, sig: int) -> None: + """Signal a shard's whole process group. + + lit runs each test in its own subprocess, and those hold the GPU contexts. + Signalling only lit leaves them behind on the node for the next job. + """ + try: + os.killpg(os.getpgid(proc.pid), sig) + except (ProcessLookupError, PermissionError): + proc.send_signal(sig) + + def run(args: argparse.Namespace) -> int: lit = args.lit or default_lit_path(args.build_dir) test_paths = args.test_paths or [os.path.join(args.build_dir, 'mlir', 'test')] - gpu_ids, _gpu_arch, gpu_msg = select_gpu_ids(args.gpus) + gpu_ids, _gpu_arch, gpu_msg = select_gpu_ids_out_of_process(args.gpus) print(f"[run_e2e_multigpu] {gpu_msg}", flush=True) # [None] means one un-pinned lit run (single-GPU / heterogeneous nodes). @@ -86,8 +138,7 @@ def run(args: argparse.Namespace) -> int: print(f"[run_e2e_multigpu] {num_shards} shard(s), {jobs_per_shard} lit workers each", flush=True) - # A single shard streams straight to the console so CI sees live progress; - # multiple shards are buffered per-GPU and dumped once they finish. + # A single shard needs no tagging, so let it inherit the console directly. if num_shards == 1: gpu_id = shard_gpus[0] cmd = build_shard_command(lit, args.lit_args, jobs_per_shard, 1, 1, test_paths) @@ -98,72 +149,69 @@ def run(args: argparse.Namespace) -> int: return 0 return subprocess.call(cmd, env=env) + labels = [] procs = [] - log_paths = [] for idx, gpu_id in enumerate(shard_gpus): cmd = build_shard_command(lit, args.lit_args, jobs_per_shard, num_shards, idx + 1, test_paths) - env = make_isolated_gpu_env(gpu_id) label = f"GPU {gpu_id}" if gpu_id is not None else "single" print(f"[run_e2e_multigpu] shard {idx + 1}/{num_shards} on {label}: {' '.join(cmd)}", flush=True) if args.dry_run: continue - log_path = os.path.join(args.build_dir, f"e2e-shard-{idx}.log") - log_paths.append((label, log_path)) - log_file = open(log_path, 'wb') - procs.append((label, log_file, - subprocess.Popen(cmd, env=env, stdout=log_file, stderr=subprocess.STDOUT))) + env = make_isolated_gpu_env(gpu_id) or os.environ.copy() + # lit is itself Python; without this it block-buffers into our pipe and + # the console would go quiet for minutes at a time. + env['PYTHONUNBUFFERED'] = '1' + labels.append(label) + procs.append( + subprocess.Popen(cmd, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + start_new_session=True)) if args.dry_run: return 0 + stdout_lock = threading.Lock() + readers = [ + threading.Thread(target=_forward_output, args=(label, proc, stdout_lock), daemon=True) + for label, proc in zip(labels, procs) + ] + for reader in readers: + reader.start() + failures = [] aborted = [] - pending = list(range(len(procs))) - # Heartbeat: keep the console alive during the buffered run so Jenkins' - # timeout(activity: true) does not fire. - start = time.time() - last_beat = start - heartbeat_secs = 30 + pending = set(range(len(procs))) while pending: time.sleep(1) - now = time.time() - if now - last_beat >= heartbeat_secs: - last_beat = now - print( - f"[run_e2e_multigpu] still running: {len(pending)}/{len(procs)} " - f"shard(s) active, {int(now - start)}s elapsed", - flush=True) - for i in list(pending): - label, log_file, proc = procs[i] - rc = proc.poll() + for i in sorted(pending): + rc = procs[i].poll() if rc is None: continue - log_file.close() - pending.remove(i) + pending.discard(i) if rc != 0: - failures.append((label, rc)) + failures.append((labels[i], rc)) # Once a shard fails, stop the rest so the run aborts promptly. if args.fail_fast and failures and pending: - for i in pending: - procs[i][2].terminate() - for i in pending: - label, log_file, proc = procs[i] + for i in sorted(pending): + _terminate_tree(procs[i], signal.SIGTERM) + for i in sorted(pending): try: - proc.wait(timeout=15) + procs[i].wait(timeout=15) except subprocess.TimeoutExpired: - proc.kill() - log_file.close() - aborted.append(label) - pending = [] + _terminate_tree(procs[i], signal.SIGKILL) + procs[i].wait() + aborted.append(labels[i]) + pending.clear() - # Surface every shard's output in the CI console. - for label, log_path in log_paths: - print(f"\n===== lit output: {label} ({log_path}) =====", flush=True) - with open(log_path, 'r', errors='replace') as f: - sys.stdout.write(f.read()) + for reader in readers: + reader.join(timeout=30) if failures: summary = ', '.join(f"{label} (exit {rc})" for label, rc in failures) @@ -200,8 +248,9 @@ def main() -> int: parser.add_argument('--max-total-jobs', type=int, default=None, - help='Upper bound on lit workers across all shards ' - '(default: the number of CPUs this process may use)') + help='Upper bound on lit workers across all shards, subject to a ' + 'floor of one worker per shard (default: the number of CPUs this ' + 'process may use)') parser.add_argument('--gpus', type=int, nargs='+', diff --git a/mlir/utils/performance/tests/test_gpu_topology.py b/mlir/utils/performance/tests/test_gpu_topology.py index e217ad2b5efe..b2fdeaadae5b 100644 --- a/mlir/utils/performance/tests/test_gpu_topology.py +++ b/mlir/utils/performance/tests/test_gpu_topology.py @@ -79,12 +79,14 @@ def test_homogeneous_multi_gpu(self, monkeypatch): assert gpu_ids == [0, 1, 2, 3, 4, 5, 6, 7] assert arch == "gfx942" - def test_mixed_archs_no_request_falls_back(self, monkeypatch): + def test_mixed_archs_pin_gpu_zero(self, monkeypatch): + # Compilation targets the reported arch, so it has to be the one that + # will actually run the kernels. _clear_visible_devices(monkeypatch) _fake_archs(monkeypatch, ["gfx942", "gfx1100"]) gpu_ids, arch, msg = gpu_topology.select_gpu_ids() - assert gpu_ids == [None] - assert arch is None + assert gpu_ids == [0] + assert arch == "gfx942" assert "mixed" in msg def test_requested_valid_homogeneous_subset(self, monkeypatch): diff --git a/mlir/utils/performance/tests/test_parameterSweeps.py b/mlir/utils/performance/tests/test_parameterSweeps.py new file mode 100644 index 000000000000..c925df4df729 --- /dev/null +++ b/mlir/utils/performance/tests/test_parameterSweeps.py @@ -0,0 +1,82 @@ +# Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions. +# See https://llvm.org/LICENSE.txt for license information. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +""" +Tests for parameterSweeps.py. + +Covers how sweep configurations are handed out to GPUs. The work itself is +stubbed out, so no GPU or build tree is needed. +""" +import asyncio +import sys +from pathlib import Path + +# Ensure we can import parameterSweeps (lives in mlir/utils/performance). +_test_dir = Path(__file__).resolve().parent +_sys_path_parent = str(_test_dir.parent) +if _sys_path_parent not in sys.path: + sys.path.insert(0, _sys_path_parent) +# Mock hip and amd_arch_db so the module imports without ROCm (CI has no GPU). +exec( + open(_test_dir / "mock_hip.py").read(), { + "__file__": str(_test_dir / "mock_hip.py"), + "sys": sys + }) + +import parameterSweeps # noqa: E402 + + +def _options(gpu_ids, concurrent_tests=2): + return parameterSweeps.Options(debug=False, + quiet=True, + debug_fails=False, + arch="gfx942", + flags=[], + concurrent_tests=concurrent_tests, + num_cu=304, + num_chiplets=8, + gpu_ids=gpu_ids) + + +def _run_sweep(monkeypatch, options, num_configs): + """Run a sweep with the actual test execution stubbed out. + + Returns the GPU id each configuration was handed to, in order. + """ + assignments = [] + + async def fake_drop_good_config(config, options, paths, gpu_id=None): + assignments.append((config, gpu_id)) + return parameterSweeps.TestResult.PASS + + monkeypatch.setattr(parameterSweeps, "drop_good_config", fake_drop_good_config) + result = asyncio.run( + parameterSweeps.sweep_parameters(range(num_configs), lambda p, o: p, options, paths=None)) + return assignments, result + + +class TestSweepGpuAssignment: + """Tests for how sweep_parameters spreads configurations over GPUs.""" + + def test_configs_cycle_through_selected_gpus(self, monkeypatch): + assignments, _ = _run_sweep(monkeypatch, _options(gpu_ids=(0, 1, 2)), num_configs=7) + assert [gpu_id for _, gpu_id in assignments] == [0, 1, 2, 0, 1, 2, 0] + + def test_every_config_is_dispatched_once_in_order(self, monkeypatch): + assignments, (passed, invalid, failing) = _run_sweep(monkeypatch, + _options(gpu_ids=(0, 1)), + num_configs=5) + assert [config for config, _ in assignments] == [0, 1, 2, 3, 4] + assert (passed, invalid, failing) == (5, 0, []) + + def test_single_gpu_leaves_every_config_unpinned(self, monkeypatch): + assignments, _ = _run_sweep(monkeypatch, _options(gpu_ids=(None,)), num_configs=3) + assert [gpu_id for _, gpu_id in assignments] == [None, None, None] + + def test_assignment_spans_batches(self, monkeypatch): + # Configurations are dispatched in batches of `concurrent_tests`; the + # rotation has to continue across batch boundaries, not restart. + assignments, _ = _run_sweep(monkeypatch, + _options(gpu_ids=(0, 1, 2), concurrent_tests=2), + num_configs=6) + assert [gpu_id for _, gpu_id in assignments] == [0, 1, 2, 0, 1, 2]