studio: tighten torchao Windows-ROCm comments and test docstrings (#6610)

studio/backend/core/rag/embeddings.py

@@ -50,14 +50,10 @@ _torchao_stub_done = False
 
 
 def _install_torchao_stub_once() -> None:
-    """Neutralize torchao before the first sentence-transformers import.
-
-    transformers.quantizers imports torchao, which loads torch's c10d
-    distributed backend at module level; the AMD Windows ROCm wheels omit it
-    (no RCCL), so the import aborts and silently drops the ST embedder to the
-    llama-server fallback. The training/export workers install this stub at
-    process start, but the embedder runs in the main backend process, which
-    otherwise never does. No-op off Windows ROCm. Runs once (under ``_lock``)."""
+    """Neutralize torchao before importing sentence-transformers. On Windows ROCm,
+    torchao (pulled in by transformers.quantizers) imports an absent c10d backend
+    and aborts, dropping the embedder to llama-server. Workers stub it too; the
+    embedder runs in the main process. No-op elsewhere; runs once under ``_lock``."""
     global _torchao_stub_done
     if _torchao_stub_done:
         return

studio/backend/tests/test_torchao_select.py

@@ -72,10 +72,9 @@ def test_default_spec_matches_table(monkeypatch):
 
 
 def test_skips_torchao_on_windows_rocm():
-    """The overrides step must skip torchao on Windows ROCm. There is no working
-    torchao build there: it loads torch's c10d distributed backend at import,
-    which the AMD Windows wheels omit, so `import torchao` raises and takes
-    transformers.quantizers with it. Studio stubs torchao at runtime instead."""
+    """The overrides step must skip torchao on Windows ROCm: no working build exists
+    there (it imports an absent c10d backend and crashes transformers.quantizers),
+    so the installer skips it and relies on the runtime stub instead."""
     source = _INSTALL_SCRIPT.read_text(encoding = "utf-8")
     # Branches on the Windows-ROCm marker set by _ensure_rocm_torch ...
     assert "elif _rocm_windows_torch_installed:" in source

studio/install_python_stack.py

@@ -2234,16 +2234,14 @@ def install_python_stack() -> int:
     # 4. Overrides (torchao) -- force-reinstall. The torchao version is chosen to
     #    match the torch installed in the venv so its C++ extensions load (see
     #    _select_torchao_spec). Skip when torch is unavailable (e.g. Intel Mac
-    #    GGUF-only mode): torchao requires torch. Also skipped on Windows ROCm,
-    #    which has no working torchao build (see below).
+    #    GGUF-only mode): torchao requires torch. Also skipped on Windows ROCm
+    #    (no working build; see below).
     if NO_TORCH:
         _progress("dependency overrides (skipped, no torch)")
     elif _rocm_windows_torch_installed:
-        # torchao has no working Windows ROCm build: it loads torch's c10d
-        # distributed backend at import, which the AMD Windows wheels omit (no
-        # RCCL), so `import torchao` raises and takes transformers.quantizers
-        # down with it. Studio stubs torchao at runtime (core/_torchao_stub.py),
-        # so installing it only ships a package that crashes on import -- skip it.
+        # No working Windows ROCm torchao build: it imports an absent c10d backend
+        # and crashes transformers.quantizers. Studio stubs it at runtime, so
+        # installing it only ships a package that crashes on import -- skip it.
         _progress("dependency overrides (skipped, Windows ROCm)")
         _safe_print("   Windows ROCm -- skipping torchao (no working build; stubbed at runtime)")
     else:

tests/studio/install/test_rocm_support.py

@@ -2332,10 +2332,9 @@ class TestWorkerWindowsRocmPatches:
         assert "install_torchao_windows_rocm_stub()" in source
 
     def test_embedder_calls_shared_torchao_stub(self):
-        """rag/embeddings.py must install the stub before the first sentence-
-        transformers import. The embedder runs in the main backend process (not a
-        stubbed worker), so without this transformers -> torchao crashes on Windows
-        ROCm and the ST embedder silently drops to the llama-server fallback."""
+        """embeddings.py must install the stub before importing sentence-transformers:
+        it runs in the main process (not a stubbed worker), so otherwise transformers
+        -> torchao crashes on Windows ROCm and the embedder drops to llama-server."""
         source = _EMBEDDINGS_PATH.read_text(encoding = "utf-8")
         assert "install_torchao_windows_rocm_stub()" in source