From ebd2b5c60fb72edb40e9799b97ee7e30c2893195 Mon Sep 17 00:00:00 2001 From: ys2025-AI Date: Thu, 8 Oct 2026 14:11:55 +0800 Subject: [PATCH] fix(mooncake): bind per-thread NPU device for buffer register/unregister Mooncake's Ascend transport is bound to a thread-local ACL context. put()/get() run register_buffer/unregister_buffer on per-call ThreadPoolExecutor workers that never called npu.set_device, so unregister_buffer failed with INVALID_PARAMS (aclrtGetCurrentContext failed) and the registered host buffers were leaked. Capture the device once on the owning thread in __init__ and bind it in each pool worker through the executor initializer. Gate on is_initialized() so a process that does not use the NPU is left untouched, and add hardware-free unit tests for the capture/bind behavior. Signed-off-by: ys2025-AI --- tests/test_mooncake_client_thread_context.py | 105 ++++++++++++++++++ .../storage/clients/mooncake_client.py | 54 ++++++++- 2 files changed, 157 insertions(+), 2 deletions(-) create mode 100644 tests/test_mooncake_client_thread_context.py diff --git a/tests/test_mooncake_client_thread_context.py b/tests/test_mooncake_client_thread_context.py new file mode 100644 index 00000000..01ce3808 --- /dev/null +++ b/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 diff --git a/transfer_queue/storage/clients/mooncake_client.py b/transfer_queue/storage/clients/mooncake_client.py index d6914902..ff38bcb7 100644 --- a/transfer_queue/storage/clients/mooncake_client.py +++ b/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]