diff --git a/tests/test_mooncake_client_thread_context.py b/tests/test_mooncake_client_thread_context.py new file mode 100644 index 00000000..f38636bd --- /dev/null +++ b/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] diff --git a/transfer_queue/storage/clients/mooncake_client.py b/transfer_queue/storage/clients/mooncake_client.py index d6914902..d7661bf5 100644 --- a/transfer_queue/storage/clients/mooncake_client.py +++ b/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]