Fix CPU simulated NUMA env parsing

Signed-off-by: Kevin-Li-2025 <2242139@qq.com>

Author: Kevin-Li-2025

tests/test_envs.py

@@ -116,6 +116,25 @@ def test_precompiled_install_flags_are_orthogonal() -> None:
         assert environment_variables["VLLM_USE_PRECOMPILED_RUST"]() is True
 
 
+@pytest.mark.parametrize(
+    ("env_value", "expected"),
+    [
+        (None, False),
+        ("0", False),
+        ("1", True),
+    ],
+)
+def test_cpu_sim_multi_numa_env(
+    monkeypatch: pytest.MonkeyPatch, env_value: str | None, expected: bool
+) -> None:
+    if env_value is None:
+        monkeypatch.delenv("VLLM_CPU_SIM_MULTI_NUMA", raising=False)
+    else:
+        monkeypatch.setenv("VLLM_CPU_SIM_MULTI_NUMA", env_value)
+
+    assert envs.VLLM_CPU_SIM_MULTI_NUMA is expected
+
+
 class TestEnvWithChoices:
     """Test cases for env_with_choices function."""
 

tests/utils_/test_cpu_resource_utils.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm.utils import cpu_resource_utils as cr_utils
+
+
+@pytest.fixture(autouse=True)
+def _clear_cpu_memory_envs(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.delenv(cr_utils.DEVICE_CONTROL_ENV_VAR, raising=False)
+    monkeypatch.delenv("VLLM_CPU_SIM_MULTI_NUMA", raising=False)
+
+
+def _stub_linux_memory_affinity(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(cr_utils.sys, "platform", "linux")
+    monkeypatch.setattr(cr_utils, "get_memory_affinity", lambda: [0, 1, 2, 3])
+
+
+@pytest.mark.parametrize("sim_multi_numa", [None, "0"])
+def test_get_visible_memory_node_uses_visible_nodes_when_not_simulating(
+    monkeypatch: pytest.MonkeyPatch, sim_multi_numa: str | None
+) -> None:
+    _stub_linux_memory_affinity(monkeypatch)
+    monkeypatch.setenv(cr_utils.DEVICE_CONTROL_ENV_VAR, "1,3")
+    if sim_multi_numa is not None:
+        monkeypatch.setenv("VLLM_CPU_SIM_MULTI_NUMA", sim_multi_numa)
+
+    assert cr_utils.get_visible_memory_node() == [1, 3]
+
+
+def test_get_visible_memory_node_ignores_visible_nodes_when_simulating(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _stub_linux_memory_affinity(monkeypatch)
+    monkeypatch.setenv(cr_utils.DEVICE_CONTROL_ENV_VAR, "1,3")
+    monkeypatch.setenv("VLLM_CPU_SIM_MULTI_NUMA", "1")
+
+    assert cr_utils.get_visible_memory_node() == [0, 1, 2, 3]

vllm/envs.py

@@ -52,6 +52,7 @@
     VLLM_CPU_KVCACHE_SPACE: int | None = 0
     VLLM_CPU_OMP_THREADS_BIND: str = "auto"
     VLLM_CPU_NUM_OF_RESERVED_CPU: int | None = None
+    VLLM_CPU_SIM_MULTI_NUMA: bool = False
     VLLM_CPU_SGL_KERNEL: bool = False
     VLLM_CPU_ATTN_SPLIT_KV: bool = True
     VLLM_ZENTORCH_WEIGHT_PREPACK: bool = True
@@ -838,6 +839,10 @@ def _resolve_rust_frontend_path() -> str | None:
         if "VLLM_CPU_NUM_OF_RESERVED_CPU" in os.environ
         else None
     ),
+    # (CPU backend only) whether to simulate multiple NUMA nodes on a single node.
+    "VLLM_CPU_SIM_MULTI_NUMA": lambda: bool(
+        int(os.getenv("VLLM_CPU_SIM_MULTI_NUMA", "0"))
+    ),
     # (CPU backend only) whether to use SGL kernels, optimized for small batch.
     "VLLM_CPU_SGL_KERNEL": lambda: bool(int(os.getenv("VLLM_CPU_SGL_KERNEL", "0"))),
     # (CPU backend only) whether to enable attention spilt KV.
@@ -2123,6 +2128,7 @@ def compile_factors() -> dict[str, object]:
         "VLLM_V1_OUTPUT_PROC_CHUNK_SIZE",
         "VLLM_CPU_KVCACHE_SPACE",
         "VLLM_CPU_MOE_PREPACK",
+        "VLLM_CPU_SIM_MULTI_NUMA",
         "VLLM_ZENTORCH_WEIGHT_PREPACK",
         "VLLM_TEST_FORCE_LOAD_FORMAT",
         "VLLM_ENABLE_CUDA_COMPATIBILITY",

vllm/utils/cpu_resource_utils.py

@@ -11,6 +11,8 @@
 import psutil
 import regex as re
 
+from vllm import envs
+
 DEVICE_CONTROL_ENV_VAR = "CPU_VISIBLE_MEMORY_NODES"
 
 
@@ -136,7 +138,7 @@ def get_visible_memory_node() -> list[int]:
 
     env_key = DEVICE_CONTROL_ENV_VAR
     if (
-        ("VLLM_CPU_SIM_MULTI_NUMA" not in os.environ)
+        not envs.VLLM_CPU_SIM_MULTI_NUMA
         and env_key in os.environ
         and os.environ[env_key] != ""
     ):

vllm/utils/ompmultiprocessing.py

@@ -33,7 +33,7 @@ def __init__(self, config: "VllmConfig"):
         # is always 1 for non-MoE models
         self.internal_dp_size = config.parallel_config._api_process_count
 
-        self.simulate_multi_node = os.environ.get("VLLM_CPU_SIM_MULTI_NUMA", "0") != "0"
+        self.simulate_multi_node = envs.VLLM_CPU_SIM_MULTI_NUMA
         ld_preload_str = os.getenv("LD_PRELOAD", "")
         self.use_iomp = "libiomp" in ld_preload_str or "libomp" in ld_preload_str
         self.use_gomp = "libgomp" in ld_preload_str