Initialize the TPU devices. Args: cluster_resolver: A tf.distribute.cluster_resolver.TPUClusterResolver, which provides information about the TPU cluster. Returns: The tf.tpu.Topology object for the topology of the TPU cluster. Raises: RuntimeError: If no TPU devices foun
(cluster_resolver=None)
| 38 | |
| 39 | @tf_export("tpu.experimental.initialize_tpu_system") |
| 40 | def initialize_tpu_system(cluster_resolver=None): |
| 41 | """Initialize the TPU devices. |
| 42 | |
| 43 | Args: |
| 44 | cluster_resolver: A tf.distribute.cluster_resolver.TPUClusterResolver, |
| 45 | which provides information about the TPU cluster. |
| 46 | Returns: |
| 47 | The tf.tpu.Topology object for the topology of the TPU cluster. |
| 48 | |
| 49 | Raises: |
| 50 | RuntimeError: If no TPU devices found for eager execution. |
| 51 | """ |
| 52 | job = None |
| 53 | if cluster_resolver is None: |
| 54 | # If no cluster resolver is specified, and running eagerly, execute the init |
| 55 | # ops in the current device scope. |
| 56 | if context.executing_eagerly(): |
| 57 | curr_device = device.DeviceSpec.from_string(context.context().device_name) |
| 58 | if curr_device.job is not None: |
| 59 | job = "{}/replica:0/task:0".format(curr_device.job) |
| 60 | |
| 61 | cluster_resolver = TPUClusterResolver("") |
| 62 | assert isinstance(cluster_resolver, TPUClusterResolver) |
| 63 | |
| 64 | tpu_name = compat.as_text(cluster_resolver._tpu) # pylint: disable=protected-access |
| 65 | if tpu_name in _INITIALIZED_TPU_SYSTEMS: |
| 66 | logging.warning("TPU system %s has already been initialized. " |
| 67 | "Reinitializing the TPU can cause previously created " |
| 68 | "variables on TPU to be lost.") |
| 69 | |
| 70 | logging.info("Initializing the TPU system: %s", tpu_name) |
| 71 | |
| 72 | if context.executing_eagerly(): |
| 73 | # This function looks as it is for the following non-intuitive reasons. |
| 74 | # tpu.initialize_system creates a dummy op whose sole purpose is to trigger |
| 75 | # DistributedTPURewritePass. This pass actually adds real ops that |
| 76 | # initialize the TPU system. Thus, we can't simply run tpu.initialize_system |
| 77 | # eagerly. We need to wrap it in defun and trigger the rewrite passes on it. |
| 78 | if tpu_name not in _LOCAL_MASTERS: |
| 79 | # Explicitly place the tpu.initialize_system in the first worker to |
| 80 | # avoid the output node match multiple devices error. |
| 81 | job = "{}/replica:0/task:0".format(cluster_resolver.get_job_name()) |
| 82 | |
| 83 | @function.defun |
| 84 | def _tpu_init_fn(): |
| 85 | return tpu.initialize_system(job=job) |
| 86 | |
| 87 | # The TPU_SYSTEM device must match the device used in tpu.initialize_system |
| 88 | # exactly, otherwise you can get errors if there are multiple TPU_SYSTEM |
| 89 | # devices available. |
| 90 | with ops.device(tpu._tpu_system_device_name(job)): # pylint: disable=protected-access |
| 91 | output = _tpu_init_fn() |
| 92 | |
| 93 | # Clear out the eager context caches since the memory is invalid now. |
| 94 | logging.info("Clearing out eager caches") |
| 95 | context.context()._clear_caches() # pylint: disable=protected-access |
| 96 | |
| 97 | serialized_topology = output.numpy() |
nothing calls this directly
no test coverage detected