[fix](mooncake): init per-thread accelerator context for buffer registe…

tests/test_mooncake_client_thread_context.py

@@ -0,0 +1,105 @@
+# Copyright 2025 The TransferQueue Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the per-thread NPU device binding used by MooncakeStoreClient.
+
+Mooncake's Ascend transport is bound to a thread-local ACL context, but
+register/unregister run on per-call ``ThreadPoolExecutor`` workers. These tests
+stay hardware-free by faking the accelerator.
+"""
+
+from transfer_queue.storage.clients import mooncake_client as mc
+
+
+class _FakeNpu:
+    def __init__(self, current=5, initialized=True):
+        self._current = current
+        self._initialized = initialized
+        self.set_devices = []
+        self.current_device_calls = 0
+
+    def is_initialized(self):
+        return self._initialized
+
+    def current_device(self):
+        self.current_device_calls += 1
+        return self._current
+
+    def set_device(self, device):
+        self.set_devices.append(device)
+
+
+def _new_client(npu_device):
+    client = object.__new__(mc.MooncakeStoreClient)
+    client._npu_device = npu_device
+    return client
+
+
+def test_resolve_npu_device_captures_owning_thread_device(monkeypatch):
+    fake_npu = _FakeNpu(current=3)
+    monkeypatch.setattr(mc, "NPU_IMPORTED", True)
+    monkeypatch.setattr(mc.torch, "npu", fake_npu, raising=False)
+
+    assert mc._resolve_npu_device() == 3
+    assert fake_npu.current_device_calls == 1
+
+
+def test_resolve_npu_device_none_without_npu(monkeypatch):
+    monkeypatch.setattr(mc, "NPU_IMPORTED", False)
+
+    assert mc._resolve_npu_device() is None
+
+
+def test_resolve_npu_device_none_when_not_initialized(monkeypatch):
+    fake_npu = _FakeNpu(initialized=False)
+    monkeypatch.setattr(mc, "NPU_IMPORTED", True)
+    monkeypatch.setattr(mc.torch, "npu", fake_npu, raising=False)
+
+    assert mc._resolve_npu_device() is None
+    assert fake_npu.current_device_calls == 0
+
+
+def test_worker_binds_device_captured_at_construction(monkeypatch):
+    # A context-less worker reports device 0; the client captured 3 in __init__.
+    fake_npu = _FakeNpu(current=0)
+    monkeypatch.setattr(mc.torch, "npu", fake_npu, raising=False)
+
+    _new_client(npu_device=3)._bind_npu_device()
+
+    assert fake_npu.set_devices == [3]
+    assert fake_npu.current_device_calls == 0
+
+
+def test_worker_bind_noop_without_npu_device(monkeypatch):
+    fake_npu = _FakeNpu()
+    monkeypatch.setattr(mc.torch, "npu", fake_npu, raising=False)
+
+    _new_client(npu_device=None)._bind_npu_device()
+
+    assert fake_npu.set_devices == []
+
+
+def test_worker_bind_logs_and_swallows_failure(monkeypatch):
+    class _FailingNpu(_FakeNpu):
+        def set_device(self, device):
+            raise RuntimeError("set_device failed")
+
+    fake_npu = _FailingNpu()
+    monkeypatch.setattr(mc.torch, "npu", fake_npu, raising=False)
+    warnings = []
+    monkeypatch.setattr(mc.logger, "warning", lambda *a, **k: warnings.append((a, k)))
+
+    _new_client(npu_device=3)._bind_npu_device()  # must not raise
+
+    assert warnings and warnings[0][1].get("exc_info") is True

transfer_queue/storage/clients/mooncake_client.py

@@ -33,6 +33,34 @@
 
 logger = get_logger(__name__)
 
+try:
+    import torch_npu  # noqa: F401
+
+    NPU_IMPORTED: bool = True
+except Exception:  # pragma: no cover - accelerator optional
+    NPU_IMPORTED = False
+
+
+def _resolve_npu_device() -> int | None:
+    """Return the NPU device this process is bound to, or ``None`` if it uses no NPU.
+
+    Ascend ACL contexts are thread-local, so this must run on the thread that owns
+    the device (``__init__``). Resolving it on a context-less pool worker would
+    always report device 0. Gating on ``is_initialized()`` (rather than
+    ``is_available()``) keeps a process that never touches the NPU from allocating
+    device memory and contending for device 0.
+    """
+    if not NPU_IMPORTED:
+        return None
+    try:
+        if not torch.npu.is_initialized():
+            return None
+        return int(torch.npu.current_device())
+    except Exception:
+        logger.warning("Failed to resolve the current NPU device.", exc_info=True)
+        return None
+
+
 MOONCAKE_STORE_IMPORTED: bool = True
 try:
     from mooncake.store import MooncakeDistributedStore, ReplicateConfig
@@ -121,6 +149,9 @@ def __init__(self, config: dict[str, Any]):
             hard_pin = not offload_enabled
         self.replica_config.with_hard_pin = bool(hard_pin)
 
+        # Capture on the owning thread; pool workers reuse it via _bind_npu_device.
+        self._npu_device = _resolve_npu_device()
+
         self._store = MooncakeDistributedStore()
         ret = self._store.setup(
             self.local_hostname,
@@ -134,6 +165,19 @@ def __init__(self, config: dict[str, Any]):
         if ret != 0:
             raise RuntimeError(f"Mooncake store setup failed with error code: {ret}")
 
+    def _bind_npu_device(self) -> None:
+        """``ThreadPoolExecutor`` initializer: bind a worker to ``self._npu_device``."""
+        if self._npu_device is None:
+            return
+        try:
+            torch.npu.set_device(self._npu_device)
+        except Exception:
+            logger.warning(
+                f"Failed to bind NPU device {self._npu_device} on a worker thread; "
+                "Mooncake buffer register/unregister may fail.",
+                exc_info=True,
+            )
+
     def put(self, keys: list[str], values: list[Any]) -> list[dict | None]:
         """Stores multiple key-value pairs to MooncakeStore.
 
@@ -173,7 +217,10 @@ def put(self, keys: list[str], values: list[Any]) -> list[dict | None]:
 
         tensor_futures: list[Future[None]] = []
         bytes_futures: list[Future[list[int]]] = []
-        with ThreadPoolExecutor(max_workers=MAX_BATCH_WORKER_THREADS) as executor:
+        with ThreadPoolExecutor(
+            max_workers=MAX_BATCH_WORKER_THREADS,
+            initializer=self._bind_npu_device,
+        ) as executor:
             if not use_gdr_path:
                 for i in range(0, len(tensor_keys), BATCH_SIZE_LIMIT):
                     batch_keys = tensor_keys[i : i + BATCH_SIZE_LIMIT]
@@ -373,7 +420,10 @@ def get(
                 results[idx] = val
 
         futures = []
-        with ThreadPoolExecutor(max_workers=MAX_BATCH_WORKER_THREADS) as executor:
+        with ThreadPoolExecutor(
+            max_workers=MAX_BATCH_WORKER_THREADS,
+            initializer=self._bind_npu_device,
+        ) as executor:
             for i in range(0, len(cpu_tensor_indices), BATCH_SIZE_LIMIT):
                 batch_keys = cpu_tensor_keys[i : i + BATCH_SIZE_LIMIT]
                 batch_shapes = cpu_tensor_shapes[i : i + BATCH_SIZE_LIMIT]