[Bugfix] Preserve user PYTORCH_CUDA_ALLOC_CONF keys in _scoped_allocator_max_split

[Sahil170595]

Purpose

_scoped_allocator_max_split (added in #41268) temporarily sets max_split_size_mb:20 around load_model to reduce allocator fragmentation. PyTorch's AcceleratorAllocatorConfig::parseArgs (c10/core/AllocatorConfig.cpp @ v2.13.0) resets max_split_size_mb, garbage_collection_threshold, and roundup_power2_divisions to defaults at the start of every call:

  // The following option will be reset to its default value if not explicitly
  // set each time.
  max_split_size_ = std::numeric_limits<size_t>::max();
  roundup_power2_divisions_.assign(kRoundUpPowerOfTwoIntervals, 0);
  garbage_collection_threshold_ = 0;

So for a user running with, e.g. PYTORCH_CUDA_ALLOC_CONF=garbage_collection_threshold:0.8,roundup_power2_divisions:4, the scoped override silently disabled both keys during weight loading, and the restore path — which passed back only max_split_size_mb — wiped them for the rest of the process lifetime. Reproduced against real torch:

$ PYTORCH_CUDA_ALLOC_CONF=garbage_collection_threshold:0.8,roundup_power2_divisions:4 python
>>> torch._C._accelerator_getAllocatorSettings()
'garbage_collection_threshold:0.8,roundup_power2_divisions:4'
>>> torch._C._accelerator_setAllocatorSettings("max_split_size_mb:64")   # entry
>>> torch._C._accelerator_getAllocatorSettings()
'max_split_size_mb:64'                                                   # user keys gone
>>> torch._C._accelerator_setAllocatorSettings("max_split_size_mb:17592186044415")  # restore
>>> torch._C._accelerator_getAllocatorSettings()
'max_split_size_mb:17592186044415'                                       # gone permanently

The same held for settings supplied via the newer PYTORCH_ALLOC_CONF or set programmatically, which the env-var-based restore missed entirely (their max_split_size_mb was also replaced by the SIZE_MAX // 2**20 approximation instead of the user's value).

The fix: snapshot the full settings string with torch._C._accelerator_getAllocatorSettings() (it returns the stored last_allocator_settings_ string, so the round-trip is exact), append the max_split_size_mb override for the scoped region (keys are parsed sequentially, so the appended override wins), and restore the snapshot on exit. parseArgs("") still performs the defaults reset, so the empty-snapshot case restores true defaults. Also drops the now-unused regex import.

Found while auditing worker-startup state for #48518 (see my comment there); this PR is independent of that issue's L2 claim, which did not reproduce. Checked open PRs for PYTORCH_CUDA_ALLOC_CONF / max_split_size_mb / allocator-settings terms: existing PRs (#43923, #42405, #48243, #49151) touch expandable_segments and CuMem pool behavior; none address this.

Test Plan

New test_scoped_allocator_max_split_preserves_alloc_conf in tests/v1/worker/test_gpu_worker.py: asserts user keys survive inside the scoped region alongside the override, the exact settings string is restored on exit, and the empty-settings case round-trips. Runs against the real torch allocator-settings bindings; no GPU allocation required.

python -m pytest tests/v1/worker/test_gpu_worker.py -v

Test Result

tests/v1/worker/test_gpu_worker.py::test_startup_plan_fingerprint_sensitivity PASSED [ 33%]
tests/v1/worker/test_gpu_worker.py::test_startup_plan_apply_gate PASSED  [ 66%]
tests/v1/worker/test_gpu_worker.py::test_scoped_allocator_max_split_preserves_alloc_conf PASSED [100%]
======================== 3 passed, 14 warnings in 2.54s ========================

(WSL2, Python 3.12.3, torch 2.13.0+cu130, editable VLLM_USE_PRECOMPILED install.) Negative control: with vllm/v1/worker/gpu_worker.py reverted to main and the new test kept, the test fails; with the fix it passes. ruff check / ruff format --check (pre-commit pin v0.14.0) pass on both changed files.

Model evaluation: not applicable — no change to model output, kernel selection, or numerics; only which allocator-config keys survive load_model.

This change was developed with AI assistance (Claude Fable 5). I have reviewed every changed line and run the tests above.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Sahil170595.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[Sahil170595]

Self-review summary before closing, since this is now superseded.

What this fixed: _scoped_allocator_max_split temporarily sets max_split_size_mb:20 around load_model to reduce allocator fragmentation. It applied the override with torch._C._accelerator_setAllocatorSettings("max_split_size_mb:20") and, on exit, restored only max_split_size_mb.

Why it was a bug: torch's parseArgs resets max_split_size_mb, garbage_collection_threshold and roundup_power2_divisions to their defaults on every call unless they appear in the string. So a user running with e.g. garbage_collection_threshold:0.8,roundup_power2_divisions:4 lost both keys during weight loading, and the restore wiped them for the rest of the process. This PR snapshots the full settings string, appends the override (keys parse in order, so it wins), and restores the snapshot on exit.

Status: #58430 landed the same snapshot / prepend / restore logic today while extending the scoped limit to the memory-profiling pass, and it covers cases this PR did not: ROCm, non-native allocator backends, the exception path (the override is applied inside the try), trailing commas in the user's string, and older torch without _accelerator_getAllocatorSettings. The trailing-comma case was a real gap here: torch accepts garbage_collection_threshold:0.8,, but this PR would then build garbage_collection_threshold:0.8,,max_split_size_mb:20, which torch rejects with ValueError: Unrecognized key ',', and since the call sat outside the try it would have failed worker startup (checked against torch 2.13). Closing as superseded.