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. 👍