| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import os |
| 5 | import shlex |
| 6 | import subprocess |
| 7 | from pathlib import Path |
| 8 | |
| 9 | REPO_ROOT = Path(__file__).resolve().parents[2] |
| 10 | HELPER = REPO_ROOT / ".buildkite" / "scripts" / "docker-build-metadata-args.sh" |
| 11 | ROCM_CI_BAKE = REPO_ROOT / ".buildkite" / "scripts" / "ci-bake-rocm.sh" |
| 12 | ROCM_IMAGE_SMOKE = REPO_ROOT / ".buildkite" / "scripts" / "rocm" / "smoke-test-image.sh" |
| 13 | |
| 14 | |
| 15 | def run_helper( |
| 16 | *args: str, |
| 17 | env: dict[str, str] | None = None, |
| 18 | path: str | None = None, |
| 19 | ) -> list[str]: |
| 20 | helper_env = {"PATH": path or os.environ["PATH"]} |
| 21 | if env: |
| 22 | helper_env.update(env) |
| 23 | result = subprocess.run( |
| 24 | ["bash", str(HELPER), *args], |
| 25 | check=True, |
| 26 | env=helper_env, |
| 27 | stdout=subprocess.PIPE, |
| 28 | text=True, |
| 29 | ) |
| 30 | return shlex.split(result.stdout) |
| 31 | |
| 32 | |
| 33 | def option_values(args: list[str], option: str) -> list[str]: |
| 34 | return [args[i + 1] for i, arg in enumerate(args[:-1]) if arg == option] |
| 35 | |
| 36 | |
| 37 | def build_args(args: list[str]) -> dict[str, str]: |
| 38 | values = {} |
| 39 | for value in option_values(args, "--build-arg"): |
| 40 | key, arg_value = value.split("=", 1) |
| 41 | values[key] = arg_value |
| 42 | return values |
| 43 | |
| 44 | |
| 45 | def test_release_metadata_args_prefer_pipeline_id() -> None: |
| 46 | args = run_helper( |
| 47 | "cu130-ubuntu2404", |
| 48 | env={ |
| 49 | "BUILDKITE": "1", |
| 50 | "BUILDKITE_COMMIT": "abc123", |
| 51 | "BUILDKITE_PIPELINE_ID": "pipe-uuid", |
| 52 | "BUILDKITE_PIPELINE_SLUG": "release", |
| 53 | "BUILDKITE_BUILD_URL": "https://buildkite.example/vllm/builds/1", |
| 54 | "RELEASE_VERSION": "v0.20.0", |
| 55 | }, |
| 56 | ) |
| 57 | |
| 58 | assert build_args(args) == { |
| 59 | "VLLM_BUILD_COMMIT": "abc123", |
| 60 | "VLLM_BUILD_PIPELINE": "pipe-uuid", |
| 61 | "VLLM_BUILD_URL": "https://buildkite.example/vllm/builds/1", |
| 62 | "VLLM_IMAGE_TAG": "vllm/vllm-openai:v0.20.0-cu130-ubuntu2404", |
| 63 | } |
| 64 | expected_tag = ( |
| 65 | "public.ecr.aws/q9t5s3a7/vllm-release-repo:" |
| 66 | f"abc123-{os.uname().machine}-cu130-ubuntu2404" |
| 67 | ) |
| 68 | assert option_values(args, "--tag") == [expected_tag] |
| 69 | |
| 70 | |
| 71 | def test_nightly_metadata_args_fall_back_to_pipeline_slug() -> None: |
| 72 | args = run_helper( |
| 73 | "ubuntu2404", |
| 74 | env={ |
| 75 | "BUILDKITE": "1", |
| 76 | "BUILDKITE_COMMIT": "def456", |
| 77 | "BUILDKITE_PIPELINE_SLUG": "release", |
| 78 | "BUILDKITE_BUILD_URL": "https://buildkite.example/vllm/builds/2", |
| 79 | "NIGHTLY": "1", |
| 80 | }, |
| 81 | ) |
| 82 | |
| 83 | assert build_args(args) == { |
| 84 | "VLLM_BUILD_COMMIT": "def456", |
| 85 | "VLLM_BUILD_PIPELINE": "release", |
| 86 | "VLLM_BUILD_URL": "https://buildkite.example/vllm/builds/2", |
| 87 | "VLLM_IMAGE_TAG": "vllm/vllm-openai:nightly-def456-ubuntu2404", |
| 88 | } |
| 89 | expected_tag = ( |
| 90 | "public.ecr.aws/q9t5s3a7/vllm-release-repo:" |
| 91 | f"def456-{os.uname().machine}-ubuntu2404" |
| 92 | ) |
| 93 | assert option_values(args, "--tag") == [expected_tag] |
| 94 | |
| 95 | |
| 96 | def test_local_metadata_args_use_local_overrides() -> None: |
| 97 | args = run_helper( |
| 98 | env={ |
| 99 | "VLLM_IMAGE_TAG": "local/test:dev", |
| 100 | "VLLM_BUILD_COMMIT": "localsha", |
| 101 | "VLLM_BUILD_PIPELINE": "local-pipeline", |
| 102 | "VLLM_BUILD_URL": "https://buildkite.example/local", |
| 103 | }, |
| 104 | ) |
| 105 | |
| 106 | assert build_args(args) == { |
| 107 | "VLLM_BUILD_COMMIT": "localsha", |
| 108 | "VLLM_BUILD_PIPELINE": "local-pipeline", |
| 109 | "VLLM_BUILD_URL": "https://buildkite.example/local", |
| 110 | "VLLM_IMAGE_TAG": "local/test:dev", |
| 111 | } |
| 112 | assert option_values(args, "--tag") == ["local/test:dev"] |
| 113 | |
| 114 | |
| 115 | def test_release_version_lookup_failure_falls_back_to_commit( |
| 116 | tmp_path: Path, |
| 117 | ) -> None: |
| 118 | fake_bin = tmp_path / "bin" |
| 119 | fake_bin.mkdir() |
| 120 | buildkite_agent = fake_bin / "buildkite-agent" |
| 121 | buildkite_agent.write_text("#!/bin/sh\nexit 1\n") |
| 122 | buildkite_agent.chmod(0o755) |
| 123 | |
| 124 | args = run_helper( |
| 125 | "cu129", |
| 126 | env={ |
| 127 | "BUILDKITE": "1", |
| 128 | "BUILDKITE_COMMIT": "fallback123", |
| 129 | "BUILDKITE_PIPELINE_SLUG": "release", |
| 130 | }, |
| 131 | path=f"{fake_bin}:{os.environ['PATH']}", |
| 132 | ) |
| 133 | |
| 134 | assert build_args(args)["VLLM_IMAGE_TAG"] == ("vllm/vllm-openai:vfallback123-cu129") |
| 135 | |
| 136 | |
| 137 | def test_vllm_openai_image_embeds_metadata_contract() -> None: |
| 138 | dockerfile = (REPO_ROOT / "docker" / "Dockerfile").read_text() |
| 139 | |
| 140 | for expected in ( |
| 141 | "ARG VLLM_BUILD_COMMIT", |
| 142 | "ARG VLLM_BUILD_PIPELINE", |
| 143 | "ARG VLLM_BUILD_URL", |
| 144 | "ARG VLLM_IMAGE_TAG", |
| 145 | "VLLM_BUILD_COMMIT=${VLLM_BUILD_COMMIT:-unknown}", |
| 146 | "VLLM_BUILD_PIPELINE=${VLLM_BUILD_PIPELINE:-local}", |
| 147 | "VLLM_BUILD_URL=${VLLM_BUILD_URL:-}", |
| 148 | "VLLM_IMAGE_TAG=${VLLM_IMAGE_TAG:-local/vllm-openai:dev}", |
| 149 | 'ai.vllm.build.commit="${VLLM_BUILD_COMMIT}"', |
| 150 | 'ai.vllm.build.pipeline="${VLLM_BUILD_PIPELINE}"', |
| 151 | 'ai.vllm.build.url="${VLLM_BUILD_URL}"', |
| 152 | 'ai.vllm.image.tag="${VLLM_IMAGE_TAG}"', |
| 153 | ): |
| 154 | assert expected in dockerfile |
| 155 | |
| 156 | |
| 157 | def test_rust_build_cache_excludes_git_metadata() -> None: |
| 158 | import torch |
| 159 | |
| 160 | from vllm.platforms import current_platform |
| 161 | |
| 162 | dockerfile_names = ["Dockerfile", "Dockerfile.cpu"] |
| 163 | # CPU jobs can reuse ROCm artifacts, which omit the XPU Dockerfile. |
| 164 | if not current_platform.is_rocm() and ( |
| 165 | not current_platform.is_cpu() or torch.version.hip is None |
| 166 | ): |
| 167 | dockerfile_names.append("Dockerfile.xpu") |
| 168 | for name in dockerfile_names: |
| 169 | dockerfile = (REPO_ROOT / "docker" / name).read_text() |
| 170 | cached_stage, exact_version_stage = dockerfile.split( |
| 171 | "FROM rust-build-cache AS rust-build", maxsplit=1 |
| 172 | ) |
| 173 | exact_version_stage = exact_version_stage.split("\nFROM ", maxsplit=1)[0] |
| 174 | cached_run = cached_stage.rsplit("RUN ", maxsplit=1)[1] |
| 175 | |
| 176 | assert 'SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+docker.cache"' in cached_run |
| 177 | assert "source=.git,target=.git" not in cached_run |
| 178 | assert "source=.git,target=.git" in exact_version_stage |
| 179 | assert 'SETUPTOOLS_SCM_PRETEND_METADATA="{dirty=false}"' in exact_version_stage |
| 180 | assert "bash tools/build_rust.sh" in exact_version_stage |
| 181 | |
| 182 | |
| 183 | def test_rocm_ci_base_bake_embeds_content_hash_label() -> None: |
| 184 | bake_file = (REPO_ROOT / "docker" / "docker-bake-rocm.hcl").read_text() |
| 185 | |
| 186 | for expected in ( |
| 187 | 'variable "CI_BASE_CONTENT_HASH"', |
| 188 | 'target "ci-base-rocm"', |
| 189 | 'target = "ci_base"', |
| 190 | '"vllm.ci_base.content_hash" = CI_BASE_CONTENT_HASH', |
| 191 | ): |
| 192 | assert expected in bake_file |
| 193 | |
| 194 | |
| 195 | def test_rocm_ci_base_metadata_inputs_cover_ci_base_files() -> None: |
| 196 | ci_bake = ROCM_CI_BAKE.read_text() |
| 197 | |
| 198 | for expected in ( |
| 199 | "requirements/common.txt", |
| 200 | "requirements/rocm.txt", |
| 201 | "requirements/test/rocm.txt", |
| 202 | "docker/Dockerfile.rocm", |
| 203 | ): |
| 204 | assert expected in ci_bake |
| 205 | |
| 206 | |
| 207 | def test_rocm_ci_smoke_runs_in_shared_buildkit_graph() -> None: |
| 208 | dockerfile = (REPO_ROOT / "docker" / "Dockerfile.rocm").read_text() |
| 209 | ci_hcl = (REPO_ROOT / "docker" / "ci-rocm.hcl").read_text() |
| 210 | full_image_group = ci_hcl.split('group "test-rocm-ci-with-wheel"', maxsplit=1)[ |
| 211 | 1 |
| 212 | ].split("}", maxsplit=1)[0] |
| 213 | |
| 214 | for expected in ( |
| 215 | "FROM test AS test_smoke", |
| 216 | "smoke-test-image.sh --inside", |
| 217 | "FROM scratch AS export_test_smoke", |
| 218 | 'target "smoke-test-rocm-ci"', |
| 219 | 'target = "export_test_smoke"', |
| 220 | 'output = ["type=local,dest=./build/rocm-smoke-export"]', |
| 221 | ): |
| 222 | assert expected in dockerfile or expected in ci_hcl |
| 223 | assert '"smoke-test-rocm-ci"' in full_image_group |
| 224 | assert 'target "smoke-test-rocm-ci"' in ROCM_CI_BAKE.read_text() |
| 225 | |
| 226 | |
| 227 | def prepare_rocm_smoke_test( |
| 228 | tmp_path: Path, |
| 229 | *, |
| 230 | marker_id: str, |
| 231 | build_id: str, |
| 232 | ) -> tuple[dict[str, str], Path, Path, Path]: |
| 233 | fake_bin = tmp_path / "bin" |
| 234 | fake_bin.mkdir() |
| 235 | docker_called = tmp_path / "docker-called" |
| 236 | docker = fake_bin / "docker" |
| 237 | docker.write_text('#!/bin/sh\ntouch "$FAKE_DOCKER_CALLED"\nexit 99\n') |
| 238 | docker.chmod(0o755) |
| 239 | |
| 240 | marker = tmp_path / "build" / "rocm-smoke-export" / "vllm-smoke-ok" |
| 241 | marker.parent.mkdir(parents=True) |
| 242 | marker.write_text(f"{marker_id}\n") |
| 243 | env = os.environ.copy() |
| 244 | env.update( |
| 245 | { |
| 246 | "BUILDKITE_BUILD_ID": build_id, |
| 247 | "FAKE_DOCKER_CALLED": str(docker_called), |
| 248 | "PATH": f"{fake_bin}:{env['PATH']}", |
| 249 | } |
| 250 | ) |
| 251 | env.pop("ROCM_CI_ARTIFACT_ONLY", None) |
| 252 | env.pop("VLLM_CI_SMOKE_IMAGE", None) |
| 253 | return env, marker, docker, docker_called |
| 254 | |
| 255 | |
| 256 | def test_rocm_smoke_marker_avoids_host_image_pull(tmp_path: Path) -> None: |
| 257 | env, marker, _, docker_called = prepare_rocm_smoke_test( |
| 258 | tmp_path, |
| 259 | marker_id="build-123", |
| 260 | build_id="build-123", |
| 261 | ) |
| 262 | |
| 263 | result = subprocess.run( |
| 264 | ["bash", str(ROCM_IMAGE_SMOKE)], |
| 265 | check=True, |
| 266 | cwd=tmp_path, |
| 267 | env=env, |
| 268 | stdout=subprocess.PIPE, |
| 269 | text=True, |
| 270 | ) |
| 271 | |
| 272 | assert "verified inside BuildKit" in result.stdout |
| 273 | assert not docker_called.exists() |
| 274 | assert not marker.exists() |
| 275 | |
| 276 | |
| 277 | def test_rocm_smoke_rejects_marker_from_another_build(tmp_path: Path) -> None: |
| 278 | env, marker, _, docker_called = prepare_rocm_smoke_test( |
| 279 | tmp_path, |
| 280 | marker_id="previous-build", |
| 281 | build_id="current-build", |
| 282 | ) |
| 283 | |
| 284 | result = subprocess.run( |
| 285 | ["bash", str(ROCM_IMAGE_SMOKE)], |
| 286 | check=False, |
| 287 | cwd=tmp_path, |
| 288 | env=env, |
| 289 | stderr=subprocess.PIPE, |
| 290 | text=True, |
| 291 | ) |
| 292 | |
| 293 | assert result.returncode == 1 |
| 294 | assert "previous-build, not current-build" in result.stderr |
| 295 | assert not docker_called.exists() |
| 296 | assert marker.exists() |
| 297 | |
| 298 | |
| 299 | def test_rocm_smoke_override_streams_current_checks_to_docker( |
| 300 | tmp_path: Path, |
| 301 | ) -> None: |
| 302 | env, marker, docker, _ = prepare_rocm_smoke_test( |
| 303 | tmp_path, |
| 304 | marker_id="build-123", |
| 305 | build_id="build-123", |
| 306 | ) |
| 307 | docker_args = tmp_path / "docker-args" |
| 308 | docker_stdin = tmp_path / "docker-stdin" |
| 309 | docker.write_text( |
| 310 | '#!/bin/sh\nprintf "%s\\n" "$@" > "$FAKE_DOCKER_ARGS"\n' |
| 311 | 'cat > "$FAKE_DOCKER_STDIN"\n' |
| 312 | ) |
| 313 | env.update( |
| 314 | { |
| 315 | "FAKE_DOCKER_ARGS": str(docker_args), |
| 316 | "FAKE_DOCKER_STDIN": str(docker_stdin), |
| 317 | "IMAGE_TAG": "rocm/vllm-ci:built", |
| 318 | "VLLM_CI_SMOKE_IMAGE": "rocm/vllm-ci:override", |
| 319 | } |
| 320 | ) |
| 321 | |
| 322 | subprocess.run( |
| 323 | ["bash", str(ROCM_IMAGE_SMOKE)], |
| 324 | check=True, |
| 325 | cwd=tmp_path, |
| 326 | env=env, |
| 327 | ) |
| 328 | |
| 329 | assert "rocm/vllm-ci:override" in docker_args.read_text().splitlines() |
| 330 | assert docker_args.read_text().splitlines()[-3:] == ["-s", "--", "--inside"] |
| 331 | assert "run_smoke_checks()" in docker_stdin.read_text() |
| 332 | assert marker.exists() |
| 333 | |
| 334 | |
| 335 | def test_rocm_git_fetch_disables_automatic_maintenance(tmp_path: Path) -> None: |
| 336 | fake_bin = tmp_path / "bin" |
| 337 | fake_bin.mkdir() |
| 338 | git = fake_bin / "git" |
| 339 | git.write_text('#!/bin/sh\nprintf "%s\\n" "$@"\n') |
| 340 | git.chmod(0o755) |
| 341 | |
| 342 | env = os.environ.copy() |
| 343 | env["PATH"] = f"{fake_bin}:{env['PATH']}" |
| 344 | result = subprocess.run( |
| 345 | [ |
| 346 | "bash", |
| 347 | "-c", |
| 348 | 'source "$1"; git_fetch_with_timeout --quiet origin HEAD', |
| 349 | "bash", |
| 350 | str(ROCM_CI_BAKE), |
| 351 | ], |
| 352 | check=True, |
| 353 | env=env, |
| 354 | stdout=subprocess.PIPE, |
| 355 | text=True, |
| 356 | ) |
| 357 | |
| 358 | assert result.stdout.splitlines() == [ |
| 359 | "fetch", |
| 360 | "--no-auto-maintenance", |
| 361 | "--quiet", |
| 362 | "origin", |
| 363 | "HEAD", |
| 364 | ] |