OptionA experimenting

[wenjung2007]

No description provided.

[datenglin]

In https://github.com/google/tpu-sync/compare/main...wenjung2007:tpu-sync:wenjung/option-a-f0818d4?expand=1#diff-7b25c8e872c18ddf25a6241a6115017db172f1a2dc3536a97fee730612bf36abR125, can you change

echo "=== [7/7] Verifying Build Artifacts ==="
echo "Verifying dynamic loader symbols on native extension..."
patchelf --print-needed "${WORKSPACE_DIR}/tpu_sync/frameworks/torch/_tpu_raiden_torch.so" | grep "libpywrap_${TORCH_GLUE_SUFFIX}_common.so"

echo "Verifying Python import test..."
python3 -c "import torch; print('PyTorch loaded:', torch.__version__); from tpu_sync.frameworks.torch import _tpu_raiden_torch; print('TPU Sync torch module verified successfully!'); from tpu_sync.api.torch import kv_cache_manager; print('KV cache manager verified successfully!')"

to

echo "Installing generated wheel for verification..."
pip install --force-reinstall "${WHEEL_PATH}"

echo "Verifying Python import test from clean directory..."
(cd /tmp && python3 -c "import torch; print('PyTorch loaded:', torch.__version__); from tpu_sync.api.torch import torch_abi, kv_cache_manager; ext = torch_abi.load_extension('tpu_sync.frameworks.torch', '_tpu_raiden_torch'); print('TPU Sync torch module verified successfully:', ext); impl = kv_cache_manager._torch_impl(); print('KV cache manager verified successfully:', impl)")

[datenglin]

For your case, can you remove line 127 and 128 and use the following

package_dirs = [pathlib.Path(p).resolve() for p in pkg.__path__]
target_dir = None
built = []
for pdir in package_dirs:
  suffixes = bundled_suffixes(pdir, stem)
  if suffixes:
    target_dir = pdir
    built = suffixes
    break