Summary
machinelearning_run() hard-crashes the process (Fatal Python error: Aborted) when libtpu is already loaded by JAX (i.e. any real JAX-on-TPU workload). libtpu_metric._initialize() does from libtpu import sdk, whose C-extension re-registers protobuf descriptors that are already in the global pool → descriptor.cc Check failed. _initialize() only catches ImportError, which cannot catch a C++ abort().
Environment
google-cloud-mldiagnostics 1.0.2 (the version shipping in current TPU images)- GKE TPU (tpu7x), plain-SPMD SFT (local libtpu present), JAX 0.10.0
Error (verbatim)
E0000 descriptor_database.cc:633 File already exists in database: google/protobuf/timestamp.proto
F0000 descriptor.cc:2236 Check failed ...
Fatal Python error: Aborted
Call path: machinelearning_run → get_software_config → get_libtpu_version → libtpu_metric._initialize → from libtpu import sdk. Reproduced deterministically — all 16 workers crashlooped (RESTART climbing 0→1→2).
Root cause
The SDK re-imports the libtpu C-extension instead of reusing the module JAX already loaded; the second registration of timestamp.proto is fatal.
Note on HEAD
HEAD commit 08d554d ("Enforce JAX-style RTLD_LOCAL and RTLD_DEEPBIND dlopen isolation during LibTPU SDK import") adds RTLD_LOCAL|RTLD_DEEPBIND and broadens the except — which likely mitigates this. But it is not in the released 1.0.2/1.0.3 wheel that ships in the images, so deployed workloads still abort.
Asks
- Cut a release containing
08d554d so deployed images stop aborting.
- Consider, in addition to dlopen isolation: if
libtpu is already in sys.modules, reuse it / read the version via importlib.metadata.version("libtpu") without importing the C-extension; and broaden the guard to except BaseException so a probe failure degrades to "n/a" rather than aborting the trainer.
— Reported via Navi on behalf of @lokic233 (Meta MRS-CE).
Summary
machinelearning_run()hard-crashes the process (Fatal Python error: Aborted) whenlibtpuis already loaded by JAX (i.e. any real JAX-on-TPU workload).libtpu_metric._initialize()doesfrom libtpu import sdk, whose C-extension re-registers protobuf descriptors that are already in the global pool →descriptor.ccCheck failed._initialize()only catchesImportError, which cannot catch a C++abort().Environment
google-cloud-mldiagnostics1.0.2 (the version shipping in current TPU images)Error (verbatim)
Call path:
machinelearning_run→get_software_config→get_libtpu_version→libtpu_metric._initialize→from libtpu import sdk. Reproduced deterministically — all 16 workers crashlooped (RESTART climbing 0→1→2).Root cause
The SDK re-imports the
libtpuC-extension instead of reusing the module JAX already loaded; the second registration oftimestamp.protois fatal.Note on HEAD
HEAD commit
08d554d("Enforce JAX-style RTLD_LOCAL and RTLD_DEEPBIND dlopen isolation during LibTPU SDK import") addsRTLD_LOCAL|RTLD_DEEPBINDand broadens the except — which likely mitigates this. But it is not in the released 1.0.2/1.0.3 wheel that ships in the images, so deployed workloads still abort.Asks
08d554dso deployed images stop aborting.libtpuis already insys.modules, reuse it / read the version viaimportlib.metadata.version("libtpu")without importing the C-extension; and broaden the guard toexcept BaseExceptionso a probe failure degrades to"n/a"rather than aborting the trainer.— Reported via Navi on behalf of @lokic233 (Meta MRS-CE).