diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 0968248679..6768f50b5c 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -955,13 +955,36 @@ def _determine_attention_impl_for_gpu_estimate(config) -> str: import types as _types if _sys.platform == "win32": + # Dummy class for any name torch.distributed tries to import from these stubs + class _Dummy: + pass + for _c10d_name in ( "torch._C._distributed_c10d", "torch._C._distributed_autograd", "torch._C._distributed_rpc", ): if _c10d_name not in _sys.modules: - _sys.modules[_c10d_name] = _types.ModuleType(_c10d_name) + _stub = _types.ModuleType(_c10d_name) + # torch.distributed imports these names from _distributed_c10d; + # provide no-op dummies so the import doesn't raise AttributeError. + for _sym in ( + "FakeProcessGroup", + "ProcessGroup", + "Work", + "Store", + "PrefixStore", + "FileStore", + "TCPStore", + "HashStore", + "Reducer", + "Logger", + "DistributedDebugLevel", + "GradBucket", + "BuiltinCommHookType", + ): + setattr(_stub, _sym, _Dummy) + _sys.modules[_c10d_name] = _stub try: import torch.distributed as _td diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 7eeb58281a..d602602c6c 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -2417,6 +2417,28 @@ class TestServerStartupRocmFixes: source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8") assert 'platform == "win32"' in source or "win32" in source + def test_hardware_py_stub_exposes_fake_process_group(self): + """hardware.py stub must set FakeProcessGroup so torch.distributed doesn't raise AttributeError.""" + source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8") + assert "FakeProcessGroup" in source + + def test_hardware_py_stub_exposes_process_group(self): + """hardware.py stub must set ProcessGroup on the c10d stub.""" + source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8") + assert "ProcessGroup" in source + + def test_hardware_py_stub_uses_setattr_for_symbols(self): + """hardware.py must use setattr to populate stub symbols dynamically.""" + source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8") + assert "setattr" in source + + def test_hardware_py_stub_all_c10d_siblings_covered(self): + """hardware.py must stub all three torch._C._distributed_* submodules.""" + source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8") + assert "_distributed_c10d" in source + assert "_distributed_autograd" in source + assert "_distributed_rpc" in source + if __name__ == "__main__": pytest.main([__file__, "-v"])