Thread that raises an exception if it encounters an error.
| 958 | start_time = time.perf_counter() |
| 959 | tpu_arrays = await asyncio.gather(*tasks) |
| 960 | |
| 961 | total_time = time.perf_counter() - start_time |
| 962 | logging.info("Pipelined colocated deserialization completed in %.2f seconds", total_time) |
| 963 | |
| 964 | # Deterministically release event loop and executor on the sidecar. |
| 965 | colocated_mgr.teardown.specialize(devices=cpu_devices)() |
| 966 | # Explicitly trigger GC to ensure deterministic memory release within the |
| 967 | # colocated-python sidecar container. |
| 968 | gc.collect() |
| 969 | # Run gc.collect() and log diagnostics on the sidecar. |
| 970 | _colocated_teardown.specialize(devices=cpu_devices)() |
| 971 | |
| 972 | return tpu_arrays |
| 973 | |
| 974 | |
| 975 | # Reference: |
| 976 | # https://github.com/google/orbax/blob/ebb3e6d75f9ccb52bf862f1740943a45b18f4dac/checkpoint/orbax/checkpoint/future.py#L49 |
| 977 | class _ThreadRaisingException(threading.Thread): |
| 978 | """Thread that raises an exception if it encounters an error.""" |