fix(worker): skip redundant virtualenv setup when replicas share the same venv

[m199369309]

Summary

Fix a race condition when multiple model replicas are deployed concurrently on the same worker and share a single virtualenv path.

Root Cause

When two replicas share the same virtualenv (e.g., two copies of the same model), each replica's setup flow acquires _exclusive_venv_path_lock sequentially. The second replica's install_packages call downgrades flashinfer from the AOT version back to the JIT version that vllm hard-pins, overwriting .so files while the first replica's child subprocess is importing torch/vllm — causing a hang and eventual subpool creation timeout (60s).

vllm 0.21.0 hard-pins flashinfer-cubin==0.6.8.post1, but the AOT post-install hook upgrades to 0.6.11.post3 for sm_120 Blackwell support. When install_packages runs for the second replica, it detects the dependency mismatch and downgrades flashinfer back to 0.6.8, overwriting .so files while the first replica's child process is loading them.

This was reproduced 5 consecutive times in production with 100% failure rate for the second replica (regardless of GPU allocation).

Fix

Add a process-local _venv_setup_done: Dict[str, int] that maps venv path to a fingerprint (hash of the sorted package list). Only the first replica sharing a given venv with the same package set runs install_packages and all post-install hooks; subsequent replicas skip the entire block. When the package requirements change (e.g., after a config update), the fingerprint mismatch invalidates the cache and setup re-runs.

The guard covers all venv-mutating operations inside the lock:

• install_packages
• ensure_sglang_inherited_packages_compatible_post_install
• apply_flashinfer_aot_post_install
• ensure_flashinfer_cubin_matches_post_install

An on-disk marker file (.xinference_setup_done) provides an additional guard against external venv deletion during the same process lifetime.

Scope

Process-local — the correct scope for the race being fixed (multiple replicas deploying concurrently within a single Worker process). Cross-worker races on a shared XINFERENCE_VIRTUAL_ENV_DIR still require the _exclusive_venv_path_lock; this change does not address that case.

Verification

• All pre-commit hooks pass: black, flake8, isort, mypy, codespell
• 8 new unit + integration tests covering: first install, same-packages skip, different-packages invalidate, missing marker after venv deletion, different path, and two _prepare_virtual_env end-to-end cases
• No new dependencies or configuration changes

[gemini-code-assist]

Code Review

This pull request introduces a process-level cache (_venv_setup_done) to track completed virtual environment setups, preventing redundant package installations and post-install hooks when multiple replicas share the same virtualenv. While this optimization is beneficial, a potential issue arises if a virtualenv is deleted and recreated during the worker's lifetime, as the process would skip setting up the newly recreated environment. To address this, it is recommended to use a marker file inside the virtualenv directory to verify that the setup is actually present on disk before skipping it.

[gemini-code-assist]

If a virtualenv is deleted (either manually or via the remove_virtual_env API) during the lifetime of the worker process, and then the same model is launched again, _create_virtual_env_manager will recreate the basic virtualenv directory. However, because venv_path is still present in _venv_setup_done, _prepare_virtual_env will skip the package installation and post-install hooks entirely, leaving the newly created virtualenv empty and causing the model launch to fail.

To make this robust against virtualenv deletion/recreation, we can write a small marker file (e.g., .xinference_setup_done) inside the virtualenv directory upon successful setup. We then skip setup only if venv_path is in _venv_setup_done and the marker file exists.

        with _exclusive_venv_path_lock(venv_path := str(virtual_env_manager.env_path)):
            marker_path = os.path.join(venv_path, ".xinference_setup_done")
            if venv_path not in _venv_setup_done or not os.path.exists(marker_path):
                if modern_sglang_kernel:
                    # SGLang 0.5.11 renamed the distribution while retaining the
                    # same import package.  Remove the cached legacy owner before
                    # the new wheel writes those files.
                    cls._uninstall_venv_package(virtual_env_manager, "sgl-kernel")
                if force_reinstall_xllamacpp:
                    cls._uninstall_venv_package(virtual_env_manager, "xllamacpp")
                virtual_env_manager.install_packages(packages, **conf, **variables)

                # Post-install: flashinfer AOT workaround for sm_120 Blackwell.
                # vllm 0.21.0 hard-pins flashinfer-cubin==0.6.8.post1 which has JIT
                # compilation failure on sm_120. Force-upgrade to AOT versions.
                # Run under the same lock — uv pip install mutates the venv and
                # must stay serialized with install_packages() and other AOT
                # upgrades when multiple replicas/workers share this venv.
                from .virtual_env_manager import (
                    apply_flashinfer_aot_post_install,
                    ensure_flashinfer_cubin_matches_post_install,
                    ensure_sglang_inherited_packages_compatible_post_install,
                )

                ensure_sglang_inherited_packages_compatible_post_install(
                    model_engine, virtual_env_manager
                )
                if XINFERENCE_VIRTUAL_ENV_OFFLINE_INSTALL:
                    logger.info(
                        "Skipping the FlashInfer AOT post-install from its public "
                        "wheel index in explicit offline-install mode"
                    )
                else:
                    apply_flashinfer_aot_post_install(
                        model_engine,
                        architectures,
                        virtual_env_manager,
                        conf,
                        cuda_version,
                    )
                ensure_flashinfer_cubin_matches_post_install(
                    model_engine,
                    virtual_env_manager,
                    allow_public_install=not XINFERENCE_VIRTUAL_ENV_OFFLINE_INSTALL,
                )

                try:
                    with open(marker_path, "w") as f:
                        f.write("done")
                except OSError:
                    pass
                _venv_setup_done.add(venv_path)
            else:
                logger.info(
                    "Virtual env %s already set up in this process; "
                    "skipping install_packages and post-install hooks",
                    venv_path,
                )

[OliverBryant]

Review

The diagnosis holds up. The venv path is built from v4/model_name/engine/python_version (worker.py:3206) with no model_uid/replica component, so replicas of the same model necessarily share one venv, and _exclusive_venv_path_lock only serializes setup — it does not make it idempotent. The second replica re-runs install_packages, and the flashinfer AOT (0.6.11) vs. vllm hard-pin (0.6.8) churn rewrites .so files while the first replica's child process is importing them. Making venv setup idempotent is the right direction.

But the skip condition as written keys only on the venv path, while the work it guards depends on inputs that are not determined by the path. That turns an intermittent 60s timeout into a deterministic silent mis-install. Details inline; summarizing the blocking one and the smaller items:

Most important: the skip key must include the resolved package set, not just the path (see inline comment on the if venv_path not in _venv_setup_done line). I'd treat this one as must-fix before merge.

Should fix:

1. The marker file's semantics are contradictory — it adds no information the in-process set doesn't already have, and it does nothing for the cross-process case it appears aimed at.
2. _venv_setup_done only ever grows; remove_virtual_env / confirm_and_remove_model don't evict.
3. The PR description's claim that process-local is "the correct scope" is too strong — _exclusive_venv_path_lock's own docstring targets "multiple replicas/workers", and a shared XINFERENCE_VIRTUAL_ENV_DIR across workers still races. Not necessary to solve here, but worth correcting in the description so this isn't read as closed.
4. The # See optimize/20260702/2026070209.md provenance pointer was dropped; please keep it.
5. No tests. Whether to skip is easily extracted into a small pure helper — three cases (first install / repeat with same package set / repeat with a different package set) would pin down the blocking issue above. CLAUDE.md also asks for tests on behavior changes.

Nit: the walrus inside with _exclusive_venv_path_lock(venv_path := ...) reads poorly; a plain assignment on the preceding line is clearer.

Also note this branch carries an unrelated change (pin_sentence_transformers_numpy_abi at worker.py:2958) that isn't mentioned in the description — if it's not intended to be part of this fix, please split it out.

One unrelated meta point: please drop the 🤖 Generated with Claude Code trailer from the PR description and the Co-Authored-By: Claude line from commit 745fb3c, per this repo's convention.

[OliverBryant]

Must-fix before merge: the skip key is only venv_path, but _prepare_virtual_env's inputs are not uniquely determined by that path.

virtual_env_packages is a per-launch user parameter (restful_api.py:846, and --virtual-env-package in cmdline.py:1016), and architectures / model_format come from this launch's resolution (worker.py:3395-3402). All of them feed the packages list computed above, while the venv path only encodes model_name/engine/python_version.

Concrete failure: launch model X with --virtual-env-package foo, then launch the same X with --virtual-env-package bar. The second launch hits this branch, skips the entire install block, and bar is never installed — the model starts and fails on import, with nothing in the log but "already set up". Same for two model_formats (awq/gptq) that resolve to different sets via get_engine_model_format_virtualenv_packages.

Suggestion: fingerprint the resolved inputs and include that in the key, e.g. sha256 over sorted(packages) plus the install-relevant bits of conf/variables. Skip only on a fingerprint match; on mismatch, fall through and install. That preserves the fix for the actual replica case (identical inputs) without swallowing genuinely different ones.

[qinxuye]

The package-list part is fixed, but this is not fully addressed yet: install-relevant settings are still omitted from the fingerprint. I reproduced two calls with the same ["demo-pkg"] requirement while changing index_url from one source to another; install_packages was invoked only once, so the second source/configuration was silently ignored. Please fingerprint the normalized package list together with conf/variables (or an equivalent canonical representation of every setup input), as suggested above.

[m199369309]

Fixed in d42cfbb — _make_fingerprint() now hashes the full tuple of (sorted(packages), sorted(conf.items()), sorted(variables.items()), sorted(architectures or [])), so a change to index_url, skip_installed, engine variables, or architectures all produce a different fingerprint and correctly trigger re-setup.

The conf dict is derived from dict(settings) with packages and inherit_pip_config popped (those are already surfaced through the packages list and the install framework). All remaining keys — including index_url, skip_installed, extra_index_url, etc. — contribute to the fingerprint.

[OliverBryant]

The marker file's semantics don't line up with either job it could do.

For the same-process case, the in-memory set already covers "we installed this"; the only extra fact the marker contributes is "the directory still exists" — and since remove_virtual_env rmtrees the whole tree, os.path.exists(venv_path) gives you that without writing a file.

For the cross-process case (worker restart), the condition is venv_path not in _venv_setup_done or not os.path.exists(marker_path). After a restart the set is empty, so the first disjunct is already true and full setup re-runs regardless of the marker. So it buys nothing there either.

Two coherent options:

• Drop the marker; use venv_path in _venv_setup_done and os.path.exists(venv_path), and evict in remove_virtual_env.
• Or make the marker actually do the cross-process work: write the package-set fingerprint from the other comment into it, and compare on read. That subsumes the in-memory set, fixes the blocking issue, and additionally avoids one redundant reinstall after every worker restart.

I'd go with the second.

[OliverBryant]

except OSError: pass swallows the write failure silently. A failed marker write means the next launch redoes the full setup, which is worth a logger.debug/logger.warning rather than being invisible during diagnosis.

[OliverBryant]

_venv_setup_done only ever grows. remove_virtual_env (worker.py:4244) and confirm_and_remove_model delete venvs without discarding the corresponding entry, so the set and the filesystem drift apart.

The marker check currently papers over this. If the marker is simplified per the other comment, a _venv_setup_done.discard(...) in the removal paths becomes mandatory — worth adding either way so the two pieces of state stay in sync.

[OliverBryant]

The # See optimize/20260702/2026070209.md provenance pointer that was here got dropped in the re-indent. Please restore it — it's the only link back to the root-cause writeup for the flashinfer AOT workaround.

[m199369309]

Thanks for the thorough review — all points are spot-on. I've pushed a fix commit (c00f683) that addresses the key findings:

Blocking issue — fixed:

The skip key now includes a package fingerprint. _venv_setup_done changed from Set[str] to Dict[str, int] mapping venv_path → hash(tuple(sorted(packages))). Two helper functions (_should_skip_venv_setup / _mark_venv_setup_done) are extracted so the three-condition check (path in dict + hash matches + marker exists) is testable. If the same model is re-launched with different packages, the hash differs and setup re-runs.

Should-fix items — fixed:

1. Tests — 8 new tests in test_utils.py covering all three skip cases + marker-missing + different-path + two _prepare_virtual_env integration tests.

2. Provenance pointer — Restored # See optimize/20260702/2026070209.md.

3. Walrus readability — venv_path extracted to a plain assignment before the with statement.

4. PR description — Dropped the Claude trailer and corrected the scope claim to acknowledge that cross-worker races remain in play.

5. Unrelated pin_sentence_transformers_numpy_abi — Verified this line already exists on origin/main (worker.py:2951). Not introduced by this PR — it just appears in diff context because it's near the changed area.

Skipped / considered:

• _venv_setup_done eviction on remove_virtual_env — The marker-file check (os.path.exists) now handles this correctly: if the venv is deleted externally, the marker is gone, the skip guard returns False, and setup re-runs. The dict entry being stale is harmless (it'll be overwritten with the new hash after the re-setup). Explicit eviction would require plumbing a callback into VirtualEnvManager.remove_virtual_env, which touches a different layer — worth a follow-up but not a correctness blocker now.

• Marker file semantics — With the package fingerprint, the marker file now serves a clear single purpose: detecting external venv deletion. Without it, the in-process dict would be the sole source of truth, and a deleted-then-recreated venv would be skipped incorrectly. The os.path.exists gate closes that gap.

[qinxuye]

One additional concurrency issue remains on the current head.

[qinxuye]

apply_vllm_patches below still mutates files inside the same shared virtualenv (for example, kv_cache_utils.py) after _exclusive_venv_path_lock has been released, and it runs for every replica even when setup was skipped. Two replicas can therefore both read/write the patch target, or one child can start importing while another replica rewrites it—the same class of race this PR is meant to eliminate. Please keep this mutation under the same deduplication/lock and include patch-relevant inputs such as architectures in the setup fingerprint.

[m199369309]

Good catch — apply_vllm_patches was indeed outside the lock, allowing the same class of race condition this PR aims to eliminate. Fixed in d42cfbb:

1. Moved apply_vllm_patches inside _exclusive_venv_path_lock and the skip guard, so it is serialized and deduplicated alongside install_packages and post-install hooks.
2. architectures is now part of the setup fingerprint via _make_fingerprint(). A change in architectures produces a different fingerprint, triggering re-setup (and re-patching) rather than incorrectly skipping.

The vllm-patches import is guarded by try/except ImportError so this path is a no-op when vllm is not installed.

[qinxuye]

One blocking fingerprint issue remains.

[qinxuye]

This fingerprint fails for valid list-valued install settings. VirtualEnvSettings.extra_index_url, find_links, and trusted_host all accept List[str], so tuple(sorted(conf.items())) still contains a list and hash(key) raises TypeError: unhashable type: 'list' before virtualenv setup starts. This also affects the automatically generated extra_index_url list for PyTorch CUDA wheels. Please canonicalize nested values first (for example, serialize the complete input with deterministic JSON or recursively freeze lists/dicts), and add a regression test using a list-valued conf.

[m199369309]

Great catch — this would have been a crash-on-launch regression. Fixed in a87d1e7:

1. Added _freeze() helper that recursively converts dict to sorted tuple of (k, freeze(v)) pairs and list to tuple of frozen values, ensuring every nested value is hashable.
2. _make_fingerprint() now passes conf and variables through _freeze() instead of tuple(sorted(.items())).
3. Added regression test test_make_fingerprint_handles_list_valued_conf with extra_index_url: ["https://...", "https://..."] verifying no TypeError, different lists produce different fingerprints, and identical inputs are stable.

[qinxuye]

One platform-specific concurrency issue remains.

[qinxuye]

On Windows, _exclusive_venv_path_lock is a no-op, while _prepare_virtual_env runs via asyncio.to_thread and the worker allows multiple concurrent launches. Two replicas can therefore both pass _should_skip_venv_setup() before either one records completion, then concurrently run install_packages and the post-install patches against the same virtualenv—the race this PR is intended to prevent. Please add a process-local per-venv lock that works on every platform (the Unix file lock can remain for cross-process serialization), and add a concurrent regression test for this check-and-setup sequence.