TPU Connectivity Issue on Google Colab
Hi @arnava13 ,

Could you please provide more context and minimal repro code snippet ? Thanks! Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/63832">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/63832">No</a>
 After ensuring training of my models is running properly on CPUs and GPUs, I am getting the following error when training on colab. I have initialised my TPU strategy as usual, and I am running all model related and training operations under strategy.scope(). 
`
err: File "/content/SeniorHonoursProject/BaCoN-II/train.py", line 737, in <module>
err: main()
err: File "/content/SeniorHonoursProject/BaCoN-II/train.py", line 697, in main
err: model, history = my_train(model, loss,
err: File "/content/SeniorHonoursProject/BaCoN-II/train.py", line 169, in my_train
err: new_history = model.fit(train_dataset.dataset, epochs=epochs,
err: File "/usr/local/lib/python3.10/dist-packages/keras/src/utils/traceback_utils.py", line 70, in error_handler
err: raise e.with_traceback(filtered_tb) from None
err: File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py", line 362, in _numpy
err: raise core._status_to_exception(e) from None  # pylint: disable=protected-access
err: tensorflow.python.framework.errors_impl.InternalError: 8 root error(s) found.
err: (0) INTERNAL: {{function_node __inference_train_function_11526}} failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:50188: Failed to connect to remote host: Connection refused
err: Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
err: :UNKNOWN:failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:50188: Failed to connect to remote host: Connection refused {grpc_status:14, created_time:"2024-03-22T23:36:04.934268533+00:00"}
err: [[{{node MultiDeviceIteratorGetNextFromShard}}]]
err: Executing non-communication op <MultiDeviceIteratorGetNextFromShard> originally returned UnavailableError, and was replaced by InternalError to avoid invoking TF network error handling logic.
err: [[RemoteCall]]
err: [[IteratorGetNextAsOptional]]
err: [[TPUReplicate/_compile/_9902494219978988908/_4/_384]]
err: (1) INTERNAL: {{function_node __inference_train_function_11526}} failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:50188: Failed to connect to remote host: Connection refused
err: Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
err: :UNKNOWN:failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:50188: Failed to connect to remote host: Connection refused {grpc_status:14, created_time:"2024-03-22T23:36:04.934268533+00:00"}
err: [[{{node MultiDeviceIteratorGetNextFromShard}}]]
err: Executing non-communication op <MultiDeviceIteratorGetNextFromShard> originally returned UnavailableError, and was replaced by InternalError to avoid invoking TF network error handling logic.
err: [[RemoteCall]]
err: [[IteratorGetNextAsOptional]]
err: [[Pad_3/_250]]
err: (2) INTERNAL: {{function_node __inference_train_function_11526}} failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:50188: Failed to connect to remote host: Connection refused
err: Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
err: :UNKNOWN:failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:50188: Failed to connect to remote host: Connection refused {grpc_status:14, created_time:"2024-03-22T23:36:04.934268533+00:00"}
err: [[{{node MultiDeviceIteratorGetNextFromShard}}]]
err: Executing non-communication op <MultiDeviceIteratorGetNextFromShard> originally returned UnavailableError, and was replaced by InternalError to avoid invoking TF network error handling logic.
err: [[RemoteCall]]
err: [[IteratorGetNextAsOptional]]
err: [[Pad_15/_466]]
err: (3) INTERNAL: {{function_node __inference_train_function_11526}} failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:50188: Failed to connect to remote host: Connection refused
err: Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
err: :UNKNOWN:failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:50188: Failed to connect to remote host: Connection refused {grpc_status:14, created_time:"2024-03-22T23:36:04.934268533+00:00"}
err: [[{{node MultiDeviceIteratorGetNextFromShard}}] ... [truncated]
err: Exception ignored in atexit callback: <function async_wait at 0x7ec0f4a74790>
err: Traceback (most recent call last):
err: File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/eager/context.py", line 2833, in async_wait
err: context().sync_executors()
err: File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/eager/context.py", line 749, in sync_executors
err: pywrap_tfe.TFE_ContextSyncExecutors(self._context_handle)
err: tensorflow.python.framework.errors_impl.InternalError: 8 root error(s) found.
err: (0) INTERNAL: {{function_node __inference_train_function_11526}} failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:50188: Failed to connect to remote host: Connection refused
err: Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
err: :UNKNOWN:failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:50188: Failed to connect to remote host: Connection refused {grpc_status:14, created_time:"2024-03-22T23:36:04.934268533+00:00"}
err: [[{{node MultiDeviceIteratorGetNextFromShard}}]]
err: Executing non-communication op <MultiDeviceIteratorGetNextFromShard> originally returned UnavailableError, and was replaced by InternalError to avoid invoking TF network error handling logic.
err: [[RemoteCall]]
err: [[IteratorGetNextAsOptional]]
err: [[TPUReplicate/_compile/_9902494219978988908/_4/_384]]
err: (1) INTERNAL: {{function_node __inference_train_function_11526}} failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:50188: Failed to connect to remote host: Connection refused
err: Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
err: :UNKNOWN:failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:50188: Failed to connect to remote host: Connection refused {grpc_status:14, created_time:"2024-03-22T23:36:04.934268533+00:00"}
err: [[{{node MultiDeviceIteratorGetNextFromShard}}]]
err: Executing non-communication op <MultiDeviceIteratorGetNextFromShard> originally returned UnavailableError, and was replaced by InternalError to avoid invoking TF network error handling logic.
err: [[RemoteCall]]
err: [[IteratorGetNextAsOptional]]
err: [[Pad_3/_250]]
err: (2) INTERNAL: {{function_node __inference_train_function_11526}} failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:50188: Failed to connect to remote host: Connection refused
err: Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
err: :UNKNOWN:failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:50188: Failed to connect to remote host: Connection refused {grpc_status:14, created_time:"2024-03-22T23:36:04.934268533+00:00"}
err: [[{{node MultiDeviceIteratorGetNextFromShard}}]]
err: Executing non-communication op <MultiDeviceIteratorGetNextFromShard> originally returned UnavailableError, and was replaced by InternalError to avoid invoking TF network error handling logic.
err: [[RemoteCall]]
err: [[IteratorGetNextAsOptional]]
err: [[Pad_15/_466]]
err: (3) INTERNAL: {{function_node __inference_train_function_11526}} failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:50188: Failed to connect to remote host: Connection refused
err: Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
err: :UNKNOWN:failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:50188: Failed to connect to remote host: Connection refused {grpc_status:14, created_time:"2024-03-22T23:36:04.934268533+00:00"}
err: [[{{node MultiDeviceIteratorGetNextFromShard}}] ... [truncated]
out: g1D)
` Note also that I have initialised the TPU with the following lines, 

`
tpu_resolver = tf.distribute.cluster_resolver.TPUClusterResolver()  # Automatically detects the TPU
tf.config.experimental_connect_to_cluster(tpu_resolver)  # Connects to the TPU cluster
tf.tpu.experimental.initialize_tpu_system(tpu_resolver)  # Initializes the TPU system
strategy = tf.distribute.TPUStrategy(tpu_resolver)
tpu_device = tpu_resolver.master()  # Retrieves the TPU device URI
print("Running on TPU:", tpu_device)
`

and got the following print, indicating that the TPU is connected.
`
Running on TPU: grpc://10.74.203.82:8470
`

 Fixed. The issue was using a generator function to build datasets in my pipeline. Using tf.data.Dataset.from_tensor_slices instead fixes the issue. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/63832">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/63832">No</a>

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

source

### TensorFlow version

2.15.0

### Custom code

Yes

### OS platform and distribution

_No response_

### Mobile device

_No response_

### Python version

3.10.12

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

Code was designed by someone else, I've just been adapting it to use TPUs and TensorFlow datasets in training. When training in TPU mode I get the attached error.

### Standalone code to reproduce the issue

```shell
https://github.com/arnava13/SeniorHonoursProject/tree/main/BaCoN-II

train.py has the training loop. 
data_generator.py has the datagenerators.
Training is run by opening training_colab.ipynb directly on colab's from github option.
```


### Relevant log output

```shell
err: Traceback (most recent call last):
err: File "/content/SeniorHonoursProject/BaCoN-II/train.py", line 763, in <module>
err: main()
err: File "/content/SeniorHonoursProject/BaCoN-II/train.py", line 733, in main
err: model, history = my_train(model, optimizer, loss,
err: File "/content/SeniorHonoursProject/BaCoN-II/train.py", line 168, in my_train
err: for IDs, x_batch_train, y_batch_train in train_generator.dataset:
err: File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/distribute/input_lib.py", line 264, in __next__
err: return self.get_next()
err: File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/distribute/input_lib.py", line 331, in get_next
err: num_replicas_with_values = _calculate_replicas_with_values(
err: File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/distribute/input_lib.py", line 196, in _calculate_replicas_with_values
err: client_has_values = math_ops.reduce_sum(worker_has_values, keepdims=True)
err: File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/ops/weak_tensor_ops.py", line 88, in wrapper
err: return op(*args, **kwargs)
err: File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
err: raise e.with_traceback(filtered_tb) from None
err: File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py", line 498, in shape
err: self._tensor_shape = tensor_shape.TensorShape(self._shape_tuple())
err: tensorflow.python.framework.errors_impl.InternalError: failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:36768: Failed to connect to remote host: Connection refused
err: Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
err: :UNKNOWN:failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:36768: Failed to connect to remote host: Connection refused {created_time:"2024-03-15T21:29:31.939700993+00:00", grpc_status:14}
err: [[{{node MultiDeviceIteratorGetNextFromShard}}]]
err: Executing non-communication op <MultiDeviceIteratorGetNextFromShard> originally returned UnavailableError, and was replaced by InternalError to avoid invoking TF network error handling logic.
err: [[RemoteCall]]
err: Exception ignored in atexit callback: <function async_wait at 0x7b4d88d84670>
err: Traceback (most recent call last):
err: File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/eager/context.py", line 2833, in async_wait
err: context().sync_executors()
err: File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/eager/context.py", line 749, in sync_executors
err: pywrap_tfe.TFE_ContextSyncExecutors(self._context_handle)
err: tensorflow.python.framework.errors_impl.InternalError: failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:36768: Failed to connect to remote host: Connection refused
err: Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
err: :UNKNOWN:failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:36768: Failed to connect to remote host: Connection refused {created_time:"2024-03-15T21:29:31.939700993+00:00", grpc_status:14}
err: [[{{node MultiDeviceIteratorGetNextFromShard}}]]
err: Executing non-communication op <MultiDeviceIteratorGetNextFromShard> originally returned UnavailableError, and was replaced by InternalError to avoid invoking TF network error handling logic.
err: [[RemoteCall]]
err: 2024-03-15 21:29:31.946346: W tensorflow/core/distributed_runtime/eager/remote_tensor_handle_data.cc:77] Unable to destroy remote tensor handles. If you are running a tf.function, it usually indicates some op in the graph gets an error: failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:36768: Failed to connect to remote host: Connection refused
err: Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
err: :UNKNOWN:failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:36768: Failed to connect to remote host: Connection refused {created_time:"2024-03-15T21:29:31.939700993+00:00", grpc_status:14}
err: [[{{node MultiDeviceIteratorGetNextFromShard}}]]
err: Executing non-communication op <MultiDeviceIteratorGetNextFromShard> originally returned UnavailableError, and was replaced by InternalError to avoid invoking TF network error handling logic.
err: [[RemoteCall]]
err: 2024-03-15 21:29:32.185465: W ./tensorflow/core/distributed_runtime/eager/destroy_tensor_handle_node.h:59] Ignoring an error encountered when deleting remote tensors handles: INVALID_ARGUMENT: Unable to find the relevant tensor remote_handle: Op ID: 2124, Output num: 0
err: Additional GRPC error information from remote target /job:worker/replica:0/task:0 while calling /tensorflow.eager.EagerService/Enqueue:
err: :{"created":"@1710538172.185353260","description":"Error received from peer ipv4:10.120.183.98:8470","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Unable to find the relevant tensor remote_handle: Op ID: 2124, Output num: 0","grpc_status":3} [type.googleapis.com/tensorflow.core.platform.ErrorSourceProto='\x08\x05']
```
