ci: qualify immediate two-GPU execution on E1 PRs

.github/scripts/e1/PR.md

@@ -0,0 +1,51 @@
+# E1 immediate GPU PR qualification
+
+`E1 Immediate GPU on PR` verifies that a pull request dispatches a real two-GPU
+test to the `colossalai-e1-h20` runner on `gpu-h20-5`.
+
+This initial rollout only accepts same-repository pull requests to `main` from
+`ci/e1-runner-bootstrap`, authored by `richardoo-707`. It reacts to opening,
+updating and reopening that PR, and only to changes in the E1 scripts/workflow.
+It deliberately supports a draft qualification PR. It does not merge the PR.
+
+The workflow uses the ordinary `pull_request` event and a read-only repository
+token. It does not use `pull_request_target`, PATs, booking credentials or website
+reservations. Resource use is authorized manually for this development trial.
+The branch/author restrictions are rollout limits, not a sandbox for untrusted
+code; do not generalize this shared-host workflow to external contributions.
+
+The job downloads the official GitHub source archive for the exact PR merge
+commit (`GITHUB_SHA`) over HTTPS, extracts only the four E1 test scripts into a
+fresh temporary directory, and records the archive SHA256. This avoids the
+node's failing connection to the Git HTTPS endpoint on `github.com`; the
+`codeload.github.com` endpoint is reachable. It does not use a mutable branch
+snapshot or a pre-existing developer checkout, and it does not require a PAT.
+This limited extraction is suitable for this infrastructure probe; project-wide
+tests will need a full source checkout. Downloading is bounded to four minutes.
+
+After source preparation, CPU-only unit tests validate the GPU selector and event guard.
+The selector considers memory, utilization and running compute processes,
+chooses two currently idle GPU UUIDs, and passes them to `run_gpu_smoke.sh`.
+That wrapper checks occupancy again immediately before starting the container.
+If fewer than two GPUs are idle, or a selected GPU becomes busy, the job fails
+without preempting another user's task. There is no automatic reservation or
+waiting queue in this immediate trial.
+
+The container exposes only the chosen pair. A successful exit requires the
+actual `E1_GPU_SMOKE_PASS` record with two NCCL workers, covering matrix
+multiplication, all-reduce, DDP gradients and three analytical SGD updates.
+The cached image is pinned to a digest and downloads no packages or models.
+This is runner qualification, not ColossalAI project/version compatibility.
+
+One E1 PR GPU job runs at a time. The job is bounded to 12 minutes and the GPU
+test to six minutes with shorter container-level deadlines. Ordinary exits and
+cancellation clean up the job's own container; host failure or SIGKILL can bypass
+cleanup, so inspect the recorded container ID when recovering from those cases.
+
+The Actions log and job summary record the PR head SHA, tested merge SHA,
+runner, chosen GPU UUIDs, timestamps and the real result. Detailed wrapper logs
+also remain in the runner's personal CI test-results directory. No GPU pass is
+reported when selection, execution or numerical checks fail.
+
+Before expanding beyond this PR trial, connect the real reservation system,
+agree the eligible GPU pool, and review the workflow's trust and trigger policy.

.github/scripts/e1/gpu_smoke.py

@@ -0,0 +1,144 @@
+"""Small, assertion-based runner qualification; this does not test ColossalAI."""
+
+import argparse
+import json
+import os
+import sys
+from datetime import timedelta
+
+import torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+
+
+def require(condition, message):
+    if not condition:
+        raise RuntimeError(message)
+
+
+def environment():
+    return {
+        "python": sys.version.split()[0],
+        "torch": torch.__version__,
+        "cuda_build": torch.version.cuda,
+        "distributed_available": dist.is_available(),
+        "nccl_available": dist.is_available() and dist.is_nccl_available(),
+    }
+
+
+def check_collectives_and_training(device, rank, local_rank):
+    # Exercise actual tensor computation, then verify the numerical result.
+    left = torch.ones((128, 128), device=device)
+    right = torch.full((128, 128), 2.0, device=device)
+    torch.testing.assert_close(left @ right, torch.full_like(left, 256.0), rtol=0, atol=0)
+
+    reduced = torch.tensor([rank + 1.0], device=device)
+    dist.all_reduce(reduced, op=dist.ReduceOp.SUM)
+    torch.testing.assert_close(reduced, torch.tensor([3.0], device=device), rtol=0, atol=0)
+
+    model = torch.nn.Linear(1, 1, bias=False).to(device)
+    with torch.no_grad():
+        model.weight.zero_()
+    ddp = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None)
+    optimizer = torch.optim.SGD(ddp.parameters(), lr=0.01)
+
+    # Different data on each rank makes missing gradient synchronization observable.
+    # Global x = [1, 2, 3, 4], y = 3*x. mean(x**2) = 7.5, so grad = 15*(w-3).
+    inputs = torch.tensor([[1.0 + 2 * rank], [2.0 + 2 * rank]], device=device)
+    targets = 3 * inputs
+    expected_weight = 0.0
+    for _ in range(3):
+        optimizer.zero_grad(set_to_none=True)
+        loss = torch.nn.functional.mse_loss(ddp(inputs), targets)
+        require(torch.isfinite(loss).item(), "Non-finite training loss")
+        loss.backward()
+
+        expected_gradient = 15.0 * (expected_weight - 3.0)
+        torch.testing.assert_close(
+            ddp.module.weight.grad,
+            torch.full_like(ddp.module.weight, expected_gradient),
+            rtol=1e-5,
+            atol=1e-5,
+        )
+        optimizer.step()
+        expected_weight -= 0.01 * expected_gradient
+        torch.testing.assert_close(
+            ddp.module.weight,
+            torch.full_like(ddp.module.weight, expected_weight),
+            rtol=1e-5,
+            atol=1e-5,
+        )
+
+    weight = ddp.module.weight.detach().contiguous()
+    gathered = [torch.empty_like(weight) for _ in range(2)]
+    dist.all_gather(gathered, weight)
+    torch.testing.assert_close(gathered[0], gathered[1], rtol=0, atol=0)
+    return float(weight.item())
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    modes = parser.add_mutually_exclusive_group()
+    modes.add_argument("--check", action="store_true", help="Inspect the build without initializing CUDA")
+    modes.add_argument("--cpu-control", action="store_true", help="Validate test math on CPU/Gloo only")
+    args = parser.parse_args()
+
+    metadata = environment()
+    if args.check:
+        require(metadata["distributed_available"], "PyTorch distributed is unavailable")
+        require(metadata["nccl_available"], "PyTorch was not built with NCCL")
+        require(metadata["cuda_build"] is not None, "PyTorch was not built with CUDA")
+        print(json.dumps({"mode": "environment_check", "gpu_tested": False, **metadata}), flush=True)
+        return
+
+    world_size = int(os.environ.get("WORLD_SIZE", "0"))
+    rank = int(os.environ.get("RANK", "-1"))
+    local_rank = int(os.environ.get("LOCAL_RANK", "-1"))
+    require(world_size == 2 and rank in (0, 1) and local_rank in (0, 1), "Launch exactly two torchrun workers")
+
+    if args.cpu_control:
+        require(torch.cuda.device_count() == 0, "CPU control must run in a container without GPU access")
+        device = torch.device("cpu")
+        backend = "gloo"
+    else:
+        require(torch.cuda.is_available(), "CUDA unavailable: GPU test cannot be skipped or passed")
+        require(torch.cuda.device_count() == 2, "Expose exactly two assigned GPUs to this container")
+        require(metadata["nccl_available"], "NCCL unavailable")
+        torch.cuda.set_device(local_rank)
+        device = torch.device("cuda", local_rank)
+        backend = "nccl"
+        properties = torch.cuda.get_device_properties(device)
+        print(
+            json.dumps({"rank": rank, "device": str(device), "name": properties.name, "uuid": str(properties.uuid)}),
+            flush=True,
+        )
+
+    torch.set_num_threads(1)
+    dist.init_process_group(backend=backend, timeout=timedelta(seconds=45))
+    try:
+        weight = check_collectives_and_training(device, rank, local_rank)
+        if device.type == "cuda":
+            torch.cuda.synchronize(device)
+        dist.barrier()
+        if rank == 0:
+            marker = "E1_CPU_CONTROL_PASS" if args.cpu_control else "E1_GPU_SMOKE_PASS"
+            print(
+                json.dumps(
+                    {
+                        "result": marker,
+                        "gpu_tested": not args.cpu_control,
+                        "backend": backend,
+                        "world_size": world_size,
+                        "checks": ["matmul", "all_reduce_sum", "ddp_gradients", "sgd_updates", "equal_weights"],
+                        "final_weight": weight,
+                        **metadata,
+                    }
+                ),
+                flush=True,
+            )
+    finally:
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

.github/scripts/e1/pr_gpu.py

@@ -0,0 +1,166 @@
+#!/usr/bin/env python3
+"""Immediate PR qualification using two currently idle GPUs on gpu-h20-5.
+
+Temporary rollout for one trusted internal PR branch, under manual resource
+authorization. This does not contact or claim a reservation from the website.
+"""
+
+import argparse
+import json
+import os
+import re
+import signal
+import socket
+import subprocess
+import sys
+from datetime import datetime, timezone
+from pathlib import Path
+
+GPU_UUID = re.compile(r"GPU-[a-fA-F0-9]{8}(?:-[a-fA-F0-9]{4}){3}-[a-fA-F0-9]{12}\Z")
+
+
+def validate_event(event):
+    pr = event["pull_request"]
+    if (
+        event["repository"]["full_name"] != "hpcaitech/ColossalAI"
+        or pr["head"]["repo"]["full_name"] != "hpcaitech/ColossalAI"
+        or pr["head"]["ref"] != "ci/e1-runner-bootstrap"
+        or pr["base"]["ref"] != "main"
+        or pr["user"]["login"] != "richardoo-707"
+    ):
+        raise RuntimeError("This qualification is restricted to the authorized internal E1 PR")
+    return pr
+
+
+def select_idle(inventory, processes):
+    busy = set()
+    for line in processes.splitlines():
+        value = line.strip()
+        if value:
+            if not GPU_UUID.fullmatch(value):
+                raise RuntimeError("GPU process information is unknown")
+            busy.add(value)
+    candidates, seen = [], set()
+    for line in inventory.splitlines():
+        fields = [part.strip() for part in line.split(",")]
+        if len(fields) != 4:
+            raise RuntimeError("Unexpected GPU inventory format")
+        index, gpu, memory, utilization = fields
+        if not GPU_UUID.fullmatch(gpu) or gpu in seen:
+            raise RuntimeError("Invalid or duplicate GPU UUID")
+        seen.add(gpu)
+        if not all(value.isdigit() for value in (index, memory, utilization)):
+            raise RuntimeError("GPU occupancy is unknown")
+        if gpu not in busy and int(memory) <= 256 and int(utilization) == 0:
+            candidates.append((int(index), gpu))
+    if len(candidates) < 2:
+        raise RuntimeError("Fewer than two idle GPUs; immediate test cannot run")
+    candidates.sort()
+    return candidates[:2]
+
+
+def query_gpus():
+    inventory = subprocess.run(
+        ["nvidia-smi", "--query-gpu=index,uuid,memory.used,utilization.gpu", "--format=csv,noheader,nounits"],
+        check=True,
+        text=True,
+        capture_output=True,
+        timeout=10,
+    ).stdout
+    processes = subprocess.run(
+        ["nvidia-smi", "--query-compute-apps=gpu_uuid", "--format=csv,noheader"],
+        check=True,
+        text=True,
+        capture_output=True,
+        timeout=10,
+    ).stdout
+    return select_idle(inventory, processes)
+
+
+def has_gpu_pass(log):
+    for line in log.splitlines():
+        try:
+            record = json.loads(line)
+        except ValueError:
+            continue
+        if isinstance(record, dict) and (
+            record.get("result") == "E1_GPU_SMOKE_PASS"
+            and record.get("gpu_tested") is True
+            and record.get("backend") == "nccl"
+            and record.get("world_size") == 2
+        ):
+            return True
+    return False
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--output", type=Path, required=True)
+    args = parser.parse_args()
+    if os.name != "posix" or socket.gethostname().split(".")[0] != "gpu-h20-5":
+        raise RuntimeError("This job must execute on gpu-h20-5")
+    if os.environ.get("GITHUB_EVENT_NAME") != "pull_request":
+        raise RuntimeError("This entry point requires a GitHub pull_request event")
+    event = json.loads(Path(os.environ["GITHUB_EVENT_PATH"]).read_text())
+    pr = validate_event(event)
+    args.output.mkdir(parents=True, mode=0o700)
+    result = {
+        "status": "failed",
+        "gpu_tested": False,
+        "website_contacted": False,
+        "allocation_source": "manual_authorization_idle_selection",
+        "host": "gpu-h20-5",
+        "pull_request": pr["number"],
+        "head_sha": pr["head"]["sha"],
+        "checkout_sha": os.environ.get("GITHUB_SHA"),
+        "source_method": "github_codeload_merge_commit",
+        "source_archive_sha256": os.environ.get("E1_SOURCE_ARCHIVE_SHA256"),
+        "runner": os.environ.get("RUNNER_NAME"),
+        "started_at": datetime.now(timezone.utc).isoformat(),
+    }
+    child = None
+
+    def interrupted(signum, frame):
+        raise RuntimeError(f"Interrupted by signal {signum}")
+
+    signal.signal(signal.SIGTERM, interrupted)
+    signal.signal(signal.SIGINT, interrupted)
+    try:
+        selected = query_gpus()
+        result["gpu_indices"] = [index for index, gpu in selected]
+        result["gpu_uuids"] = [gpu for index, gpu in selected]
+        print("Selected idle GPUs: " + json.dumps(selected), flush=True)
+        command = ["bash", str(Path(__file__).with_name("run_gpu_smoke.sh")), "--run", *result["gpu_uuids"]]
+        env = {key: value for key, value in os.environ.items() if not key.startswith("E1_BOOKING_")}
+        with (args.output / "test.log").open("w") as log:
+            child = subprocess.Popen(command, stdout=log, stderr=subprocess.STDOUT, env=env, start_new_session=True)
+            code = child.wait(timeout=360)
+        log_text = (args.output / "test.log").read_text()
+        print(log_text, flush=True)
+        if code != 0 or not has_gpu_pass(log_text):
+            raise RuntimeError(f"GPU test failed (exit={code}); no successful two-GPU qualification")
+        result.update(status="passed", gpu_tested=True)
+    except (RuntimeError, OSError, ValueError, subprocess.SubprocessError) as error:
+        result["error"] = str(error)
+    finally:
+        signal.signal(signal.SIGTERM, signal.SIG_IGN)
+        signal.signal(signal.SIGINT, signal.SIG_IGN)
+        if child is not None and child.poll() is None:
+            os.killpg(child.pid, signal.SIGTERM)
+            try:
+                child.wait(timeout=25)
+            except subprocess.TimeoutExpired:
+                os.killpg(child.pid, signal.SIGKILL)
+                child.wait(timeout=5)
+        result["finished_at"] = datetime.now(timezone.utc).isoformat()
+        (args.output / "result.json").write_text(json.dumps(result, indent=2) + "\n")
+    print(json.dumps(result, indent=2), flush=True)
+    return 0 if result["status"] == "passed" else 1
+
+
+if __name__ == "__main__":
+    try:
+        sys.exit(main())
+    except (RuntimeError, OSError, ValueError, KeyError) as error:
+        print(f"ERROR: {error}", file=sys.stderr)
+        sys.exit(1)