fix(runtime): bind NPU device before Mooncake transport setup

[curnane-lab]

Motivation

Launching an online disaggregated training run on Ascend NPU (e.g. specforge train -c examples/configs/qwen3.5-4b-domino-online-npu.yaml) aborts during startup with:

ascend_direct_transport.cpp:528] Call aclrtGetDevice failed, ret: 107002
ascend_direct_transport.cpp:96] AscendDirectTransport: cannot allocate local seg
transfer_engine_impl.cpp:199] Failed to install Ascend transport
RuntimeError: Mooncake setup failed (status -600)

Ascend error 107002 is ACL_ERROR_RT_CONTEXT_NULL — no device context has been established for the calling process. The failing process is the disaggregated capture producer (specforge/cli.py:126), which deliberately skips init_distributed() so it never initializes an accelerator "merely to publish feature refs". On Ascend, however, MooncakeDistributedStore.setup() installs the AscendDirectTransport, which calls aclrtGetDevice; with no NPU bound to the process, that call fails and the whole run dies before training starts.

(The trainer/consumer role already binds its NPU in init_distributed before building the store, so it was unaffected — but the producer failure tears the job down.)

Modifications

• Add _bind_transport_device() in specforge/runtime/data_plane/mooncake_store.py, called inside _connect_store() immediately before MooncakeDistributedStore.setup(), so the local NPU context exists before the Ascend transport is installed.
  • Binds via the existing rank-aware specforge.distributed._bind_local_device (reads LOCAL_RANK/RANK), so every rank pins the transport to its own NPU and the call is idempotent with the trainer's earlier binding.
  • NPU only. CUDA's transport defaults to device 0 without an explicit set_device, and the producer intentionally avoids initializing CUDA, so CUDA/CPU paths are left untouched.
• Add _ascend_runtime_available() to robustly detect Ascend on a producer that has not yet imported torch_npu (so get_device_type() still reports cpu). It only activates torch_npu when Ascend visible-device env vars (ASCEND_RT_VISIBLE_DEVICES / ASCEND_VISIBLE_DEVICES) are set, avoiding a spurious torch_npu import on CUDA/CPU hosts, and degrades gracefully if the import or availability check fails.
• Add unit tests covering: no-op on plain CPU and on CUDA, binding on NPU, the producer path where torch_npu is not yet active, the env-gated torch_npu activation, and that _connect_store binds the device before setup().

Related Issues

Fixes the Ascend NPU online disaggregated launch failure (Mooncake setup failed (status -600) / aclrtGetDevice 107002).

Accuracy Test

Not applicable — no model/kernel/architecture changes; this only fixes device-context binding ordering in the data-plane transport setup.

Benchmark & Profiling

Not applicable — no steady-state performance impact; the added binding runs once during store construction.

Checklist

• Format your code according to the Code Formatting with Pre-Commit. (black, isort, autoflake clean)
• Add unit tests as outlined in Running Unit Tests. (tests/test_runtime/test_mooncake_store.py::TestMooncakeTransportDeviceBinding)
• Update documentation / docstrings / example tutorials as needed. (Behavior documented in module docstrings.)
• Provide throughput / latency benchmark results and accuracy evaluation results as needed. (Not applicable.)