Implement flock-based bitmap GPU allocator for CI parallelization

[Copilot]

Summary

Implements a lightweight, flock-based GPU allocator for CI parallelization on our 8-GPU runner. Enables dynamic GPU allocation with variable requests (1, 2, 4, 8 GPUs), eliminating hardcoded assignments and enabling true parallel execution.

Key Features

GPU Allocator Core (.github/scripts/gpu_allocator.sh)

• Bitmap-based allocation: 8-bit bitmap (0-255) tracks GPU state atomically via flock on /tmp/iris_gpu_state
• Non-contiguous allocation: Any N free GPUs can be used (no contiguous requirement)
• Out-of-order release safe: Jobs can finish in any order without deadlocks or GPU waste
• Retry loop: 1-minute backoff, up to 3 hours total wait time (180 retries × 60s)
• State validation: Detects and resets corrupted or out-of-range bitmap values
• Handles set -e: Works correctly in scripts with errexit enabled
• Diagnostic logging: Configuration and allocation status for troubleshooting
• Throughput-oriented: First-fit scheduling (non-FIFO acceptable)

Workflow-Level GPU Management

• acquire_gpus.sh: Acquires GPUs and exports GPU_DEVICES to GITHUB_ENV
• release_gpus.sh: Releases GPUs via bitmap clearing with if: always() guarantee
• All workflows have dedicated acquire/release GPU steps with 3-hour timeouts
• Release always runs regardless of test outcome

Distributed Test Runner (tests/run_tests_distributed.py)

• Simple worker script that runs pytest under torchrun
• Invoked directly by torchrun (no recursive launcher pattern)
• Eliminates EADDRINUSE port conflicts through torchrun's automatic port management

Performance Benchmarks (examples/)

• 3 benchmarks support both torchrun and mp.spawn modes:
  • examples/07_gemm_all_scatter/benchmark.py
  • examples/11_gemm_all_scatter_producer_consumer/benchmark.py
  • examples/12_gemm_all_scatter_bulk_synchronous/benchmark.py
• Auto-detects execution mode via RANK and LOCAL_RANK environment variables
• Maintains backward compatibility for standalone execution

External Validation Tests

• Invoke torchrun directly (avoids nested process issues)
• Both tests use: torchrun --nproc_per_node=2 <test_file>

Test Runner Integration

• Reads GPU_DEVICES from environment (set by workflow)
• Invokes torchrun directly for distributed tests
• No internal GPU allocation logic

Workflow Simplification

• Removed hardcoded gpu_devices from all workflow matrices
• Removed port cleanup steps (obsolete with torchrun)
• Jobs acquire GPUs dynamically for true parallel execution
• All jobs have 3-hour timeouts matching GPU allocator wait time

Implementation Details

Bitmap Allocator

• State file: 8-bit integer (0-255) representing GPU allocation bitmap
• Bit position N = GPU N (bit 0 = GPU 0, bit 1 = GPU 1, etc.)
• Bit value: 1 = allocated, 0 = free
• Allocation: finds N free GPUs (bits=0), sets bits to 1
• Release: clears bits using bitwise AND with inverse mask
• ALLOCATED_GPU_BITMAP env var stores allocation mask for cleanup

Example bitmap states:

• 0b00000000 (0): All GPUs free
• 0b00000011 (3): GPUs 0,1 allocated
• 0b00011111 (31): GPUs 0-4 allocated
• 0b11111111 (255): All 8 GPUs allocated

Advantages over index-based allocation:

• Out-of-order job completion safe (no deadlocks)
• No GPU waste from out-of-order releases
• Immediate GPU reuse when freed
• Better utilization with variable GPU requests

Impact

Before: Sequential job execution, static GPU assignment, port conflicts, potential deadlocks from out-of-order completion

After: Parallel execution constrained only by GPU availability. A 1-GPU job and 4-GPU job can run concurrently with proper isolation. Jobs complete in any order without blocking. All distributed components use torchrun for automatic port management.

Original prompt

---

This section details on the original issue you should resolve

<issue_title>Implement Lightweight GPU Allocator for CI (Supports Variable GPU Requests)</issue_title>
<issue_description>We need a lightweight GPU isolation mechanism for CI jobs running on our single-node 8-GPU machine.

Requirements:

• Support variable GPU requests (e.g., 1, 2, 4, 8 GPUs)
• Ensure no overlapping GPU usage
• Maximize machine utilization
• Avoid heavy infrastructure (e.g., Slurm)
• Throughput-oriented scheduling (non-FIFO is acceptable)

Design:

• Use a single global flock to protect allocation state
• Maintain a shared state file tracking free GPUs
• When a job requests N GPUs:
  • If enough GPUs are free → allocate immediately
  • Otherwise → release lock, sleep, and retry
• On job exit, GPUs are released back into the pool

This approach provides deterministic isolation and high utilization without introducing a full scheduler.</issue_description>

Comments on the Issue (you are @copilot in this section)

• Fixes Implement Lightweight GPU Allocator for CI (Supports Variable GPU Requests) #369

---

💬 We'd love your input! Share your thoughts on Copilot coding agent in our 2 minute survey.

[Copilot]

Pull request overview

Introduces a lightweight, flock-based GPU allocation mechanism for CI so matrix jobs can dynamically acquire the number of GPUs they need, rather than relying on hardcoded GPU assignments in the workflow.

Changes:

• Added a shared-state GPU allocator script that blocks/retries until enough GPUs are available and exports GPU_DEVICES for HIP_VISIBLE_DEVICES.
• Updated the test runner script to automatically acquire/release GPUs when not explicitly provided.
• Simplified the CI workflow matrix by removing static gpu_devices entries and always letting the runner auto-allocate.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 6 comments.

File	Description
.github/workflows/iris-tests.yml	Removes hardcoded GPU device mappings from the test matrices and relies on dynamic allocation.
.github/scripts/run_tests.sh	Sources the allocator and acquires/releases GPUs automatically when devices aren’t provided.
.github/scripts/gpu_allocator.sh	Implements the flock-protected shared state allocator and cleanup logic.

[mawad-amd]

@copilot apply changes based on the comments in this thread

[Copilot]

@copilot apply changes based on the comments in this thread

Applied all review feedback. Changes in commit eb6bae9:

1. Removed set -e from sourced file (no side effects)
2. Added state validation in init_gpu_state (detects/resets invalid/out-of-range values)
3. Fixed race condition in acquire_gpus (captures start index atomically via temp file)
4. Changed out-of-order release to preserve state instead of resetting to 0 (prevents GPU overlap)
5. Made cleanup trap optional via enable_gpu_cleanup_trap() (composes with existing traps)
6. Fixed GPU_DEVICES env variable handling in run_tests.sh (honors environment, arg overrides)

All changes tested and verified.

[Copilot]

Pull request overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated 12 comments.

[Copilot]

run_tests_distributed.py now assumes it is always launched via torchrun (env:// rendezvous). Running it directly (e.g., python tests/run_tests_distributed.py ...) will fail because MASTER_ADDR/MASTER_PORT are not set and there's no fallback. Consider adding a guard that either (a) errors with a clear message telling users to run via torchrun, or (b) retains the previous mp.spawn fallback when RANK/WORLD_SIZE/LOCAL_RANK are missing.

[Copilot]

The process-group backend is hard-coded to nccl. If this script is ever executed in an environment without CUDA/HIP available, dist.init_process_group(backend="nccl") will fail. Consider selecting backend="nccl" if torch.cuda.is_available() else "gloo" (matching the pattern used elsewhere in the repo).

[Copilot]

The header comment says this script provides "automatic GPU allocation" / "will use allocator" when gpu_devices isn't provided, but the implementation only warns when GPU_DEVICES is empty and never calls the allocator. Suggest updating the header/usage text to reflect the actual behavior (workflow-level acquire/release), or implement the documented fallback.

[Copilot]

GPU_DEVICES can be overridden via the environment, but the benchmark always runs torchrun --nproc_per_node=8. If GPU_DEVICES is set to fewer than 8 devices, this will oversubscribe and fail. Consider deriving --nproc_per_node from the requested/available GPU count (e.g., count of comma-separated IDs), or enforce that GPU_DEVICES must contain 8 GPUs when using this script.

[Copilot]

In torchrun mode, this passes rank=int(os.environ["RANK"]) into _worker, but _worker uses that parameter as local_rank for rank= and device_id=torch.device(f"cuda:{local_rank}"). That’s only correct when global rank == local rank. Please use LOCAL_RANK for the first _worker argument / device selection, and keep RANK for the global rank.

Suggested change

	world_size = int(os.environ.get("WORLD_SIZE", 1))
	init_url = os.environ.get("MASTER_ADDR", "127.0.0.1") + ":" + os.environ.get("MASTER_PORT", "29500")
	_worker(rank, world_size, f"tcp://{init_url}", args)
	local_rank = int(os.environ.get("LOCAL_RANK", 0))
	world_size = int(os.environ.get("WORLD_SIZE", 1))
	init_url = os.environ.get("MASTER_ADDR", "127.0.0.1") + ":" + os.environ.get("MASTER_PORT", "29500")
	_worker(local_rank, world_size, f"tcp://{init_url}", args)

[Copilot]

State validation hard-codes the max bitmap value to 255, but the allocator also supports overriding MAX_GPUS. If MAX_GPUS is changed, this range check becomes incorrect. Consider computing the max as (1 << MAX_GPUS) - 1 to keep validation consistent with configuration.

[Copilot]

This job acquires GPUs (sets GPU_DEVICES), but the subsequent container invocation in this workflow does not pass --gpus "$GPU_DEVICES", so the container will still see all GPUs and can overlap with other jobs. Update the container step to consume the allocated GPU_DEVICES (same pattern as external-gluon-validation-test).

[Copilot]

In torchrun mode, this uses RANK as the first argument to _worker, but _worker treats that parameter as the per-process GPU index (local_rank) and uses it for device_id. For multi-node torchrun (or any case where global rank != local rank), this will select the wrong device. Use LOCAL_RANK for device selection / _worker's first argument, and use RANK only for the global rank passed to the process group.

Suggested change

	rank = int(os.environ.get("RANK", 0))
	world_size = int(os.environ.get("WORLD_SIZE", 1))
	init_url = os.environ.get("MASTER_ADDR", "127.0.0.1") + ":" + os.environ.get("MASTER_PORT", "29500")
	_worker(rank, world_size, f"tcp://{init_url}", args)
	local_rank = int(os.environ.get("LOCAL_RANK", 0))
	world_size = int(os.environ.get("WORLD_SIZE", 1))
	init_url = os.environ.get("MASTER_ADDR", "127.0.0.1") + ":" + os.environ.get("MASTER_PORT", "29500")
	_worker(local_rank, world_size, f"tcp://{init_url}", args)

[Copilot]

The GPU state file is created in /tmp via echo "0" > "$GPU_STATE_FILE" without setting permissions. On multi-user runners this can allow unintended modification/corruption of allocator state. Consider creating it with restrictive permissions (e.g., umask 077 for the write, or chmod 600 "$GPU_STATE_FILE" after creation).

[Copilot]

release_gpus.sh proceeds even if ALLOCATED_GPU_BITMAP is missing (as long as GPU_DEVICES is set), but release_gpus uses only the bitmap mask. In that scenario, the script will log details and then effectively release nothing, leaking allocations in the shared bitmap. Consider either requiring ALLOCATED_GPU_BITMAP to be set (fail fast), or reconstructing the mask from GPU_DEVICES as a fallback.

Suggested change

	# Fail fast if GPU_DEVICES is set but the bitmap mask is missing
	if [ -n "$GPU_DEVICES" ] && [ -z "$ALLOCATED_GPU_BITMAP" ]; then
	echo "[RELEASE-GPUS] ERROR: GPU_DEVICES is set but ALLOCATED_GPU_BITMAP is missing; cannot safely release GPUs"
	exit 1
	fi