diff --git a/.github/scripts/ci/container_batch.py b/.github/scripts/ci/container_batch.py new file mode 100644 index 000000000000..81d746cee2d9 --- /dev/null +++ b/.github/scripts/ci/container_batch.py @@ -0,0 +1,249 @@ +"""Run an allocated GPU batch in a disposable dependency-image container. + +The caller must hold the host GPU lock until this function (including cleanup) +returns. Docker-client process cleanup alone does not stop a container. +""" + +import argparse +import json +import os +import re +import signal +import subprocess +import uuid +from pathlib import Path + +OWNER_LABEL = "org.colossalai.ci.owner" +SCOPE_LABEL = "org.colossalai.ci.scope" +OWNER = "gpu-batches" +IMAGE_REF = re.compile(r"[a-zA-Z0-9][a-zA-Z0-9._:/-]*@sha256:[0-9a-f]{64}\Z") +IMAGE_ID = re.compile(r"sha256:[0-9a-f]{64}\Z") +SCOPE = re.compile(r"[a-zA-Z0-9][a-zA-Z0-9_.-]{0,100}\Z") +IMAGE_INPUTS = ( + "Dockerfile", + "Dockerfile.dockerignore", + "requirements.lock", + "requirements-torch.lock", + "verify_image.py", +) + + +def check_image_recipe(source, image_inputs): + for name in IMAGE_INPUTS: + if (source / name).read_bytes() != (image_inputs / name).read_bytes(): + raise RuntimeError(f"CI image does not match {name}; build and qualify a new dependency image") + + +def docker(*args): + try: + return subprocess.run(["docker", *args], check=True, capture_output=True, text=True, timeout=60).stdout.strip() + except subprocess.CalledProcessError as error: + raise RuntimeError(f"Docker {args[0]} failed: {error.stderr.strip()}") from error + + +def resolve_image(reference): + if not IMAGE_ID.fullmatch(reference) and not IMAGE_REF.fullmatch(reference): + raise RuntimeError( + "COLOSSAL_CI_IMAGE must be a full local sha256:<64 hex digits> image ID " + "or an already-loaded registry image@sha256:digest" + ) + # Inspect never pulls. Build/load the image on the runner's Docker daemon + # before testing; no registry or upload is required for local image IDs. + image_id = docker("image", "inspect", "--format", "{{.Id}}", reference) + if not IMAGE_ID.fullmatch(image_id): + raise RuntimeError("Docker returned an invalid image ID") + if IMAGE_ID.fullmatch(reference) and image_id != reference: + raise RuntimeError("Docker returned a different local image ID") + return image_id + + +def validate_scope(scope): + if not SCOPE.fullmatch(scope): + raise RuntimeError("Invalid CI_CONTAINER_SCOPE") + return scope + + +def create_command(name, scope, image_id, source, output, gpu_uuids, batch, env): + validate_scope(scope) + if not IMAGE_ID.fullmatch(image_id): + raise RuntimeError("Container launch requires a resolved image ID") + if ( + not gpu_uuids + or len(set(gpu_uuids)) != len(gpu_uuids) + or any(not re.fullmatch(r"GPU-[a-fA-F0-9]{8}(?:-[a-fA-F0-9]{4}){3}-[a-fA-F0-9]{12}", gpu) for gpu in gpu_uuids) + ): + raise RuntimeError("Container launch requires allocated GPU UUIDs") + source, output = source.resolve(strict=True), output.resolve(strict=True) + if not (source / ".github/scripts/ci/run_container_batch.sh").is_file(): + raise RuntimeError("Container entry script is missing from the selected source") + if any("," in str(path) or "\n" in str(path) for path in (source, output)): + raise RuntimeError("Docker bind paths must not contain commas or newlines") + + cpus = env.get("COLOSSAL_CI_CPUS", "8") + memory = env.get("COLOSSAL_CI_MEMORY", "64g") + shm = env.get("COLOSSAL_CI_SHM", "8g") + if not re.fullmatch(r"[1-9][0-9]*(?:\.[0-9]+)?", cpus): + raise RuntimeError("COLOSSAL_CI_CPUS must be positive") + if any(not re.fullmatch(r"[1-9][0-9]*[mg]", value) for value in (memory, shm)): + raise RuntimeError("COLOSSAL_CI_MEMORY and COLOSSAL_CI_SHM must be positive m/g sizes") + + # No host venv, CUDA toolkit, home, Docker socket, credentials or compiler + # cache is mounted. Reports are written as the runner UID/GID, not root. + command = [ + "docker", + "create", + "--name", + name, + "--pull=never", + "--init", + "--label", + f"{OWNER_LABEL}={OWNER}", + "--label", + f"{SCOPE_LABEL}={scope}", + "--user", + f"{os.getuid()}:{os.getgid()}", + "--cap-drop=ALL", + "--security-opt=no-new-privileges", + "--cpus", + cpus, + "--memory", + memory, + "--memory-swap", + memory, + "--shm-size", + shm, + "--pids-limit=4096", + "--ulimit=memlock=-1:-1", + "--gpus", + '"device=' + ",".join(gpu_uuids) + '"', + "--mount", + f"type=bind,src={source},dst=/ci-source,readonly", + "--mount", + f"type=bind,src={output},dst=/ci-results", + "--workdir", + "/tmp", + ] + container_env = { + "NVIDIA_VISIBLE_DEVICES": ",".join(gpu_uuids), + "NVIDIA_DRIVER_CAPABILITIES": "compute,utility", + # Do not reuse the host indices (e.g. 4,7) after device filtering. + "CUDA_VISIBLE_DEVICES": ",".join(str(index) for index in range(len(gpu_uuids))), + "FAST_MODE": env.get("FAST_MODE", "1"), + "CI_SOURCE_SHA": env.get("CI_SOURCE_SHA", ""), + "CI_IMAGE_ID": image_id, + "GITHUB_RUN_ID": env.get("GITHUB_RUN_ID", "local"), + "COLOSSAL_HF_OFFLINE": env.get("COLOSSAL_HF_OFFLINE", "1"), + "MAX_JOBS": env.get("MAX_JOBS", "4"), + "TORCH_CUDA_ARCH_LIST": env.get("TORCH_CUDA_ARCH_LIST", "9.0"), + "PYTHONDONTWRITEBYTECODE": "1", + } + for key, value in container_env.items(): + command.extend(["--env", f"{key}={value}"]) + command.extend( + [ + "--entrypoint", + "/bin/bash", + image_id, + "/ci-source/.github/scripts/ci/run_container_batch.sh", + batch, + container_env["CUDA_VISIBLE_DEVICES"], + ] + ) + return command + + +def remove_owned_container(name, scope): + """Remove only an existing container with this exact job's ownership labels.""" + # Listing first distinguishes 'absent' from a daemon/permission failure. + ids = docker( + "ps", + "--all", + "--quiet", + "--no-trunc", + "--filter", + f"label={OWNER_LABEL}={OWNER}", + "--filter", + f"label={SCOPE_LABEL}={validate_scope(scope)}", + "--filter", + f"name=^/{name}$", + ).splitlines() + for container_id in ids: + info = json.loads(docker("inspect", container_id))[0] + labels = info["Config"].get("Labels") or {} + if info["Name"] != f"/{name}" or labels.get(OWNER_LABEL) != OWNER or labels.get(SCOPE_LABEL) != scope: + raise RuntimeError("Refusing to remove a container with different ownership") + docker("rm", "--force", container_id) + + +def run_container(batch, selected, source, output, image_id, env, stream, result): + scope = validate_scope(env.get("CI_CONTAINER_SCOPE") or f"local-{uuid.uuid4().hex}") + name = f"colossalai-ci-{scope}-{batch}-{uuid.uuid4().hex[:12]}" + command = create_command(name, scope, image_id, source, output, [gpu for _, gpu in selected], batch, env) + result.update({"container_name": name, "container_scope": scope, "image_id": image_id}) + + def interrupted(signum, frame): + raise RuntimeError(f"Container batch interrupted by signal {signum}") + + old_handlers = {sig: signal.signal(sig, interrupted) for sig in (signal.SIGTERM, signal.SIGINT)} + try: + try: + subprocess.run(command, check=True, capture_output=True, text=True, timeout=60) + except subprocess.CalledProcessError as error: + result["container_create_error"] = error.stderr + raise + client_code = stream(["docker", "start", "--attach", name], output / "pytest.log", env) + state = json.loads(docker("inspect", "--format", "{{json .State}}", name)) + result["container_state"] = state + if state.get("Running") or state.get("Status") != "exited": + raise RuntimeError("Docker client exited before the container finished") + if state.get("OOMKilled"): + raise RuntimeError("Container exceeded its configured memory limit") + if state.get("Error"): + raise RuntimeError(f"Container runtime error: {state['Error']}") + return state["ExitCode"] or client_code + finally: + # Also runs when create times out: the daemon might have created the + # container despite losing the client's response. The job's always() + # cleanup is a second safety net for runner cancellation. + for sig in old_handlers: + signal.signal(sig, signal.SIG_IGN) + try: + remove_owned_container(name, scope) + result["container_cleanup"] = "completed" + except (OSError, RuntimeError, ValueError, subprocess.SubprocessError) as error: + result["container_cleanup"] = f"failed: {error}" + raise + finally: + for sig, handler in old_handlers.items(): + signal.signal(sig, handler) + + +def cleanup_scope(scope): + """Job finalizer; never prune a daemon or remove another job's containers.""" + validate_scope(scope) + ids = docker( + "ps", + "--all", + "--quiet", + "--no-trunc", + "--filter", + f"label={OWNER_LABEL}={OWNER}", + "--filter", + f"label={SCOPE_LABEL}={scope}", + ).splitlines() + for container_id in ids: + info = json.loads(docker("inspect", container_id))[0] + remove_owned_container(info["Name"].removeprefix("/"), scope) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + mode = parser.add_mutually_exclusive_group(required=True) + mode.add_argument("--cleanup-scope") + mode.add_argument("--check-recipe", type=Path, help="Source docker/torch213 directory") + parser.add_argument("--image-inputs", type=Path, default=Path("/opt/ci")) + args = parser.parse_args() + if args.cleanup_scope: + cleanup_scope(args.cleanup_scope) + else: + check_image_recipe(args.check_recipe, args.image_inputs) diff --git a/.github/scripts/ci/run_container_batch.sh b/.github/scripts/ci/run_container_batch.sh new file mode 100644 index 000000000000..5ff98b6d367f --- /dev/null +++ b/.github/scripts/ci/run_container_batch.sh @@ -0,0 +1,79 @@ +#!/usr/bin/env bash +# Runs inside the dependency image, as the runner's non-root UID/GID. +set -euo pipefail + +BATCH="${1:?batch is required}" +GPU_LIST="${2:?container-local GPU list is required}" +# A changed recipe/lock must not accidentally pass CI using an older image. +/opt/venv/bin/python /ci-source/.github/scripts/ci/container_batch.py \ + --check-recipe /ci-source/docker/torch213 --image-inputs /opt/ci +WORK_ROOT="$(mktemp -d /tmp/colossalai-batch.XXXXXX)" +export HOME="${WORK_ROOT}/home" +export USER=colossalai-ci +export LOGNAME=colossalai-ci +export TMPDIR="${WORK_ROOT}/tmp" +export RUNNER_TEMP="${WORK_ROOT}/cache" +export COLOSSAL_CACHE_ROOT="${RUNNER_TEMP}" +export TORCH_EXTENSIONS_DIR="${RUNNER_TEMP}/torch-extensions" +export COLOSSAL_RESULTS=/ci-results +export COLOSSAL_REPO="${WORK_ROOT}/repo" +export COLOSSAL_CUDA_HOME=/usr/local/cuda +export CUDA_HOME=/usr/local/cuda +export COLOSSAL_EXPECT_TORCH_PREFIX=2.13. +export CUDA_VISIBLE_DEVICES="${GPU_LIST}" +export PIP_NO_INDEX=1 +export PIP_DISABLE_PIP_VERSION_CHECK=1 +unset COLOSSAL_VENV PYTHONPATH PYTEST_ADDOPTS +mkdir -p "${HOME}" "${TMPDIR}" "${RUNNER_TEMP}" "${COLOSSAL_REPO}" + +# The read-only bind is only an input. Builds and tests can mutate this private +# copy without leaving root-owned artifacts or stale .so files in the checkout. +cp -R --no-preserve=ownership /ci-source/. "${COLOSSAL_REPO}/" +chmod -R u+w "${COLOSSAL_REPO}" + +# A writable per-container venv reuses the image's read-only dependencies via a +# .pth file. No multi-GB Torch copy, network install or host venv is necessary. +IMAGE_PYTHON=/opt/venv/bin/python +"${IMAGE_PYTHON}" -m venv --without-pip "${WORK_ROOT}/venv" +export COLOSSAL_PYTHON="${WORK_ROOT}/venv/bin/python" +export PATH="${WORK_ROOT}/venv/bin:/usr/local/cuda/bin:${PATH}" +IMAGE_SITE="$("${IMAGE_PYTHON}" -c 'import sysconfig; print(sysconfig.get_path("purelib"))')" +LOCAL_SITE="$("${COLOSSAL_PYTHON}" -c 'import sysconfig; print(sysconfig.get_path("purelib"))')" +printf '%s\n' "${IMAGE_SITE}" > "${LOCAL_SITE}/ci-image-dependencies.pth" + +cd "${COLOSSAL_REPO}" +"${COLOSSAL_PYTHON}" - <<'PY' +import json +import os +import sys +from pathlib import Path + +import torch + +assert torch.__version__.startswith("2.13."), torch.__version__ +assert torch.version.cuda == "13.0", torch.version.cuda +assert torch.cuda.is_available(), "CUDA is not available inside the container" +assert torch.cuda.device_count() == len(os.environ["CUDA_VISIBLE_DEVICES"].split(",")) +profile = { + "source_sha": os.environ["CI_SOURCE_SHA"], + "image_id": os.environ["CI_IMAGE_ID"], + "python": sys.version, + "torch": torch.__version__, + "cuda": torch.version.cuda, + "gpu_names": [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())], +} +print(json.dumps(profile, indent=2), flush=True) +Path("/ci-results/runtime.json").write_text(json.dumps(profile, indent=2) + "\n") +PY + +# Batch 3 explicitly validates the AOT build. Other batches compile the exact +# source's kernels on demand (JIT) in disposable, batch-local caches. +BUILD_EXT=0 +if [[ "${BATCH}" == "3" ]]; then + BUILD_EXT=1 +fi +export BUILD_EXT +printf 'Installing source %s with BUILD_EXT=%s\n' "${CI_SOURCE_SHA}" "${BUILD_EXT}" +"${COLOSSAL_PYTHON}" -m pip install --no-deps --no-build-isolation -e . +"${COLOSSAL_PYTHON}" -m pip check +exec bash .github/scripts/ci/run_colossalai_batch.sh "${BATCH}" "${GPU_LIST}" diff --git a/.github/scripts/ci/run_gpu_batch.py b/.github/scripts/ci/run_gpu_batch.py index 435493faa14a..b46cb1e624ed 100644 --- a/.github/scripts/ci/run_gpu_batch.py +++ b/.github/scripts/ci/run_gpu_batch.py @@ -15,6 +15,8 @@ from datetime import datetime, timezone from pathlib import Path +from container_batch import resolve_image, run_container + GPU_UUID = re.compile(r"GPU-[a-fA-F0-9]{8}(?:-[a-fA-F0-9]{4}){3}-[a-fA-F0-9]{12}\Z") GPU_COUNTS = { "2": 1, @@ -88,7 +90,11 @@ def select_idle(inventory, processes, count, allowed=None): def query_gpus(): inventory = subprocess.run( - ["nvidia-smi", "--query-gpu=index,uuid,memory.used,utilization.gpu", "--format=csv,noheader,nounits"], + [ + "nvidia-smi", + "--query-gpu=index,uuid,memory.used,utilization.gpu", + "--format=csv,noheader,nounits", + ], check=True, text=True, capture_output=True, @@ -184,6 +190,10 @@ def main(): parser.add_argument("--batch", required=True, choices=sorted(GPU_COUNTS)) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--lock", type=Path, default=Path("/tmp/colossalai-ci-gpu.lock")) + parser.add_argument( + "--image", + help="Preloaded dependency image: full sha256:ID or registry@sha256:digest; omit for host mode", + ) args = parser.parse_args() if os.name != "posix": @@ -202,6 +212,9 @@ def main(): } try: + image_id = resolve_image(args.image) if args.image is not None else None + result["executor"] = "docker" if image_id else "host" + result["image_reference"] = args.image with args.lock.open("w") as lock: fcntl.flock(lock, fcntl.LOCK_EX) inventory, processes = query_gpus() @@ -223,13 +236,31 @@ def main(): GPU_COUNTS[args.batch], {index for index, _ in selected}, ) - if [index for index, _ in still_idle] != result["gpu_indices"]: + if still_idle != selected: raise RuntimeError("GPU occupancy changed before the batch could start") script = Path(__file__).with_name("run_colossalai_batch.sh") gpu_list = ",".join(str(index) for index, _ in selected) env = os.environ.copy() - code = stream_process(["bash", str(script), args.batch, gpu_list], args.output / "pytest.log", env) + if image_id: + # Cleanup is part of run_container and completes BEFORE this + # lock context exits. Killing the Docker client is not enough. + code = run_container( + args.batch, + selected, + Path(__file__).resolve().parents[3], + args.output, + image_id, + env, + stream_process, + result, + ) + else: + code = stream_process( + ["bash", str(script), args.batch, gpu_list], + args.output / "pytest.log", + env, + ) if code != 0: raise RuntimeError(f"Batch {args.batch} failed with exit code {code}") result["status"] = "passed" diff --git a/.github/scripts/ci/test_container_batch.py b/.github/scripts/ci/test_container_batch.py new file mode 100644 index 000000000000..a551027c4a7d --- /dev/null +++ b/.github/scripts/ci/test_container_batch.py @@ -0,0 +1,341 @@ +import io +import json +import os +import signal +import subprocess +import tempfile +import unittest +from contextlib import redirect_stdout +from pathlib import Path +from unittest.mock import Mock, patch + +import container_batch as container +import run_gpu_batch as runner + +IMAGE = "sha256:" + "a" * 64 +REFERENCE = "ghcr.io/example/ci@sha256:" + "b" * 64 +GPU_A = "GPU-00000000-0000-0000-0000-000000000001" +GPU_B = "GPU-00000000-0000-0000-0000-000000000002" +SCOPE = "gh-123-1-layer-2" + + +class ContainerCommandTests(unittest.TestCase): + def setUp(self): + self.directory = tempfile.TemporaryDirectory() + self.addCleanup(self.directory.cleanup) + self.root = Path(self.directory.name) + script = self.root / ".github/scripts/ci/run_container_batch.sh" + script.parent.mkdir(parents=True) + script.touch() + self.output = self.root / "reports" + self.output.mkdir() + + def command(self, **env): + return container.create_command( + "test-container", + SCOPE, + IMAGE, + self.root, + self.output, + [GPU_A, GPU_B], + "5", + env, + ) + + def test_uuid_filter_and_logical_device_mapping(self): + command = self.command(CUDA_VISIBLE_DEVICES="4,7") + self.assertEqual(command[command.index("--gpus") + 1], f'"device={GPU_A},{GPU_B}"') + self.assertIn("CUDA_VISIBLE_DEVICES=0,1", command) + self.assertIn(f"NVIDIA_VISIBLE_DEVICES={GPU_A},{GPU_B}", command) + self.assertEqual(command[-2:], ["5", "0,1"]) + + def test_only_source_and_reports_are_mounted(self): + command = self.command(COLOSSAL_VENV="/host/private/venv", SECRET_TOKEN="not-for-container") + mounts = [command[index + 1] for index, value in enumerate(command) if value == "--mount"] + self.assertEqual( + mounts, + [ + f"type=bind,src={self.root.resolve()},dst=/ci-source,readonly", + f"type=bind,src={self.output.resolve()},dst=/ci-results", + ], + ) + self.assertNotIn("SECRET_TOKEN=not-for-container", command) + self.assertNotIn("COLOSSAL_VENV=/host/private/venv", command) + self.assertNotIn("--privileged", command) + self.assertNotIn("--ipc=host", command) + self.assertIn(f"{os.getuid()}:{os.getgid()}", command) + self.assertIn("--pull=never", command) + + def test_resource_limits(self): + command = self.command(COLOSSAL_CI_CPUS="4", COLOSSAL_CI_MEMORY="32g", COLOSSAL_CI_SHM="4g") + for flag, expected in ( + ("--cpus", "4"), + ("--memory", "32g"), + ("--memory-swap", "32g"), + ("--shm-size", "4g"), + ): + self.assertEqual(command[command.index(flag) + 1], expected) + for env in ( + {"COLOSSAL_CI_CPUS": "0"}, + {"COLOSSAL_CI_MEMORY": "unlimited"}, + {"COLOSSAL_CI_SHM": "0g"}, + ): + with self.subTest(env=env), self.assertRaises(RuntimeError): + self.command(**env) + + def test_mutable_or_empty_image_is_rejected_before_docker(self): + with patch.object(container, "docker") as docker: + for reference in ("", "ci:latest", "ci:torch213", "ci@sha256:invalid", "sha256:abc"): + with self.subTest(reference=reference), self.assertRaises(RuntimeError): + container.resolve_image(reference) + docker.assert_not_called() + + def test_registry_digest_resolves_to_local_content_id(self): + with patch.object(container, "docker", return_value=IMAGE) as docker: + self.assertEqual(container.resolve_image(REFERENCE), IMAGE) + docker.assert_called_once_with("image", "inspect", "--format", "{{.Id}}", REFERENCE) + + def test_local_image_id_needs_no_registry(self): + with patch.object(container, "docker", return_value=IMAGE) as docker: + self.assertEqual(container.resolve_image(IMAGE), IMAGE) + docker.assert_called_once_with("image", "inspect", "--format", "{{.Id}}", IMAGE) + + def test_missing_local_image_fails_without_pulling(self): + with patch.object(container, "docker", side_effect=RuntimeError("No such image")) as docker: + with self.assertRaisesRegex(RuntimeError, "No such image"): + container.resolve_image(IMAGE) + docker.assert_called_once_with("image", "inspect", "--format", "{{.Id}}", IMAGE) + + def test_wrong_local_image_id_fails_closed(self): + with patch.object(container, "docker", return_value="sha256:" + "c" * 64): + with self.assertRaisesRegex(RuntimeError, "different local image ID"): + container.resolve_image(IMAGE) + + def run_batch(self, stream): + result = {} + code = container.run_container( + "5", + [(4, GPU_A), (7, GPU_B)], + self.root, + self.output, + IMAGE, + {"CI_CONTAINER_SCOPE": SCOPE}, + stream, + result, + ) + return code, result + + def test_success_and_test_failure_preserve_exit_and_cleanup(self): + for exit_code in (0, 1, 5, 124): + with ( + self.subTest(exit_code=exit_code), + patch.object(container.subprocess, "run"), + patch.object( + container, + "docker", + return_value=json.dumps( + { + "Running": False, + "Status": "exited", + "ExitCode": exit_code, + } + ), + ), + patch.object(container, "remove_owned_container") as cleanup, + ): + code, result = self.run_batch(Mock(return_value=exit_code)) + self.assertEqual(code, exit_code) + self.assertEqual(result["container_cleanup"], "completed") + cleanup.assert_called_once_with(result["container_name"], SCOPE) + + def test_create_failure_still_attempts_cleanup(self): + for error in ( + subprocess.CalledProcessError(125, ["docker", "create"]), + subprocess.TimeoutExpired("docker", 60), + ): + with ( + self.subTest(error=error), + patch.object(container.subprocess, "run", side_effect=error), + patch.object(container, "remove_owned_container") as cleanup, + ): + with self.assertRaises(subprocess.SubprocessError): + self.run_batch(Mock()) + cleanup.assert_called_once() + + def test_interruption_cleans_container_and_restores_signal_handlers(self): + previous = signal.getsignal(signal.SIGTERM) + with patch.object(container.subprocess, "run"), patch.object(container, "remove_owned_container") as cleanup: + with self.assertRaisesRegex(RuntimeError, "cancelled"): + self.run_batch(Mock(side_effect=RuntimeError("cancelled"))) + cleanup.assert_called_once() + self.assertEqual(signal.getsignal(signal.SIGTERM), previous) + + def test_running_or_oom_container_is_not_a_pass(self): + for state in ( + {"Running": True}, + {"Running": False, "Status": "exited", "OOMKilled": True, "ExitCode": 0}, + ): + with ( + self.subTest(state=state), + patch.object(container.subprocess, "run"), + patch.object(container, "docker", return_value=json.dumps(state)), + patch.object(container, "remove_owned_container") as cleanup, + ): + with self.assertRaises(RuntimeError): + self.run_batch(Mock(return_value=0)) + cleanup.assert_called_once() + + def test_cleanup_failure_cannot_report_success(self): + with ( + patch.object(container.subprocess, "run"), + patch.object( + container, + "docker", + return_value=json.dumps( + { + "Running": False, + "Status": "exited", + "ExitCode": 0, + } + ), + ), + patch.object( + container, + "remove_owned_container", + side_effect=RuntimeError("daemon unavailable"), + ), + ): + with self.assertRaisesRegex(RuntimeError, "daemon unavailable"): + self.run_batch(Mock(return_value=0)) + + +class ContainerOwnershipTests(unittest.TestCase): + def test_absent_container_needs_no_removal(self): + with patch.object(container, "docker", return_value="") as docker: + container.remove_owned_container("owned", SCOPE) + self.assertEqual(docker.call_count, 1) + + def test_removal_checks_exact_name_and_both_labels(self): + good = { + "Name": "/owned", + "Config": { + "Labels": { + container.OWNER_LABEL: container.OWNER, + container.SCOPE_LABEL: SCOPE, + } + }, + } + with patch.object(container, "docker", side_effect=["container-id", json.dumps([good]), ""]) as docker: + container.remove_owned_container("owned", SCOPE) + self.assertEqual(docker.call_args.args, ("rm", "--force", "container-id")) + for info in ( + {"Name": "/someone-else", "Config": good["Config"]}, + { + "Name": "/owned", + "Config": {"Labels": {container.OWNER_LABEL: container.OWNER}}, + }, + {"Name": "/owned", "Config": {"Labels": {container.SCOPE_LABEL: SCOPE}}}, + ): + with ( + self.subTest(info=info), + patch.object(container, "docker", side_effect=["container-id", json.dumps([info])]) as docker, + ): + with self.assertRaises(RuntimeError): + container.remove_owned_container("owned", SCOPE) + self.assertEqual(docker.call_count, 2) + + def test_invalid_cleanup_scope_is_rejected(self): + with patch.object(container, "docker") as docker: + for scope in ("", "*", "other\njob", "x" * 102): + with self.subTest(scope=scope), self.assertRaises(RuntimeError): + container.cleanup_scope(scope) + docker.assert_not_called() + + +class ImageRecipeTests(unittest.TestCase): + def test_recipe_must_match_image_including_locks_and_verifier(self): + with tempfile.TemporaryDirectory() as directory: + source, image = Path(directory) / "source", Path(directory) / "image" + source.mkdir() + image.mkdir() + for name in container.IMAGE_INPUTS: + (source / name).write_text("same\n") + (image / name).write_text("same\n") + container.check_image_recipe(source, image) + for name in container.IMAGE_INPUTS: + (source / name).write_text("changed\n") + with self.subTest(name=name), self.assertRaisesRegex(RuntimeError, "does not match"): + container.check_image_recipe(source, image) + (source / name).write_text("same\n") + (image / "requirements.lock").unlink() + with self.assertRaises(FileNotFoundError): + container.check_image_recipe(source, image) + + +@unittest.skipUnless(os.name == "posix", "flock requires POSIX") +class LockedContainerTests(unittest.TestCase): + def test_container_lifecycle_is_inside_gpu_lock(self): + import fcntl + + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + lock_path = root / "gpu.lock" + command = [ + "run_gpu_batch.py", + "--batch", + "5", + "--output", + str(root / "out"), + "--lock", + str(lock_path), + "--image", + REFERENCE, + ] + + def fake_container(*args): + with lock_path.open("w") as second: + with self.assertRaises(BlockingIOError): + fcntl.flock(second, fcntl.LOCK_EX | fcntl.LOCK_NB) + self.assertEqual(args[1], [(4, GPU_A), (7, GPU_B)]) + return 0 + + with ( + patch("sys.argv", command), + patch.dict(os.environ, {"COLOSSAL_GPU_POOL": "4,7"}), + patch.object(runner, "resolve_image", return_value=IMAGE), + patch.object(runner, "query_gpus", return_value=(f"4,{GPU_A},4,0\n7,{GPU_B},4,0", "")), + patch.object(runner, "run_container", side_effect=fake_container) as launch, + redirect_stdout(io.StringIO()), + ): + self.assertEqual(runner.main(), 0) + launch.assert_called_once() + with lock_path.open("w") as unlocked: + fcntl.flock(unlocked, fcntl.LOCK_EX | fcntl.LOCK_NB) + + def test_empty_explicit_image_never_falls_back_to_host(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + command = [ + "run_gpu_batch.py", + "--batch", + "4a", + "--output", + str(root / "out"), + "--lock", + str(root / "gpu.lock"), + "--image", + "", + ] + with ( + patch("sys.argv", command), + patch.object(runner, "query_gpus") as query, + patch.object(runner, "stream_process") as stream, + redirect_stdout(io.StringIO()), + ): + self.assertEqual(runner.main(), 1) + query.assert_not_called() + stream.assert_not_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/scripts/ci/test_draft_pr_workflows.py b/.github/scripts/ci/test_draft_pr_workflows.py new file mode 100644 index 000000000000..b9ec1243a017 --- /dev/null +++ b/.github/scripts/ci/test_draft_pr_workflows.py @@ -0,0 +1,301 @@ +"""CPU-only regression checks for Draft admission and GPU-pool isolation.""" + +import ast +import json +import os +import subprocess +import tempfile +import unittest +from pathlib import Path + +from check_pr_runner_policy import event_enabled, is_github_hosted, load_workflow, workflow_paths + +ROOT = Path(__file__).resolve().parents[3] +WORKFLOWS = ROOT / ".github/workflows" +REUSABLE = "./.github/workflows/ci_reusable_gpu_batches.yml" +UPSTREAM = "hpcaitech/ColossalAI" + + +def evaluate_condition(expression, context): + """Evaluate only the typed boolean/string subset used by these job guards. + + This is a test helper, not a general GitHub expression engine. Unknown + syntax fails the test; actual workflow syntax is also checked by actionlint. + """ + tree = ast.parse(expression.replace("&&", " and ").replace("||", " or ").strip(), mode="eval") + + def folded(value): + return value.casefold() if isinstance(value, str) else value + + def visit(node): + if isinstance(node, ast.Constant): + return node.value + if isinstance(node, ast.Name) and node.id in {"true", "false"}: + return node.id == "true" + if isinstance(node, ast.Attribute): + parts = [] + while isinstance(node, ast.Attribute): + parts.append(node.attr) + node = node.value + if isinstance(node, ast.Name): + return context[".".join([node.id, *reversed(parts)])] + if isinstance(node, ast.BoolOp): + values = (visit(value) for value in node.values) + return all(values) if isinstance(node.op, ast.And) else any(values) + if isinstance(node, ast.Compare) and len(node.ops) == 1: + left, right = folded(visit(node.left)), folded(visit(node.comparators[0])) + if isinstance(node.ops[0], ast.Eq): + return left == right + if isinstance(node.ops[0], ast.NotEq): + return left != right + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name): + values = [visit(value) for value in node.args] + if node.func.id == "fromJSON" and len(values) == 1: + return json.loads(values[0]) + if node.func.id == "contains" and len(values) == 2: + return folded(values[1]) in [folded(value) for value in values[0]] + raise AssertionError(f"Unsupported workflow condition syntax: {ast.dump(node)}") + + return bool(visit(tree.body)) + + +def github_context(draft=False, association="OWNER", head=UPSTREAM, repository=UPSTREAM, event="pull_request"): + return { + "github.repository": repository, + "github.event_name": event, + "github.event.pull_request.draft": draft, + "github.event.pull_request.head.repo.full_name": head, + "github.event.pull_request.author_association": association, + } + + +class DraftPullRequestRoutingTests(unittest.TestCase): + def setUp(self): + self.layer2 = load_workflow(WORKFLOWS / "ci_layer_2_gpu_pr.yml") + + def test_trusted_draft_and_ready_prs_have_distinct_jobs(self): + for association in ("OWNER", "MEMBER", "COLLABORATOR"): + for draft in (True, False): + with self.subTest(association=association, draft=draft): + context = github_context(draft=draft, association=association) + jobs = self.layer2["jobs"] + self.assertEqual(evaluate_condition(jobs["container-smoke"]["if"], context), draft) + self.assertEqual(evaluate_condition(jobs["gpu-regression"]["if"], context), not draft) + + def test_forks_untrusted_authors_and_other_repositories_never_admit_gpu_jobs(self): + for draft in (True, False): + for overrides in ( + {"head": "contributor/ColossalAI"}, + {"association": "CONTRIBUTOR"}, + {"association": "NONE"}, + {"repository": "contributor/ColossalAI", "head": "contributor/ColossalAI"}, + {"event": "push"}, + ): + with self.subTest(draft=draft, overrides=overrides): + context = github_context(draft=draft, **overrides) + for job in self.layer2["jobs"].values(): + self.assertFalse(evaluate_condition(job["if"], context)) + + def test_manual_dispatch_retains_full_regression(self): + context = github_context(event="workflow_dispatch") + self.assertFalse(evaluate_condition(self.layer2["jobs"]["container-smoke"]["if"], context)) + self.assertTrue(evaluate_condition(self.layer2["jobs"]["gpu-regression"]["if"], context)) + + def test_smoke_membership_check_name_and_fast_mode_are_explicit(self): + smoke, core = (self.layer2["jobs"][key] for key in ("container-smoke", "gpu-regression")) + self.assertNotEqual(smoke["name"], core["name"]) + self.assertEqual(core["name"], "Torch 2.13 GPU core (batches 2 and 5)") + self.assertEqual(smoke["uses"], REUSABLE) + self.assertEqual(core["uses"], REUSABLE) + self.assertEqual(smoke["with"]["source_sha"], "${{ github.sha }}") + self.assertEqual(smoke["with"]["batches"], "4a") + self.assertEqual(smoke["with"]["gpu_pool_variable"], "COLOSSAL_CI_SMOKE_GPU_POOL") + self.assertEqual(smoke["with"]["fast_mode"], "false") + self.assertLessEqual(int(smoke["with"]["timeout_minutes"]), 15) + self.assertEqual(core["with"]["batches"], "2 4a 5") + self.assertEqual(core["with"]["fast_mode"], "true") + self.assertNotIn("gpu_pool_variable", core["with"]) + + def test_ready_and_converted_to_draft_events_reconcile_layer2(self): + types = self.layer2["on"]["pull_request"]["types"] + self.assertIn("ready_for_review", types) + self.assertIn("converted_to_draft", types) + self.assertEqual(self.layer2["concurrency"]["cancel-in-progress"], "true") + + def test_retired_e1_host_workflow_cannot_run_alongside_containers(self): + self.assertFalse((WORKFLOWS / "e1-gpu-on-pr.yml").exists()) + self.assertFalse((WORKFLOWS / "e1-gpu-on-pr.yaml").exists()) + for path in workflow_paths(WORKFLOWS): + source = path.read_text(encoding="utf-8") + for legacy_entry in ( + "E1_COLOSSALAI_PYTHON", + ".github/scripts/e1/pr_gpu.py", + ".github/scripts/e1/colossalai_suite.py", + ): + with self.subTest(workflow=path.name, entry=legacy_entry): + self.assertNotIn(legacy_entry, source) + + def test_retired_e1_tests_remain_in_container_core_batches(self): + manifest = (ROOT / ".github/scripts/ci/run_colossalai_batch.sh").read_text(encoding="utf-8") + for case, path, test_name in ( + ("2|02", "tests/test_booster/test_accelerator.py", "test_accelerator"), + ("5|05", "tests/test_booster/test_plugin/test_dp_plugin_base.py", "test_dp_plugin_dataloader"), + ): + with self.subTest(batch=case, test=test_name): + batch = manifest.split(f" {case})\n", 1)[1].split(" ;;", 1)[0] + self.assertIn(path, batch) + module = ast.parse((ROOT / path).read_text(encoding="utf-8")) + self.assertIn(test_name, [node.name for node in module.body if isinstance(node, ast.FunctionDef)]) + + def test_legacy_gpu_entrypoints_skip_drafts_but_keep_ready_prs(self): + for filename in ("compatibility_test_on_pr.yml", "cuda_ext_check_before_merge.yml"): + workflow = load_workflow(WORKFLOWS / filename) + self.assertIn("ready_for_review", workflow["on"]["pull_request"]["types"]) + for job_id, job in workflow["jobs"].items(): + with self.subTest(workflow=filename, job=job_id): + self.assertFalse(evaluate_condition(job["if"], github_context(draft=True))) + self.assertTrue(evaluate_condition(job["if"], github_context(draft=False))) + self.assertFalse(evaluate_condition(job["if"], github_context(head="contributor/ColossalAI"))) + + def test_cuda_extension_manual_dispatch_remains_available(self): + workflow = load_workflow(WORKFLOWS / "cuda_ext_check_before_merge.yml") + for job in workflow["jobs"].values(): + self.assertTrue(evaluate_condition(job["if"], github_context(event="workflow_dispatch"))) + + def test_all_other_private_pr_jobs_have_a_draft_or_post_merge_gate(self): + for path in workflow_paths(WORKFLOWS): + workflow = load_workflow(path) + if not event_enabled(workflow, "pull_request"): + continue + jobs = workflow["jobs"] + + def gated(job_id): + job = jobs[job_id] + condition = " ".join(job.get("if", "").split()) + if any( + guard in condition + for guard in ( + "github.event.pull_request.draft == false", + "github.event.pull_request.merged == true", + ) + ): + return True + needs = job.get("needs", []) + needs = [needs] if isinstance(needs, str) else needs + return "always()" not in condition and any(gated(parent) for parent in needs) + + for job_id, job in jobs.items(): + if is_github_hosted(job.get("runs-on")): + continue + if path.name == "ci_layer_2_gpu_pr.yml" and job_id == "container-smoke": + continue + with self.subTest(workflow=path.name, job=job_id): + self.assertTrue(gated(job_id), "Private Draft PR jobs must be gated before runner allocation") + + +class LegacyPrBuildRetirementTests(unittest.TestCase): + def test_legacy_pr_build_and_coverage_workflows_are_removed(self): + for name in ("build_on_pr", "report_test_coverage"): + for suffix in (".yml", ".yaml"): + with self.subTest(workflow=name + suffix): + self.assertFalse((WORKFLOWS / (name + suffix)).exists()) + + def test_active_workflows_do_not_reference_retired_pr_builds(self): + for path in workflow_paths(WORKFLOWS): + source = path.read_text(encoding="utf-8") + for entry in ("Build on PR", "Report Test Coverage", "build_on_pr.yml", "report_test_coverage.yml"): + with self.subTest(workflow=path.name, entry=entry): + self.assertNotIn(entry, source) + + def test_layer3_retains_batch_three_aot_compilation(self): + workflow = load_workflow(WORKFLOWS / "ci_layer_3_full.yml") + default = workflow["on"]["workflow_dispatch"]["inputs"]["batches"]["default"] + self.assertIn("3", default.split()) + self.assertIn(f"|| '{default}'", workflow["jobs"]["torch-213"]["with"]["batches"]) + entrypoint = (ROOT / ".github/scripts/ci/run_container_batch.sh").read_text(encoding="utf-8") + build_mode = "BUILD_EXT=0\n" + entrypoint.split("\nBUILD_EXT=0\n", 1)[1].split("\nexport BUILD_EXT\n", 1)[0] + for batch, expected in (("3", "1"), ("2", "0"), ("4a", "0"), ("5", "0")): + with self.subTest(batch=batch): + result = subprocess.run( + ["bash", "-eu", "-c", build_mode + '\nprintf "%s" "$BUILD_EXT"'], + env={"PATH": os.defpath, "BATCH": batch}, + capture_output=True, + text=True, + check=True, + timeout=10, + ) + self.assertEqual(result.stdout, expected) + + +class GpuPoolConfigurationTests(unittest.TestCase): + def setUp(self): + self.workflow = load_workflow(WORKFLOWS / "ci_reusable_gpu_batches.yml") + self.job = self.workflow["jobs"]["run-batches"] + + def test_pool_lookup_has_no_truthy_fallback_and_core_default_is_unchanged(self): + inputs = self.workflow["on"]["workflow_call"]["inputs"] + self.assertEqual(inputs["gpu_pool_variable"]["default"], "COLOSSAL_GPU_POOL") + self.assertEqual(self.job["env"]["COLOSSAL_GPU_POOL"], "${{ vars[inputs.gpu_pool_variable] }}") + self.assertEqual(self.job["env"]["CI_GPU_LOCK"], "${{ vars.COLOSSAL_GPU_LOCK }}") + for filename in ("ci_layer_3_full.yml", "ci_gpu_idle_monitor.yml"): + for job in load_workflow(WORKFLOWS / filename)["jobs"].values(): + if job.get("uses") == REUSABLE: + self.assertNotIn("gpu_pool_variable", job["with"]) + + def validate_configuration(self, pool_variable, pool, batches="4a"): + script = self.job["steps"][0]["run"] + for key, value in (("inputs.layer", "draft-smoke"), ("github.run_id", "123"), ("github.run_attempt", "1")): + script = script.replace("${{ " + key + " }}", value) + self.assertNotIn("${{", script) + with tempfile.TemporaryDirectory() as directory: + env_file = Path(directory) / "github-env" + result = subprocess.run( + ["bash", "-c", script], + env={ + "PATH": os.defpath, + "COLOSSAL_CI_IMAGE": "sha256:" + "a" * 64, + "CI_GPU_LOCK": str(Path(directory) / "shared.lock"), + "GPU_POOL_VARIABLE": pool_variable, + "COLOSSAL_GPU_POOL": pool, + "REQUESTED_BATCHES": batches, + "RUNNER_TEMP": directory, + "GITHUB_ENV": str(env_file), + }, + capture_output=True, + text=True, + timeout=10, + ) + self.assertEqual(env_file.exists(), result.returncode == 0) + return result + + def test_empty_smoke_pool_fails_before_results_or_gpu_work(self): + result = self.validate_configuration("COLOSSAL_CI_SMOKE_GPU_POOL", "") + self.assertNotEqual(result.returncode, 0) + self.assertIn("COLOSSAL_CI_SMOKE_GPU_POOL", result.stderr) + self.assertIn("no pool fallback", result.stderr) + + def test_smoke_accepts_an_explicit_pool_including_gpu_zero(self): + for pool in ("0", "3", "3,7"): + with self.subTest(pool=pool): + result = self.validate_configuration("COLOSSAL_CI_SMOKE_GPU_POOL", pool) + self.assertEqual(result.returncode, 0, result.stderr) + + def test_main_pool_still_accepts_multicard_batches(self): + result = self.validate_configuration("COLOSSAL_GPU_POOL", "2,3", batches="2 4a 5") + self.assertEqual(result.returncode, 0, result.stderr) + + def test_invalid_pool_or_profile_and_smoke_scope_expansion_fail(self): + for variable, pool, batches in ( + ("OTHER_POOL", "3", "4a"), + ("COLOSSAL_GPU_POOL", "", "2 4a 5"), + ("COLOSSAL_CI_SMOKE_GPU_POOL", "3", "2 4a 5"), + ("COLOSSAL_CI_SMOKE_GPU_POOL", "all", "4a"), + ("COLOSSAL_CI_SMOKE_GPU_POOL", "3,", "4a"), + ("COLOSSAL_CI_SMOKE_GPU_POOL", "3\n4", "4a"), + ): + with self.subTest(variable=variable, pool=pool, batches=batches): + self.assertNotEqual(self.validate_configuration(variable, pool, batches).returncode, 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/scripts/ci/test_optional_apex_guards.py b/.github/scripts/ci/test_optional_apex_guards.py new file mode 100644 index 000000000000..f66a62375134 --- /dev/null +++ b/.github/scripts/ci/test_optional_apex_guards.py @@ -0,0 +1,146 @@ +"""Check optional Apex probes without importing the GPU-dependent test modules. + +Execute the actual importlib imports and module-level guard from each test file. +This lets the dependency-free CI helper suite catch unused-import autofixes that +silently turn an optional dependency probe into an unconditional True. +""" + +import ast +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +ROOT = Path(__file__).resolve().parents[3] +OPTIMIZER = "tests/test_zero/test_gemini/test_optim.py" +CHECKPOINT = "tests/test_checkpoint_io/test_gemini_checkpoint_io.py" +GUARDS = {OPTIMIZER: "APEX_AMP_AVAILABLE", CHECKPOINT: "APEX_AVAILABLE"} + + +def read_module(path): + return ast.parse((ROOT / path).read_text(encoding="utf-8"), filename=path) + + +def probe(path, modules): + tree = read_module(path) + imports = [node for node in tree.body if isinstance(node, ast.ImportFrom) and node.module == "importlib"] + guard = next( + node + for node in tree.body + if isinstance(node, ast.Try) + and any(isinstance(child, ast.Name) and child.id == GUARDS[path] for child in ast.walk(node)) + ) + code = compile(ast.Module(body=[*imports, guard], type_ignores=[]), path, "exec") + + def import_optional(name): + if name not in modules: + raise ModuleNotFoundError(f"No module named {name!r}") + value = modules[name] + if isinstance(value, Exception): + raise value + return value + + namespace = {} + with patch("importlib.import_module", side_effect=import_optional) as importer: + exec(code, namespace) + return namespace[GUARDS[path]], [call.args[0] for call in importer.call_args_list] + + +def normalization_modules(): + return { + "apex.normalization": SimpleNamespace(FusedLayerNorm=object, FusedRMSNorm=object), + "fused_layer_norm_cuda": SimpleNamespace(), + } + + +class OptionalApexGuardTests(unittest.TestCase): + def test_missing_apex_is_unavailable(self): + for path in GUARDS: + with self.subTest(path=path): + available, imports = probe(path, {}) + self.assertFalse(available) + self.assertTrue(imports, "The guard must attempt an import, not assign a constant") + + def test_normalization_does_not_substitute_for_legacy_amp(self): + available, imports = probe(OPTIMIZER, normalization_modules()) + self.assertFalse(available) + self.assertEqual(imports, ["apex.amp"]) + + def test_importable_legacy_amp_is_available(self): + available, imports = probe(OPTIMIZER, {"apex.amp": SimpleNamespace()}) + self.assertTrue(available) + self.assertEqual(imports, ["apex.amp"]) + + def test_fused_normalization_does_not_require_legacy_amp(self): + available, imports = probe(CHECKPOINT, normalization_modules()) + self.assertTrue(available) + self.assertEqual(imports, ["apex.normalization", "fused_layer_norm_cuda"]) + + def test_python_only_apex_does_not_enable_fused_kernels(self): + modules = normalization_modules() + del modules["fused_layer_norm_cuda"] + available, imports = probe(CHECKPOINT, modules) + self.assertFalse(available) + self.assertIn("fused_layer_norm_cuda", imports) + + def test_both_normalization_classes_are_required(self): + for missing in ("FusedLayerNorm", "FusedRMSNorm"): + with self.subTest(missing=missing): + modules = normalization_modules() + delattr(modules["apex.normalization"], missing) + self.assertFalse(probe(CHECKPOINT, modules)[0]) + + def test_unloadable_cuda_extension_is_unavailable(self): + modules = normalization_modules() + modules["fused_layer_norm_cuda"] = ImportError("undefined symbol") + self.assertFalse(probe(CHECKPOINT, modules)[0]) + + def test_unexpected_initialization_errors_are_not_silenced(self): + for path, module in ((OPTIMIZER, "apex.amp"), (CHECKPOINT, "apex.normalization")): + with self.subTest(path=path), self.assertRaisesRegex(RuntimeError, "unexpected failure"): + probe(path, {module: RuntimeError("unexpected failure")}) + + def test_optimizer_skip_uses_the_probe(self): + function = next( + node + for node in read_module(OPTIMIZER).body + if isinstance(node, ast.FunctionDef) and node.name == "test_optim" + ) + marker = next( + node + for node in function.decorator_list + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr == "skipif" + ) + condition = compile(ast.Expression(marker.args[0]), OPTIMIZER, "eval") + for available in (True, False): + with self.subTest(available=available): + self.assertEqual(eval(condition, {GUARDS[OPTIMIZER]: available}), not available) + + def test_checkpoint_keeps_native_normalization_without_apex(self): + expected = {"exam_state_dict_with_origin", "exam_state_dict"} + functions = [ + node for node in read_module(CHECKPOINT).body if isinstance(node, ast.FunctionDef) and node.name in expected + ] + self.assertEqual({node.name for node in functions}, expected) + for node in functions: + setting = next( + child + for child in node.body + if isinstance(child, ast.Assign) + and any( + isinstance(target, ast.Name) and target.id == "enable_fused_normalization" + for target in child.targets + ) + ) + condition = compile(ast.Expression(setting.value), CHECKPOINT, "eval") + for tp_size in (1, 2): + for available in (True, False): + with self.subTest(function=node.name, tp_size=tp_size, available=available): + self.assertEqual( + eval(condition, {GUARDS[CHECKPOINT]: available, "tp_size": tp_size}), + tp_size > 1 and available, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/scripts/e1/PR.md b/.github/scripts/e1/PR.md index f0d725c549a2..947821103ef1 100644 --- a/.github/scripts/e1/PR.md +++ b/.github/scripts/e1/PR.md @@ -1,6 +1,16 @@ -# E1 ColossalAI PR regression +# E1 ColossalAI PR regression (retired) -The `ColossalAI GPU on internal PR` workflow runs two existing ColossalAI tests: +> Historical reference only. The `e1-gpu-on-pr.yml` workflow has been removed; +> no active workflow invokes this host-venv PR suite. PR GPU regression now uses +> the container-based Layer 2 workflow: batch 2 covers the accelerator test and +> batch 5 covers the two-rank DP dataloader test. See the +> [current CI workflow documentation](../../workflows/README.md#three-layer-ci-rollout). +> The scripts are retained for history, not as an automatic CI fallback. +> The runner registration and infrastructure bootstrap probe are unchanged. + +The sections below describe the retired implementation, not the current PR flow. + +The `ColossalAI GPU on internal PR` workflow ran two existing ColossalAI tests: - `tests/test_booster/test_accelerator.py::test_accelerator`: verifies model placement on CPU and CUDA. - `tests/test_booster/test_plugin/test_dp_plugin_base.py::test_dp_plugin_dataloader`: launches two NCCL workers through `colossalai.launch`, prepares a DPPlugin dataloader, and checks that ranks receive different data. diff --git a/.github/workflows/README.md b/.github/workflows/README.md index 4b85cb0d4fa9..5e2d23f6e0bf 100644 --- a/.github/workflows/README.md +++ b/.github/workflows/README.md @@ -12,25 +12,98 @@ by the release workflows. ## Three-layer CI rollout The replacement test pipeline is being qualified in three resource tiers. The -legacy build, example, and compatibility workflows remain enabled while the new -tiers are validated. Retiring a legacy workflow is a separate change that must -also migrate required checks, notifications, and repository rules. +legacy example and compatibility workflows remain enabled for non-Draft PRs +while the new tiers are validated. Retiring any additional legacy workflow +must also account for required checks, notifications, and repository rules. + +The initial `e1-gpu-on-pr.yml` host-venv workflow has been retired in favor of +the container-based Layer 2 workflow. Its accelerator test is covered by batch +2 and its two-rank DP dataloader test by batch 5. Do not restore the personal +Python-environment path as a fallback. The `colossalai-e1-h20` runner label and +the infrastructure-only `e1-runner-bootstrap.yml` probe are retained; retiring +a workflow does not unregister or reconfigure the runner. + +Before merging this retirement, check branch protection and rulesets for the +old `gpu-smoke` check from `ColossalAI GPU on internal PR`. If it is required, +migrate that specific requirement to the qualified full Layer 2 check +`Torch 2.13 GPU core (batches 2 and 5) / run-batches`; do not substitute the +Draft-only smoke check or remove unrelated requirements. Removing a workflow +file does not update repository settings or cancel existing runs. Other PRs +will stop using E1 once their workflow snapshot includes this removal. + +The legacy `Build on PR` workflow (`build_on_pr.yml`) and its dedicated +`report_test_coverage.yml` consumer have also been retired. CPU checks and GPU +regressions use the three-layer pipeline; container batch 3 retains the +`BUILD_EXT=1` AOT compilation check. The batch manifest, not the old broad +`pytest tests/` invocation, defines the replacement's coverage. Dependency, +toolchain, and kernel changes still need the relevant Layer 3 batches; a Layer +2 pass alone does not replace that validation. + +Before merging, also check for required statuses named `Detect file change` +or `Build and Test Colossal-AI` from the old PR build and migrate their policy +to qualified layer checks. Deleting these files does not change branch +protection, cancel queued/running jobs, or unregister any runner. Old runs +retain their workflow snapshots. The old coverage-comment feature is removed, +not migrated: the new pytest reports do not imply code-coverage measurement. +Legacy scheduled builds, example/compatibility checks, and release workflows +are outside this retirement and remain enabled. | Layer | Workflow file | Trigger | Coverage and resources | | --- | --- | --- | --- | | 1 | `ci_layer_1_pr.yml` | Every pull request to `main` | Formatting and genuinely CPU-only unit tests on GitHub-hosted runners. | -| 2 | `ci_layer_2_gpu_pr.yml` | Trusted same-repository pull requests and manual dispatch | Core batches 2 and 5 on the dedicated `colossalai-e1-h20` runner. | +| Draft smoke | `ci_layer_2_gpu_pr.yml` | Trusted same-repository Draft pull requests | Only batch 4a, in a container on `colossalai-e1-h20`, using the separate smoke GPU pool. | +| 2 | `ci_layer_2_gpu_pr.yml` | Trusted same-repository Ready pull requests and manual dispatch | Core batches 2, 4a and 5 in containers on the `colossalai-e1-h20` runner. | | 3 | `ci_layer_3_full.yml` | Weekly schedule and manual dispatch | The complete GPU batch set, including four- and eight-GPU coverage. | | Monitor | `ci_gpu_idle_monitor.yml` | Hourly schedule and manual dispatch | Probes the authorized H20 pool; runs layer 2 with at least two idle GPUs and layer 3 with all eight idle GPUs. | | Shared | `ci_reusable_gpu_batches.yml` | Called by layers 2 and 3 | Fetches the pinned source, runs selected batches, uploads reports, and writes the run summary. | -The execution flow is: +Layer 2 retains its historical check name, `Torch 2.13 GPU core (batches 2 and 5)`, +so existing required-check rules do not break; its actual batch list now also +includes `4a`. Rename that check only together with branch-protection settings. + +### Draft PR container validation + +A trusted same-repository Draft PR runs the separate check +`Torch 2.13 container smoke (batch 4a)`. It calls the same container executor as +the full GPU jobs, tests the PR merge SHA, uploads its own reports, and removes +only its own containers. Batch 4a requires one GPU, runs with `FAST_MODE=0`, and +is unaffected by `FAST_TEST`. The entire smoke job has a 15-minute limit, +including source download and waiting for the shared GPU lock. + +Set `COLOSSAL_CI_SMOKE_GPU_POOL` to GPUs explicitly authorized for the smoke +test window (for example `3` for physical GPU 3 on the runner host). It is a +separate repository variable, **not** a replacement for `COLOSSAL_GPU_POOL`. +An unset/empty smoke pool fails before starting containers or allocating GPUs; +it never falls back to the main pool. Clear the smoke variable when temporary +authorization expires. Both profiles use the same `COLOSSAL_GPU_LOCK` and +`COLOSSAL_CI_IMAGE`; a smoke-only lock would not coordinate with scheduled CI. + +When the PR becomes Ready, `ready_for_review` starts the full `2 4a 5` check +using the main pool. Converting it back to Draft cancels the previous Layer 2 +run through the existing per-PR concurrency group and starts the smoke route. +The two jobs have different check names: a green smoke is not evidence that +the skipped full regression passed. Do not bypass required full checks to merge. +External forks and authors outside OWNER/MEMBER/COLLABORATOR do not run either +GPU route. Manual Layer 2 dispatch remains a full regression, not a smoke run. + +Remaining legacy PR GPU/private-runner jobs are gated out for Draft PRs, +including compatibility matrices and CUDA extension builds. Their Ready-PR +behavior and existing explicit manual triggers are retained. These admission conditions do +not cancel legacy runs already started while a PR was Ready; inspect those runs +separately before assuming that converting a PR to Draft has freed every GPU. +Scheduled jobs are unchanged and may still hold the shared lock. CPU/lint jobs +remain available for Drafts and forks. + +Execution path: ```text layer trigger -> reusable GPU workflow - -> run_gpu_batch.py - -> run_colossalai_batch.sh + -> verify the preloaded dependency image by immutable image ID + -> run_gpu_batch.py (host GPU lock + idle-device checks) + -> container_batch.py (one disposable container per batch) + -> run_container_batch.sh (install this source SHA) + -> run_colossalai_batch.sh (existing test membership) ``` `.github/scripts/ci/run_colossalai_batch.sh` is the source of truth for batch @@ -46,18 +119,24 @@ devices. The host lock complements rather than replaces the team's external reservation policy. The hourly monitor is deliberately capacity-aware. It skips layer 2 unless -both of its batches can be admitted, and skips the full layer 3 matrix unless +all of its batches can be admitted, and skips the full layer 3 matrix unless all eight GPUs are idle. The normal per-batch selector checks occupancy again immediately before every launch, so a stale hourly probe still fails closed. ### Rollout configuration 1. Keep the `colossalai-e1-h20` runner label restricted to this repository. -2. Set `COLOSSAL_GPU_POOL` when CI may use only a subset of physical GPU - indices. -3. Optionally set `COLOSSAL_TORCH213_VENV`, `COLOSSAL_CUDA130_HOME`, and - `COLOSSAL_GPU_LOCK`. The reusable workflow contains the current E1 paths as - rollout defaults. +2. Set `COLOSSAL_GPU_POOL` to the **authorized** comma-separated physical GPU + indices, and `COLOSSAL_GPU_LOCK` to the shared lock file used by every CI + executor on this machine. Both variables are required; there is no personal + directory default. The monitor and GPU jobs must use the same lock. +3. Build and validate the [dependency image](../../docker/torch213/README.md) + on the runner's Docker daemon (or import it once with `docker load`). Set + `COLOSSAL_CI_IMAGE` to its full local `sha256:<64 hex digits>` image ID. + The runner needs Docker access, NVIDIA Container Toolkit, a compatible driver, + and host Python 3.10+ with PyYAML. Publishing the image is **not required**. + The Docker daemon/context used for provisioning and CI must match. CI no longer needs + `COLOSSAL_TORCH213_VENV` or `COLOSSAL_CUDA130_HOME`. Layer 1 remains CPU-only. 4. Make `Formatting` and `CPU unit tests` required only after layer 1 has passed repeatedly. Make layer 2 required only after GPU capacity and the trusted-PR policy have been agreed. @@ -67,6 +146,87 @@ immediately before every launch, so a stale hourly probe still fails closed. 6. Add another Torch/CUDA profile to layer 3 only after its isolated environment or pinned image has been provisioned. +### Container execution and image rollout + +The image holds Python dependencies and CUDA build tools, **not ColossalAI**. +The workflow fetches the exact source SHA and mounts it read-only. Each batch +copies it into its private container layer, creates a small writable venv that +references the image dependencies, and installs the checkout with `--no-deps +--no-build-isolation`. Batch 3 sets `BUILD_EXT=1` to check AOT compilation; +other batches use source-specific JIT kernels as needed. Compiler caches are +fresh per container. The initial rollout intentionally has no shared cache or +testmon incremental selection. + +The container rejects a mismatch between the source's image recipe/locks and +the files embedded in the image. Updating a lock file alone cannot make CI +silently validate the old dependency environment. + +The host selector passes **GPU UUIDs** to Docker, not `--gpus all`. Inside the +container, the allocated devices become logical `0,1,...`; host indices must +not be reused there. No personal venv, host CUDA toolkit, home directory, +credentials or Docker socket is mounted. Only the report bind is writable; +containers run as the runner UID/GID so artifact ownership stays correct. +Containers are not a security boundary for hostile pull requests: the host +scripts still have Docker access. Keep the trusted-ref policy and runner-group +restrictions. + +| Repository variable | Default / purpose | +| --- | --- | +| `COLOSSAL_CI_IMAGE` | Required, full local `sha256:image-ID` (or an already-loaded registry digest). No tag, automatic pull or host-environment fallback. | +| `COLOSSAL_GPU_POOL` | Required, GPUs reserved/authorized for CI. Idle does not mean authorized. | +| `COLOSSAL_CI_SMOKE_GPU_POOL` | Required for Draft smoke only. Authorized candidate GPUs for its one-GPU batch 4a; never falls back to the main pool. | +| `COLOSSAL_GPU_LOCK` | Required, agreed shared host lock. Use the same path as legacy executors. | +| `COLOSSAL_CI_CPUS` | `8`, CPU quota per active container. | +| `COLOSSAL_CI_MEMORY` | `64g`, RAM limit; swap does not add to this limit. | +| `COLOSSAL_CI_SHM` | `8g`, private shared memory for multiprocessing/NCCL. | + +Tune RAM/CPU limits with the administrator before the full eight-GPU matrix; +CPU-offload tests can need more RAM than the initial smoke profile. These limits +do not reserve GPU memory or cap image/layer disk usage. Docker build/load/run +write to daemon storage, so keeping the checkout and exported image `.tar` in a +personal directory is **not** equivalent to confining all writes there. Do not +move the shared daemon's data root or restart it to achieve personal isolation. + +Bootstrap with **one PR**: first build the candidate from the local branch on +an approved builder; verify dependencies; build/load it on the test daemon; +qualify batches `4a`, `3`, `4b`, then `2 5` and cancellation cleanup on allocated +GPUs using `run_gpu_batch.py --image sha256:...`. None of these steps requires +opening or merging a PR, publishing the image, or upgrading the host venv. +Save the qualified image ID, source SHA and test reports, optionally archive the +image with `docker save`, then configure the runner/repository variables and +open the single combined PR. GitHub event triggering, permissions and artifact +upload still need verification on that PR; fix them on the same branch/PR. +Open it as a **Draft** first to qualify only the container smoke route. Configure +the smoke pool independently; do not shrink the main pool to one GPU, which +would prevent normal two-/four-/eight-GPU jobs from running. After the Draft +check succeeds, arrange resources for the full Layer 2 job and mark the same PR +Ready. The GPU profile is not fully qualified by one FP8 smoke test alone. +Merge only after reviewing failures and skipped coverage. +Do not merge the workflow switch before the image and runner are ready. +Subsequent code-only PRs reuse the image and install their own source. Rebuild +the image when the lock files, base/toolchain or Dockerfile change, not on every +PR or `version.txt` change. This CI image is separate from the user release image. + +An image tar is a portable backup, not an executable container directory. +Load it once into the actual runner daemon when restoring or moving machines; +ordinary runs reuse the local image without repeatedly loading the tar. If +several machines share the runner label, preload the same image on each one +or restrict the job to the provisioned runner. No registry credentials are needed. + +For rollback, select a previously **qualified** image ID compatible with +the source under test, or revert the executor switch. Do not silently fall back +to a developer's venv. The E1 host-venv PR workflow and legacy `Build on PR` +workflow are retired. Legacy scheduled builds, example/compatibility checks, +and the release pipeline are retained. The remaining legacy GPU/private-runner +PR jobs are gated out for Drafts. + +`FAST_MODE=1` continues to reduce the model matrices of batches 2 and 5; batch +4a is unaffected. Layer 3 defaults to `FAST_MODE=0`, which unsets `FAST_TEST`. +Installing testmon does not enable selection: a later change must validate +coverage recording, image/batch-scoped persistent data, multiprocessing and +CUDA/C++ change invalidation before opting in. Core smoke and full regressions +must remain unconditional. + ### Failure diagnosis and targeted validation Layer 3 accepts an explicit `batches` input. After changing code, rerun the @@ -86,9 +246,13 @@ The batch wrapper uses `run_with_timeout.py` to supervise a dedicated test process group and clean its workers before the logging `tee` waits for EOF. The GPU runner also checks whether the batch process exited while streaming logs; it must not wait indefinitely for EOF held open by -orphaned workers. Its finalizer terminates only that batch's process group, -including on timeout, before releasing the CI GPU lock. This does not cancel -unrelated GPU jobs. +orphaned workers. In container mode it also removes the uniquely named container +while still holding the lock; both ownership labels and the exact name must +match. An `always()` step retries cleanup for this job's scope on cancellation. +It never runs `docker system prune`, daemon-wide cleanup or unrelated-process +termination. A runner crash, forced kill or unavailable Docker daemon can still +prevent cleanup: stop admitting jobs and have the administrator inspect the +reported container name/scope before returning the runner to service. ## Active workflow inventory @@ -118,7 +282,6 @@ runner groups to approved workflows, in GitHub settings. | Workflow file | Trigger | Purpose | | --- | --- | --- | -| `build_on_pr.yml` | Pull request and branch lifecycle events | Legacy Colossal-AI build and unit-test workflow. | | `build_on_schedule.yml` | Schedule or manual dispatch | Legacy scheduled full unit tests. | | `example_check_on_pr.yml` | Pull request | Runs tests for examples affected by a pull request. | | `example_check_on_schedule.yml` | Schedule or manual dispatch | Runs the scheduled example suite. | @@ -127,7 +290,6 @@ runner groups to approved workflows, in GitHub settings. | `compatibility_test_on_schedule.yml` | Schedule or manual dispatch | Runs the compatibility matrix from `.compatibility`. | | `compatibility_test_on_dispatch.yml` | Manual dispatch | Runs the required Torch and CUDA version inputs. Neither input has an implicit `all` default. | | `cuda_ext_check_before_merge.yml` | Release pull request or manual dispatch | Checks CUDA extension builds before release. | -| `report_test_coverage.yml` | Completion of the legacy PR build | Reports test coverage back to the pull request. | | `requirements_compatibility.yml` | Scoped requirements changes or manual dispatch | Validates dependency installation and compatibility. | | `run_chatgpt_examples.yml` | Pull request | Runs ChatGPT example checks when their scoped files change. | | `run_chatgpt_unit_tests.yml` | Pull request | Runs ChatGPT unit tests when their scoped files change. | @@ -157,7 +319,6 @@ runner groups to approved workflows, in GitHub settings. | Workflow file | Trigger | Purpose | | --- | --- | --- | | `e1-runner-bootstrap.yml` | Scoped push or manual dispatch | Qualifies the E1 self-hosted runner and its runtime. | -| `e1-gpu-on-pr.yml` | Trusted internal pull request | Runs the existing E1 GPU regression while the layered pipeline is qualified. | ### Repository maintenance diff --git a/.github/workflows/build_on_pr.yml b/.github/workflows/build_on_pr.yml deleted file mode 100644 index 20f0b91912a6..000000000000 --- a/.github/workflows/build_on_pr.yml +++ /dev/null @@ -1,211 +0,0 @@ -name: Build on PR - -on: - pull_request: - types: [synchronize, opened, reopened, ready_for_review, closed] - branches: - - "main" - - "develop" - - "feature/**" - paths: - - ".github/workflows/build_on_pr.yml" # run command & env variables change - - "colossalai/**" # source code change - - "!colossalai/**.md" # ignore doc change - - "extensions/**" # cuda extension change - - "!extensions/**.md" # ignore doc change - - "requirements/**" # requirements change - - "tests/**" # test change - - "!tests/**.md" # ignore doc change - - "pytest.ini" # test config change - - "setup.py" # install command change - create: - delete: - -jobs: - detect: - name: Detect file change - if: | - github.event_name == 'pull_request' && - (github.event.action == 'synchronize' || github.event.action == 'opened' || github.event.action == 'reopened' || github.event.action == 'ready_for_review') && - github.event.pull_request.draft == false && - github.event.pull_request.base.repo.full_name == 'hpcaitech/ColossalAI' && - github.event.pull_request.head.repo.full_name == github.repository - outputs: - changedExtenisonFiles: ${{ steps.find-extension-change.outputs.all_changed_files }} - anyExtensionFileChanged: ${{ steps.find-extension-change.outputs.any_changed }} - changedLibraryFiles: ${{ steps.find-lib-change.outputs.all_changed_files }} - anyLibraryFileChanged: ${{ steps.find-lib-change.outputs.any_changed }} - runs-on: [self-hosted, ubuntu-latest] - concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}-detect-change - cancel-in-progress: true - steps: - - uses: actions/checkout@v2 - with: - fetch-depth: 0 - ref: ${{ github.event.pull_request.head.sha }} - - - name: Locate base commit - id: locate-base-sha - run: | - curBranch=$(git rev-parse --abbrev-ref HEAD) - commonCommit=$(git merge-base origin/main $curBranch) - echo $commonCommit - echo "baseSHA=$commonCommit" >> $GITHUB_OUTPUT - - - name: Find the changed extension-related files - id: find-extension-change - uses: tj-actions/changed-files@v35 - with: - base_sha: ${{ steps.locate-base-sha.outputs.baseSHA }} - files: | - extensions/** - colossalai/kernel/** - setup.py - - - name: Find the changed library-related files - id: find-lib-change - uses: tj-actions/changed-files@v35 - with: - base_sha: ${{ steps.locate-base-sha.outputs.baseSHA }} - files: | - **/*.py - **/*.h - **/*.cpp - **/*.cu - **/*.txt - - - name: List changed files - run: | - for file in ${{ steps.find-extension-change.outputs.all_changed_files }}; do - echo "$file was changed" - done - for file in ${{ steps.find-lib-change.outputs.all_changed_files }}; do - echo "$file was changed" - done - - build: - name: Build and Test Colossal-AI - needs: detect - if: >- - needs.detect.outputs.anyLibraryFileChanged == 'true' && - github.event.pull_request.head.repo.full_name == github.repository - runs-on: [self-hosted, ubuntu-latest] - container: - image: image-cloud.luchentech.com/hpcaitech/pytorch-cuda:2.2.2-12.1.0 - options: --gpus all --shm-size=2g --rm -v /dev/shm -v /data/scratch:/data/scratch - timeout-minutes: 90 - defaults: - run: - shell: bash - concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}-run-test - cancel-in-progress: true - steps: - - name: Checkout TensorNVMe - uses: actions/checkout@v2 - with: - repository: hpcaitech/TensorNVMe - ssh-key: ${{ secrets.SSH_KEY_FOR_CI }} - path: TensorNVMe - - - name: Restore TensorNVMe Cache - run: | - if [ -d /github/home/tensornvme_cache ] && [ ! -z "$(ls -A /github/home/tensornvme_cache/)" ]; then - cp -p -r /github/home/tensornvme_cache/* /__w/ColossalAI/ColossalAI/TensorNVMe - fi - - - name: Install TensorNVMe - run: | - cd TensorNVMe - conda install cmake - pip install -r requirements.txt - DISABLE_URING=1 pip install -v --no-cache-dir . - - - name: Store TensorNVMe Cache - run: | - cd TensorNVMe - cp -p -r ./build /github/home/tensornvme_cache/ - cp -p -r ./cmake-build /github/home/tensornvme_cache/ - - - name: Checkout Colossal-AI - uses: actions/checkout@v2 - with: - ssh-key: ${{ secrets.SSH_KEY_FOR_CI }} - - - name: Restore Colossal-AI Cache - if: needs.detect.outputs.anyExtensionFileChanged != 'true' - run: | - # -p flag is required to preserve the file timestamp to avoid ninja rebuild - if [ -d /github/home/cuda_ext_cache ] && [ ! -z "$(ls -A /github/home/cuda_ext_cache/)" ]; then - cp -p -r /github/home/cuda_ext_cache/* /__w/ColossalAI/ColossalAI/ - fi - - - name: Install flash-attention - run: | - pip install flash-attn==2.7.4.post1 --no-build-isolation - - - name: Install Colossal-AI - run: | - BUILD_EXT=1 pip install -v -e . - pip install --no-cache-dir -r requirements/requirements-test.txt - - - name: Store Colossal-AI Cache - run: | - # -p flag is required to preserve the file timestamp to avoid ninja rebuild - cp -p -r /__w/ColossalAI/ColossalAI/build /github/home/cuda_ext_cache/ - - - name: Execute Unit Testing - run: | - CURL_CA_BUNDLE="" PYTHONPATH=$PWD FAST_TEST=1 pytest \ - -m "not largedist" \ - --durations=0 \ - --ignore tests/test_analyzer \ - --ignore tests/test_auto_parallel \ - --ignore tests/test_fx \ - --ignore tests/test_autochunk \ - --ignore tests/test_legacy \ - --ignore tests/test_smoothquant \ - tests/ - env: - LD_LIBRARY_PATH: /github/home/.tensornvme/lib:/usr/local/nvidia/lib:/usr/local/nvidia/lib64 - LLAMA_PATH: /data/scratch/llama-tiny - MOE_TENSOR_PATH: /data/scratch/moe_tensors - HF_ENDPOINT: https://hf-mirror.com - - - name: Collate artifact - env: - PR_NUMBER: ${{ github.event.number }} - changedLibraryFiles: ${{ needs.detect.outputs.changedLibraryFiles }} - anyLibraryFileChanged: ${{ needs.detect.outputs.anyLibraryFileChanged }} - changedExtenisonFiles: ${{ needs.detect.outputs.changedExtenisonFiles }} - run: | - mkdir report - echo $PR_NUMBER > ./report/pr_number - - # generate coverage.xml if any - if [ "$anyLibraryFileChanged" == "true" ] && [ -e .coverage ]; then - allFiles="" - for file in $changedLibraryFiles; do - if [ "$allFiles" == "" ]; then - allFiles=$file - else - allFiles=$allFiles,$file - fi - done - - coverage report --data-file .coverage --include $allFiles > ./coverage.txt - - covPercentage=$(tail -n 1 coverage.txt | grep -o '[1-9]*%$') - covNum=${covPercentage::-1} - mv coverage.txt ./report - echo $covNum > ./report/cov_number - else - echo "No coverage report is generated" - fi - - - name: Upload test coverage artifact - uses: actions/upload-artifact@v4 - with: - name: report - path: report/ diff --git a/.github/workflows/ci_gpu_idle_monitor.yml b/.github/workflows/ci_gpu_idle_monitor.yml index fae69c9ab005..0dde64f11661 100644 --- a/.github/workflows/ci_gpu_idle_monitor.yml +++ b/.github/workflows/ci_gpu_idle_monitor.yml @@ -25,7 +25,7 @@ jobs: layer3_ready: ${{ steps.capacity.outputs.layer3_ready }} env: COLOSSAL_GPU_POOL: ${{ vars.COLOSSAL_GPU_POOL }} - CI_GPU_LOCK: ${{ vars.COLOSSAL_GPU_LOCK || '/mnt/beegfs/ricardoo/ci/gpu-h20-5/locks/ci-gpu.lock' }} + CI_GPU_LOCK: ${{ vars.COLOSSAL_GPU_LOCK }} PYTHONDONTWRITEBYTECODE: "1" steps: # Keep this gate self-contained. The H20 runners can reach the Actions @@ -37,6 +37,8 @@ jobs: shell: bash run: | set -euo pipefail + : "${CI_GPU_LOCK:?Set COLOSSAL_GPU_LOCK to the shared CI GPU lock path}" + : "${COLOSSAL_GPU_POOL:?Set COLOSSAL_GPU_POOL to the authorized physical GPU indices}" python3 -B - <<'PY' import fcntl import json @@ -154,7 +156,7 @@ jobs: with: layer: layer-2-hourly source_sha: ${{ github.sha }} - batches: "2 5" + batches: "2 4a 5" fast_mode: true timeout_minutes: 240 artifact_retention_days: 14 diff --git a/.github/workflows/ci_layer_2_gpu_pr.yml b/.github/workflows/ci_layer_2_gpu_pr.yml index 65ffb431e483..f39619d6c4f9 100644 --- a/.github/workflows/ci_layer_2_gpu_pr.yml +++ b/.github/workflows/ci_layer_2_gpu_pr.yml @@ -2,7 +2,7 @@ name: CI / Layer 2 / Trusted PR GPU on: pull_request: - types: [opened, synchronize, reopened, ready_for_review] + types: [opened, synchronize, reopened, ready_for_review, converted_to_draft] branches: [main] paths: - .github/scripts/ci/** @@ -10,6 +10,7 @@ on: - .github/workflows/ci_reusable_gpu_batches.yml - colossalai/** - extensions/** + - docker/torch213/** - requirements/** - tests/** - pytest.ini @@ -24,18 +25,41 @@ concurrency: cancel-in-progress: true jobs: + container-smoke: + # This is a separate check, not a successful full Layer 2 regression. + name: Torch 2.13 container smoke (batch 4a) + if: >- + github.repository == 'hpcaitech/ColossalAI' && + github.event_name == 'pull_request' && + github.event.pull_request.draft == true && + github.event.pull_request.head.repo.full_name == github.repository && + contains(fromJSON('["OWNER", "MEMBER", "COLLABORATOR"]'), github.event.pull_request.author_association) + uses: ./.github/workflows/ci_reusable_gpu_batches.yml + with: + layer: draft-smoke + source_sha: ${{ github.sha }} + batches: "4a" + gpu_pool_variable: COLOSSAL_CI_SMOKE_GPU_POOL + fast_mode: false + timeout_minutes: 15 + artifact_retention_days: 14 + gpu-regression: + # Preserve the existing check name until branch-protection rules migrate. + # Batch 4a is also included below for the FP8 compatibility regression. name: Torch 2.13 GPU core (batches 2 and 5) if: >- github.repository == 'hpcaitech/ColossalAI' && (github.event_name == 'workflow_dispatch' || - (github.event.pull_request.head.repo.full_name == github.repository && + (github.event_name == 'pull_request' && + github.event.pull_request.draft == false && + github.event.pull_request.head.repo.full_name == github.repository && contains(fromJSON('["OWNER", "MEMBER", "COLLABORATOR"]'), github.event.pull_request.author_association))) uses: ./.github/workflows/ci_reusable_gpu_batches.yml with: layer: layer-2 source_sha: ${{ github.sha }} - batches: "2 5" + batches: "2 4a 5" fast_mode: true timeout_minutes: 240 artifact_retention_days: 14 diff --git a/.github/workflows/ci_reusable_gpu_batches.yml b/.github/workflows/ci_reusable_gpu_batches.yml index db92b9e2e456..3abe38968fec 100644 --- a/.github/workflows/ci_reusable_gpu_batches.yml +++ b/.github/workflows/ci_reusable_gpu_batches.yml @@ -15,6 +15,11 @@ on: description: Space-separated GPU batch identifiers required: true type: string + gpu_pool_variable: + description: Repository variable holding the authorized GPU pool (core or Draft smoke) + required: false + default: COLOSSAL_GPU_POOL + type: string fast_mode: description: Use the reduced FAST_TEST model matrix required: false @@ -40,21 +45,58 @@ jobs: timeout-minutes: ${{ inputs.timeout_minutes }} env: CI_SOURCE_SHA: ${{ inputs.source_sha }} - COLOSSAL_VENV: ${{ vars.COLOSSAL_TORCH213_VENV || '/mnt/beegfs/ColossalAI/wangzhijian/envs/colossalai-torch213' }} - COLOSSAL_CUDA_HOME: ${{ vars.COLOSSAL_CUDA130_HOME || '/mnt/beegfs/ColossalAI/wangzhijian/toolchains/cuda-13.0' }} - COLOSSAL_EXPECT_TORCH_PREFIX: "2.13." - COLOSSAL_GPU_POOL: ${{ vars.COLOSSAL_GPU_POOL }} - CI_GPU_LOCK: ${{ vars.COLOSSAL_GPU_LOCK || '/mnt/beegfs/ricardoo/ci/gpu-h20-5/locks/ci-gpu.lock' }} + COLOSSAL_CI_IMAGE: ${{ vars.COLOSSAL_CI_IMAGE }} + COLOSSAL_CI_CPUS: ${{ vars.COLOSSAL_CI_CPUS || '8' }} + COLOSSAL_CI_MEMORY: ${{ vars.COLOSSAL_CI_MEMORY || '64g' }} + COLOSSAL_CI_SHM: ${{ vars.COLOSSAL_CI_SHM || '8g' }} + CI_CONTAINER_SCOPE: gh-${{ github.run_id }}-${{ github.run_attempt }}-${{ inputs.layer }} + # Direct lookup is intentional: an unset smoke pool must never fall back + # to the main pool used by Ready PRs, schedules and manual regressions. + GPU_POOL_VARIABLE: ${{ inputs.gpu_pool_variable }} + COLOSSAL_GPU_POOL: ${{ vars[inputs.gpu_pool_variable] }} + CI_GPU_LOCK: ${{ vars.COLOSSAL_GPU_LOCK }} REQUESTED_BATCHES: ${{ inputs.batches }} REQUESTED_FAST_MODE: ${{ inputs.fast_mode }} PYTHONDONTWRITEBYTECODE: "1" steps: - - name: Initialize result directory + - name: Validate configuration and initialize results shell: bash run: | + set -euo pipefail + : "${COLOSSAL_CI_IMAGE:?Set COLOSSAL_CI_IMAGE to the validated local sha256:image-ID}" + [[ "$COLOSSAL_CI_IMAGE" =~ ^(sha256:[0-9a-f]{64}|[a-zA-Z0-9][a-zA-Z0-9._:/-]*@sha256:[0-9a-f]{64})$ ]] + : "${CI_GPU_LOCK:?Set COLOSSAL_GPU_LOCK to the shared CI GPU lock path}" + case "$GPU_POOL_VARIABLE" in + COLOSSAL_GPU_POOL) ;; + COLOSSAL_CI_SMOKE_GPU_POOL) + [[ "$REQUESTED_BATCHES" == "4a" ]] || { + printf 'The Draft smoke pool may only run batch 4a.\n' >&2 + exit 1 + } + ;; + *) printf 'Unsupported GPU pool variable: %s\n' "$GPU_POOL_VARIABLE" >&2; exit 1 ;; + esac + if [[ -z "$COLOSSAL_GPU_POOL" ]]; then + printf 'Set %s to the authorized physical GPU indices; no pool fallback is allowed.\n' "$GPU_POOL_VARIABLE" >&2 + exit 1 + fi + [[ "$COLOSSAL_GPU_POOL" =~ ^[0-9]+(,[0-9]+)*$ ]] || { + printf '%s must contain comma-separated physical GPU indices.\n' "$GPU_POOL_VARIABLE" >&2 + exit 1 + } result_dir="$RUNNER_TEMP/colossalai-${{ inputs.layer }}-${{ github.run_id }}-${{ github.run_attempt }}" + mkdir -p "$result_dir" printf 'CI_RESULTS_DIR=%s\n' "$result_dir" >> "$GITHUB_ENV" - printf 'COLOSSAL_RESULTS=%s\n' "$result_dir" >> "$GITHUB_ENV" + + - name: Verify the preloaded local dependency image + timeout-minutes: 2 + shell: bash + run: | + set -euo pipefail + # Build on this Docker daemon, or load the approved local image tar + # once beforehand. Missing images fail; CI never pulls or pushes. + docker image inspect --format '{{.Id}}' "$COLOSSAL_CI_IMAGE" \ + | tee "$CI_RESULTS_DIR/image-id.txt" - name: Fetch the exact GitHub source archive timeout-minutes: 5 @@ -80,6 +122,7 @@ jobs: run: | set -euo pipefail bash -n "$COLOSSAL_REPO/.github/scripts/ci/run_colossalai_batch.sh" + bash -n "$COLOSSAL_REPO/.github/scripts/ci/run_container_batch.sh" python3 -B -m unittest discover \ -s "$COLOSSAL_REPO/.github/scripts/ci" -p 'test_*.py' -v @@ -99,13 +142,22 @@ jobs: if ! python3 -B "$COLOSSAL_REPO/.github/scripts/ci/run_gpu_batch.py" \ --batch "$batch" \ --output "$CI_RESULTS_DIR/batch-$batch" \ - --lock "$CI_GPU_LOCK"; then + --lock "$CI_GPU_LOCK" \ + --image "$COLOSSAL_CI_IMAGE"; then status=1 printf 'Batch %s failed; continuing to collect the remaining results.\n' "$batch" >&2 fi done exit "$status" + - name: Clean up only this job's containers + if: always() && env.COLOSSAL_REPO != '' + timeout-minutes: 3 + shell: bash + run: | + python3 -B "$COLOSSAL_REPO/.github/scripts/ci/container_batch.py" \ + --cleanup-scope "$CI_CONTAINER_SCOPE" + - name: Upload GPU test reports if: always() uses: actions/upload-artifact@v7 @@ -124,6 +176,8 @@ jobs: printf -- '- Source SHA: `%s`\n' "$CI_SOURCE_SHA" printf -- '- Runner: `%s`\n' "$RUNNER_NAME" printf -- '- Runtime profile: Torch 2.13 / CUDA 13.0\n' + printf -- '- Dependency image: `%s`\n' "$COLOSSAL_CI_IMAGE" + printf -- '- GPU pool variable: `%s`\n' "$GPU_POOL_VARIABLE" printf -- '- Batches: `%s`\n\n' "$REQUESTED_BATCHES" find "$CI_RESULTS_DIR" -name result.json -print -exec cat {} \; 2>/dev/null || true } >> "$GITHUB_STEP_SUMMARY" diff --git a/.github/workflows/compatibility_test_on_pr.yml b/.github/workflows/compatibility_test_on_pr.yml index e93dfc53733e..23c0ccd81133 100644 --- a/.github/workflows/compatibility_test_on_pr.yml +++ b/.github/workflows/compatibility_test_on_pr.yml @@ -2,6 +2,7 @@ name: Compatibility Test on PR on: pull_request: + types: [opened, synchronize, reopened, ready_for_review] paths: - "version.txt" - ".compatibility" @@ -9,7 +10,9 @@ on: jobs: matrix_preparation: name: Prepare Container List - if: github.event.pull_request.head.repo.full_name == github.repository + if: >- + github.event.pull_request.draft == false && + github.event.pull_request.head.repo.full_name == github.repository runs-on: [self-hosted, ubuntu-latest] outputs: matrix: ${{ steps.set-matrix.outputs.matrix }} @@ -37,6 +40,7 @@ jobs: needs: matrix_preparation if: >- github.repository == 'hpcaitech/ColossalAI' && + github.event.pull_request.draft == false && github.event.pull_request.head.repo.full_name == github.repository runs-on: [self-hosted, ubuntu-latest] strategy: diff --git a/.github/workflows/cuda_ext_check_before_merge.yml b/.github/workflows/cuda_ext_check_before_merge.yml index 9e93fa69b442..cfd967ba9e10 100644 --- a/.github/workflows/cuda_ext_check_before_merge.yml +++ b/.github/workflows/cuda_ext_check_before_merge.yml @@ -3,6 +3,7 @@ name: Check CUDA Extension Build Before Merge on: workflow_dispatch: pull_request: + types: [opened, synchronize, reopened, ready_for_review] paths: - 'version.txt' @@ -12,7 +13,8 @@ jobs: if: >- github.repository == 'hpcaitech/ColossalAI' && (github.event_name != 'pull_request' || - github.event.pull_request.head.repo.full_name == github.repository) + (github.event.pull_request.draft == false && + github.event.pull_request.head.repo.full_name == github.repository)) runs-on: [self-hosted, ubuntu-latest] outputs: matrix: ${{ steps.set-matrix.outputs.matrix }} @@ -29,7 +31,8 @@ jobs: needs: matrix_preparation if: >- github.event_name != 'pull_request' || - github.event.pull_request.head.repo.full_name == github.repository + (github.event.pull_request.draft == false && + github.event.pull_request.head.repo.full_name == github.repository) runs-on: [self-hosted, ubuntu-latest] strategy: fail-fast: false diff --git a/.github/workflows/e1-gpu-on-pr.yml b/.github/workflows/e1-gpu-on-pr.yml deleted file mode 100644 index c42f1a7e63c9..000000000000 --- a/.github/workflows/e1-gpu-on-pr.yml +++ /dev/null @@ -1,66 +0,0 @@ -name: ColossalAI GPU on internal PR - -on: - pull_request: - types: [opened, synchronize, reopened, ready_for_review] - -permissions: - contents: read - -concurrency: - group: e1-colossalai-pr-${{ github.event.pull_request.number }} - cancel-in-progress: false - -jobs: - gpu-smoke: - if: >- - github.repository == 'hpcaitech/ColossalAI' && - github.event.pull_request.head.repo.full_name == 'hpcaitech/ColossalAI' - runs-on: colossalai-e1-h20 - timeout-minutes: 12 - env: - E1_COLOSSALAI_PYTHON: /mnt/beegfs/ColossalAI/wangzhijian/envs/colossalai-torch213/bin/python - E1_PR_RESULTS: /mnt/beegfs/ricardoo/ci/gpu-h20-5/test-results/pr-${{ github.run_id }}-${{ github.run_attempt }} - steps: - - name: Fetch full source from the exact PR merge commit - timeout-minutes: 4 - shell: bash - run: | - set -euo pipefail - [[ "$GITHUB_SHA" =~ ^[0-9a-f]{40}$ ]] - source_dir="$(mktemp -d "$RUNNER_TEMP/e1-source.XXXXXX")" - archive="$source_dir/source.tar.gz" - curl --fail --location --proto '=https' --proto-redir '=https' \ - --connect-timeout 10 --max-time 180 --retry 1 --retry-max-time 200 \ - --output "$archive" \ - "https://codeload.github.com/hpcaitech/ColossalAI/tar.gz/$GITHUB_SHA" - tar --extract --gzip --file "$archive" --directory "$source_dir" \ - --strip-components=1 --no-same-owner --no-same-permissions - archive_sha="$(sha256sum "$archive" | cut -d ' ' -f 1)" - printf 'Source commit: %s\nArchive SHA256: %s\n' "$GITHUB_SHA" "$archive_sha" - printf 'E1_SOURCE_DIR=%s\nE1_SOURCE_ARCHIVE_SHA256=%s\n' \ - "$source_dir" "$archive_sha" >> "$GITHUB_ENV" - - name: Test resource selection without using GPUs - run: python3 -B -m unittest discover -s "$E1_SOURCE_DIR/.github/scripts/e1" -p test_pr_gpu.py -v - - name: Check ColossalAI dependencies and collect selected tests without GPUs - timeout-minutes: 2 - run: >- - "$E1_COLOSSALAI_PYTHON" -B "$E1_SOURCE_DIR/.github/scripts/e1/colossalai_suite.py" - --collect-only --output "$E1_PR_RESULTS-preflight" - - name: Run ColossalAI tests on two idle GPUs - run: python3 -B "$E1_SOURCE_DIR/.github/scripts/e1/pr_gpu.py" --suite colossalai --output "$E1_PR_RESULTS" - - name: Record result - if: always() - shell: bash - run: | - { - printf 'E1 ColossalAI two-GPU regression on gpu-h20-5\n\n' - printf 'Manual resource authorization; reservation website not used.\n\n' - if test -f "$E1_PR_RESULTS/result.json"; then - printf '```json\n' - cat "$E1_PR_RESULTS/result.json" - printf '\n```\n' - else - printf 'GPU test did not produce a result; inspect the failed step.\n' - fi - } >> "$GITHUB_STEP_SUMMARY" diff --git a/.github/workflows/report_test_coverage.yml b/.github/workflows/report_test_coverage.yml deleted file mode 100644 index 2f4c58bf9ad7..000000000000 --- a/.github/workflows/report_test_coverage.yml +++ /dev/null @@ -1,81 +0,0 @@ -name: Report Test Coverage - -on: - workflow_run: - workflows: [Build on PR] - types: - - completed - -jobs: - report-test-coverage: - runs-on: [self-hosted, ubuntu-latest] - if: ${{ github.event.workflow_run.conclusion == 'success' }} - steps: - - name: "Download artifact" - uses: actions/github-script@v6 - with: - script: | - let allArtifacts = await github.rest.actions.listWorkflowRunArtifacts({ - owner: context.repo.owner, - repo: context.repo.repo, - run_id: context.payload.workflow_run.id, - }); - let matchArtifact = allArtifacts.data.artifacts.filter((artifact) => { - return artifact.name == "report" - })[0]; - let download = await github.rest.actions.downloadArtifact({ - owner: context.repo.owner, - repo: context.repo.repo, - artifact_id: matchArtifact.id, - archive_format: 'zip', - }); - let fs = require('fs'); - fs.writeFileSync(`${process.env.GITHUB_WORKSPACE}/report.zip`, Buffer.from(download.data)); - - - name: "Unzip artifact" - id: unzip - run: | - unzip report.zip - if [ -f "coverage.txt" ]; then - echo "hasReport=true" >> $GITHUB_OUTPUT - else - echo "hasReport=false" >> $GITHUB_OUTPUT - fi - - - name: Make Coverage Report Collapsable - if: steps.unzip.outputs.hasReport == 'true' - run: | - covNum=$(cat cov_number) - title="The code coverage for the changed files is ${covNum}%." - touch coverage_report.txt - echo $title >> coverage_report.txt - echo " " >> coverage_report.txt - echo "
" >> coverage_report.txt - echo "Click me to view the complete report" >> coverage_report.txt - echo " " >> coverage_report.txt - echo "\`\`\`" >> coverage_report.txt - cat coverage.txt >> coverage_report.txt - echo "\`\`\`" >> coverage_report.txt - echo "
" >> coverage_report.txt - mv coverage_report.txt coverage.txt - - - name: "Comment on PR" - if: steps.unzip.outputs.hasReport == 'true' - uses: actions/github-script@v6 - with: - github-token: ${{ secrets.GITHUB_TOKEN }} - script: | - let fs = require('fs'); - let issue_number = Number(fs.readFileSync('./pr_number')); - let owner = context.repo.owner; - let repo = context.repo.repo; - let run_id = context.payload.workflow_run.id; - let run_url = `https://github.com/${owner}/${repo}/actions/runs/${run_id}` - let body = fs.readFileSync('./coverage.txt', {encoding:'utf8', flag:'r'}) - - await github.rest.issues.createComment({ - owner: owner, - repo: repo, - issue_number: issue_number, - body: body - }); diff --git a/.github/workflows/requirements_compatibility.yml b/.github/workflows/requirements_compatibility.yml index aaa365687396..0ef6e4015ff3 100644 --- a/.github/workflows/requirements_compatibility.yml +++ b/.github/workflows/requirements_compatibility.yml @@ -63,14 +63,24 @@ jobs: assert all(Version(version) in python_versions for version in ("3.10", "3.11", "3.12")) assert Version("3.13") not in python_versions - requirements = {Requirement(item).name.lower() for item in requires("colossalai") or []} + requirements = { + Requirement(item).name.lower(): Requirement(item) + for item in requires("colossalai") or [] + } assert "pre-commit" not in requirements + # Keep the CPU baseline installable while allowing the GPU CI profile. + torch_versions = requirements["torch"].specifier + assert Version("2.5.1") in torch_versions + assert Version("2.13.0+cu130") in torch_versions + assert Version("2.14.0") not in torch_versions PY - name: Check removed NumPy APIs run: ruff check --select NPY201 colossalai tests applications examples - name: Run compatibility smoke tests + working-directory: ${{ github.workspace }} run: | - pytest -m cpu tests/test_cluster/test_process_group_mesh.py - pytest applications/ColossalChat/tests/test_code_reward_testing_util.py + # Keep the repository root on sys.path for applications.* imports. + python -m pytest -m cpu tests/test_cluster/test_process_group_mesh.py + python -m pytest applications/ColossalChat/tests/test_code_reward_testing_util.py diff --git a/docker/torch213/.gitignore b/docker/torch213/.gitignore new file mode 100644 index 000000000000..b0895b3620c2 --- /dev/null +++ b/docker/torch213/.gitignore @@ -0,0 +1,2 @@ +wheelhouse/* +!wheelhouse/.gitkeep diff --git a/docker/torch213/Dockerfile b/docker/torch213/Dockerfile new file mode 100644 index 000000000000..c977b552e2da --- /dev/null +++ b/docker/torch213/Dockerfile @@ -0,0 +1,54 @@ +# Use BuildKit's bundled frontend; do not fetch docker/dockerfile from Docker Hub. +# Build from the repository root. No GPU is needed to build this dependency image. +ARG CUDA_IMAGE=nvidia/cuda:13.0.3-cudnn-devel-ubuntu22.04 +FROM ${CUDA_IMAGE} AS base + +# Shared hosts may default to the NVIDIA runtime. GPU CI overrides this with +# the allocated UUIDs; dependency builds and plain containers need no devices. +ENV NVIDIA_VISIBLE_DEVICES=void +ARG DEBIAN_FRONTEND=noninteractive +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3.10 python3.10-dev python3.10-venv \ + build-essential cmake ninja-build git curl ca-certificates \ + libaio-dev libibverbs1 libibverbs-dev ibverbs-providers librdmacm1 rdma-core \ + && rm -rf /var/lib/apt/lists/* +RUN python3.10 -m venv /opt/venv +ENV PATH=/opt/venv/bin:/usr/local/cuda/bin:${PATH} \ + CUDA_HOME=/usr/local/cuda \ + PYTHONUNBUFFERED=1 \ + PIP_DEFAULT_TIMEOUT=180 \ + MAX_JOBS=4 + +FROM base AS dependencies +ARG PIP_INDEX_URL=https://repo.huaweicloud.com/repository/pypi/simple +ARG TORCH_INDEX_URL=https://download.pytorch.org/whl/cu130 +COPY docker/torch213/requirements*.lock /opt/ci/ +# The optional wheelhouse is a read-only build input, never a final-image layer. +# Core wheels are either supplied by the builder or fetched from the Torch index. +RUN --mount=type=bind,source=docker/torch213/wheelhouse,target=/wheelhouse,readonly \ + --mount=type=cache,id=colossalai-torch213-pip,target=/root/.cache/pip,sharing=locked \ + if find /wheelhouse -maxdepth 1 -name '*.whl' -print -quit | grep -q .; then \ + python -m pip install --no-deps --no-index --find-links=/wheelhouse \ + --require-hashes -r /opt/ci/requirements-torch.lock; \ + else \ + python -m pip install --no-deps --index-url "${TORCH_INDEX_URL}" \ + --require-hashes -r /opt/ci/requirements-torch.lock; \ + fi && \ + python -m pip install --no-deps --index-url "${PIP_INDEX_URL}" \ + -r /opt/ci/requirements.lock && \ + python -m pip check + +# Inherit the installed environment, including removal of the bootstrap pip and +# setuptools metadata. COPY over the base venv would merge directories and +# resurrect their old .dist-info entries alongside the locked versions. +FROM dependencies AS ci +ARG SOURCE_COMMIT=unknown +LABEL org.opencontainers.image.title="ColossalAI Torch 2.13 CI dependencies" \ + org.opencontainers.image.source="https://github.com/hpcaitech/ColossalAI" \ + org.opencontainers.image.revision="${SOURCE_COMMIT}" +COPY docker/torch213/requirements*.lock docker/torch213/verify_image.py \ + docker/torch213/Dockerfile docker/torch213/Dockerfile.dockerignore /opt/ci/ +WORKDIR /workspace +# This checks metadata, imports and CPU operations; CUDA remains hidden. +RUN NVIDIA_VISIBLE_DEVICES=void CUDA_VISIBLE_DEVICES="" python /opt/ci/verify_image.py +CMD ["/bin/bash"] diff --git a/docker/torch213/Dockerfile.dockerignore b/docker/torch213/Dockerfile.dockerignore new file mode 100644 index 000000000000..ee63802b4a64 --- /dev/null +++ b/docker/torch213/Dockerfile.dockerignore @@ -0,0 +1,13 @@ +# Only dependency build inputs belong in this context. In particular, do not +# copy the checkout, prebuilt extensions, personal environments or credentials. +** +!docker/ +!docker/torch213/ +!docker/torch213/Dockerfile +!docker/torch213/Dockerfile.dockerignore +!docker/torch213/requirements.lock +!docker/torch213/requirements-torch.lock +!docker/torch213/verify_image.py +!docker/torch213/wheelhouse/ +!docker/torch213/wheelhouse/.gitkeep +!docker/torch213/wheelhouse/*.whl diff --git a/docker/torch213/README.md b/docker/torch213/README.md new file mode 100644 index 000000000000..e5b607315517 --- /dev/null +++ b/docker/torch213/README.md @@ -0,0 +1,251 @@ +# Torch 2.13 / CUDA 13.0 CI dependency image + +This image provides an isolated Python environment and CUDA build tools for the +existing layered GPU tests. It contains **no ColossalAI checkout, installed +ColossalAI distribution or prebuilt ColossalAI extensions**. Install the exact +CI source commit after starting a container. This prevents a new PR's Python +code from using an older image's CUDA kernels. + +The reusable GPU workflow uses this dependency-image profile while retaining +host-side GPU allocation and locking. Runner labels, required checks, the user +release image and `version.txt` are unchanged. The first CPU dependency check +and actual container/GPU qualification are separate acceptance steps. + +## Dependency profile + +| Component | Pinned profile | +| --- | --- | +| Platform | Linux x86_64, Ubuntu 22.04, Python 3.10 | +| CUDA Toolkit | 13.0.3 (nvcc reports 13.0.88) | +| PyTorch | 2.13.0+cu130 | +| torchvision | 0.28.0+cu130 | +| torchaudio | 2.11.0+cu130 (existing development artifact; see below) | +| Triton | 3.7.1 | +| Transformers / PEFT | 4.51.3 / 0.13.2 | +| pytest / pytest-testmon | 7.4.4 / 2.2.0 | +| setuptools | 84.0.0 | +| PyTorch's Python cuDNN package | nvidia-cudnn-cu13==9.20.0.48 | + +The base image also contains a system cuDNN library (9.14 in the original image). +It is distinct from the Python wheel used by PyTorch. Use +`torch.backends.cudnn.version()` to inspect the PyTorch runtime; the base image's +labels alone do not describe that runtime. + +`requirements-torch.lock` pins the four core wheel artifacts, including SHA256 +hashes. `requirements.lock` pins the rest of the existing development environment, +including transitive dependencies and packaging tools. All entries are installed +with `--no-deps`, then `pip check` verifies that the snapshot is complete. This is +a Python 3.10/Linux profile, not a universal lock for all supported platforms. +Non-core packages are version-pinned but not hash-locked; Ubuntu packages still +come from apt repositories. For published images record the resolved base digest +and use the final local image by its full immutable image ID. A registry is optional. + +The existing torchaudio wheel has no `Requires-Dist: torch` declaration. Its +successful `pip check` is therefore **not** evidence of general TorchAudio 2.13 +compatibility. Image validation covers importing it and CPU resampling. Keep this +artifact pinned for reproducing the current environment; separately validate any +audio workload or replacement wheel before changing it. + +Optional backends **Apex, flash-attn, vLLM, TensorNVMe and SmoothQuant extensions** +are not included. Tests that require them may skip. In particular, the Apex batch +requires the legacy `apex.amp` API, not merely an arbitrary Apex installation. +Report skipped coverage separately; this image is not a claim of passing every +optional test. The general `requirements/requirements-test.txt` still includes +flash-attn, so do not reinstall it blindly on top of this locked core profile. + +The library permits `torch>=2.2.0,<2.14`, retaining the 2.5.1 CPU baseline. That +installation range is distinct from this exact GPU profile and does not imply +that every intermediate Torch/CUDA combination has been tested. Transformers +remains at 4.51.3. On Torch 2.13 the FP8 implementation currently uses its eager +fallback; passing FP8 tests does not establish `torch.compile` support. + +## Build + +Use a dedicated or explicitly approved build host. A normal Docker build/load/run +writes to the shared daemon's storage outside the checkout; keeping this recipe +in a personal directory does **not** confine those writes to that directory. +Do not run these commands on a shared server when only personal-directory writes +are permitted. The runtime example below also requires explicit GPU allocation. + +Run from the repository root. Docker with BuildKit is required; GPUs are not +needed during the dependency image build. The Dockerfile-specific ignore file +restricts the context to the lock files, verification script and optional wheels. +The default package mirror is Huawei Cloud; core wheels use the PyTorch cu130 +index. An approved registry mirror or base digest can be supplied via `CUDA_IMAGE`. +The recipe uses BuildKit's bundled Dockerfile frontend so it does not need an +extra `docker/dockerfile` download from Docker Hub. The inspected build host +uses BuildKit 0.27.1; use a recent BuildKit with support for bind/cache mounts. +GPU devices are hidden by default, and the CI executor overrides that setting +only with its allocated UUIDs. + +```bash +# Build a candidate from the committed recipe; use a fresh tag for each revision. +docker build --platform linux/amd64 --progress=plain \ + --build-arg SOURCE_COMMIT="$(git rev-parse HEAD)" \ + -f docker/torch213/Dockerfile \ + -t colossalai-ci:torch2.13-cu130-candidate . + +# For a mirror/pinned base, add these build arguments to the command above. +# --build-arg CUDA_IMAGE=your-approved-registry/cuda@sha256:BASE_DIGEST +# --build-arg PIP_INDEX_URL=https://repo.huaweicloud.com/repository/pypi/simple +``` + +If core downloads are slow, place the four matching wheels in +`docker/torch213/wheelhouse/` first. They are git-ignored and mounted only during +the build; they are not copied into the final image. A nonempty wheelhouse must +contain **all four** pinned wheels. Pip verifies the hashes and fails on missing +or different artifacts rather than silently downloading replacements. Otherwise +the builder downloads the pinned artifacts from `TORCH_INDEX_URL`. + +```bash +# Optional: download the pinned core artifacts before building. +python -m pip download --no-deps --require-hashes \ + --only-binary=:all: --platform manylinux_2_28_x86_64 \ + --platform manylinux_2_27_x86_64 --python-version 310 --implementation cp --abi cp310 \ + --index-url https://download.pytorch.org/whl/cu130 \ + -r docker/torch213/requirements-torch.lock \ + --dest docker/torch213/wheelhouse + +# Verify dependencies and CPU extension operations without exposing any GPUs. +# Explicit runc avoids inheriting a host's default NVIDIA runtime. +docker run --rm --runtime=runc \ + -e NVIDIA_VISIBLE_DEVICES=void -e CUDA_VISIBLE_DEVICES= \ + colossalai-ci:torch2.13-cu130-candidate \ + python /opt/ci/verify_image.py +``` + +## Keep the image on the SSH server (no registry required) + +Building with `docker build` already puts the image into that Docker daemon's +local image store. Record its full image ID and use that same daemon for tests +and the runner. A tag is convenient for humans but can be overwritten; CI pins +the content ID. No image upload, registry account or publishing workflow is +needed. The initial build still downloads base images and dependencies unless +they are already cached or supplied locally. + +```bash +# Record the full immutable ID from the local Docker store; this does not pull. +export COLOSSAL_CI_IMAGE="$(docker image inspect --format '{{.Id}}' colossalai-ci:torch2.13-cu130-candidate)" +printf 'CI image ID: %s\n' "$COLOSSAL_CI_IMAGE" +``` + +Optionally save an image archive in an approved personal directory. Set +`CI_IMAGE_ARCHIVE_ROOT`, for example, to +`/mnt/beegfs/ColossalAI/wangzhijian/images`. Each export below gets a new directory +so it does not overwrite an existing archive. Do not commit large image tars to Git. + +```bash +# Export a local backup and record both its checksum and Docker image ID. +: "${CI_IMAGE_ARCHIVE_ROOT:?Set an approved directory for image archives}" +: "${COLOSSAL_CI_IMAGE:?Set the full local sha256:image-ID}" +mkdir -p "$CI_IMAGE_ARCHIVE_ROOT" +CI_IMAGE_BUNDLE="$(mktemp -d "$CI_IMAGE_ARCHIVE_ROOT/torch213-cu130.XXXXXX")" +docker image save --output "$CI_IMAGE_BUNDLE/image.tar" "$COLOSSAL_CI_IMAGE" +printf '%s\n' "$COLOSSAL_CI_IMAGE" > "$CI_IMAGE_BUNDLE/image.id" +(cd "$CI_IMAGE_BUNDLE" && sha256sum image.tar > image.tar.sha256) +printf 'Saved image bundle: %s\n' "$CI_IMAGE_BUNDLE" + +# Only when restoring or provisioning another daemon: check the archive first, +# then load it. Do not repeat this on every CI run if the image is already present. +(cd "$CI_IMAGE_BUNDLE" && sha256sum --check image.tar.sha256) && \ + docker image load --input "$CI_IMAGE_BUNDLE/image.tar" +export COLOSSAL_CI_IMAGE="$(tr -d '\n' < "$CI_IMAGE_BUNDLE/image.id")" +docker image inspect --format '{{.Id}}' "$COLOSSAL_CI_IMAGE" +``` + +The tar file and Docker's active storage are different things. Docker does not +run directly from `image.tar`: loading it writes layers/metadata to daemon +storage, and starting containers writes there too. A personal archive path does +not bypass shared-server restrictions. Do not change the shared daemon's data +root, restart it, prune unrelated data or claim this is zero-impact on the host. +If all writes must remain private, arrange an administrator-approved isolated +runtime first; this change does not install one automatically. + +## Validate before opening the single PR + +Use the unmerged branch's scripts directly; they are the same entry points that +the workflow calls. Local Git commits are fine and do not create a PR. This +proves the container execution path before GitHub needs the new workflow. + +Qualify the immutable candidate on an approved GPU machine. The host needs a +compatible NVIDIA driver, Docker, NVIDIA Container Toolkit and Python 3.10+; +it does not need your Python environment or host CUDA toolkit. Reserve GPUs +with the team **before** adding them to `COLOSSAL_GPU_POOL`. An idle GPU is not +automatically available for CI use. All executors must use the same agreed lock. + +```bash +# Required pre-agreed settings. Missing values stop the example rather than +# using every idle GPU or inventing a second, ineffective lock. +: "${COLOSSAL_CI_IMAGE:?Set the candidate's full local sha256:image-ID}" +: "${COLOSSAL_GPU_POOL:?Set authorized physical GPU indices, e.g. 4 for one allocated GPU}" +: "${COLOSSAL_GPU_LOCK:?Set the existing shared CI GPU lock file}" +: "${CI_QUALIFICATION_ROOT:?Set an approved personal directory for source and reports}" +export COLOSSAL_GPU_POOL +export CI_SOURCE_SHA="$(git rev-parse HEAD)" +export FAST_MODE=0 + +# Test committed source only. No local .so, editable install or .git credentials +# are carried into the container. Commit the proposed changes before this step. +mkdir -p "$CI_QUALIFICATION_ROOT" +CI_SOURCE_DIR="$(mktemp -d "$CI_QUALIFICATION_ROOT/source.XXXXXX")" +git archive "$CI_SOURCE_SHA" | tar -x -C "$CI_SOURCE_DIR" +docker image inspect --format '{{.Id}}' "$COLOSSAL_CI_IMAGE" + +# First run FP8 batch 4a through the same allocation, container and cleanup code +# used by Actions. Logs, JUnit XML and runtime/result JSON go into this directory. +python3 -B "$CI_SOURCE_DIR/.github/scripts/ci/run_gpu_batch.py" \ + --image "$COLOSSAL_CI_IMAGE" --batch 4a \ + --lock "$COLOSSAL_GPU_LOCK" \ + --output "$CI_QUALIFICATION_ROOT/results-$CI_SOURCE_SHA/batch-4a" +``` + +Repeat the last command with batch `3` (AOT build + Adam), `4b` (inference CUDA), +then `2` and `5` (single-/two-GPU core); batch 5 requires two authorized GPUs. +For cancellation qualification, interrupt **only this test command**, then +check its reported container name is absent from `docker ps -a` and that no +test-owned GPU processes remain. Do not kill or prune unrelated jobs. + +Batch 4a is the FP8 cast/linear/hook regression and is unaffected by `FAST_TEST`. +Batch 3 (Adam kernels) and 4b (inference CUDA kernels) are also unaffected. Layer 2 +selects batches `2 4a 5` with `FAST_MODE=1`, which reduces model matrices in 2/5. +Layer 3 defaults to `FAST_MODE=0`; the batch script **unsets** `FAST_TEST` instead +of setting it to zero. Testmon selection is not enabled. + +The executor holds the host lock through named-container cleanup, filters +devices by UUID and supplies logical `0,1,...` inside the container. Source is +mounted read-only and copied into the container; reports are the only writable +bind, owned by the runner UID/GID. A lightweight private venv references image +dependencies without copying Torch or modifying the base venv. All compiler +caches are initially disposable. Batch 3 compiles AOT with `BUILD_EXT=1`; other +batches exercise JIT as needed. No old ColossalAI binary is reused. + +Before installation, the container compares its embedded Dockerfile, context +ignore file, locks and verifier with the selected source. A mismatch fails with +an explicit rebuild message rather than testing a dependency-upgrade PR against +the old environment. Documentation-only and normal source changes reuse the +image. When changing the recipe, qualify the candidate before pointing CI at it. + +After candidate qualification, set repository variables documented in the +[workflow README](../../.github/workflows/README.md#container-execution-and-image-rollout), +then open the **one combined PR** and exercise its workflow before merging the +switch. Start with a Draft PR: its separate container-smoke check runs only 4a +using `COLOSSAL_CI_SMOKE_GPU_POOL`, not the main `COLOSSAL_GPU_POOL`. Missing +smoke authorization fails closed. Mark the same PR Ready only after arranging +the resources for the full Layer 2 batches `2 4a 5`; smoke success does not +replace those checks. No dependency-image rebuild is needed for routing-only +workflow/script changes because none of the embedded image inputs changes. +The runner checks the preloaded local image before acquiring GPUs; it +does not pull, push or repeatedly reload archives. A local image on another +daemon is insufficient, even when a tar is present on a shared filesystem. +Preload every machine eligible for the runner label or restrict that label to +the provisioned host. Source download and artifact upload still use GitHub. + +Pre-PR script runs do not prove GitHub event dispatch, permissions, Checks or +artifact upload. Verify those on the single PR, and make any further fixes on +that same branch. Do not merge first just to discover whether containers work. + +This recipe does not automatically publish images, configure Docker on the +shared server, or change GitHub required checks. Build/load once per dependency +profile change; code-only PRs reuse the image. Keep the local image ID, optional +archive checksum, base digest, source SHA, build log and qualification reports together. A recipe +or successful CPU check alone is not proof that the GPU CI profile passes. diff --git a/docker/torch213/requirements-torch.lock b/docker/torch213/requirements-torch.lock new file mode 100644 index 000000000000..1b59c7ce4227 --- /dev/null +++ b/docker/torch213/requirements-torch.lock @@ -0,0 +1,7 @@ +# Linux x86_64 / CPython 3.10 / CUDA 13.0. Hashes identify the tested artifacts. +# torchaudio is the existing 2.11 build; its metadata has no torch requirement. +# verify_image.py checks import and resampling, not general audio compatibility. +torch==2.13.0+cu130 --hash=sha256:83767298109148f10de124287b392bc58d74f9938e9430f26f9fd669efb8af9c +torchvision==0.28.0+cu130 --hash=sha256:19c2d18bbc8c3b1a666f0d41f91d8f624f912faffad421371061dddda52ba8ea +torchaudio==2.11.0+cu130 --hash=sha256:91497b6525569dcd0871991d32bb4beda90775bcbe9050328ea7b8022106e2fa +triton==3.7.1 --hash=sha256:3d62b1df358f49fcc229d5b77e7a346d66a93e93c19f6872936eadba9aac5f01 diff --git a/docker/torch213/requirements.lock b/docker/torch213/requirements.lock new file mode 100644 index 000000000000..e235e03a08ef --- /dev/null +++ b/docker/torch213/requirements.lock @@ -0,0 +1,136 @@ +# Dependency snapshot for Linux x86_64, CPython 3.10 and requirements-torch.lock. +# Captured from the existing Torch 2.13 development environment. No ColossalAI, +# personal paths, VCS dependencies or optional GPU backends are included. +# All entries are installed with --no-deps and checked with pip check. +pip==26.2.1 +setuptools==84.0.0 +wheel==0.48.0 +accelerate==1.15.0 +aiohappyeyeballs==2.7.1 +aiohttp==3.14.3 +aiosignal==1.4.0 +annotated-doc==0.0.5 +annotated-types==0.8.0 +anyio==4.15.1 +async-timeout==5.0.1 +attrs==26.1.0 +bcrypt==5.0.0 +beautifulsoup4==4.15.0 +bitsandbytes==0.50.2 +certifi==2026.7.22 +cffi==2.1.1 +cfgv==3.5.0 +charset-normalizer==3.5.1 +click==8.5.0 +contexttimer==0.3.3 +coverage==7.2.3 +cryptography==50.0.1 +cuda-bindings==13.4.1 +cuda-pathfinder==1.8.1 +cuda-toolkit==13.0.3.0 +datasets==5.0.1 +decorator==5.3.1 +Deprecated==1.3.1 +diffusers==0.29.0 +dill==0.4.1 +distlib==0.4.3 +einops==0.8.2 +exceptiongroup==1.3.1 +fabric==3.2.3 +fastapi==0.141.1 +filelock==3.32.7 +frozenlist==1.8.0 +fsspec==2026.6.0 +galore-torch==1.0 +google==3.0.0 +h11==0.16.0 +hf-xet==1.6.0 +httpcore==1.0.9 +httpx==0.28.1 +huggingface_hub==0.36.2 +identify==2.6.19 +idna==3.19 +importlib_metadata==9.0.1 +iniconfig==2.3.0 +invoke==2.2.1 +Jinja2==3.1.6 +jsonschema==4.26.0 +jsonschema-specifications==2025.9.1 +markdown-it-py==4.2.0 +MarkupSafe==3.0.3 +mdurl==0.1.2 +mpmath==1.3.0 +msgpack==1.2.2 +multidict==6.8.0 +multiprocess==0.70.19 +networkx==3.4.2 +ninja==1.13.2 +nodeenv==1.10.0 +numpy==2.2.6 +nvidia-cublas==13.1.1.3 +nvidia-cuda-cupti==13.0.85 +nvidia-cuda-nvrtc==13.0.88 +nvidia-cuda-runtime==13.0.96 +nvidia-cudnn-cu13==9.20.0.48 +nvidia-cufft==12.0.0.61 +nvidia-cufile==1.15.1.6 +nvidia-curand==10.4.0.35 +nvidia-cusolver==12.0.4.66 +nvidia-cusparse==12.6.3.3 +nvidia-cusparselt-cu13==0.8.1 +nvidia-nccl-cu13==2.29.7 +nvidia-nvjitlink==13.4.52 +nvidia-nvshmem-cu13==3.4.5 +nvidia-nvtx==13.0.85 +packaging==26.3 +pandas==2.3.3 +paramiko==5.0.0 +peft==0.13.2 +pillow==12.3.0 +platformdirs==4.11.8 +pluggy==1.6.0 +plumbum==2.0.2 +pre_commit==4.6.2 +propcache==0.5.2 +protobuf==7.36.1 +psutil==7.2.2 +pyarrow==25.0.1 +pycparser==3.0 +pydantic==2.13.5 +pydantic_core==2.46.5 +Pygments==2.21.0 +PyNaCl==1.6.2 +pytest==7.4.4 +pytest-testmon==2.2.0 +python-dateutil==2.9.0.post0 +python-discovery==1.6.0 +pytz==2026.3.post1 +PyYAML==6.0.3 +ray==2.58.0 +referencing==0.37.0 +regex==2026.9.10 +requests==2.34.2 +rich==15.0.0 +rpds-py==0.30.0 +rpyc==6.0.0 +safetensors==0.8.0 +sentencepiece==0.2.2 +six==1.17.0 +soupsieve==2.9.2 +starlette==1.6.0 +sympy==1.14.0 +timm==1.0.29 +tokenizers==0.21.4 +tomli==2.4.1 +tqdm==4.70.1 +transformers==4.51.3 +typing-inspection==0.4.4 +typing_extensions==4.16.0 +tzdata==2026.4 +urllib3==2.8.0 +uvicorn==0.29.0 +virtualenv==21.7.9 +wrapt==2.4.1 +xxhash==4.0.1 +yarl==1.25.1 +zipp==4.1.0 diff --git a/docker/torch213/verify_image.py b/docker/torch213/verify_image.py new file mode 100644 index 000000000000..01ed90e15fd8 --- /dev/null +++ b/docker/torch213/verify_image.py @@ -0,0 +1,62 @@ +"""Validate the dependency image without importing a checkout or accessing GPUs.""" + +import importlib +import importlib.metadata as metadata +import json +import shutil +import subprocess +import sys +from pathlib import Path + +from packaging.requirements import Requirement +from packaging.version import Version + + +def main(): + assert sys.version_info[:2] == (3, 10), sys.version + root = Path(__file__).resolve().parent + versions = {} + for lock_name in ("requirements.lock", "requirements-torch.lock"): + for line in (root / lock_name).read_text().splitlines(): + line = line.split("#", 1)[0].strip() + if not line: + continue + requirement = Requirement(line.split(" --hash=", 1)[0]) + actual = metadata.version(requirement.name) + assert Version(actual) in requirement.specifier, (requirement.name, actual, str(requirement.specifier)) + versions[requirement.name] = actual + + subprocess.run([sys.executable, "-m", "pip", "check"], check=True) + for executable in ("nvcc", "g++", "cmake", "ninja", "git"): + assert shutil.which(executable), f"Missing build tool: {executable}" + nvcc = subprocess.check_output(["nvcc", "--version"], text=True) + assert "release 13.0" in nvcc, nvcc + + # A base image must not silently supply a stale ColossalAI package or AOT .so. + try: + metadata.distribution("colossalai") + except metadata.PackageNotFoundError: + pass + else: + raise AssertionError("Install ColossalAI from the CI source SHA at runtime, not in this image") + + for module in ("transformers", "peft", "diffusers", "triton", "testmon.pytest_testmon"): + importlib.import_module(module) + + import torch + import torchaudio + import torchvision + + assert torch.version.cuda == "13.0", torch.version.cuda + # CPU extension smoke checks catch some ABI mismatches that pip check cannot. + boxes = torch.tensor([[0.0, 0.0, 1.0, 1.0], [0.0, 0.0, 1.0, 1.0]]) + assert torchvision.ops.nms(boxes, torch.tensor([0.9, 0.8]), 0.5).tolist() == [0] + waveform = torch.zeros(1, 160) + assert torchaudio.functional.resample(waveform, 16000, 8000).shape == (1, 80) + # pytest --help loads installed entry points, including pytest-testmon. + subprocess.run([sys.executable, "-m", "pytest", "--help"], check=True, stdout=subprocess.DEVNULL) + print(json.dumps({"python": sys.version, "torch_cuda": torch.version.cuda, "packages": versions}, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/docker/torch213/wheelhouse/.gitkeep b/docker/torch213/wheelhouse/.gitkeep new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/requirements/requirements-test.txt b/requirements/requirements-test.txt index fa08ef3e61aa..c0f31c9406c4 100644 --- a/requirements/requirements-test.txt +++ b/requirements/requirements-test.txt @@ -1,6 +1,6 @@ pytest coverage==7.2.3 -git+https://github.com/hpcaitech/pytest-testmon +pytest-testmon==2.2.0 torchvision timm torchaudio>=0.13.1 diff --git a/requirements/requirements.txt b/requirements/requirements.txt index 1c8d78719497..c84a28969484 100644 --- a/requirements/requirements.txt +++ b/requirements/requirements.txt @@ -7,7 +7,7 @@ click fabric contexttimer ninja -torch>=2.2.0,<=2.5.1 +torch>=2.2.0,<2.14 safetensors einops pydantic diff --git a/tests/test_checkpoint_io/test_gemini_checkpoint_io.py b/tests/test_checkpoint_io/test_gemini_checkpoint_io.py index ba60de0bd1c1..716a56e65bc8 100644 --- a/tests/test_checkpoint_io/test_gemini_checkpoint_io.py +++ b/tests/test_checkpoint_io/test_gemini_checkpoint_io.py @@ -1,4 +1,5 @@ import os +from importlib import import_module import pytest import torch @@ -23,9 +24,10 @@ ASYNC_MODES = [False, True] if HAS_TENSORNVME else [False] try: - import apex.amp - - APEX_AVAILABLE = True + apex_normalization = import_module("apex.normalization") + # Apex's Python-only installation lacks the CUDA kernel used by both norms. + import_module("fused_layer_norm_cuda") + APEX_AVAILABLE = all(hasattr(apex_normalization, name) for name in ("FusedLayerNorm", "FusedRMSNorm")) except ImportError: APEX_AVAILABLE = False diff --git a/tests/test_fp8/test_fp8_linear.py b/tests/test_fp8/test_fp8_linear.py index d035957f2a31..25bd29a2cc19 100644 --- a/tests/test_fp8/test_fp8_linear.py +++ b/tests/test_fp8/test_fp8_linear.py @@ -4,35 +4,32 @@ from torch.testing import assert_close from colossalai.accelerator import get_accelerator -from colossalai.quantization.fp8 import linear_fp8 +from colossalai.quantization.fp8 import _LinearFp8, linear_fp8 from colossalai.utils import get_current_device D_IN, D_OUT = 16, 32 B, S = 2, 64 -DTYPE = torch.bfloat16 -@pytest.mark.skipif(get_accelerator().get_device_capability()[0] < 9, reason="Test requires device capability >= 9.0") -@pytest.mark.parametrize("use_bias", [True, False]) -@pytest.mark.parametrize("use_batch", [True, False]) -def test_fp8_linear(use_bias: bool, use_batch: bool): +def check_fp8_linear(linear, use_bias: bool, use_batch: bool, dtype: torch.dtype): + torch.manual_seed(42) # create tensors - w = torch.rand(D_OUT, D_IN, device=get_current_device(), dtype=DTYPE, requires_grad=True) + w = torch.rand(D_OUT, D_IN, device=get_current_device(), dtype=dtype, requires_grad=True) ref_w = w.clone().detach().requires_grad_() if use_batch: x_shape = (B, S, D_IN) else: x_shape = (S, D_IN) - x = torch.rand(x_shape, device=get_current_device(), dtype=DTYPE, requires_grad=True) + x = torch.rand(x_shape, device=get_current_device(), dtype=dtype, requires_grad=True) ref_x = x.clone().detach().requires_grad_() if use_bias: - bias = torch.rand(D_OUT, device=get_current_device(), dtype=DTYPE, requires_grad=True) + bias = torch.rand(D_OUT, device=get_current_device(), dtype=dtype, requires_grad=True) ref_bias = bias.clone().detach().requires_grad_() else: bias = None ref_bias = None - out = linear_fp8(x, w, bias) + out = linear(x, w, bias) assert out.shape == x_shape[:-1] + (D_OUT,) out.sum().backward() ref_out = F.linear(ref_x, ref_w, ref_bias) @@ -43,3 +40,34 @@ def test_fp8_linear(use_bias: bool, use_batch: bool): assert_close(w.grad, ref_w.grad, rtol=0.2, atol=0.1) if use_bias: assert_close(bias.grad, ref_bias.grad, rtol=0.2, atol=0.1) + + +@pytest.mark.skipif(get_accelerator().get_device_capability()[0] < 9, reason="Test requires device capability >= 9.0") +@pytest.mark.parametrize("use_bias", [True, False]) +@pytest.mark.parametrize("use_batch", [True, False]) +def test_fp8_linear(use_bias: bool, use_batch: bool): + # Test the public entry point, including the existing 2.13 eager fallback. + check_fp8_linear(linear_fp8, use_bias, use_batch, torch.bfloat16) + + +@pytest.mark.skipif(get_accelerator().get_device_capability()[0] < 9, reason="Test requires device capability >= 9.0") +@pytest.mark.parametrize("legacy_return", [False, True], ids=["tensor-return", "tuple-return"]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("use_bias", [False, True]) +@pytest.mark.parametrize("use_batch", [False, True]) +def test_fp8_linear_scaled_mm_return(monkeypatch, legacy_return, dtype, use_bias, use_batch): + scaled_mm = torch._scaled_mm + calls = [] + + def scaled_mm_with_return_contract(*args, **kwargs): + result = scaled_mm(*args, **kwargs) + output = result[0] if isinstance(result, tuple) else result + calls.append(output.shape) + # amax is unused by _LinearFp8; exercise both historical return contracts. + return (output, output.abs().max()) if legacy_return else output + + monkeypatch.setattr(torch, "_scaled_mm", scaled_mm_with_return_contract) + # Test the actual forward and backward kernels without compiler caching of + # the patched function. This does not claim torch.compile coverage on 2.13. + check_fp8_linear(_LinearFp8.apply, use_bias, use_batch, dtype) + assert len(calls) == 3, "Expected forward, input-gradient and weight-gradient matmuls" diff --git a/tests/test_zero/test_gemini/test_optim.py b/tests/test_zero/test_gemini/test_optim.py index b89c591882d6..7d94d01a4cf5 100644 --- a/tests/test_zero/test_gemini/test_optim.py +++ b/tests/test_zero/test_gemini/test_optim.py @@ -1,3 +1,5 @@ +from importlib import import_module + import pytest import torch import torch.distributed as dist @@ -15,8 +17,8 @@ from tests.kit.model_zoo import model_zoo, run_fwd_bwd try: - import apex.amp - + # Keep the capability probe visible to unused-import cleanup tools. + import_module("apex.amp") APEX_AMP_AVAILABLE = True except ImportError: APEX_AMP_AVAILABLE = False