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

[ys2025-AI]

What

Bind the Ascend ACL device context on every thread that calls Mooncake register_buffer/unregister_buffer, i.e. the worker threads of the per-call ThreadPoolExecutor in put()/get().

Why

MooncakeStoreClient.put()/get() create a fresh ThreadPoolExecutor per call and register/unregister transient buffers inside those worker threads. ACL contexts are thread-local, so the workers had no context:

• register_buffer silently fell back to host memory;
• unregister_buffer failed with INVALID_PARAMS, leaking the registration and flooding logs with aclrtGetCurrentContext failed.

Changes

transfer_queue/storage/clients/mooncake_client.py

• add _ensure_accelerator_context(): no-op when torch.npu is unavailable (CUDA/CPU/generic unaffected), otherwise npu.set_device(...) for the calling thread with a one-time context warmup;
• call it once before store.setup(...) (the client thread needs a context too);
• pass it as initializer=_ensure_accelerator_context to both ThreadPoolExecutors (put() and get()).

_register_all_buffers/_unregister_all_buffers run only inside those pools, so no other call sites are needed.

Tests

tests/test_mooncake_client_thread_context.py (hardware-free):

• _ensure_accelerator_context sets the device and warms up only once;
• it is a no-op when no NPU is available;
• put() and get() build their executor with a non-None initializer that binds the context (executor faked, store faked).

python -m pytest -q tests/test_mooncake_client_thread_context.py   # 4 passed
python -m compileall -q transfer_queue tutorial tests
ruff check / ruff format --check                                   # clean

Verification on hardware (Ascend 910B + mooncake-transfer-engine 0.3.11.post1)

Before: Unregister buffer failed / aclrtGetCurrentContext failed 1500+ times per short run.
After: 0 occurrences; register/unregister succeed; no leak.

[ascend-robot]

CLA Signature Pass

ys2025-AI, thanks for your pull request. All authors of the commits have signed the CLA. 👍

[0oshowero0]

[AI Review] Duplicate submission (tooling retry) — superseded by the review below, which carries the same findings as inline comments.

[0oshowero0]

[AI Review]

The root-cause analysis is right: ACL contexts are thread-local, put()/get() build a fresh ThreadPoolExecutor per call, and binding the device in the pool initializer is the cheapest correct hook. The claim that _register_all_buffers/_unregister_all_buffers only run inside those two pools also checks out — the GDR path does register on the caller thread, but it is gated by torch.cuda.is_initialized() and is unreachable on Ascend.

Two things I think should be fixed before merge:

1. The device index is resolved inside the context-less worker thread, so npu.current_device() always reports 0 there and the workers never bind the device the client actually uses. Details in the inline comment on set_device.
2. The once-per-process torch.zeros warmup cannot be doing what its comment claims, and it hides the fact that aclrtSetDevice alone is already sufficient.

The rest is reducible. The torch_npu probe can be resolved once at module import (MOONCAKE_STORE_IMPORTED already shows the pattern in this file), the module-level mutable flag can disappear with the warmup, and roughly two thirds of the new test file either asserts that a line of code was written or is never exercised at all.

A shape that keeps the fix while dropping the global state, the per-call import, the swallowed failure and the wrong device (~10 lines instead of 33):

try:
    import torch_npu  # noqa: F401

    NPU_IMPORTED = True
except ImportError:
    NPU_IMPORTED = False

        # ACL contexts are thread-local: the per-call pool workers must bind the
        # same device as the client, otherwise unregister_buffer fails with
        # INVALID_PARAMS and leaks the registration.
        self._npu_device = (
            torch.npu.current_device() if NPU_IMPORTED and torch.npu.is_initialized() else None
        )

    def _bind_npu_device(self) -> None:
        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.", exc_info=True)

with initializer=self._bind_npu_device on both executors, plus one test asserting that the worker binds the device captured at construction time rather than 0.

[ys2025-AI]

[AI Review]

The root-cause analysis is right: ACL contexts are thread-local, put()/get() build a fresh ThreadPoolExecutor per call, and binding the device in the pool initializer is the cheapest correct hook. The claim that _register_all_buffers/_unregister_all_buffers only run inside those two pools also checks out — the GDR path does register on the caller thread, but it is gated by torch.cuda.is_initialized() and is unreachable on Ascend.

Two things I think should be fixed before merge:

1. The device index is resolved inside the context-less worker thread, so npu.current_device() always reports 0 there and the workers never bind the device the client actually uses. Details in the inline comment on set_device.
2. The once-per-process torch.zeros warmup cannot be doing what its comment claims, and it hides the fact that aclrtSetDevice alone is already sufficient.

The rest is reducible. The torch_npu probe can be resolved once at module import (MOONCAKE_STORE_IMPORTED already shows the pattern in this file), the module-level mutable flag can disappear with the warmup, and roughly two thirds of the new test file either asserts that a line of code was written or is never exercised at all.

A shape that keeps the fix while dropping the global state, the per-call import, the swallowed failure and the wrong device (~10 lines instead of 33):

try:
    import torch_npu  # noqa: F401

    NPU_IMPORTED = True
except ImportError:
    NPU_IMPORTED = False

        # ACL contexts are thread-local: the per-call pool workers must bind the
        # same device as the client, otherwise unregister_buffer fails with
        # INVALID_PARAMS and leaks the registration.
        self._npu_device = (
            torch.npu.current_device() if NPU_IMPORTED and torch.npu.is_initialized() else None
        )

    def _bind_npu_device(self) -> None:
        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.", exc_info=True)

with initializer=self._bind_npu_device on both executors, plus one test asserting that the worker binds the device captured at construction time rather than 0.

Thanks for the review — all points addressed in ebd2b5c:

• Wrong device on workers: device is now captured once on the owning thread in init (_resolve_npu_device); workers bind that value via _bind_npu_device, never calling current_device() on a context-less thread.
• Misleading warmup: removed torch.zeros and the module-level flag entirely; set_device alone is relied upon.
• Swallowed failure: bind/resolve failures now logger.warning(..., exc_info=True).
• Import probe: torch_npu is resolved once at import (NPU_IMPORTED); is_available() replaced with is_initialized() so non-NPU processes are untouched.
• Tests: dropped _FakeStore and the two tautological executor tests; new tests cover device capture, the not-initialized no-op, binding the captured device (not 0), and logged-and-swallowed bind failure.

[Chase-Rong]

LGTM. Thanks for your fix. Validated this PR on Ascend A3 hardware with mooncake-transfer-engine-npu 0.3.11.post1.

Results:

• Worker device binding: With the owning thread on npu:1, the base version’s workers used device 0 and produced 24 registration failures and 24 unregistration failures (-600). With this PR, all 24 workers had a valid ACL context on device 1; all 159 registrations and 159 unregistrations returned 0.

• Real data transfers: CPU tensor/object batches completed three put/get/clear rounds with 802 keys per round. NPU tensor/object transfers also completed three rounds using the ascend protocol. All data-integrity and key-removal checks passed.

• Without diagnostic instrumentation: The base version logged 22 occurrences each of aclrtGetCurrentContext failed and Unregister buffer failed; both counts dropped to zero with this PR.