Skip to content

Commit 2979f7a

Browse files
tjtanaatjtanaa
committed
collect the wheel dependencies after the whole environment has been setup, this will be the most portable and it should have taken into account the custom torch, torchvision, triton, amdsmi and vllm's dependencies
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
1 parent 0d6e14c commit 2979f7a

File tree

3 files changed

+64
-82
lines changed

3 files changed

+64
-82
lines changed

‎.github/workflows/build-rocm-wheel.yml‎

Lines changed: 28 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,14 @@ on:
3434
description: 'Run ID to reuse artifacts from (required if debug_mode or skip_to_job4 is true)'
3535
required: false
3636
default: ''
37+
upload_wheels:
38+
description: 'Upload wheels to S3 (false for nightly builds, true for releases)'
39+
required: false
40+
default: 'false'
41+
type: choice
42+
options:
43+
- 'false'
44+
- 'true'
3745
permissions:
3846
contents: write
3947
pages: write
@@ -104,85 +112,7 @@ jobs:
104112
name: rocm-base-image
105113
path: artifacts/docker-image/rocm-base-image.tar.gz
106114
retention-days: 7
107-
collect-dependency-wheels:
108-
name: Collect Dependency Wheels from PyPI
109-
needs: build-base-wheels
110-
if: github.event.inputs.debug_mode != 'true' && github.event.inputs.skip_to_job4 != 'true'
111-
runs-on: ubuntu-latest
112-
steps:
113-
- name: Checkout repository
114-
uses: actions/checkout@v4
115-
with:
116-
fetch-depth: 0 # Fetch full history for correct version detection
117-
- name: Set up Python
118-
uses: actions/setup-python@v5
119-
with:
120-
python-version: ${{ github.event.inputs.python_version }}
121-
- name: Install pip-tools
122-
run: |
123-
python -m pip install --upgrade pip
124-
pip install pip-tools packaging requests
125-
- name: Download base wheels artifact
126-
uses: actions/download-artifact@v4
127-
with:
128-
name: rocm-base-wheels
129-
path: base-wheels
130-
- name: List downloaded base wheels
131-
run: |
132-
echo "Base wheels that will be used for dependency resolution:"
133-
ls -lh base-wheels/
134-
- name: Download dependency wheels with version ranges
135-
run: |
136-
python .github/scripts/download_dependency_wheels.py \
137-
--requirements requirements/rocm.txt \
138-
--output-dir artifacts/dependency-wheels \
139-
--python-version ${{ github.event.inputs.python_version }} \
140-
--max-versions 5 \
141-
--base-wheels-dir base-wheels
142-
- name: List downloaded wheels
143-
run: |
144-
echo "Downloaded wheels:"
145-
ls -lh artifacts/dependency-wheels/
146-
echo "Total size:"
147-
du -sh artifacts/dependency-wheels/
148-
echo "Total count: $(ls artifacts/dependency-wheels/*.whl 2>/dev/null | wc -l) wheels"
149115

150-
- name: Validate critical dependencies
151-
run: |
152-
WHEEL_COUNT=$(ls artifacts/dependency-wheels/*.whl 2>/dev/null | wc -l)
153-
echo "Downloaded $WHEEL_COUNT wheels"
154-
155-
# With transitive dependencies, we should have 150+ wheels
156-
if [ "$WHEEL_COUNT" -lt 150 ]; then
157-
echo "::error::Too few wheels downloaded ($WHEEL_COUNT). Expected at least 150 (including transitive dependencies)."
158-
echo "::error::Check the download logs above for failures."
159-
exit 1
160-
fi
161-
162-
# Check for specific critical packages (including transitive deps)
163-
CRITICAL_PKGS="regex numpy transformers tokenizers protobuf pydantic aiohttp requests tqdm fastapi anyio"
164-
MISSING=""
165-
166-
for pkg in $CRITICAL_PKGS; do
167-
if ! ls artifacts/dependency-wheels/${pkg}-*.whl 2>/dev/null | grep -q .; then
168-
MISSING="$MISSING $pkg"
169-
fi
170-
done
171-
172-
if [ -n "$MISSING" ]; then
173-
echo "::error::Missing critical packages:$MISSING"
174-
echo "::error::These packages are required for vllm to function properly."
175-
exit 1
176-
fi
177-
178-
echo "✅ All critical packages present!"
179-
180-
- name: Upload dependency wheels
181-
uses: actions/upload-artifact@v4
182-
with:
183-
name: dependency-wheels
184-
path: artifacts/dependency-wheels/*.whl
185-
retention-days: 7
186116
build-vllm-wheel:
187117
name: Build vLLM ROCm Wheel
188118
needs: build-base-wheels
@@ -243,6 +173,16 @@ jobs:
243173
.
244174
ls -lh dist/
245175
ls -lh dist/*.whl
176+
177+
echo ""
178+
echo "Extracted dependency wheels:"
179+
ls -lh dist/dependency-wheels/
180+
echo "Total dependency wheels: $(ls dist/dependency-wheels/*.whl 2>/dev/null | wc -l)"
181+
182+
echo ""
183+
echo "Dependency list:"
184+
cat dist/all-dependencies.txt
185+
246186
- name: Install auditwheel and repair wheel
247187
run: |
248188
python -m pip install auditwheel patchelf
@@ -269,10 +209,18 @@ jobs:
269209
name: vllm-wheel
270210
path: artifacts/vllm-wheel/*.whl
271211
retention-days: 7
212+
213+
- name: Upload dependency wheels
214+
uses: actions/upload-artifact@v4
215+
with:
216+
name: dependency-wheels
217+
path: dist/dependency-wheels/*.whl
218+
retention-days: 7
219+
272220
create-pypi-repository:
273221
name: Create PyPI Repository and Publish
274-
needs: [build-base-wheels, collect-dependency-wheels, build-vllm-wheel]
275-
if: always() && (needs.build-vllm-wheel.result == 'success' || github.event.inputs.skip_to_job4 == 'true') && (github.event.inputs.debug_mode != 'true' || github.event.inputs.skip_to_job4 == 'true' || (needs.build-base-wheels.result == 'skipped' && needs.collect-dependency-wheels.result == 'skipped'))
222+
needs: [build-base-wheels, build-vllm-wheel]
223+
if: github.event.inputs.upload_wheels == 'true' && always() && (needs.build-vllm-wheel.result == 'success' || github.event.inputs.skip_to_job4 == 'true')
276224
runs-on: ubuntu-latest-l
277225
steps:
278226
- name: Checkout repository

‎docker/Dockerfile.rocm‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,33 @@ RUN cd vllm \
5454
&& python3 -m pip install -r requirements/rocm.txt \
5555
&& python3 setup.py clean --all \
5656
&& python3 setup.py bdist_wheel --dist-dir=dist
57+
FROM build_vllm AS collect_dependencies
58+
ARG COMMON_WORKDIR
59+
# Install pipdeptree for dependency analysis
60+
RUN python3 -m pip install pipdeptree
61+
62+
# Install all built wheels into the environment
63+
RUN cd ${COMMON_WORKDIR}/vllm && python3 -m pip install dist/*.whl
64+
65+
# Export dependency tree
66+
RUN pipdeptree --freeze > ${COMMON_WORKDIR}/all-dependencies.txt
67+
68+
# Show what we collected
69+
RUN echo "Collected dependencies:" && cat ${COMMON_WORKDIR}/all-dependencies.txt
70+
71+
# Verify custom ROCm wheels are available at /install
72+
RUN echo "Custom ROCm wheels at /install:" && ls -lh /install/
73+
74+
# Download all dependency wheels
75+
# pip will use wheels from /install (referenced in pipdeptree output)
76+
RUN mkdir -p ${COMMON_WORKDIR}/dependency-wheels && \
77+
pip download -r ${COMMON_WORKDIR}/all-dependencies.txt \
78+
-d ${COMMON_WORKDIR}/dependency-wheels/ \
79+
--only-binary=:all:
80+
81+
# Show downloaded wheels
82+
RUN echo "Downloaded dependency wheels:" && ls -lh ${COMMON_WORKDIR}/dependency-wheels/
83+
5784
FROM scratch AS export_vllm
5885
ARG COMMON_WORKDIR
5986
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/dist/*.whl /
@@ -63,6 +90,8 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests
6390
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples
6491
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
6592
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
93+
COPY --from=collect_dependencies ${COMMON_WORKDIR}/dependency-wheels /dependency-wheels
94+
COPY --from=collect_dependencies ${COMMON_WORKDIR}/all-dependencies.txt /all-dependencies.txt
6695

6796
# -----------------------
6897
# Test vLLM image

‎docker/Dockerfile.rocm_base.minimal‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -87,8 +87,13 @@ RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
8787
cp /install/*.whl /app/debs
8888

8989
FROM base AS final
90-
RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \
91-
pip install /install/*.whl
90+
# Copy wheels to permanent location first (so pip metadata references remain valid)
91+
RUN --mount=type=bind,from=debs,src=/app/debs,target=/tmp/debs \
92+
mkdir -p /install && \
93+
cp /tmp/debs/*.whl /install/
94+
95+
# Install from permanent location
96+
RUN pip install /install/*.whl
9297

9398
ARG BASE_IMAGE
9499
ARG TRITON_BRANCH

0 commit comments

Comments
 (0)