fix: stub is_torchelastic_launched on torch.distributed for Windows ROCm

resolve_attention_implementation calls is_torchelastic_launched() which
does not exist in the incomplete torch.distributed shipped with the
Windows ROCm wheel, causing a warning on every model config load in the
server process. Add it to the stub table alongside the four helpers
already patched in _determine_attention_impl_for_gpu_estimate.

Also adds two tests: one confirming the new stub and one confirming all
five core distributed helpers are covered.
This commit is contained in:
LeoBorcherding 2026-05-15 15:01:32 -05:00
commit 75ef5993ce
2 changed files with 23 additions and 0 deletions

View file

@ -958,6 +958,7 @@ def _determine_attention_impl_for_gpu_estimate(config) -> str:
("is_available", lambda: False),
("get_rank", lambda: 0),
("get_world_size", lambda: 1),
("is_torchelastic_launched", lambda: False),
):
if not hasattr(_td, _attr):
setattr(_td, _attr, _stub)

View file

@ -846,6 +846,28 @@ class TestHardwareRocmFlag:
assert '"cuda"' in func_body
assert '"rocm"' in func_body
def test_distributed_stubs_cover_is_torchelastic_launched(self):
"""_determine_attention_impl_for_gpu_estimate must stub is_torchelastic_launched.
resolve_attention_implementation calls is_torchelastic_launched() on
Windows ROCm where torch.distributed ships without that helper, causing
a warning: 'module torch.distributed has no attribute is_torchelastic_launched'.
"""
hw_path = (
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
)
source = hw_path.read_text(encoding = "utf-8")
assert "is_torchelastic_launched" in source
def test_distributed_stubs_cover_core_helpers(self):
"""_determine_attention_impl_for_gpu_estimate must stub the four core distributed helpers."""
hw_path = (
PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
)
source = hw_path.read_text(encoding = "utf-8")
for attr in ("is_initialized", "is_available", "get_rank", "get_world_size"):
assert attr in source, f"distributed stub for '{attr}' missing from hardware.py"
# =============================================================================
# TEST: tokenizer_utils.py -- error message