| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Helpers for reading and editing PyTorch's allocator configuration. |
| 4 | |
| 5 | Torch's allocator config is parsed as a whole: ``parseArgs`` resets every |
| 6 | option that is not explicitly present in the string it is handed. Writing a |
| 7 | single field therefore silently drops the rest of the user's configuration, |
| 8 | for the remainder of the process: |
| 9 | |
| 10 | >>> # PYTORCH_ALLOC_CONF=expandable_segments:True,max_split_size_mb:512 |
| 11 | >>> torch._C._accelerator_setAllocatorSettings("expandable_segments:False") |
| 12 | >>> # max_split_size is now SIZE_MAX, not 512 MiB -- and writing |
| 13 | >>> # "expandable_segments:True" back does not bring it home either. |
| 14 | |
| 15 | So anything that wants to toggle one option has to read the current config, |
| 16 | flip just that field, and write the whole string back. |
| 17 | """ |
| 18 | |
| 19 | import os |
| 20 | |
| 21 | import torch |
| 22 | |
| 23 | # Checked in torch's own precedence order: the accelerator-agnostic variable |
| 24 | # wins, then the CUDA one, then the HIP one (which is what ROCm users set). |
| 25 | ALLOC_CONF_ENV_VARS = ( |
| 26 | "PYTORCH_ALLOC_CONF", |
| 27 | "PYTORCH_CUDA_ALLOC_CONF", |
| 28 | "PYTORCH_HIP_ALLOC_CONF", |
| 29 | ) |
| 30 | |
| 31 | EXPANDABLE_SEGMENTS = "expandable_segments" |
| 32 | |
| 33 | |
| 34 | def alloc_conf_from_env( |
| 35 | environ: "os._Environ[str] | dict[str, str] | None" = None, |
| 36 | ) -> str: |
| 37 | """The allocator config string the process was started with, if any.""" |
| 38 | env = os.environ if environ is None else environ |
| 39 | for name in ALLOC_CONF_ENV_VARS: |
| 40 | conf = env.get(name) |
| 41 | if conf: |
| 42 | return conf |
| 43 | return "" |
| 44 | |
| 45 | |
| 46 | def live_alloc_conf() -> str | None: |
| 47 | """The allocator config string currently in effect, or ``None``. |
| 48 | |
| 49 | Read from the live allocator rather than the environment, so that a |
| 50 | *runtime* write by another component is visible. |
| 51 | """ |
| 52 | try: |
| 53 | settings = torch.cuda.memory._snapshot().get("allocator_settings") |
| 54 | if settings is not None: |
| 55 | for name in ALLOC_CONF_ENV_VARS: |
| 56 | conf = settings.get(name) |
| 57 | if conf is not None: |
| 58 | return conf |
| 59 | except Exception: # noqa: BLE001 |
| 60 | pass |
| 61 | return None |
| 62 | |
| 63 | |
| 64 | def current_alloc_conf() -> str: |
| 65 | """The live allocator config, falling back to the environment.""" |
| 66 | live = live_alloc_conf() |
| 67 | return live if live is not None else alloc_conf_from_env() |
| 68 | |
| 69 | |
| 70 | def conf_flag_enabled(conf: str, key: str) -> bool: |
| 71 | """Parse ``<key>:<bool>`` out of an allocator config string.""" |
| 72 | for field in conf.split(","): |
| 73 | name, _, value = field.partition(":") |
| 74 | if name.strip() == key: |
| 75 | return value.strip().lower() in ("true", "1") |
| 76 | return False |
| 77 | |
| 78 | |
| 79 | def with_conf_flag(conf: str, key: str, enabled: bool) -> str: |
| 80 | """Return ``conf`` with only ``key`` flipped, every other field verbatim.""" |
| 81 | want = f"{key}:{'True' if enabled else 'False'}" |
| 82 | fields = [f for f in (f.strip() for f in conf.split(",")) if f] |
| 83 | for i, field in enumerate(fields): |
| 84 | if field.partition(":")[0].strip() == key: |
| 85 | fields[i] = want |
| 86 | return ",".join(fields) |
| 87 | fields.append(want) |
| 88 | return ",".join(fields) |
| 89 | |
| 90 | |
| 91 | def expandable_segments_enabled_from_env( |
| 92 | environ: "os._Environ[str] | dict[str, str] | None" = None, |
| 93 | ) -> bool: |
| 94 | return conf_flag_enabled(alloc_conf_from_env(environ), EXPANDABLE_SEGMENTS) |
| 95 | |
| 96 | |
| 97 | def expandable_segments_enabled() -> bool | None: |
| 98 | """Live ``expandable_segments`` state, or ``None`` if it cannot be read. |
| 99 | |
| 100 | ``None`` is deliberately distinct from ``False``: the environment can never |
| 101 | reflect a runtime write, so a caller verifying a write it just issued must |
| 102 | not read "cannot tell" as "the write did not take". |
| 103 | """ |
| 104 | try: |
| 105 | settings = torch.cuda.memory._snapshot().get("allocator_settings") |
| 106 | if settings is not None and EXPANDABLE_SEGMENTS in settings: |
| 107 | return bool(settings[EXPANDABLE_SEGMENTS]) |
| 108 | except Exception: # noqa: BLE001 |
| 109 | pass |
| 110 | return None |
| 111 | |
| 112 | |
| 113 | def set_alloc_conf(conf: str) -> None: |
| 114 | """Write a complete allocator config string.""" |
| 115 | setter = getattr(torch._C, "_accelerator_setAllocatorSettings", None) |
| 116 | if setter is not None: |
| 117 | setter(conf) |
| 118 | return |
| 119 | torch.cuda.memory._set_allocator_settings(conf) |