Retrieves TPU system metadata given a TPUClusterResolver.
(tpu_cluster_resolver)
| 57 | |
| 58 | |
| 59 | def get_tpu_system_metadata(tpu_cluster_resolver): |
| 60 | """Retrieves TPU system metadata given a TPUClusterResolver.""" |
| 61 | master = tpu_cluster_resolver.master() |
| 62 | |
| 63 | # pylint: disable=protected-access |
| 64 | cluster_spec = tpu_cluster_resolver.cluster_spec() |
| 65 | cluster_def = cluster_spec.as_cluster_def() if cluster_spec else None |
| 66 | tpu_system_metadata = ( |
| 67 | tpu_system_metadata_lib._query_tpu_system_metadata( |
| 68 | master, |
| 69 | cluster_def=cluster_def, |
| 70 | query_topology=False)) |
| 71 | |
| 72 | return tpu_system_metadata |
| 73 | |
| 74 | |
| 75 | @contextlib.contextmanager |
no test coverage detected