From e55d076f89fd01a0538a3e496d8ff20bf7980100 Mon Sep 17 00:00:00 2001 From: Oxana Korzh Date: Fri, 25 Sep 2026 16:34:19 -0600 Subject: [PATCH] [Bugfix] Don't drop the rest of the allocator config when toggling expandable segments (#57982) Signed-off-by: Oxana Korzh --- .buildkite/test_areas/misc.yaml | 3 + tests/device_allocator/test_alloc_conf.py | 130 ++++++++++++++++++++++ vllm/device_allocator/alloc_conf.py | 119 ++++++++++++++++++++ vllm/device_allocator/cumem.py | 26 +++-- 4 files changed, 270 insertions(+), 8 deletions(-) create mode 100644 tests/device_allocator/test_alloc_conf.py create mode 100644 vllm/device_allocator/alloc_conf.py diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index e71102ebf1..f4b07ffd97 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -495,6 +495,7 @@ steps: source_file_dependencies: &cpu-async-engine-deps - vllm/assets/ - vllm/config/ + - vllm/device_allocator/ - vllm/engine/arg_utils.py - vllm/entrypoints/chat_utils.py - vllm/entrypoints/mcp/ @@ -536,9 +537,11 @@ steps: - tests/parser - tests/transformers_utils - tests/config + - tests/device_allocator device: cpu-small commands: - python3 standalone_tests/lazy_imports.py + - pytest -v -s device_allocator - pytest -v -s test_envs.py - pytest -v -s test_outputs.py - pytest -v -s test_pcp_dp.py diff --git a/tests/device_allocator/test_alloc_conf.py b/tests/device_allocator/test_alloc_conf.py new file mode 100644 index 0000000000..8a9128c5f0 --- /dev/null +++ b/tests/device_allocator/test_alloc_conf.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.device_allocator import alloc_conf + +pytestmark = pytest.mark.cpu_test + +ES = alloc_conf.EXPANDABLE_SEGMENTS + + +@pytest.mark.parametrize( + "environ,expected", + [ + ({}, False), + ({"PYTORCH_ALLOC_CONF": "expandable_segments:True"}, True), + ({"PYTORCH_ALLOC_CONF": "expandable_segments:true"}, True), + ({"PYTORCH_ALLOC_CONF": "expandable_segments:1"}, True), + ({"PYTORCH_ALLOC_CONF": "expandable_segments:False"}, False), + ({"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}, True), + # ROCm users set this one; it was previously not looked at. + ({"PYTORCH_HIP_ALLOC_CONF": "expandable_segments:True"}, True), + ( + {"PYTORCH_ALLOC_CONF": "max_split_size_mb:512,expandable_segments:True"}, + True, + ), + ({"PYTORCH_ALLOC_CONF": "max_split_size_mb:512"}, False), + ({"PYTORCH_ALLOC_CONF": "not_expandable_segments_really:True"}, False), + ({"PYTORCH_ALLOC_CONF": ""}, False), + # torch's precedence: the accelerator-agnostic variable wins. + ( + { + "PYTORCH_ALLOC_CONF": "expandable_segments:False", + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + }, + False, + ), + ], +) +def test_expandable_segments_from_env(environ, expected): + assert alloc_conf.expandable_segments_enabled_from_env(environ) is expected + + +@pytest.mark.parametrize( + "conf,key,expected", + [ + ("", ES, False), + ("expandable_segments:True", ES, True), + ("expandable_segments:TRUE", ES, True), + ("expandable_segments:1", ES, True), + ("expandable_segments:0", ES, False), + (" expandable_segments : True ", ES, True), + ("max_split_size_mb:512", ES, False), + ("max_split_size_mb:512", "max_split_size_mb", False), + ("garbage_collection_threshold:0.9", ES, False), + ], +) +def test_conf_flag_enabled(conf, key, expected): + assert alloc_conf.conf_flag_enabled(conf, key) is expected + + +@pytest.mark.parametrize( + "conf,enabled,expected", + [ + ("", False, "expandable_segments:False"), + ("expandable_segments:True", False, "expandable_segments:False"), + ("expandable_segments:False", True, "expandable_segments:True"), + ( + "max_split_size_mb:512,expandable_segments:True", + False, + "max_split_size_mb:512,expandable_segments:False", + ), + ( + "garbage_collection_threshold:0.9,max_split_size_mb:512", + False, + "garbage_collection_threshold:0.9,max_split_size_mb:512," + "expandable_segments:False", + ), + ], +) +def test_with_conf_flag_preserves_other_fields(conf, enabled, expected): + assert alloc_conf.with_conf_flag(conf, ES, enabled) == expected + + +def test_with_conf_flag_round_trips(): + """The regression this guards: torch resets every option absent from the + string it is handed, so a one-field write drops the rest permanently.""" + conf = ( + "max_split_size_mb:512,expandable_segments:True," + "garbage_collection_threshold:0.9" + ) + off = alloc_conf.with_conf_flag(conf, ES, False) + assert "max_split_size_mb:512" in off + assert "garbage_collection_threshold:0.9" in off + assert alloc_conf.with_conf_flag(off, ES, True) == conf + + +def test_alloc_conf_from_env_precedence(): + assert ( + alloc_conf.alloc_conf_from_env( + { + "PYTORCH_CUDA_ALLOC_CONF": "max_split_size_mb:256", + "PYTORCH_HIP_ALLOC_CONF": "max_split_size_mb:128", + } + ) + == "max_split_size_mb:256" + ) + assert alloc_conf.alloc_conf_from_env({}) == "" + + +def test_current_alloc_conf_prefers_live_state(monkeypatch): + monkeypatch.setattr(alloc_conf, "live_alloc_conf", lambda: "max_split_size_mb:64") + monkeypatch.setattr(alloc_conf, "alloc_conf_from_env", lambda: "stale:True") + assert alloc_conf.current_alloc_conf() == "max_split_size_mb:64" + + monkeypatch.setattr(alloc_conf, "live_alloc_conf", lambda: None) + assert alloc_conf.current_alloc_conf() == "stale:True" + + +def test_expandable_segments_enabled_returns_none_when_unreadable(monkeypatch): + """'cannot read' must stay distinct from 'disabled'.""" + + class _NoSnapshot: + @staticmethod + def _snapshot(): + raise RuntimeError("no allocator here") + + monkeypatch.setattr(alloc_conf.torch.cuda, "memory", _NoSnapshot) + assert alloc_conf.expandable_segments_enabled() is None diff --git a/vllm/device_allocator/alloc_conf.py b/vllm/device_allocator/alloc_conf.py new file mode 100644 index 0000000000..405a4ecac3 --- /dev/null +++ b/vllm/device_allocator/alloc_conf.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Helpers for reading and editing PyTorch's allocator configuration. + +Torch's allocator config is parsed as a whole: ``parseArgs`` resets every +option that is not explicitly present in the string it is handed. Writing a +single field therefore silently drops the rest of the user's configuration, +for the remainder of the process: + + >>> # PYTORCH_ALLOC_CONF=expandable_segments:True,max_split_size_mb:512 + >>> torch._C._accelerator_setAllocatorSettings("expandable_segments:False") + >>> # max_split_size is now SIZE_MAX, not 512 MiB -- and writing + >>> # "expandable_segments:True" back does not bring it home either. + +So anything that wants to toggle one option has to read the current config, +flip just that field, and write the whole string back. +""" + +import os + +import torch + +# Checked in torch's own precedence order: the accelerator-agnostic variable +# wins, then the CUDA one, then the HIP one (which is what ROCm users set). +ALLOC_CONF_ENV_VARS = ( + "PYTORCH_ALLOC_CONF", + "PYTORCH_CUDA_ALLOC_CONF", + "PYTORCH_HIP_ALLOC_CONF", +) + +EXPANDABLE_SEGMENTS = "expandable_segments" + + +def alloc_conf_from_env( + environ: "os._Environ[str] | dict[str, str] | None" = None, +) -> str: + """The allocator config string the process was started with, if any.""" + env = os.environ if environ is None else environ + for name in ALLOC_CONF_ENV_VARS: + conf = env.get(name) + if conf: + return conf + return "" + + +def live_alloc_conf() -> str | None: + """The allocator config string currently in effect, or ``None``. + + Read from the live allocator rather than the environment, so that a + *runtime* write by another component is visible. + """ + try: + settings = torch.cuda.memory._snapshot().get("allocator_settings") + if settings is not None: + for name in ALLOC_CONF_ENV_VARS: + conf = settings.get(name) + if conf is not None: + return conf + except Exception: # noqa: BLE001 + pass + return None + + +def current_alloc_conf() -> str: + """The live allocator config, falling back to the environment.""" + live = live_alloc_conf() + return live if live is not None else alloc_conf_from_env() + + +def conf_flag_enabled(conf: str, key: str) -> bool: + """Parse ``:`` out of an allocator config string.""" + for field in conf.split(","): + name, _, value = field.partition(":") + if name.strip() == key: + return value.strip().lower() in ("true", "1") + return False + + +def with_conf_flag(conf: str, key: str, enabled: bool) -> str: + """Return ``conf`` with only ``key`` flipped, every other field verbatim.""" + want = f"{key}:{'True' if enabled else 'False'}" + fields = [f for f in (f.strip() for f in conf.split(",")) if f] + for i, field in enumerate(fields): + if field.partition(":")[0].strip() == key: + fields[i] = want + return ",".join(fields) + fields.append(want) + return ",".join(fields) + + +def expandable_segments_enabled_from_env( + environ: "os._Environ[str] | dict[str, str] | None" = None, +) -> bool: + return conf_flag_enabled(alloc_conf_from_env(environ), EXPANDABLE_SEGMENTS) + + +def expandable_segments_enabled() -> bool | None: + """Live ``expandable_segments`` state, or ``None`` if it cannot be read. + + ``None`` is deliberately distinct from ``False``: the environment can never + reflect a runtime write, so a caller verifying a write it just issued must + not read "cannot tell" as "the write did not take". + """ + try: + settings = torch.cuda.memory._snapshot().get("allocator_settings") + if settings is not None and EXPANDABLE_SEGMENTS in settings: + return bool(settings[EXPANDABLE_SEGMENTS]) + except Exception: # noqa: BLE001 + pass + return None + + +def set_alloc_conf(conf: str) -> None: + """Write a complete allocator config string.""" + setter = getattr(torch._C, "_accelerator_setAllocatorSettings", None) + if setter is not None: + setter(conf) + return + torch.cuda.memory._set_allocator_settings(conf) diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index 63d4a1f4e0..2e5612e5fa 100644 --- a/vllm/device_allocator/cumem.py +++ b/vllm/device_allocator/cumem.py @@ -10,7 +10,6 @@ # the only successful approach is to call cuda driver API in C. import atexit import gc -import os from collections.abc import Callable, Iterator from contextlib import contextmanager from typing import Any @@ -18,6 +17,13 @@ from typing import Any import torch from vllm.device_allocator import AllocationData, HandleType +from vllm.device_allocator.alloc_conf import ( + EXPANDABLE_SEGMENTS, + conf_flag_enabled, + current_alloc_conf, + set_alloc_conf, + with_conf_flag, +) from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.system_utils import find_loaded_library @@ -371,13 +377,17 @@ class CuMemAllocator: # Expandable segments are incompatible with the memory pool used for # sleep mode (see https://github.com/pytorch/pytorch/issues/147851). - # If the user has enabled expandable segments via - # PYTORCH_CUDA_ALLOC_CONF, temporarily disable them for the duration - # of the memory pool context and restore on exit. - conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "") - expandable_was_enabled = "expandable_segments:True" in conf + # If the user has enabled expandable segments, temporarily disable + # them for the duration of the memory pool context and restore on + # exit. The whole config string is rewritten, not just this one + # field: torch resets every option that is absent from the string it + # is handed, so a bare "expandable_segments:False" would also drop + # max_split_size_mb, garbage_collection_threshold and + # roundup_power2_divisions for the rest of the process. + prev_conf = current_alloc_conf() + expandable_was_enabled = conf_flag_enabled(prev_conf, EXPANDABLE_SEGMENTS) if expandable_was_enabled: - torch.cuda.memory._set_allocator_settings("expandable_segments:False") + set_alloc_conf(with_conf_flag(prev_conf, EXPANDABLE_SEGMENTS, False)) old_tag = self.current_tag self.current_tag = tag @@ -411,7 +421,7 @@ class CuMemAllocator: finally: self.current_tag = old_tag if expandable_was_enabled: - torch.cuda.memory._set_allocator_settings("expandable_segments:True") + set_alloc_conf(prev_conf) def get_current_usage(self) -> int: """Get the total number of bytes allocated in the memory pool."""