diff --git a/tests/utils_/test_numa_utils.py b/tests/utils_/test_numa_utils.py index 0f615fb8c4..9f718703a7 100644 --- a/tests/utils_/test_numa_utils.py +++ b/tests/utils_/test_numa_utils.py @@ -464,3 +464,49 @@ def test_parallel_config_validates_numa_bind_nodes(): def test_parallel_config_rejects_invalid_numa_bind_cpus(cpuset): with pytest.raises(ValueError, match="numa_bind_cpus"): ParallelConfig(numa_bind_cpus=[cpuset]) + + +def _fake_numactl_run(rejected_args): + """Fake ``numactl`` that fails when any of ``rejected_args`` is present.""" + + def run(cmd, *args, **kwargs): + arg_str = " ".join(cmd[1:-1]) + ok = not any(bad in arg_str for bad in rejected_args) + return SimpleNamespace(returncode=0 if ok else 1) + + return run + + +def test_configure_subprocess_numa_fallback(monkeypatch): + import multiprocessing + + monkeypatch.setattr("shutil.which", lambda name: "/usr/bin/numactl") + monkeypatch.setattr(numa_utils.envs, "VLLM_WORKER_MULTIPROC_METHOD", "spawn") + node_config = _make_config(numa_bind=True, numa_bind_nodes=[0]) + + monkeypatch.setattr(numa_utils.subprocess, "run", _fake_numactl_run([])) + with numa_utils.configure_subprocess(node_config, local_rank=0): + assert os.environ[numa_utils._NUMACTL_ARGS_ENV] == "--cpunodebind=0 --membind=0" + + membind_fails = _fake_numactl_run(["--membind="]) + monkeypatch.setattr(numa_utils.subprocess, "run", membind_fails) + with numa_utils.configure_subprocess(node_config, local_rank=0): + assert os.environ[numa_utils._NUMACTL_ARGS_ENV] == "--cpunodebind=0" + + cpu_config = _make_config( + numa_bind=True, + numa_bind_nodes=[0], + numa_bind_cpus=["0-3"], + ) + with numa_utils.configure_subprocess(cpu_config, local_rank=0): + assert os.environ[numa_utils._NUMACTL_ARGS_ENV] == "--physcpubind=0-3" + + before = multiprocessing.spawn.get_executable() + monkeypatch.setattr( + numa_utils.subprocess, + "run", + _fake_numactl_run(["--cpunodebind=", "--membind="]), + ) + with numa_utils.configure_subprocess(node_config, local_rank=0): + assert multiprocessing.spawn.get_executable() == before + assert numa_utils._NUMACTL_ARGS_ENV not in os.environ diff --git a/vllm/utils/numa_utils.py b/vllm/utils/numa_utils.py index 6e4b4b471c..2e52935ea6 100644 --- a/vllm/utils/numa_utils.py +++ b/vllm/utils/numa_utils.py @@ -473,6 +473,37 @@ def log_current_affinity_state(label: str) -> None: _log_numactl_show(label) +def _probe_numactl_args(numactl_args: str) -> bool: + """Whether ``numactl true`` succeeds in this (parent) environment.""" + try: + result = subprocess.run( + ["numactl", *numactl_args.split(), "true"], + capture_output=True, + timeout=10, + ) + except (OSError, subprocess.SubprocessError): + return False + return result.returncode == 0 + + +def _resolve_numactl_args(numactl_args: str) -> str: + """Drop ``--membind`` if the container rejects it, keeping CPU binding.""" + cpu_only = " ".join( + t for t in numactl_args.split() if not t.startswith("--membind=") + ) + for candidate in (numactl_args, cpu_only, ""): + if _probe_numactl_args(candidate): + if candidate != numactl_args: + logger.warning( + "numactl args %r rejected; falling back to %r. Add " + "--cap-add SYS_NICE for full NUMA binding.", + numactl_args, + candidate or "no binding", + ) + return candidate + return "" + + @contextmanager def configure_subprocess( vllm_config: "VllmConfig", @@ -500,6 +531,11 @@ def configure_subprocess( ) executable, debug_str = _get_numactl_executable() + numactl_args = _resolve_numactl_args(numactl_args) + if not numactl_args: + # No NUMA binding possible here; launch without the wrapper. + yield + return python_executable = os.fsdecode(multiprocessing.spawn.get_executable()) with ( _set_numa_wrapper_env(numactl_args, python_executable),