[CI/Docker] Clean up redundant flashinfer cubin downloads (#22491)
This commit is contained in:
+1
-4
@@ -302,9 +302,7 @@ FROM torch_deps AS flashinfer_cache
|
||||
ARG CUDA_VERSION
|
||||
ARG INSTALL_FLASHINFER_JIT_CACHE
|
||||
ARG FLASHINFER_VERSION
|
||||
ARG BUILD_AND_DOWNLOAD_PARALLEL
|
||||
|
||||
# Download flashinfer cubins
|
||||
# Stage jit-cache artifacts into /flashinfer_jit_output for clean COPY later
|
||||
RUN --mount=type=cache,target=/root/.cache/pip \
|
||||
case "$CUDA_VERSION" in \
|
||||
@@ -319,8 +317,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
|
||||
python3 -m pip install flashinfer-jit-cache==${FLASHINFER_VERSION} --index-url https://flashinfer.ai/whl/cu${CUINDEX} \
|
||||
&& cp -r /usr/local/lib/python3.12/dist-packages/flashinfer_jit_cache /flashinfer_jit_output/ \
|
||||
&& cp -r /usr/local/lib/python3.12/dist-packages/flashinfer_jit_cache-*.dist-info /flashinfer_jit_output/ ; \
|
||||
fi \
|
||||
&& FLASHINFER_CUBIN_DOWNLOAD_THREADS=${BUILD_AND_DOWNLOAD_PARALLEL} FLASHINFER_LOGGING_LEVEL=warning python3 -m flashinfer --download-cubin
|
||||
fi
|
||||
|
||||
########################################################
|
||||
# PARALLEL STAGE 4: Dev Tools Builder (starts from base)
|
||||
|
||||
Reference in New Issue
Block a user