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

tests/test_mooncake_client_thread_context.py

@@ -0,0 +1,119 @@
+# 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 accelerator context 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 and the executor.
+"""
+
+import sys
+
+from transfer_queue.storage.clients import mooncake_client as mc
+
+
+class _FakeStore:
+    def __init__(self):
+        self.registered = []
+        self.unregistered = []
+
+    def register_buffer(self, ptr, size):
+        self.registered.append((ptr, size))
+
+    def unregister_buffer(self, ptr):
+        self.unregistered.append(ptr)
+
+
+class _FakeNpu:
+    def __init__(self, available=True, current=5):
+        self._available = available
+        self._current = current
+        self.set_devices = []
+
+    def is_available(self):
+        return self._available
+
+    def current_device(self):
+        return self._current
+
+    def set_device(self, device):
+        self.set_devices.append(device)
+
+
+def _new_client():
+    client = object.__new__(mc.MooncakeStoreClient)
+    client._store = _FakeStore()
+    client.use_gdr = False
+    client._gdr_staging = None
+    return client
+
+
+def test_ensure_accelerator_context_sets_device_and_warms_once(monkeypatch):
+    fake_npu = _FakeNpu(current=5)
+    monkeypatch.setattr(mc.torch, "npu", fake_npu, raising=False)
+    warmups = []
+    monkeypatch.setattr(mc.torch, "zeros", lambda *a, **k: warmups.append(k))
+
+    mc._ensure_accelerator_context._initialized = False
+    mc._ensure_accelerator_context()
+    mc._ensure_accelerator_context()
+
+    assert fake_npu.set_devices == [5, 5]
+    assert len(warmups) == 1  # context warmup runs once per process
+
+
+def test_ensure_accelerator_context_noop_without_npu(monkeypatch):
+    monkeypatch.setattr(mc.torch, "npu", None, raising=False)
+    # Force the optional torch_npu import to fail so torch.npu stays absent.
+    monkeypatch.setitem(sys.modules, "torch_npu", None)
+
+    mc._ensure_accelerator_context()  # must not raise or touch a device
+
+
+def _patch_executor(monkeypatch):
+    captured = []
+    real_executor = mc.ThreadPoolExecutor
+
+    class _RecordingExecutor(real_executor):
+        def __init__(self, *args, **kwargs):
+            captured.append(kwargs.get("initializer"))
+            super().__init__(*args, **kwargs)
+
+    monkeypatch.setattr(mc, "ThreadPoolExecutor", _RecordingExecutor)
+    return captured
+
+
+def test_put_pool_initializer_binds_thread_context(monkeypatch):
+    calls = []
+    monkeypatch.setattr(mc, "_ensure_accelerator_context", lambda: calls.append(1))
+    captured = _patch_executor(monkeypatch)
+
+    assert _new_client().put([], []) == []
+    assert captured and all(fn is not None for fn in captured)
+
+    captured[-1]()
+    assert calls == [1]
+
+
+def test_get_pool_initializer_binds_thread_context(monkeypatch):
+    calls = []
+    monkeypatch.setattr(mc, "_ensure_accelerator_context", lambda: calls.append(1))
+    captured = _patch_executor(monkeypatch)
+
+    assert _new_client().get([], [], [], []) == []
+    assert captured and all(fn is not None for fn in captured)
+
+    captured[-1]()
+    assert calls == [1]

transfer_queue/storage/clients/mooncake_client.py

@@ -33,6 +33,37 @@
 
 logger = get_logger(__name__)
 
+
+def _ensure_accelerator_context() -> None:
+    """Bind the calling thread to an accelerator device context.
+
+    Ascend ACL contexts are thread-local, so Mooncake register/unregister fails
+    with INVALID_PARAMS on a ThreadPoolExecutor worker that never set a device.
+    """
+    npu = getattr(torch, "npu", None)
+    if npu is None:
+        try:
+            import torch_npu  # noqa: F401
+        except Exception:  # pragma: no cover - accelerator optional
+            return
+        npu = getattr(torch, "npu", None)
+    if npu is None or not npu.is_available():
+        return
+
+    try:
+        npu.set_device(npu.current_device())
+    except Exception:
+        pass
+
+    # Force context creation once per process.
+    if not getattr(_ensure_accelerator_context, "_initialized", False):
+        try:
+            torch.zeros(1, device="npu")
+        except Exception:  # pragma: no cover - best effort
+            pass
+        _ensure_accelerator_context._initialized = True
+
+
 MOONCAKE_STORE_IMPORTED: bool = True
 try:
     from mooncake.store import MooncakeDistributedStore, ReplicateConfig
@@ -121,6 +152,8 @@ def __init__(self, config: dict[str, Any]):
             hard_pin = not offload_enabled
         self.replica_config.with_hard_pin = bool(hard_pin)
 
+        _ensure_accelerator_context()
+
         self._store = MooncakeDistributedStore()
         ret = self._store.setup(
             self.local_hostname,
@@ -173,7 +206,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=_ensure_accelerator_context,
+        ) 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 +409,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=_ensure_accelerator_context,
+        ) 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]