fix(runtime): bind NPU device before Mooncake transport setup

specforge/runtime/data_plane/mooncake_store.py

@@ -54,6 +54,7 @@
 from __future__ import annotations
 
 import logging
+import os
 import threading
 import time
 import uuid
@@ -90,6 +91,64 @@ def __init__(
         self.with_soft_pin = with_soft_pin
 
 
+# Ascend selects visible NPUs through these env vars, mirroring
+# ``CUDA_VISIBLE_DEVICES``. Their presence marks a host as Ascend even before
+# ``torch_npu`` has been imported (e.g. in a capture producer that never runs
+# ``init_distributed``).
+_ASCEND_VISIBLE_DEVICE_ENVS = ("ASCEND_RT_VISIBLE_DEVICES", "ASCEND_VISIBLE_DEVICES")
+
+
+def _ascend_runtime_available() -> bool:
+    """Report whether a usable Ascend NPU runtime is present.
+
+    ``torch.npu`` only exists once ``torch_npu`` is imported, which the canonical
+    trainer does lazily inside ``init_distributed``. A disaggregated capture
+    producer skips that path, so activate ``torch_npu`` here (only on a host that
+    actually selects Ascend devices) to detect the runtime without forcing the
+    import on CUDA/CPU hosts.
+    """
+    if getattr(torch, "npu", None) is None:
+        if not any(os.environ.get(name) for name in _ASCEND_VISIBLE_DEVICE_ENVS):
+            return False
+        try:
+            import torch_npu  # noqa: F401
+        except Exception:
+            return False
+    npu = getattr(torch, "npu", None)
+    try:
+        return npu is not None and bool(npu.is_available())
+    except Exception:  # pragma: no cover - defensive against driver faults
+        return False
+
+
+def _bind_transport_device() -> None:
+    """Bind this process's local NPU before Mooncake installs its transport.
+
+    On Ascend, ``MooncakeDistributedStore.setup()`` installs the
+    ``AscendDirectTransport``, which calls ``aclrtGetDevice`` and fails with
+    ``ACL_ERROR_RT_CONTEXT_NULL`` (107002) when no device context exists for the
+    calling process. The store is constructed at run-assembly time; a trainer
+    (consumer) has already bound its NPU in ``init_distributed``, but a capture
+    *producer* deliberately never initializes an accelerator, so the transport
+    would have no device to allocate its local segment against.
+
+    Bind only for NPU: CUDA's transport defaults to device 0 without an explicit
+    ``set_device`` and the producer intentionally avoids initializing CUDA.
+    ``_bind_local_device`` reads ``LOCAL_RANK``/``RANK`` (set by ``torchrun``), so
+    every rank pins the transport to its own NPU, and it is idempotent with the
+    trainer's earlier binding.
+    """
+    from specforge.utils import get_device_type
+
+    device_type = get_device_type()
+    if device_type == "cuda":
+        return
+    if device_type == "npu" or (device_type == "cpu" and _ascend_runtime_available()):
+        from specforge.distributed import _bind_local_device
+
+        _bind_local_device("npu")
+
+
 def _connect_store(setup_kwargs: Dict[str, Any]) -> Tuple[Any, Any]:
     """Construct a real store and return its required config type."""
     try:
@@ -104,6 +163,10 @@ def _connect_store(setup_kwargs: Dict[str, Any]) -> Tuple[Any, Any]:
             "official wheel (`mooncake-transfer-engine` for CUDA < 13, or "
             "`mooncake-transfer-engine-cuda13` for CUDA >= 13)."
         ) from e
+    # Ascend's transport is installed inside setup() and needs a bound device
+    # context; bind the local accelerator first so the transfer engine can
+    # allocate its local segment (see _bind_transport_device).
+    _bind_transport_device()
     store = MooncakeDistributedStore()
     rc = store.setup(**setup_kwargs)
     if rc is not None and int(rc) != 0:

tests/test_runtime/test_mooncake_store.py

@@ -551,6 +551,181 @@ def test_cross_process_abort_under_lease_defer_is_known_gap(self):
             consumer.get(ref)  # SHOULD raise; currently returns stale -> xfail
 
 
+class TestMooncakeTransportDeviceBinding(unittest.TestCase):
+    """The Ascend transport needs a bound device context before setup().
+
+    Mooncake installs its ``AscendDirectTransport`` inside ``setup()``, which
+    calls ``aclrtGetDevice`` and fails (107002 ACL_ERROR_RT_CONTEXT_NULL) unless
+    the process has already bound its NPU. The store is built at run-assembly
+    time, before ``init_distributed`` binds the rank, so ``_connect_store`` must
+    bind the local accelerator first.
+    """
+
+    def test_bind_is_noop_on_plain_cpu(self):
+        from unittest import mock
+
+        import specforge.distributed as sf_dist
+        import specforge.runtime.data_plane.mooncake_store as mc
+        import specforge.utils as sf_utils
+
+        with (
+            mock.patch.object(sf_utils, "get_device_type", return_value="cpu"),
+            mock.patch.object(mc, "_ascend_runtime_available", return_value=False),
+            mock.patch.object(sf_dist, "_bind_local_device") as bind,
+        ):
+            mc._bind_transport_device()
+        bind.assert_not_called()
+
+    def test_bind_is_noop_on_cuda(self):
+        # CUDA's transport needs no explicit device context, and the producer
+        # deliberately avoids initializing CUDA.
+        from unittest import mock
+
+        import specforge.distributed as sf_dist
+        import specforge.runtime.data_plane.mooncake_store as mc
+        import specforge.utils as sf_utils
+
+        with (
+            mock.patch.object(sf_utils, "get_device_type", return_value="cuda"),
+            mock.patch.object(sf_dist, "_bind_local_device") as bind,
+        ):
+            mc._bind_transport_device()
+        bind.assert_not_called()
+
+    def test_bind_pins_local_npu(self):
+        from unittest import mock
+
+        import specforge.distributed as sf_dist
+        import specforge.runtime.data_plane.mooncake_store as mc
+        import specforge.utils as sf_utils
+
+        with (
+            mock.patch.object(sf_utils, "get_device_type", return_value="npu"),
+            mock.patch.object(sf_dist, "_bind_local_device") as bind,
+        ):
+            mc._bind_transport_device()
+        bind.assert_called_once_with("npu")
+
+    def test_producer_binds_npu_when_torch_npu_not_yet_active(self):
+        # The disaggregated producer skips init_distributed, so get_device_type()
+        # reports "cpu" until torch_npu is imported. On an Ascend host the
+        # transport still needs a bound device.
+        from unittest import mock
+
+        import specforge.distributed as sf_dist
+        import specforge.runtime.data_plane.mooncake_store as mc
+        import specforge.utils as sf_utils
+
+        with (
+            mock.patch.object(sf_utils, "get_device_type", return_value="cpu"),
+            mock.patch.object(mc, "_ascend_runtime_available", return_value=True),
+            mock.patch.object(sf_dist, "_bind_local_device") as bind,
+        ):
+            mc._bind_transport_device()
+        bind.assert_called_once_with("npu")
+
+    def test_ascend_runtime_unavailable_without_env_or_module(self):
+        from unittest import mock
+
+        import specforge.runtime.data_plane.mooncake_store as mc
+
+        # torch.npu absent and no Ascend device env -> do not import torch_npu.
+        with (
+            mock.patch.object(mc.torch, "npu", None, create=True),
+            mock.patch.dict(mc.os.environ, {}, clear=True),
+        ):
+            self.assertFalse(mc._ascend_runtime_available())
+
+    def test_ascend_runtime_available_when_torch_npu_live(self):
+        from unittest import mock
+
+        import specforge.runtime.data_plane.mooncake_store as mc
+
+        fake_npu = mock.Mock()
+        fake_npu.is_available.return_value = True
+        with mock.patch.object(mc.torch, "npu", fake_npu, create=True):
+            self.assertTrue(mc._ascend_runtime_available())
+
+        fake_npu.is_available.return_value = False
+        with mock.patch.object(mc.torch, "npu", fake_npu, create=True):
+            self.assertFalse(mc._ascend_runtime_available())
+
+    def test_ascend_activation_gated_on_visible_devices(self):
+        # torch.npu absent: only attempt to import torch_npu when Ascend devices
+        # are actually selected, and tolerate a failed/absent torch_npu import.
+        import builtins
+        from unittest import mock
+
+        import specforge.runtime.data_plane.mooncake_store as mc
+
+        real_import = builtins.__import__
+
+        def _no_torch_npu(name, *args, **kwargs):
+            if name == "torch_npu":
+                raise ImportError("torch_npu not installed")
+            return real_import(name, *args, **kwargs)
+
+        # Without the Ascend device env, torch_npu must not even be imported.
+        import_calls = []
+
+        def _record_import(name, *args, **kwargs):
+            import_calls.append(name)
+            return _no_torch_npu(name, *args, **kwargs)
+
+        with (
+            mock.patch.object(mc.torch, "npu", None, create=True),
+            mock.patch.dict(mc.os.environ, {}, clear=True),
+            mock.patch.object(builtins, "__import__", side_effect=_record_import),
+        ):
+            self.assertFalse(mc._ascend_runtime_available())
+        self.assertNotIn("torch_npu", import_calls)
+
+        # With the env set but torch_npu unavailable, degrade gracefully.
+        with (
+            mock.patch.object(mc.torch, "npu", None, create=True),
+            mock.patch.dict(
+                mc.os.environ, {"ASCEND_RT_VISIBLE_DEVICES": "0,1"}, clear=True
+            ),
+            mock.patch.object(builtins, "__import__", side_effect=_no_torch_npu),
+        ):
+            self.assertFalse(mc._ascend_runtime_available())
+
+    def test_connect_store_binds_device_before_setup(self):
+        import sys
+        from unittest import mock
+
+        import specforge.runtime.data_plane.mooncake_store as mc
+
+        events = []
+
+        class _FakeDistributedStore:
+            def setup(self, **kwargs):
+                events.append(("setup", kwargs))
+                return 0
+
+        class _FakeReplicateConfig:
+            pass
+
+        fake_module = mock.Mock()
+        fake_module.MooncakeDistributedStore = _FakeDistributedStore
+        fake_module.ReplicateConfig = _FakeReplicateConfig
+
+        with (
+            mock.patch.dict(sys.modules, {"mooncake.store": fake_module}),
+            mock.patch.object(
+                mc,
+                "_bind_transport_device",
+                side_effect=lambda: events.append(("bind", None)),
+            ),
+        ):
+            store, config_type = mc._connect_store({"protocol": "tcp"})
+
+        self.assertIsInstance(store, _FakeDistributedStore)
+        self.assertIs(config_type, _FakeReplicateConfig)
+        # The device must be bound before the transport is installed by setup().
+        self.assertEqual([name for name, _ in events], ["bind", "setup"])
+
+
 @unittest.skipUnless(
     importlib.util.find_spec("mooncake") is not None,
     "mooncake package not installed; real end-to-end store test skipped",