Obtains TPU fabric topology.
(master_address, cluster_def)
| 158 | |
| 159 | |
| 160 | def _obtain_topology(master_address, cluster_def): |
| 161 | """Obtains TPU fabric topology.""" |
| 162 | try: |
| 163 | logging.info('Initializing TPU system (master: %s) to fetch topology ' |
| 164 | 'for model parallelism. This might take a while.', |
| 165 | master_address) |
| 166 | with ops.Graph().as_default(): |
| 167 | session_config = get_session_config_with_timeout( |
| 168 | _INITIAL_TPU_SYSTEM_TIMEOUT_IN_MS, cluster_def) |
| 169 | with session_lib.Session( |
| 170 | master_address, config=session_config) as sess: |
| 171 | topology = sess.run(tpu.initialize_system()) |
| 172 | return topology |
| 173 | except errors.DeadlineExceededError: |
| 174 | raise ValueError( |
| 175 | 'Fail to initialize TPU system with master (%s). ' |
| 176 | 'Please double check the TPU system is functional.' % ( |
| 177 | master_address)) |
| 178 | |
| 179 | |
| 180 | def get_session_config_with_timeout(timeout_in_secs, cluster_def): |
no test coverage detected