Fix ReshardStore to accept the native RaidenId binding.

[majunze2001]

Summary

tpu_sync.api.torch.reshard_store.ReshardStore.__init__ could not be called at
all: it unwrapped its raiden_id argument with raiden_id._impl, but
kv_cache_store.RaidenId has not had an _impl attribute since Aug 27. Every
construction raised AttributeError.

This change unwraps with getattr(raiden_id, "_impl", raiden_id), matching the
guard the other call sites in kv_cache_store.py already use, and adds the
first test for the class.

Root cause

kv_cache_store.RaidenId is resolved as

RaidenId = getattr(_impl, "RaidenId", common.RaidenId)

i.e. the nanobind class exported by _tpu_raiden_torch when present, else the
pure-Python dataclass in tpu_sync/api/common.py. Neither carries _impl.

Timeline:

• 4b2b466 (Aug 12) added reshard_store.py with raiden_id._impl. At that
  time kv_cache_store.RaidenId was a Python wrapper class holding the C++
  object in _impl, so the line was correct.
• 7efa546 (Aug 27, "fix RaidenId binding") deleted that wrapper and set
  RaidenId = _impl.RaidenId. It updated the unwrap sites in
  kv_cache_store.py with hasattr(x, "_impl") guards but did not touch
  reshard_store.py. From here on the constructor raises.
• a0f5ea1 (Aug 30) changed the alias to the current getattr(...) form with
  the dataclass fallback. Still no _impl on either branch.

Nothing tested ReshardStore, so the breakage went unnoticed.

Fix

tpu_sync/api/torch/reshard_store.py: one-line change from raiden_id._impl
to getattr(raiden_id, "_impl", raiden_id). Native ids pass through; a Python
wrapper that does carry _impl is still unwrapped.

Tests

New tpu_sync/api/torch/reshard_store_test.py (absltest, registered as a
py_test in tpu_sync/api/torch/BUILD and added to run_tests.sh):

• test_accepts_native_raiden_id: asserts the native RaidenId has no _impl
  (the premise), mocks create_reshard_store, and checks the native id is
  forwarded verbatim. Fails on the old code with AttributeError.
• test_unwraps_python_wrapper_with_impl: an object exposing _impl is
  unwrapped before reaching the factory.
• test_constructs_real_store_on_localhost: builds a real in-engine store on
  127.0.0.1 and checks the controller address and reshard service port.

Testing performed

Before the fix, the old expression on a native id:

RaidenId class: <class 'tpu_sync.frameworks.torch._tpu_raiden_torch.RaidenId'>
hasattr _impl: False
AttributeError - 'tpu_sync.frameworks.torch._tpu_raiden_torch.RaidenId' object has no attribute '_impl'

After the fix:

$ PYTHONPATH=$PWD python -u tpu_sync/api/torch/reshard_store_test.py
[       OK ] ReshardStoreTest.test_accepts_native_raiden_id
[       OK ] ReshardStoreTest.test_constructs_real_store_on_localhost
[       OK ] ReshardStoreTest.test_unwraps_python_wrapper_with_impl
Ran 3 tests in 0.332s
OK

Notes

• Python-only change; no rebuild of the extension is required.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[majunze2001]

closed since #889 is merged