A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 02:45 UTC
[ROCm] [CI] [Release] Rocm wheel pipeline with sccache (#32264)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
TJian committed
Jan 16, 2026 at 02:56 UTC
41c544f78a4d8264adec86dcf8d338c62193bf68
10 files changed
+1273
-16
.buildkite/release-pipeline.yaml
+362
@@ -214,3 +214,365 @@ steps:
214
env:
215
DOCKER_BUILDKIT: "1"
216
DOCKERHUB_USERNAME: "vllmbot"
217
+
218
+ # =============================================================================
219
+ # ROCm Release Pipeline (x86_64 only)
220
+ # =============================================================================
221
+ #
222
+ # vLLM version is determined by the Buildkite checkout (like CUDA pipeline).
223
+ # To build a specific version, trigger the build from that branch/tag.
224
+ #
225
+ # Environment variables for ROCm builds (set via Buildkite UI or schedule):
226
+ # ROCM_PYTHON_VERSION: Python version (default: 3.12)
227
+ # PYTORCH_ROCM_ARCH: GPU architectures (default: gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151)
228
+ # ROCM_UPLOAD_WHEELS: Upload to S3 (default: false for nightly, true for releases)
229
+ # ROCM_FORCE_REBUILD: Force rebuild base wheels, ignore S3 cache (default: false)
230
+ #
231
+ # Note: ROCm version is determined by BASE_IMAGE in docker/Dockerfile.rocm_base
232
+ # (currently rocm/dev-ubuntu-22.04:7.1-complete)
233
+ #
234
+ # =============================================================================
235
+
236
+ # ROCm Input Step - Collect build configuration (manual trigger only)
237
+ - input: "ROCm Wheel Release Build Configuration"
238
+ key: input-rocm-config
239
+ depends_on: ~
240
+ if: build.source == "ui"
241
+ fields:
242
+ - text: "Python Version"
243
+ key: "rocm-python-version"
244
+ default: "3.12"
245
+ hint: "Python version (e.g., 3.12)"
246
+ - text: "GPU Architectures"
247
+ key: "rocm-pytorch-rocm-arch"
248
+ default: "gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151"
249
+ hint: "Semicolon-separated GPU architectures"
250
+ - select: "Upload Wheels to S3"
251
+ key: "rocm-upload-wheels"
252
+ default: "true"
253
+ options:
254
+ - label: "No - Build only (nightly/dev)"
255
+ value: "false"
256
+ - label: "Yes - Upload to S3 (release)"
257
+ value: "true"
258
+ - select: "Force Rebuild Base Wheels"
259
+ key: "rocm-force-rebuild"
260
+ default: "false"
261
+ hint: "Ignore S3 cache and rebuild base wheels from scratch"
262
+ options:
263
+ - label: "No - Use cached wheels if available"
264
+ value: "false"
265
+ - label: "Yes - Rebuild even if cache exists"
266
+ value: "true"
267
+
268
+ # ROCm Job 1: Build ROCm Base Wheels (with S3 caching)
269
+ - label: ":rocm: Build ROCm Base Wheels"
270
+ id: build-rocm-base-wheels
271
+ depends_on:
272
+ - step: input-rocm-config
273
+ allow_failure: true # Allow failure so non-UI builds can proceed (input step is skipped)
274
+ agents:
275
+ queue: cpu_queue_postmerge
276
+ commands:
277
+ # Set configuration and check cache
278
+ - |
279
+ set -euo pipefail
280
+
281
+ # Get values from meta-data (set by input step) or use defaults
282
+ PYTHON_VERSION="$$(buildkite-agent meta-data get rocm-python-version 2>/dev/null || echo '')"
283
+ export PYTHON_VERSION="$${PYTHON_VERSION:-3.12}"
284
+
285
+ PYTORCH_ROCM_ARCH="$$(buildkite-agent meta-data get rocm-pytorch-rocm-arch 2>/dev/null || echo '')"
286
+ export PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH:-gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151}"
287
+
288
+ # Check for force rebuild flag
289
+ ROCM_FORCE_REBUILD="$${ROCM_FORCE_REBUILD:-}"
290
+ if [ -z "$${ROCM_FORCE_REBUILD}" ]; then
291
+ ROCM_FORCE_REBUILD="$$(buildkite-agent meta-data get rocm-force-rebuild 2>/dev/null || echo '')"
292
+ fi
293
+
294
+ echo "========================================"
295
+ echo "ROCm Base Wheels Build Configuration"
296
+ echo "========================================"
297
+ echo " PYTHON_VERSION: $${PYTHON_VERSION}"
298
+ echo " PYTORCH_ROCM_ARCH: $${PYTORCH_ROCM_ARCH}"
299
+ echo " ROCM_FORCE_REBUILD: $${ROCM_FORCE_REBUILD:-false}"
300
+ echo "========================================"
301
+
302
+ # Save resolved config for later jobs
303
+ buildkite-agent meta-data set "rocm-python-version" "$${PYTHON_VERSION}"
304
+ buildkite-agent meta-data set "rocm-pytorch-rocm-arch" "$${PYTORCH_ROCM_ARCH}"
305
+
306
+ # Check S3 cache for pre-built wheels
307
+ CACHE_KEY=$$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
308
+ CACHE_PATH=$$(.buildkite/scripts/cache-rocm-base-wheels.sh path)
309
+ echo ""
310
+ echo "Cache key: $${CACHE_KEY}"
311
+ echo "Cache path: $${CACHE_PATH}"
312
+
313
+ # Save cache key for downstream jobs
314
+ buildkite-agent meta-data set "rocm-cache-key" "$${CACHE_KEY}"
315
+
316
+ CACHE_STATUS="miss"
317
+ if [ "$${ROCM_FORCE_REBUILD}" != "true" ]; then
318
+ CACHE_STATUS=$$(.buildkite/scripts/cache-rocm-base-wheels.sh check)
319
+ else
320
+ echo "Force rebuild requested, skipping cache check"
321
+ fi
322
+
323
+ if [ "$${CACHE_STATUS}" = "hit" ]; then
324
+ echo ""
325
+ echo "CACHE HIT! Downloading pre-built wheels..."
326
+ echo ""
327
+ .buildkite/scripts/cache-rocm-base-wheels.sh download
328
+
329
+ # Set the S3 path for the cached Docker image (for Job 2 to download)
330
+ S3_ARTIFACT_PATH="s3://$${S3_BUCKET}/rocm/cache/$${CACHE_KEY}"
331
+ buildkite-agent meta-data set "rocm-docker-image-s3-path" "$${S3_ARTIFACT_PATH}/rocm-base-image.tar.gz"
332
+
333
+ # Mark that we used cache (for Docker image handling)
334
+ buildkite-agent meta-data set "rocm-used-cache" "true"
335
+
336
+ echo ""
337
+ echo "Cache download complete. Skipping Docker build."
338
+ echo "Docker image will be downloaded from: $${S3_ARTIFACT_PATH}/rocm-base-image.tar.gz"
339
+ else
340
+ echo ""
341
+ echo "CACHE MISS. Building from scratch..."
342
+ echo ""
343
+
344
+ # Build full base image (for later vLLM build)
345
+ DOCKER_BUILDKIT=1 docker buildx build \
346
+ --file docker/Dockerfile.rocm_base \
347
+ --tag rocm/vllm-dev:base-$${BUILDKITE_BUILD_NUMBER} \
348
+ --build-arg PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH}" \
349
+ --build-arg PYTHON_VERSION="$${PYTHON_VERSION}" \
350
+ --build-arg USE_SCCACHE=1 \
351
+ --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
352
+ --build-arg SCCACHE_REGION_NAME=us-west-2 \
353
+ --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
354
+ --load \
355
+ .
356
+
357
+ # Build debs_wheel_release stage for wheel extraction
358
+ DOCKER_BUILDKIT=1 docker buildx build \
359
+ --file docker/Dockerfile.rocm_base \
360
+ --tag rocm-base-debs:$${BUILDKITE_BUILD_NUMBER} \
361
+ --target debs_wheel_release \
362
+ --build-arg PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH}" \
363
+ --build-arg PYTHON_VERSION="$${PYTHON_VERSION}" \
364
+ --build-arg USE_SCCACHE=1 \
365
+ --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
366
+ --build-arg SCCACHE_REGION_NAME=us-west-2 \
367
+ --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
368
+ --load \
369
+ .
370
+
371
+ # Extract wheels from Docker image
372
+ mkdir -p artifacts/rocm-base-wheels
373
+ container_id=$$(docker create rocm-base-debs:$${BUILDKITE_BUILD_NUMBER})
374
+ docker cp $${container_id}:/app/debs/. artifacts/rocm-base-wheels/
375
+ docker rm $${container_id}
376
+ echo "Extracted base wheels:"
377
+ ls -lh artifacts/rocm-base-wheels/
378
+
379
+ # Upload wheels to S3 cache for future builds
380
+ echo ""
381
+ echo "Uploading wheels to S3 cache..."
382
+ .buildkite/scripts/cache-rocm-base-wheels.sh upload
383
+
384
+ # Export base Docker image for reuse in vLLM build
385
+ mkdir -p artifacts/rocm-docker-image
386
+ docker save rocm/vllm-dev:base-$${BUILDKITE_BUILD_NUMBER} | gzip > artifacts/rocm-docker-image/rocm-base-image.tar.gz
387
+ echo "Docker image size:"
388
+ ls -lh artifacts/rocm-docker-image/
389
+
390
+ # Upload large Docker image to S3 (also cached by cache key)
391
+ S3_ARTIFACT_PATH="s3://$${S3_BUCKET}/rocm/cache/$${CACHE_KEY}"
392
+ echo "Uploading Docker image to $${S3_ARTIFACT_PATH}/"
393
+ aws s3 cp artifacts/rocm-docker-image/rocm-base-image.tar.gz "$${S3_ARTIFACT_PATH}/rocm-base-image.tar.gz"
394
+
395
+ # Save the S3 path for downstream jobs
396
+ buildkite-agent meta-data set "rocm-docker-image-s3-path" "$${S3_ARTIFACT_PATH}/rocm-base-image.tar.gz"
397
+
398
+ # Mark that we did NOT use cache
399
+ buildkite-agent meta-data set "rocm-used-cache" "false"
400
+
401
+ echo ""
402
+ echo "Build complete. Wheels cached for future builds."
403
+ fi
404
+ artifact_paths:
405
+ - "artifacts/rocm-base-wheels/*.whl"
406
+ env:
407
+ DOCKER_BUILDKIT: "1"
408
+ S3_BUCKET: "vllm-wheels"
409
+
410
+ # ROCm Job 2: Build vLLM ROCm Wheel
411
+ - label: ":python: Build vLLM ROCm Wheel"
412
+ id: build-rocm-vllm-wheel
413
+ depends_on:
414
+ - step: build-rocm-base-wheels
415
+ allow_failure: false
416
+ agents:
417
+ queue: cpu_queue_postmerge
418
+ timeout_in_minutes: 180
419
+ commands:
420
+ # Download artifacts and prepare Docker image
421
+ - |
422
+ set -euo pipefail
423
+
424
+ # Ensure git tags are up-to-date (Buildkite's default fetch doesn't update tags)
425
+ # This fixes version detection when tags are moved/force-pushed
426
+ echo "Fetching latest tags from origin..."
427
+ git fetch --tags --force origin
428
+
429
+ # Log tag information for debugging version detection
430
+ echo "========================================"
431
+ echo "Git Tag Verification"
432
+ echo "========================================"
433
+ echo "Current HEAD: $(git rev-parse HEAD)"
434
+ echo "git describe --tags: $(git describe --tags 2>/dev/null || echo 'No tags found')"
435
+ echo ""
436
+ echo "Recent tags (pointing to commits near HEAD):"
437
+ git tag -l --sort=-creatordate | head -5
438
+ echo "setuptools_scm version detection:"
439
+ pip install -q setuptools_scm 2>/dev/null || true
440
+ python3 -c "import setuptools_scm; print(' Detected version:', setuptools_scm.get_version())" 2>/dev/null || echo " (setuptools_scm not available in this environment)"
441
+ echo "========================================"
442
+
443
+ # Download wheel artifacts from current build
444
+ echo "Downloading wheel artifacts from current build"
445
+ buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
446
+
447
+ # Download Docker image from S3 (too large for Buildkite artifacts)
448
+ DOCKER_IMAGE_S3_PATH="$$(buildkite-agent meta-data get rocm-docker-image-s3-path 2>/dev/null || echo '')"
449
+ if [ -z "$${DOCKER_IMAGE_S3_PATH}" ]; then
450
+ echo "ERROR: rocm-docker-image-s3-path metadata not found"
451
+ echo "This should have been set by the build-rocm-base-wheels job"
452
+ exit 1
453
+ fi
454
+ echo "Downloading Docker image from $${DOCKER_IMAGE_S3_PATH}"
455
+ mkdir -p artifacts/rocm-docker-image
456
+ aws s3 cp "$${DOCKER_IMAGE_S3_PATH}" artifacts/rocm-docker-image/rocm-base-image.tar.gz
457
+
458
+ # Load base Docker image and capture the tag
459
+ echo "Loading base Docker image..."
460
+ LOAD_OUTPUT=$$(gunzip -c artifacts/rocm-docker-image/rocm-base-image.tar.gz | docker load)
461
+ echo "$${LOAD_OUTPUT}"
462
+ # Extract the actual loaded image tag from "Loaded image: <tag>" output
463
+ # This avoids picking up stale images (like rocm/vllm-dev:nightly) already on the agent
464
+ BASE_IMAGE_TAG=$$(echo "$${LOAD_OUTPUT}" | grep "Loaded image:" | sed 's/Loaded image: //')
465
+ if [ -z "$${BASE_IMAGE_TAG}" ]; then
466
+ echo "ERROR: Failed to extract image tag from docker load output"
467
+ echo "Load output was: $${LOAD_OUTPUT}"
468
+ exit 1
469
+ fi
470
+ echo "Loaded base image: $${BASE_IMAGE_TAG}"
471
+
472
+ # Prepare base wheels for Docker build context
473
+ mkdir -p docker/context/base-wheels
474
+ touch docker/context/base-wheels/.keep
475
+ cp artifacts/rocm-base-wheels/*.whl docker/context/base-wheels/
476
+ echo "Base wheels for vLLM build:"
477
+ ls -lh docker/context/base-wheels/
478
+
479
+ # Get GPU architectures from meta-data
480
+ PYTORCH_ROCM_ARCH="$$(buildkite-agent meta-data get rocm-pytorch-rocm-arch 2>/dev/null || echo '')"
481
+ PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH:-gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151}"
482
+
483
+ echo "========================================"
484
+ echo "Building vLLM wheel with:"
485
+ echo " BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
486
+ echo " BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}"
487
+ echo " PYTORCH_ROCM_ARCH: $${PYTORCH_ROCM_ARCH}"
488
+ echo " BASE_IMAGE: $${BASE_IMAGE_TAG}"
489
+ echo "========================================"
490
+
491
+ # Build vLLM wheel using local checkout (REMOTE_VLLM=0)
492
+ DOCKER_BUILDKIT=1 docker build \
493
+ --file docker/Dockerfile.rocm \
494
+ --target export_vllm_wheel_release \
495
+ --output type=local,dest=rocm-dist \
496
+ --build-arg BASE_IMAGE="$${BASE_IMAGE_TAG}" \
497
+ --build-arg ARG_PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH}" \
498
+ --build-arg REMOTE_VLLM=0 \
499
+ --build-arg GIT_REPO_CHECK=1 \
500
+ --build-arg USE_SCCACHE=1 \
501
+ --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
502
+ --build-arg SCCACHE_REGION_NAME=us-west-2 \
503
+ --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
504
+ .
505
+
506
+ echo "Built vLLM wheel:"
507
+ ls -lh rocm-dist/*.whl
508
+
509
+ # Copy wheel to artifacts directory
510
+ mkdir -p artifacts/rocm-vllm-wheel
511
+ cp rocm-dist/*.whl artifacts/rocm-vllm-wheel/
512
+ echo "Final vLLM wheel:"
513
+ ls -lh artifacts/rocm-vllm-wheel/
514
+ artifact_paths:
515
+ - "artifacts/rocm-vllm-wheel/*.whl"
516
+ env:
517
+ DOCKER_BUILDKIT: "1"
518
+ S3_BUCKET: "vllm-wheels"
519
+
520
+ # ROCm Job 3: Upload Wheels to S3
521
+ - label: ":s3: Upload ROCm Wheels to S3"
522
+ id: upload-rocm-wheels
523
+ depends_on:
524
+ - step: build-rocm-vllm-wheel
525
+ allow_failure: false
526
+ agents:
527
+ queue: cpu_queue_postmerge
528
+ timeout_in_minutes: 60
529
+ commands:
530
+ # Download all wheel artifacts and run upload
531
+ - |
532
+ set -euo pipefail
533
+
534
+ # Check if upload is enabled (from env var, meta-data, or release branch)
535
+ ROCM_UPLOAD_WHEELS="$${ROCM_UPLOAD_WHEELS:-}"
536
+ if [ -z "$${ROCM_UPLOAD_WHEELS}" ]; then
537
+ # Try to get from meta-data (input form)
538
+ ROCM_UPLOAD_WHEELS="$$(buildkite-agent meta-data get rocm-upload-wheels 2>/dev/null || echo '')"
539
+ fi
540
+
541
+ echo "========================================"
542
+ echo "Upload check:"
543
+ echo " ROCM_UPLOAD_WHEELS: $${ROCM_UPLOAD_WHEELS}"
544
+ echo " BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}"
545
+ echo "========================================"
546
+
547
+ # Skip upload if not enabled
548
+ if [ "$${ROCM_UPLOAD_WHEELS}" != "true" ]; then
549
+ echo "Skipping S3 upload (ROCM_UPLOAD_WHEELS != true, NIGHTLY != 1, not a release branch)"
550
+ echo "To enable upload, set 'Upload Wheels to S3' to 'Yes' in the build configuration"
551
+ exit 0
552
+ fi
553
+
554
+ echo "Upload enabled, proceeding..."
555
+
556
+ # Download artifacts from current build
557
+ echo "Downloading artifacts from current build"
558
+ buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
559
+ buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
560
+
561
+ # Run upload script
562
+ bash .buildkite/scripts/upload-rocm-wheels.sh
563
+ env:
564
+ DOCKER_BUILDKIT: "1"
565
+ S3_BUCKET: "vllm-wheels"
566
+
567
+ # ROCm Job 4: Annotate ROCm Wheel Release
568
+ - label: ":memo: Annotate ROCm wheel release"
569
+ id: annotate-rocm-release
570
+ depends_on:
571
+ - step: upload-rocm-wheels
572
+ allow_failure: true
573
+ agents:
574
+ queue: cpu_queue_postmerge
575
+ commands:
576
+ - "bash .buildkite/scripts/annotate-rocm-release.sh"
577
+ env:
578
+ S3_BUCKET: "vllm-wheels"
.buildkite/scripts/annotate-rocm-release.sh
new
+74
@@ -0,0 +1,74 @@
1
+#!/bin/bash
2
+# SPDX-License-Identifier: Apache-2.0
3
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+#
5
+# Generate Buildkite annotation for ROCm wheel release
6
+
7
+set -ex
8
+
9
+# Get build configuration from meta-data
10
+# Extract ROCm version dynamically from Dockerfile.rocm_base
11
+# BASE_IMAGE format: rocm/dev-ubuntu-22.04:7.1-complete -> extracts "7.1"
12
+ROCM_VERSION=$(grep -E '^ARG BASE_IMAGE=' docker/Dockerfile.rocm_base | sed -E 's/.*:([0-9]+\.[0-9]+).*/\1/' || echo "unknown")
13
+PYTHON_VERSION=$(buildkite-agent meta-data get rocm-python-version 2>/dev/null || echo "3.12")
14
+PYTORCH_ROCM_ARCH=$(buildkite-agent meta-data get rocm-pytorch-rocm-arch 2>/dev/null || echo "gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151")
15
+
16
+# S3 URLs
17
+S3_BUCKET="${S3_BUCKET:-vllm-wheels}"
18
+S3_REGION="${AWS_DEFAULT_REGION:-us-west-2}"
19
+S3_URL="https://${S3_BUCKET}.s3.${S3_REGION}.amazonaws.com"
20
+ROCM_PATH="rocm/${BUILDKITE_COMMIT}"
21
+
22
+buildkite-agent annotate --style 'success' --context 'rocm-release-workflow' << EOF
23
+## :rocm: ROCm Wheel Release
24
+
25
+### Build Configuration
26
+| Setting | Value |
27
+|---------|-------|
28
+| **ROCm Version** | ${ROCM_VERSION} |
29
+| **Python Version** | ${PYTHON_VERSION} |
30
+| **GPU Architectures** | ${PYTORCH_ROCM_ARCH} |
31
+| **Branch** | \`${BUILDKITE_BRANCH}\` |
32
+| **Commit** | \`${BUILDKITE_COMMIT}\` |
33
+
34
+### :package: Installation
35
+
36
+**Install from this build (by commit):**
37
+\`\`\`bash
38
+uv pip install vllm --extra-index-url ${S3_URL}/${ROCM_PATH}/{rocm_variant}/
39
+
40
+# Example:
41
+uv pip install vllm --extra-index-url ${S3_URL}/${ROCM_PATH}/rocm700/
42
+\`\`\`
43
+
44
+**Install from nightly (if published):**
45
+\`\`\`bash
46
+uv pip install vllm --extra-index-url ${S3_URL}/rocm/nightly/
47
+\`\`\`
48
+
49
+### :floppy_disk: Download Wheels Directly
50
+
51
+\`\`\`bash
52
+# List all ROCm wheels
53
+aws s3 ls s3://${S3_BUCKET}/${ROCM_PATH}/
54
+
55
+# Download specific wheels
56
+aws s3 cp s3://${S3_BUCKET}/${ROCM_PATH}/vllm-*.whl .
57
+aws s3 cp s3://${S3_BUCKET}/${ROCM_PATH}/torch-*.whl .
58
+aws s3 cp s3://${S3_BUCKET}/${ROCM_PATH}/triton_rocm-*.whl .
59
+aws s3 cp s3://${S3_BUCKET}/${ROCM_PATH}/torchvision-*.whl .
60
+aws s3 cp s3://${S3_BUCKET}/${ROCM_PATH}/amdsmi-*.whl .
61
+\`\`\`
62
+
63
+### :gear: Included Packages
64
+- **vllm**: vLLM with ROCm support
65
+- **torch**: PyTorch built for ROCm ${ROCM_VERSION}
66
+- **triton_rocm**: Triton built for ROCm
67
+- **torchvision**: TorchVision for ROCm PyTorch
68
+- **amdsmi**: AMD SMI Python bindings
69
+
70
+### :warning: Notes
71
+- These wheels are built for **ROCm ${ROCM_VERSION}** and will NOT work with CUDA GPUs
72
+- Supported GPU architectures: ${PYTORCH_ROCM_ARCH}
73
+- Platform: Linux x86_64 only
74
+EOF
.buildkite/scripts/cache-rocm-base-wheels.sh
new
+140
@@ -0,0 +1,140 @@
1
+#!/usr/bin/env bash
2
+# SPDX-License-Identifier: Apache-2.0
3
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+#
5
+# Cache helper for ROCm base wheels
6
+#
7
+# This script manages caching of pre-built ROCm base wheels (torch, triton, etc.)
8
+# to avoid rebuilding them when Dockerfile.rocm_base hasn't changed.
9
+#
10
+# Usage:
11
+# cache-rocm-base-wheels.sh check - Check if cache exists, outputs "hit" or "miss"
12
+# cache-rocm-base-wheels.sh upload - Upload wheels to cache
13
+# cache-rocm-base-wheels.sh download - Download wheels from cache
14
+# cache-rocm-base-wheels.sh key - Output the cache key
15
+#
16
+# Environment variables:
17
+# S3_BUCKET - S3 bucket name (default: vllm-wheels)
18
+# PYTHON_VERSION - Python version (affects cache key)
19
+# PYTORCH_ROCM_ARCH - GPU architectures (affects cache key)
20
+#
21
+# Note: ROCm version is determined by BASE_IMAGE in Dockerfile.rocm_base,
22
+# so changes to ROCm version are captured by the Dockerfile hash.
23
+
24
+set -euo pipefail
25
+
26
+BUCKET="${S3_BUCKET:-vllm-wheels}"
27
+DOCKERFILE="docker/Dockerfile.rocm_base"
28
+CACHE_PREFIX="rocm/cache"
29
+
30
+# Generate hash from Dockerfile content + build args
31
+generate_cache_key() {
32
+ # Include Dockerfile content
33
+ if [[ ! -f "$DOCKERFILE" ]]; then
34
+ echo "ERROR: Dockerfile not found: $DOCKERFILE" >&2
35
+ exit 1
36
+ fi
37
+ local dockerfile_hash=$(sha256sum "$DOCKERFILE" | cut -c1-16)
38
+
39
+ # Include key build args that affect the output
40
+ # These should match the ARGs in Dockerfile.rocm_base that change the build output
41
+ # Note: ROCm version is determined by BASE_IMAGE in the Dockerfile, so it's captured by dockerfile_hash
42
+ local args_string="${PYTHON_VERSION:-}|${PYTORCH_ROCM_ARCH:-}"
43
+ local args_hash=$(echo "$args_string" | sha256sum | cut -c1-8)
44
+
45
+ echo "${dockerfile_hash}-${args_hash}"
46
+}
47
+
48
+CACHE_KEY=$(generate_cache_key)
49
+CACHE_PATH="s3://${BUCKET}/${CACHE_PREFIX}/${CACHE_KEY}/"
50
+
51
+case "${1:-}" in
52
+ check)
53
+ echo "Checking cache for key: ${CACHE_KEY}" >&2
54
+ echo "Cache path: ${CACHE_PATH}" >&2
55
+ echo "Variables used in cache key:" >&2
56
+ echo " PYTHON_VERSION: ${PYTHON_VERSION:-<not set>}" >&2
57
+ echo " PYTORCH_ROCM_ARCH: ${PYTORCH_ROCM_ARCH:-<not set>}" >&2
58
+
59
+ # Check if cache exists by listing objects
60
+ # We look for at least one .whl file
61
+ echo "Running: aws s3 ls ${CACHE_PATH}" >&2
62
+ S3_OUTPUT=$(aws s3 ls "${CACHE_PATH}" 2>&1) || true
63
+ echo "S3 ls output:" >&2
64
+ echo "$S3_OUTPUT" | head -5 >&2
65
+
66
+ if echo "$S3_OUTPUT" | grep -q "\.whl"; then
67
+ echo "hit"
68
+ else
69
+ echo "miss"
70
+ fi
71
+ ;;
72
+
73
+ upload)
74
+ echo "========================================"
75
+ echo "Uploading wheels to cache"
76
+ echo "========================================"
77
+ echo "Cache key: ${CACHE_KEY}"
78
+ echo "Cache path: ${CACHE_PATH}"
79
+ echo ""
80
+
81
+ if [[ ! -d "artifacts/rocm-base-wheels" ]]; then
82
+ echo "ERROR: artifacts/rocm-base-wheels directory not found" >&2
83
+ exit 1
84
+ fi
85
+
86
+ WHEEL_COUNT=$(ls artifacts/rocm-base-wheels/*.whl 2>/dev/null | wc -l)
87
+ if [[ "$WHEEL_COUNT" -eq 0 ]]; then
88
+ echo "ERROR: No wheels found in artifacts/rocm-base-wheels/" >&2
89
+ exit 1
90
+ fi
91
+
92
+ echo "Uploading $WHEEL_COUNT wheels..."
93
+ aws s3 cp --recursive artifacts/rocm-base-wheels/ "${CACHE_PATH}"
94
+
95
+ echo ""
96
+ echo "Cache upload complete!"
97
+ echo "========================================"
98
+ ;;
99
+
100
+ download)
101
+ echo "========================================"
102
+ echo "Downloading wheels from cache"
103
+ echo "========================================"
104
+ echo "Cache key: ${CACHE_KEY}"
105
+ echo "Cache path: ${CACHE_PATH}"
106
+ echo ""
107
+
108
+ mkdir -p artifacts/rocm-base-wheels
109
+ aws s3 cp --recursive "${CACHE_PATH}" artifacts/rocm-base-wheels/
110
+
111
+ echo ""
112
+ echo "Downloaded wheels:"
113
+ ls -lh artifacts/rocm-base-wheels/
114
+
115
+ WHEEL_COUNT=$(ls artifacts/rocm-base-wheels/*.whl 2>/dev/null | wc -l)
116
+ echo ""
117
+ echo "Total: $WHEEL_COUNT wheels"
118
+ echo "========================================"
119
+ ;;
120
+
121
+ key)
122
+ echo "${CACHE_KEY}"
123
+ ;;
124
+
125
+ path)
126
+ echo "${CACHE_PATH}"
127
+ ;;
128
+
129
+ *)
130
+ echo "Usage: $0 {check|upload|download|key|path}" >&2
131
+ echo "" >&2
132
+ echo "Commands:" >&2
133
+ echo " check - Check if cache exists, outputs 'hit' or 'miss'" >&2
134
+ echo " upload - Upload wheels from artifacts/rocm-base-wheels/ to cache" >&2
135
+ echo " download - Download wheels from cache to artifacts/rocm-base-wheels/" >&2
136
+ echo " key - Output the cache key" >&2
137
+ echo " path - Output the full S3 cache path" >&2
138
+ exit 1
139
+ ;;
140
+esac
.buildkite/scripts/generate-nightly-index.py
+69
-9
@@ -16,6 +16,18 @@ from urllib.parse import quote
16
17
import regex as re
18
19
+
20
+def normalize_package_name(name: str) -> str:
21
+ """
22
+ Normalize package name according to PEP 503.
23
+ https://peps.python.org/pep-0503/#normalized-names
24
+
25
+ Replace runs of underscores, hyphens, and periods with a single hyphen,
26
+ and lowercase the result.
27
+ """
28
+ return re.sub(r"[-_.]+", "-", name).lower()
29
+
30
+
31
if not sys.version_info >= (3, 12):
32
raise RuntimeError("This script requires Python 3.12 or higher.")
33
@@ -78,7 +90,13 @@ def parse_from_filename(file: str) -> WheelFileInfo:
90
version = version.removesuffix("." + variant)
91
else:
92
if "+" in version:
81
- version, variant = version.split("+")
93
+ version_part, suffix = version.split("+", 1)
94
+ # Only treat known patterns as variants (rocmXXX, cuXXX, cpu)
95
+ # Git hashes and other suffixes are NOT variants
96
+ if suffix.startswith(("rocm", "cu", "cpu")):
97
+ variant = suffix
98
+ version = version_part
99
+ # Otherwise keep the full version string (variant stays None)
100
101
return WheelFileInfo(
102
package_name=package_name,
@@ -206,6 +224,26 @@ def generate_index_and_metadata(
224
print("No wheel files found, skipping index generation.")
225
return
226
227
+ # For ROCm builds: inherit variant from vllm wheel
228
+ # All ROCm wheels should share the same variant as vllm
229
+ rocm_variant = None
230
+ for file in parsed_files:
231
+ if (
232
+ file.package_name == "vllm"
233
+ and file.variant
234
+ and file.variant.startswith("rocm")
235
+ ):
236
+ rocm_variant = file.variant
237
+ print(f"Detected ROCm variant from vllm: {rocm_variant}")
238
+ break
239
+
240
+ # Apply ROCm variant to all wheels without a variant
241
+ if rocm_variant:
242
+ for file in parsed_files:
243
+ if file.variant is None:
244
+ file.variant = rocm_variant
245
+ print(f"Inherited variant '{rocm_variant}' for {file.filename}")
246
+
247
# Group by variant
248
variant_to_files: dict[str, list[WheelFileInfo]] = {}
249
for file in parsed_files:
@@ -256,8 +294,8 @@ def generate_index_and_metadata(
294
295
variant_dir.mkdir(parents=True, exist_ok=True)
296
259
- # gather all package names in this variant
260
- packages = set(f.package_name for f in files)
297
+ # gather all package names in this variant (normalized per PEP 503)
298
+ packages = set(normalize_package_name(f.package_name) for f in files)
299
if variant == "default":
300
# these packages should also appear in the "project list"
301
# generate after all variants are processed
@@ -269,8 +307,10 @@ def generate_index_and_metadata(
307
f.write(project_list_str)
308
309
for package in packages:
272
- # filter files belonging to this package only
273
- package_files = [f for f in files if f.package_name == package]
310
+ # filter files belonging to this package only (compare normalized names)
311
+ package_files = [
312
+ f for f in files if normalize_package_name(f.package_name) == package
313
+ ]
314
package_dir = variant_dir / package
315
package_dir.mkdir(parents=True, exist_ok=True)
316
index_str, metadata_str = generate_package_index_and_metadata(
@@ -341,8 +381,13 @@ if __name__ == "__main__":
381
args = parser.parse_args()
382
383
version = args.version
344
- if "/" in version or "\\" in version:
345
- raise ValueError("Version string must not contain slashes.")
384
+ # Allow rocm/ prefix, reject other slashes and all backslashes
385
+ if "\\" in version:
386
+ raise ValueError("Version string must not contain backslashes.")
387
+ if "/" in version and not version.startswith("rocm/"):
388
+ raise ValueError(
389
+ "Version string must not contain slashes (except for 'rocm/' prefix)."
390
+ )
391
current_objects_path = Path(args.current_objects)
392
output_dir = Path(args.output_dir)
393
if not output_dir.exists():
@@ -393,8 +438,23 @@ if __name__ == "__main__":
438
# Generate index and metadata, assuming wheels and indices are stored as:
439
# s3://vllm-wheels/{wheel_dir}/<wheel files>
440
# s3://vllm-wheels/<anything>/<index files>
396
- wheel_dir = args.wheel_dir or version
397
- wheel_base_dir = Path(output_dir).parent / wheel_dir.strip().rstrip("/")
441
+ #
442
+ # For ROCm builds, version is "rocm/{commit}" and indices are uploaded to:
443
+ # - rocm/{commit}/ (same as wheels)
444
+ # - rocm/nightly/
445
+ # - rocm/{version}/
446
+ # All these are under the "rocm/" prefix, so relative paths should be
447
+ # relative to "rocm/", not the bucket root.
448
+ if args.wheel_dir:
449
+ # Explicit wheel-dir provided (e.g., for version-specific indices pointing to commit dir)
450
+ wheel_dir = args.wheel_dir.strip().rstrip("/")
451
+ elif version.startswith("rocm/"):
452
+ # For rocm/commit, wheel_base_dir should be just the commit part
453
+ # so relative path from rocm/0.12.0/rocm710/vllm/ -> ../../../{commit}/
454
+ wheel_dir = version.split("/", 1)[1]
455
+ else:
456
+ wheel_dir = version
457
+ wheel_base_dir = Path(output_dir).parent / wheel_dir
458
index_base_dir = Path(output_dir)
459
460
generate_index_and_metadata(
.buildkite/scripts/upload-rocm-wheels.sh
new
+151
@@ -0,0 +1,151 @@
1
+#!/usr/bin/env bash
2
+# SPDX-License-Identifier: Apache-2.0
3
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+#
5
+# Upload ROCm wheels to S3 with proper index generation
6
+#
7
+# Required environment variables:
8
+# AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY (or IAM role)
9
+# S3_BUCKET (default: vllm-wheels)
10
+#
11
+# S3 path structure:
12
+# s3://vllm-wheels/rocm/{commit}/ - All wheels for this commit
13
+# s3://vllm-wheels/rocm/nightly/ - Index pointing to latest nightly
14
+# s3://vllm-wheels/rocm/{version}/ - Index for release versions
15
+
16
+set -ex
17
+
18
+# ======== Configuration ========
19
+BUCKET="${S3_BUCKET:-vllm-wheels}"
20
+ROCM_SUBPATH="rocm/${BUILDKITE_COMMIT}"
21
+S3_COMMIT_PREFIX="s3://$BUCKET/$ROCM_SUBPATH/"
22
+INDICES_OUTPUT_DIR="rocm-indices"
23
+PYTHON="${PYTHON_PROG:-python3}"
24
+
25
+# ROCm uses manylinux_2_35 (Ubuntu 22.04 based)
26
+MANYLINUX_VERSION="manylinux_2_35"
27
+
28
+echo "========================================"
29
+echo "ROCm Wheel Upload Configuration"
30
+echo "========================================"
31
+echo "S3 Bucket: $BUCKET"
32
+echo "S3 Path: $ROCM_SUBPATH"
33
+echo "Commit: $BUILDKITE_COMMIT"
34
+echo "Branch: $BUILDKITE_BRANCH"
35
+echo "========================================"
36
+
37
+# ======== Part 0: Setup Python ========
38
+
39
+# Detect if python3.12+ is available
40
+has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) else 0)" 2>/dev/null || echo 0)
41
+if [[ "$has_new_python" -eq 0 ]]; then
42
+ # Use new python from docker
43
+ # Use --user to ensure files are created with correct ownership (not root)
44
+ docker pull python:3-slim
45
+ PYTHON="docker run --rm --user $(id -u):$(id -g) -v $(pwd):/app -w /app python:3-slim python3"
46
+fi
47
+
48
+echo "Using python interpreter: $PYTHON"
49
+echo "Python version: $($PYTHON --version)"
50
+
51
+# ======== Part 1: Collect and prepare wheels ========
52
+
53
+# Collect all wheels
54
+mkdir -p all-rocm-wheels
55
+cp artifacts/rocm-base-wheels/*.whl all-rocm-wheels/ 2>/dev/null || true
56
+cp artifacts/rocm-vllm-wheel/*.whl all-rocm-wheels/ 2>/dev/null || true
57
+
58
+WHEEL_COUNT=$(ls all-rocm-wheels/*.whl 2>/dev/null | wc -l)
59
+echo "Total wheels to upload: $WHEEL_COUNT"
60
+
61
+if [ "$WHEEL_COUNT" -eq 0 ]; then
62
+ echo "ERROR: No wheels found to upload!"
63
+ exit 1
64
+fi
65
+
66
+# Rename linux to manylinux in wheel filenames
67
+for wheel in all-rocm-wheels/*.whl; do
68
+ if [[ "$wheel" == *"linux"* ]] && [[ "$wheel" != *"manylinux"* ]]; then
69
+ new_wheel="${wheel/linux/$MANYLINUX_VERSION}"
70
+ mv -- "$wheel" "$new_wheel"
71
+ echo "Renamed: $(basename "$wheel") -> $(basename "$new_wheel")"
72
+ fi
73
+done
74
+
75
+echo ""
76
+echo "Wheels to upload:"
77
+ls -lh all-rocm-wheels/
78
+
79
+# ======== Part 2: Upload wheels to S3 ========
80
+
81
+echo ""
82
+echo "Uploading wheels to $S3_COMMIT_PREFIX"
83
+for wheel in all-rocm-wheels/*.whl; do
84
+ aws s3 cp "$wheel" "$S3_COMMIT_PREFIX"
85
+done
86
+
87
+# ======== Part 3: Generate and upload indices ========
88
+
89
+# List existing wheels in commit directory
90
+echo ""
91
+echo "Generating indices..."
92
+obj_json="rocm-objects.json"
93
+aws s3api list-objects-v2 --bucket "$BUCKET" --prefix "$ROCM_SUBPATH/" --delimiter / --output json > "$obj_json"
94
+
95
+mkdir -p "$INDICES_OUTPUT_DIR"
96
+
97
+# Use the existing generate-nightly-index.py
98
+# HACK: Replace regex module with stdlib re (same as CUDA script)
99
+sed -i 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py
100
+
101
+$PYTHON .buildkite/scripts/generate-nightly-index.py \
102
+ --version "$ROCM_SUBPATH" \
103
+ --current-objects "$obj_json" \
104
+ --output-dir "$INDICES_OUTPUT_DIR" \
105
+ --comment "ROCm commit $BUILDKITE_COMMIT"
106
+
107
+# Upload indices to commit directory
108
+echo "Uploading indices to $S3_COMMIT_PREFIX"
109
+aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"
110
+
111
+# Update rocm/nightly/ if on main branch and not a PR
112
+if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]] || [[ "$NIGHTLY" == "1" ]]; then
113
+ echo "Updating rocm/nightly/ index..."
114
+ aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/rocm/nightly/"
115
+fi
116
+
117
+# Extract version from vLLM wheel and update version-specific index
118
+VLLM_WHEEL=$(ls all-rocm-wheels/vllm*.whl 2>/dev/null | head -1)
119
+if [ -n "$VLLM_WHEEL" ]; then
120
+ VERSION=$(unzip -p "$VLLM_WHEEL" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
121
+ echo "Version in wheel: $VERSION"
122
+ PURE_VERSION="${VERSION%%+*}"
123
+ PURE_VERSION="${PURE_VERSION%%.rocm}"
124
+ echo "Pure version: $PURE_VERSION"
125
+
126
+ if [[ "$VERSION" != *"dev"* ]]; then
127
+ echo "Updating rocm/$PURE_VERSION/ index..."
128
+ aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/rocm/$PURE_VERSION/"
129
+ fi
130
+fi
131
+
132
+# ======== Part 4: Summary ========
133
+
134
+echo ""
135
+echo "========================================"
136
+echo "ROCm Wheel Upload Complete!"
137
+echo "========================================"
138
+echo ""
139
+echo "Wheels available at:"
140
+echo " s3://$BUCKET/$ROCM_SUBPATH/"
141
+echo ""
142
+echo "Install command (by commit):"
143
+echo " pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/$ROCM_SUBPATH/"
144
+echo ""
145
+if [[ "$BUILDKITE_BRANCH" == "main" ]] || [[ "$NIGHTLY" == "1" ]]; then
146
+ echo "Install command (nightly):"
147
+ echo " pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/rocm/nightly/"
148
+fi
149
+echo ""
150
+echo "Wheel count: $WHEEL_COUNT"
151
+echo "========================================"
docker/Dockerfile.rocm
+150
-4
@@ -3,6 +3,14 @@ ARG REMOTE_VLLM="0"
3
ARG COMMON_WORKDIR=/app
4
ARG BASE_IMAGE=rocm/vllm-dev:base
5
6
+# Sccache configuration (only used in release pipeline)
7
+ARG USE_SCCACHE
8
+ARG SCCACHE_DOWNLOAD_URL
9
+ARG SCCACHE_ENDPOINT
10
+ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
11
+ARG SCCACHE_REGION_NAME=us-west-2
12
+ARG SCCACHE_S3_NO_CREDENTIALS=0
13
+
14
FROM ${BASE_IMAGE} AS base
15
16
ARG ARG_PYTORCH_ROCM_ARCH
@@ -14,9 +22,14 @@ ENV RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=1
22
RUN apt-get update -q -y && apt-get install -q -y \
23
sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \
24
apt-transport-https ca-certificates wget curl
17
-# Remove sccache
25
RUN python3 -m pip install --upgrade pip
19
-RUN apt-get purge -y sccache; python3 -m pip uninstall -y sccache; rm -f "$(which sccache)"
26
+# Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_base)
27
+ARG USE_SCCACHE
28
+RUN if [ "$USE_SCCACHE" != "1" ]; then \
29
+ apt-get purge -y sccache || true; \
30
+ python3 -m pip uninstall -y sccache || true; \
31
+ rm -f "$(which sccache)" || true; \
32
+ fi
33
34
# Install UV
35
RUN curl -LsSf https://astral.sh/uv/install.sh | env UV_INSTALL_DIR="/usr/local/bin" sh
@@ -28,6 +41,39 @@ ENV UV_INDEX_STRATEGY="unsafe-best-match"
41
# Use copy mode to avoid hardlink failures with Docker cache mounts
42
ENV UV_LINK_MODE=copy
43
44
+# Install sccache if USE_SCCACHE is enabled (for release builds)
45
+ARG USE_SCCACHE
46
+ARG SCCACHE_DOWNLOAD_URL
47
+ARG SCCACHE_ENDPOINT
48
+ARG SCCACHE_BUCKET_NAME
49
+ARG SCCACHE_REGION_NAME
50
+ARG SCCACHE_S3_NO_CREDENTIALS
51
+RUN if [ "$USE_SCCACHE" = "1" ]; then \
52
+ if command -v sccache >/dev/null 2>&1; then \
53
+ echo "sccache already installed, skipping installation"; \
54
+ sccache --version; \
55
+ else \
56
+ echo "Installing sccache..." \
57
+ && SCCACHE_ARCH="x86_64" \
58
+ && SCCACHE_VERSION="v0.8.1" \
59
+ && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \
60
+ && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \
61
+ && tar -xzf /tmp/sccache.tar.gz -C /tmp \
62
+ && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \
63
+ && chmod +x /usr/bin/sccache \
64
+ && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \
65
+ && sccache --version; \
66
+ fi; \
67
+ fi
68
+
69
+# Set sccache environment variables only when USE_SCCACHE=1
70
+# This prevents S3 config from leaking into images when sccache is not used
71
+ARG USE_SCCACHE
72
+ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}}
73
+ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
74
+ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
75
+ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
76
+
77
ARG COMMON_WORKDIR
78
WORKDIR ${COMMON_WORKDIR}
79
@@ -53,7 +99,7 @@ FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm
99
# -----------------------
100
# vLLM build stages
101
FROM fetch_vllm AS build_vllm
56
-# Build vLLM
102
+# Build vLLM (setup.py auto-detects sccache in PATH)
103
RUN cd vllm \
104
&& python3 -m pip install -r requirements/rocm.txt \
105
&& python3 setup.py clean --all \
@@ -69,7 +115,6 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
115
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
116
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
117
72
-# -----------------------
118
# RIXL/UCX build stages
119
FROM base AS build_rixl
120
ARG RIXL_BRANCH="f33a5599"
@@ -141,6 +186,107 @@ RUN cd /opt/rixl && mkdir -p /app/install && \
186
--ucx-plugins-dir ${UCX_HOME}/lib/ucx \
187
--nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
188
189
+
190
+# -----------------------
191
+# vLLM wheel release build stage (for building distributable wheels)
192
+# This stage pins dependencies to custom ROCm wheel versions and handles version detection
193
+FROM fetch_vllm AS build_vllm_wheel_release
194
+
195
+ARG COMMON_WORKDIR
196
+
197
+# Create /install directory for custom wheels
198
+RUN mkdir -p /install
199
+
200
+# Copy custom ROCm wheels from docker/context if they exist
201
+# COPY ensures Docker cache is invalidated when wheels change
202
+# .keep file ensures directory always exists for COPY to work
203
+COPY docker/context/base-wheels/ /tmp/base-wheels/
204
+# This is how we know if we are building for a wheel release or not.
205
+# If there are not wheels found there, we are not building for a wheel release.
206
+# So we exit with an error. To skip this stage.
207
+RUN if [ -n "$(ls /tmp/base-wheels/*.whl 2>/dev/null)" ]; then \
208
+ echo "Found custom wheels - copying to /install"; \
209
+ cp /tmp/base-wheels/*.whl /install/ && \
210
+ echo "Copied custom wheels:"; \
211
+ ls -lh /install/; \
212
+ else \
213
+ echo "ERROR: No custom wheels found in docker/context/base-wheels/"; \
214
+ echo "Wheel releases require pre-built ROCm wheels."; \
215
+ exit 1; \
216
+ fi
217
+
218
+# GIT_REPO_CHECK: Verify repo is clean and tags are available (for release builds)
219
+# This matches CUDA's Dockerfile behavior for proper version detection via setuptools_scm
220
+ARG GIT_REPO_CHECK=0
221
+RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \
222
+ echo "Running repository checks..."; \
223
+ cd vllm && bash tools/check_repo.sh; \
224
+ fi
225
+
226
+# Extract version from git BEFORE any modifications (pin_rocm_dependencies.py modifies requirements/rocm.txt)
227
+# This ensures setuptools_scm sees clean repo state for version detection
228
+RUN --mount=type=bind,source=.git,target=vllm/.git \
229
+ cd vllm \
230
+ && pip install setuptools_scm \
231
+ && VLLM_VERSION=$(python3 -c "import setuptools_scm; print(setuptools_scm.get_version())") \
232
+ && echo "Detected vLLM version: ${VLLM_VERSION}" \
233
+ && echo "${VLLM_VERSION}" > /tmp/vllm_version.txt
234
+
235
+# Fail if git-based package dependencies are found in requirements files
236
+# (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI)
237
+# Extra notes: pip install is able to handle git+ URLs, but uv doesn't.
238
+RUN echo "Checking for git-based packages in requirements files..." \
239
+ && echo "Checking common.txt for git-based packages:" \
240
+ && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \
241
+ echo "ERROR: Git-based packages found in common.txt:"; \
242
+ grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \
243
+ echo "Please publish these packages to PyPI instead of using git dependencies."; \
244
+ exit 1; \
245
+ else \
246
+ echo " ✓ No git-based packages found in common.txt"; \
247
+ fi \
248
+ && echo "Checking rocm.txt for git-based packages:" \
249
+ && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \
250
+ echo "ERROR: Git-based packages found in rocm.txt:"; \
251
+ grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \
252
+ echo "Please publish these packages to PyPI instead of using git dependencies."; \
253
+ exit 1; \
254
+ else \
255
+ echo " ✓ No git-based packages found in rocm.txt"; \
256
+ fi \
257
+ && echo "All requirements files are clean - no git-based packages found"
258
+
259
+# Pin vLLM dependencies to exact versions of custom ROCm wheels
260
+# This ensures 'pip install vllm' automatically installs correct torch/triton/torchvision/amdsmi
261
+COPY tools/vllm-rocm/pin_rocm_dependencies.py /tmp/pin_rocm_dependencies.py
262
+RUN echo "Pinning vLLM dependencies to custom wheel versions..." \
263
+ && python3 /tmp/pin_rocm_dependencies.py /install ${COMMON_WORKDIR}/vllm/requirements/rocm.txt
264
+
265
+# Install dependencies using custom wheels from /install
266
+RUN cd vllm \
267
+ && echo "Building vLLM with custom wheels from /install" \
268
+ && python3 -m pip install --find-links /install -r requirements/rocm.txt \
269
+ && python3 setup.py clean --all
270
+
271
+# Build wheel using pre-extracted version to avoid dirty state from modified requirements/rocm.txt
272
+# (setup.py auto-detects sccache in PATH)
273
+RUN --mount=type=bind,source=.git,target=vllm/.git \
274
+ cd vllm \
275
+ && export SETUPTOOLS_SCM_PRETEND_VERSION=$(cat /tmp/vllm_version.txt) \
276
+ && echo "Building wheel with version: ${SETUPTOOLS_SCM_PRETEND_VERSION}" \
277
+ && python3 setup.py bdist_wheel --dist-dir=dist
278
+
279
+FROM scratch AS export_vllm_wheel_release
280
+ARG COMMON_WORKDIR
281
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/dist/*.whl /
282
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/requirements /requirements
283
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks
284
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests
285
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples
286
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
287
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
288
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
289
+
290
# -----------------------
291
# Test vLLM image
292
FROM base AS test
docker/Dockerfile.rocm_base
+104
-2
@@ -14,6 +14,14 @@ ARG AITER_REPO="https://github.com/ROCm/aiter.git"
14
ARG MORI_BRANCH="2d02c6a9"
15
ARG MORI_REPO="https://github.com/ROCm/mori.git"
16
17
+# Sccache configuration (only used in release pipeline)
18
+ARG USE_SCCACHE
19
+ARG SCCACHE_DOWNLOAD_URL
20
+ARG SCCACHE_ENDPOINT
21
+ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
22
+ARG SCCACHE_REGION_NAME=us-west-2
23
+ARG SCCACHE_S3_NO_CREDENTIALS=0
24
+
25
FROM ${BASE_IMAGE} AS base
26
27
ENV PATH=/opt/rocm/llvm/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
@@ -53,6 +61,49 @@ RUN apt-get update -y \
61
RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 Cython
62
RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/*
63
64
+# Install sccache if USE_SCCACHE is enabled (for release builds)
65
+ARG USE_SCCACHE
66
+ARG SCCACHE_DOWNLOAD_URL
67
+ARG SCCACHE_ENDPOINT
68
+ARG SCCACHE_BUCKET_NAME
69
+ARG SCCACHE_REGION_NAME
70
+ARG SCCACHE_S3_NO_CREDENTIALS
71
+RUN if [ "$USE_SCCACHE" = "1" ]; then \
72
+ echo "Installing sccache..." \
73
+ && SCCACHE_ARCH="x86_64" \
74
+ && SCCACHE_VERSION="v0.8.1" \
75
+ && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \
76
+ && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \
77
+ && tar -xzf /tmp/sccache.tar.gz -C /tmp \
78
+ && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \
79
+ && chmod +x /usr/bin/sccache \
80
+ && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \
81
+ && sccache --version; \
82
+ fi
83
+
84
+# Setup sccache for HIP compilation via HIP_CLANG_PATH
85
+# This creates wrapper scripts in a separate directory and points HIP to use them
86
+# This avoids modifying the original ROCm binaries which can break detection
87
+# NOTE: HIP_CLANG_PATH is NOT set as ENV to avoid affecting downstream images (Dockerfile.rocm)
88
+# Instead, each build stage should export HIP_CLANG_PATH=/opt/sccache-wrappers if USE_SCCACHE=1
89
+RUN if [ "$USE_SCCACHE" = "1" ]; then \
90
+ echo "Setting up sccache wrappers for HIP compilation..." \
91
+ && mkdir -p /opt/sccache-wrappers \
92
+ && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \
93
+ && chmod +x /opt/sccache-wrappers/clang++ \
94
+ && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \
95
+ && chmod +x /opt/sccache-wrappers/clang \
96
+ && echo "sccache wrappers created in /opt/sccache-wrappers"; \
97
+ fi
98
+
99
+# Set sccache environment variables only when USE_SCCACHE=1
100
+# This prevents S3 config from leaking into images when sccache is not used
101
+ARG USE_SCCACHE
102
+ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}}
103
+ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
104
+ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
105
+ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
106
+
107
108
###
109
### Triton Build
@@ -89,22 +140,42 @@ ARG PYTORCH_AUDIO_BRANCH
140
ARG PYTORCH_REPO
141
ARG PYTORCH_VISION_REPO
142
ARG PYTORCH_AUDIO_REPO
143
+ARG USE_SCCACHE
144
145
RUN git clone ${PYTORCH_REPO} pytorch
146
RUN cd pytorch && git checkout ${PYTORCH_BRANCH} \
147
&& pip install -r requirements.txt && git submodule update --init --recursive \
148
&& python3 tools/amd_build/build_amd.py \
149
+ && if [ "$USE_SCCACHE" = "1" ]; then \
150
+ export HIP_CLANG_PATH=/opt/sccache-wrappers \
151
+ && export CMAKE_C_COMPILER_LAUNCHER=sccache \
152
+ && export CMAKE_CXX_COMPILER_LAUNCHER=sccache \
153
+ && sccache --show-stats; \
154
+ fi \
155
&& CMAKE_PREFIX_PATH=$(python3 -c 'import sys; print(sys.prefix)') python3 setup.py bdist_wheel --dist-dir=dist \
156
+ && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \
157
&& pip install dist/*.whl
158
RUN git clone ${PYTORCH_VISION_REPO} vision
159
RUN cd vision && git checkout ${PYTORCH_VISION_BRANCH} \
160
+ && if [ "$USE_SCCACHE" = "1" ]; then \
161
+ export HIP_CLANG_PATH=/opt/sccache-wrappers \
162
+ && export CMAKE_C_COMPILER_LAUNCHER=sccache \
163
+ && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \
164
+ fi \
165
&& python3 setup.py bdist_wheel --dist-dir=dist \
166
+ && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \
167
&& pip install dist/*.whl
168
RUN git clone ${PYTORCH_AUDIO_REPO} audio
169
RUN cd audio && git checkout ${PYTORCH_AUDIO_BRANCH} \
170
&& git submodule update --init --recursive \
171
&& pip install -r requirements.txt \
172
+ && if [ "$USE_SCCACHE" = "1" ]; then \
173
+ export HIP_CLANG_PATH=/opt/sccache-wrappers \
174
+ && export CMAKE_C_COMPILER_LAUNCHER=sccache \
175
+ && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \
176
+ fi \
177
&& python3 setup.py bdist_wheel --dist-dir=dist \
178
+ && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \
179
&& pip install dist/*.whl
180
RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install \
181
&& cp /app/vision/dist/*.whl /app/install \
@@ -133,13 +204,19 @@ RUN mkdir -p /app/install && cp /app/mori/dist/*.whl /app/install
204
FROM base AS build_fa
205
ARG FA_BRANCH
206
ARG FA_REPO
207
+ARG USE_SCCACHE
208
RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
209
pip install /install/*.whl
210
RUN git clone ${FA_REPO}
211
RUN cd flash-attention \
212
&& git checkout ${FA_BRANCH} \
213
&& git submodule update --init \
142
- && GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist
214
+ && if [ "$USE_SCCACHE" = "1" ]; then \
215
+ export HIP_CLANG_PATH=/opt/sccache-wrappers \
216
+ && sccache --show-stats; \
217
+ fi \
218
+ && GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \
219
+ && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi
220
RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install
221
222
@@ -149,6 +226,7 @@ RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install
226
FROM base AS build_aiter
227
ARG AITER_BRANCH
228
ARG AITER_REPO
229
+ARG USE_SCCACHE
230
RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
231
pip install /install/*.whl
232
RUN git clone --recursive ${AITER_REPO}
@@ -156,13 +234,37 @@ RUN cd aiter \
234
&& git checkout ${AITER_BRANCH} \
235
&& git submodule update --init --recursive \
236
&& pip install -r requirements.txt
159
-RUN pip install pyyaml && cd aiter && PREBUILD_KERNELS=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist && ls /app/aiter/dist/*.whl
237
+RUN pip install pyyaml && cd aiter \
238
+ && if [ "$USE_SCCACHE" = "1" ]; then \
239
+ export HIP_CLANG_PATH=/opt/sccache-wrappers \
240
+ && sccache --show-stats; \
241
+ fi \
242
+ && PREBUILD_KERNELS=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
243
+ && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \
244
+ && ls /app/aiter/dist/*.whl
245
RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install
246
247
248
###
249
### Final Build
250
###
251
+
252
+# Wheel release stage -
253
+# only includes dependencies used by wheel release pipeline
254
+FROM base AS debs_wheel_release
255
+RUN mkdir /app/debs
256
+RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
257
+ cp /install/*.whl /app/debs
258
+RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \
259
+ cp /install/*.whl /app/debs
260
+RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
261
+ cp /install/*.whl /app/debs
262
+RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
263
+ cp /install/*.whl /app/debs
264
+RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
265
+ cp /install/*.whl /app/debs
266
+
267
+# Full debs stage - includes Mori (used by Docker releases)
268
FROM base AS debs
269
RUN mkdir /app/debs
270
RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
requirements/rocm-test.txt
+2
@@ -80,6 +80,8 @@ num2words==0.5.14
80
pqdm==0.2.0
81
# via lm-eval
82
83
+# Required for fastsafetensors test
84
+fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@d6f998a03432b2452f8de2bb5cefb5af9795d459
85
# Required for suffix decoding test
86
arctic-inference == 0.1.1
87
# Required for Nemotron test
requirements/rocm.txt
-1
@@ -15,5 +15,4 @@ setuptools-scm>=8
15
runai-model-streamer[s3,gcs]==0.15.3
16
conch-triton-kernels==1.2.1
17
timm>=1.0.17
18
-fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@d6f998a03432b2452f8de2bb5cefb5af9795d459
18
grpcio-tools>=1.76.0
\ No newline at end of file
tools/vllm-rocm/pin_rocm_dependencies.py
new
+221
@@ -0,0 +1,221 @@
1
+#!/usr/bin/env python3
2
+# SPDX-License-Identifier: Apache-2.0
3
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+"""
5
+Pin vLLM dependencies to exact versions of custom ROCm wheels.
6
+
7
+This script modifies vLLM's requirements files to replace version constraints
8
+with exact versions of custom-built ROCm wheels (torch, triton, torchvision, amdsmi).
9
+
10
+This ensures that 'pip install vllm' automatically installs the correct custom wheels
11
+instead of allowing pip to download different versions from PyPI.
12
+"""
13
+
14
+import re
15
+import sys
16
+from pathlib import Path
17
+
18
+
19
+def extract_version_from_wheel(wheel_name: str) -> str:
20
+ """
21
+ Extract version from wheel filename.
22
+
23
+ Example:
24
+ torch-2.9.0a0+git1c57644-cp312-cp312-linux_x86_64.whl -> 2.9.0a0+git1c57644
25
+ triton-3.4.0-cp312-cp312-linux_x86_64.whl -> 3.4.0
26
+ """
27
+ # Wheel format:
28
+ # {distribution}-{version}(-{build tag})?-{python}-{abi}-{platform}.whl
29
+ parts = wheel_name.replace(".whl", "").split("-")
30
+
31
+ if len(parts) < 5:
32
+ raise ValueError(f"Invalid wheel filename format: {wheel_name}")
33
+
34
+ # Version is the second part
35
+ version = parts[1]
36
+ return version
37
+
38
+
39
+def get_custom_wheel_versions(install_dir: str) -> dict[str, str]:
40
+ """
41
+ Read /install directory and extract versions of custom wheels.
42
+
43
+ Returns:
44
+ Dict mapping package names to exact versions
45
+ """
46
+ install_path = Path(install_dir)
47
+ if not install_path.exists():
48
+ print(f"ERROR: Install directory not found: {install_dir}", file=sys.stderr)
49
+ sys.exit(1)
50
+
51
+ versions = {}
52
+
53
+ # Map wheel prefixes to package names
54
+ # IMPORTANT: Use dashes to avoid matching substrings
55
+ # (e.g., 'torch' would match 'torchvision')
56
+ # ORDER MATTERS: This order is preserved when pinning dependencies
57
+ # in requirements files
58
+ package_mapping = [
59
+ ("torch-", "torch"), # Match torch- (not torchvision)
60
+ ("triton-", "triton"), # Match triton- (not triton_kernels)
61
+ ("triton_kernels-", "triton-kernels"), # Match triton_kernels-
62
+ ("torchvision-", "torchvision"), # Match torchvision-
63
+ ("torchaudio-", "torchaudio"), # Match torchaudio-
64
+ ("amdsmi-", "amdsmi"), # Match amdsmi-
65
+ ("flash_attn-", "flash-attn"), # Match flash_attn-
66
+ ("aiter-", "aiter"), # Match aiter-
67
+ ]
68
+
69
+ for wheel_file in install_path.glob("*.whl"):
70
+ wheel_name = wheel_file.name
71
+
72
+ for prefix, package_name in package_mapping:
73
+ if wheel_name.startswith(prefix):
74
+ try:
75
+ version = extract_version_from_wheel(wheel_name)
76
+ versions[package_name] = version
77
+ print(f"Found {package_name}=={version}", file=sys.stderr)
78
+ except Exception as e:
79
+ print(
80
+ f"WARNING: Could not extract version from {wheel_name}: {e}",
81
+ file=sys.stderr,
82
+ )
83
+ break
84
+
85
+ # Return versions in the order defined by package_mapping
86
+ ordered_versions = {}
87
+ for _, package_name in package_mapping:
88
+ if package_name in versions:
89
+ ordered_versions[package_name] = versions[package_name]
90
+ return ordered_versions
91
+
92
+
93
+def pin_dependencies_in_requirements(requirements_path: str, versions: dict[str, str]):
94
+ """
95
+ Insert custom wheel pins at the TOP of requirements file.
96
+
97
+ This ensures that when setup.py processes the file line-by-line,
98
+ custom wheels (torch, triton, etc.) are encountered FIRST, before
99
+ any `-r common.txt` includes that might pull in other dependencies.
100
+
101
+ Creates:
102
+ # Custom ROCm wheel pins (auto-generated)
103
+ torch==2.9.0a0+git1c57644
104
+ triton==3.4.0
105
+ torchvision==0.23.0a0+824e8c8
106
+ amdsmi==26.1.0+5df6c765
107
+
108
+ -r common.txt
109
+ ... rest of file ...
110
+ """
111
+ requirements_file = Path(requirements_path)
112
+
113
+ if not requirements_file.exists():
114
+ print(
115
+ f"ERROR: Requirements file not found: {requirements_path}", file=sys.stderr
116
+ )
117
+ sys.exit(1)
118
+
119
+ # Backup original file
120
+ backup_file = requirements_file.with_suffix(requirements_file.suffix + ".bak")
121
+ with open(requirements_file) as f:
122
+ original_lines = f.readlines()
123
+
124
+ # Write backup
125
+ with open(backup_file, "w") as f:
126
+ f.writelines(original_lines)
127
+
128
+ # Build header with pinned custom wheels
129
+ header_lines = [
130
+ "# Custom ROCm wheel pins (auto-generated by pin_rocm_dependencies.py)\n",
131
+ "# These must come FIRST to ensure correct dependency resolution\n",
132
+ ]
133
+
134
+ for package_name, exact_version in versions.items():
135
+ header_lines.append(f"{package_name}=={exact_version}\n")
136
+
137
+ header_lines.append("\n") # Blank line separator
138
+
139
+ # Filter out any existing entries for custom packages from original file
140
+ filtered_lines = []
141
+ removed_packages = []
142
+
143
+ for line in original_lines:
144
+ stripped = line.strip()
145
+ should_keep = True
146
+
147
+ # Check if this line is for one of our custom packages
148
+ if stripped and not stripped.startswith("#") and not stripped.startswith("-"):
149
+ for package_name in versions:
150
+ # Handle both hyphen and underscore variations
151
+ pattern_name = package_name.replace("-", "[-_]")
152
+ pattern = rf"^{pattern_name}\s*[=<>]=?\s*[\d.a-zA-Z+]+"
153
+
154
+ if re.match(pattern, stripped, re.IGNORECASE):
155
+ removed_packages.append(f"{package_name}: {stripped}")
156
+ should_keep = False
157
+ break
158
+
159
+ if should_keep:
160
+ filtered_lines.append(line)
161
+
162
+ # Combine: header + filtered original content
163
+ final_lines = header_lines + filtered_lines
164
+
165
+ # Write modified content
166
+ with open(requirements_file, "w") as f:
167
+ f.writelines(final_lines)
168
+
169
+ # Print summary
170
+ print("\n✓ Inserted custom wheel pins at TOP of requirements:", file=sys.stderr)
171
+ for package_name, exact_version in versions.items():
172
+ print(f" - {package_name}=={exact_version}", file=sys.stderr)
173
+
174
+ if removed_packages:
175
+ print("\n✓ Removed old package entries:", file=sys.stderr)
176
+ for pkg in removed_packages:
177
+ print(f" - {pkg}", file=sys.stderr)
178
+
179
+ print(f"\n✓ Patched requirements file: {requirements_path}", file=sys.stderr)
180
+ print(f" Backup saved: {backup_file}", file=sys.stderr)
181
+
182
+
183
+def main():
184
+ if len(sys.argv) != 3:
185
+ print(
186
+ f"Usage: {sys.argv[0]} <install_dir> <requirements_file>", file=sys.stderr
187
+ )
188
+ print(
189
+ f"Example: {sys.argv[0]} /install /app/vllm/requirements/rocm.txt",
190
+ file=sys.stderr,
191
+ )
192
+ sys.exit(1)
193
+
194
+ install_dir = sys.argv[1]
195
+ requirements_path = sys.argv[2]
196
+
197
+ print("=" * 70, file=sys.stderr)
198
+ print("Pinning vLLM dependencies to custom ROCm wheel versions", file=sys.stderr)
199
+ print("=" * 70, file=sys.stderr)
200
+
201
+ # Get versions from custom wheels
202
+ print(f"\nScanning {install_dir} for custom wheels...", file=sys.stderr)
203
+ versions = get_custom_wheel_versions(install_dir)
204
+
205
+ if not versions:
206
+ print("\nERROR: No custom wheels found in /install!", file=sys.stderr)
207
+ sys.exit(1)
208
+
209
+ # Pin dependencies in requirements file
210
+ print(f"\nPatching {requirements_path}...", file=sys.stderr)
211
+ pin_dependencies_in_requirements(requirements_path, versions)
212
+
213
+ print("\n" + "=" * 70, file=sys.stderr)
214
+ print("✓ Dependency pinning complete!", file=sys.stderr)
215
+ print("=" * 70, file=sys.stderr)
216
+
217
+ sys.exit(0)
218
+
219
+
220
+if __name__ == "__main__":
221
+ main()