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 11, 17:30 UTC
[Perf] Enable environment cache in EngineCore to enable the feature for UniProcExecutor as well (#29289)
Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>
Jialin Ouyang committed
Dec 10, 2025 at 11:13 UTC
9f042ba26b59e1bfc9bef031165033fa931f3457
4 files changed
+63
-4
tests/test_envs.py
+38
@@ -8,6 +8,7 @@ import pytest
8
9
import vllm.envs as envs
10
from vllm.envs import (
11
+ disable_envs_cache,
12
enable_envs_cache,
13
env_list_with_choices,
14
env_set_with_choices,
@@ -57,6 +58,43 @@ def test_getattr_with_cache(monkeypatch: pytest.MonkeyPatch):
58
envs.__getattr__ = envs.__getattr__.__wrapped__
59
60
61
+def test_getattr_with_reset(monkeypatch: pytest.MonkeyPatch) -> None:
62
+ monkeypatch.setenv("VLLM_HOST_IP", "1.1.1.1")
63
+ # __getattr__ is not decorated with functools.cache
64
+ assert not hasattr(envs.__getattr__, "cache_info")
65
+
66
+ # Enable envs cache and ignore ongoing environment changes
67
+ enable_envs_cache()
68
+ assert envs.VLLM_HOST_IP == "1.1.1.1"
69
+ # With cache enabled, the environment variable value is cached and unchanged
70
+ monkeypatch.setenv("VLLM_HOST_IP", "2.2.2.2")
71
+ assert envs.VLLM_HOST_IP == "1.1.1.1"
72
+
73
+ disable_envs_cache()
74
+ assert envs.VLLM_HOST_IP == "2.2.2.2"
75
+ # After cache disabled, the environment variable value would be synced
76
+ # with os.environ
77
+ monkeypatch.setenv("VLLM_HOST_IP", "3.3.3.3")
78
+ assert envs.VLLM_HOST_IP == "3.3.3.3"
79
+
80
+
81
+def test_is_envs_cache_enabled() -> None:
82
+ assert not envs._is_envs_cache_enabled()
83
+ enable_envs_cache()
84
+ assert envs._is_envs_cache_enabled()
85
+
86
+ # Only wrap one-layer of cache, so we only need to
87
+ # call disable once to reset.
88
+ enable_envs_cache()
89
+ enable_envs_cache()
90
+ enable_envs_cache()
91
+ disable_envs_cache()
92
+ assert not envs._is_envs_cache_enabled()
93
+
94
+ disable_envs_cache()
95
+ assert not envs._is_envs_cache_enabled()
96
+
97
+
98
class TestEnvWithChoices:
99
"""Test cases for env_with_choices function."""
100
vllm/distributed/parallel_state.py
+2
@@ -1586,6 +1586,8 @@ def destroy_distributed_environment():
1586
1587
1588
def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
1589
+ # Reset environment variable cache
1590
+ envs.disable_envs_cache()
1591
# Ensure all objects are not frozen before cleanup
1592
gc.unfreeze()
1593
vllm/envs.py
+20
@@ -1580,6 +1580,12 @@ def __getattr__(name: str):
1580
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
1581
1582
1583
+def _is_envs_cache_enabled() -> bool:
1584
+ """Checked if __getattr__ is wrapped with functools.cache"""
1585
+ global __getattr__
1586
+ return hasattr(__getattr__, "cache_clear")
1587
+
1588
+
1589
def enable_envs_cache() -> None:
1590
"""
1591
Enables caching of environment variables. This is useful for performance
@@ -1590,6 +1596,9 @@ def enable_envs_cache() -> None:
1596
runtime overhead. This also means that environment variables should NOT
1597
be updated after the service is initialized.
1598
"""
1599
+ if _is_envs_cache_enabled():
1600
+ # Avoid wrapping functools.cache multiple times
1601
+ return
1602
# Tag __getattr__ with functools.cache
1603
global __getattr__
1604
__getattr__ = functools.cache(__getattr__)
@@ -1599,6 +1608,17 @@ def enable_envs_cache() -> None:
1608
__getattr__(key)
1609
1610
1611
+def disable_envs_cache() -> None:
1612
+ """
1613
+ Resets the environment variables cache. It could be used to isolate environments
1614
+ between unit tests.
1615
+ """
1616
+ global __getattr__
1617
+ # If __getattr__ is wrapped by functions.cache, unwrap the caching layer.
1618
+ if _is_envs_cache_enabled():
1619
+ __getattr__ = __getattr__.__wrapped__
1620
+
1621
+
1622
def __dir__():
1623
return list(environment_variables.keys())
1624
vllm/v1/engine/core.py
+3
-4
@@ -211,6 +211,9 @@ class EngineCore:
211
freeze_gc_heap()
212
# If enable, attach GC debugger after static variable freeze.
213
maybe_attach_gc_debug_callback()
214
+ # Enable environment variable cache (e.g. assume no more
215
+ # environment variable overrides after this point)
216
+ enable_envs_cache()
217
218
def _initialize_kv_caches(
219
self, vllm_config: VllmConfig
@@ -672,10 +675,6 @@ class EngineCoreProc(EngineCore):
675
assert addresses.coordinator_input is not None
676
logger.info("Waiting for READY message from DP Coordinator...")
677
675
- # Enable environment variable cache (e.g. assume no more
676
- # environment variable overrides after this point)
677
- enable_envs_cache()
678
-
678
@contextmanager
679
def _perform_handshakes(
680
self,