diff --git a/specforge/runtime/data_plane/mooncake_store.py b/specforge/runtime/data_plane/mooncake_store.py index 5ef4e75d3..fe5e1535d 100644 --- a/specforge/runtime/data_plane/mooncake_store.py +++ b/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: diff --git a/tests/test_runtime/test_mooncake_store.py b/tests/test_runtime/test_mooncake_store.py index 3c1590b9f..1f0a94b05 100644 --- a/tests/test_runtime/test_mooncake_store.py +++ b/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",