[Bugfix] Don't drop the rest of the allocator config when toggling expandable segments (#57982)

Signed-off-by: Oxana Korzh <okorzh@amd.com>

.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

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

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 ``<key>:<bool>`` 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)

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."""