Skip to content

Commit 16b2d4c

Browse files
authored
fix(python-backend): make JIT subprocesses work on hosts of any size (#9679)
Two related runtime fixes for Python backends that JIT-compile CUDA kernels at first model load (FlashInfer, PyTorch inductor, triton): 1. libbackend.sh: replace `source ${EDIR}/venv/bin/activate` with a minimal manual setup (_activateVenv: export VIRTUAL_ENV, prepend PATH, unset PYTHONHOME) computed from $EDIR at runtime. `uv venv` and `python -m venv` both bake the create-time absolute path into bin/activate (e.g. VIRTUAL_ENV='/vllm/venv' from the Docker build stage), so sourcing activate on a relocated venv — copied out of the build container and unpacked at an arbitrary backend dir — prepends a stale, non-existent path to $PATH. Pip-installed CLI tools (e.g. ninja, used by FlashInfer's NVFP4 GEMM JIT) are then never found and the load aborts with FileNotFoundError. Doing the env setup ourselves matches what `uv run` does internally and sidesteps the relocation problem entirely. Generic — every Python backend benefits. 2. vllm/run.sh: replace ninja's default -j$(nproc)+2 with an adaptive MAX_JOBS = min(nproc, (MemAvailable-4)/4). Each concurrent nvcc/cudafe++ peaks at multiple GiB; the default OOM-kills on memory-tight hosts (e.g. a 16 GiB desktop loading a 27B NVFP4 model) but underutilises 100-core / 1 TB boxes. User-set MAX_JOBS still wins. Also pin NVCC_THREADS=2 unless overridden. Refs: vllm-project/vllm#20079 Assisted-by: Claude:claude-opus-4-7 [Edit] [Bash]
1 parent 8e43842 commit 16b2d4c

2 files changed

Lines changed: 41 additions & 2 deletions

File tree

backend/python/common/libbackend.sh

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -318,6 +318,21 @@ _makeVenvPortable() {
318318
}
319319

320320

321+
# Apply the venv to the current process: VIRTUAL_ENV, PATH, PYTHONHOME hygiene.
322+
# Equivalent to the runtime portion of `source bin/activate`, but computed from
323+
# $EDIR (resolved at runtime via realpath) instead of the path baked into
324+
# bin/activate at venv-create time. `uv venv` (and `python -m venv`) both bake
325+
# the create-time absolute path in, so sourcing activate on a relocated venv —
326+
# e.g. one built at /vllm/venv inside a Docker stage and unpacked under
327+
# /backends/cuda13-vllm-development/venv at runtime — silently prepends a
328+
# stale, non-existent path to $PATH. Doing the setup ourselves sidesteps that;
329+
# this is the same approach `uv run` takes internally.
330+
_activateVenv() {
331+
export VIRTUAL_ENV="${EDIR}/venv"
332+
export PATH="${EDIR}/venv/bin:${PATH}"
333+
unset PYTHONHOME
334+
}
335+
321336
# ensureVenv makes sure that the venv for the backend both exists, and is activated.
322337
#
323338
# This function is idempotent, so you can call it as many times as you want and it will
@@ -354,7 +369,7 @@ function ensureVenv() {
354369
venv_args="--copies"
355370
fi
356371
"${interpreter}" -m venv ${venv_args} "${EDIR}/venv"
357-
source "${EDIR}/venv/bin/activate"
372+
_activateVenv
358373
"${interpreter}" -m pip install --upgrade pip
359374
else
360375
if [ "x${PORTABLE_PYTHON}" == "xtrue" ]; then
@@ -375,7 +390,7 @@ function ensureVenv() {
375390
fi
376391

377392
if [ "x${VIRTUAL_ENV:-}" != "x${EDIR}/venv" ]; then
378-
source "${EDIR}/venv/bin/activate"
393+
_activateVenv
379394
fi
380395
}
381396

backend/python/vllm/run.sh

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,30 @@ set -x
33

44
backend_dir=$(dirname $0)
55

6+
# FlashInfer / PyTorch JIT-compile CUDA kernels at first model load (e.g.
7+
# the NVFP4 GEMM kernel for Blackwell SM120). Each concurrent nvcc /
8+
# cudafe++ peaks at multiple GiB during compilation; ninja's default
9+
# (-j$(nproc)+2) OOM-kills on memory-tight hosts but underutilises
10+
# 100-core / 1 TB boxes. Default MAX_JOBS to the smaller of the CPU count
11+
# and an available-memory budget at ~4 GiB per job. User-set MAX_JOBS in
12+
# the environment wins.
13+
# https://github.com/vllm-project/vllm/issues/20079
14+
if [ -z "${MAX_JOBS:-}" ]; then
15+
_ncpus=$(nproc 2>/dev/null || echo 1)
16+
_mem_avail_kb=$(awk '/^MemAvailable:/ {print $2; exit}' /proc/meminfo 2>/dev/null || echo 0)
17+
_mem_avail_gb=$(( _mem_avail_kb / 1024 / 1024 ))
18+
# Reserve ~4 GiB for the rest of the system; budget ~4 GiB per job.
19+
if [ "${_mem_avail_gb}" -gt 8 ]; then
20+
_mem_jobs=$(( (_mem_avail_gb - 4) / 4 ))
21+
else
22+
_mem_jobs=1
23+
fi
24+
[ "${_mem_jobs}" -lt 1 ] && _mem_jobs=1
25+
[ "${_mem_jobs}" -gt "${_ncpus}" ] && _mem_jobs=${_ncpus}
26+
export MAX_JOBS="${_mem_jobs}"
27+
fi
28+
export NVCC_THREADS="${NVCC_THREADS:-2}"
29+
630
if [ -d $backend_dir/common ]; then
731
source $backend_dir/common/libbackend.sh
832
else

0 commit comments

Comments
 (0)