Repository: tensorflow
Issue: Seg Fault when iterate dataset created from data service
Labels: type:bug,comp:ops,TF 2.18

### Issue type

Bug

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

No

### Source

binary

### TensorFlow version

2.18.0

### Custom code

Yes

### OS platform and distribution

Linux Ubuntu 22.04

### Mobile device

_No response_

### Python version

3.10

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

Segfault when trying to iterate dataset get from data service.

### Standalone code to reproduce the issue

```shell
# start the data service file start_dataservice.py

import tensorflow as tf

dispatcher = tf.data.experimental.service.DispatchServer(
    tf.data.experimental.service.DispatcherConfig(port=50050), start=True
)
dispatcher_address = dispatcher.target.split("://")[1]
worker = tf.data.experimental.service.WorkerServer(
    tf.data.experimental.service.WorkerConfig(dispatcher_address=dispatcher_address), start=True
)
print("Starting Worker")
worker.join()

# test file test_dataset_service.py
import tensorflow as tf
import numpy as np


flags = tf.compat.v1.app.flags

flags.DEFINE_bool("local", False, "Run data service in process")
flags.DEFINE_bool("distribute", False, "Run data service in distributed_epoch mode")
FLAGS = flags.FLAGS


def local_service():
    print("Starting Local Service")
    dispatcher = tf.data.experimental.service.DispatchServer(
        tf.data.experimental.service.DispatcherConfig(port=50050), start=True
    )
    dispatcher_address = dispatcher.target.split("://")[1]
    worker = tf.data.experimental.service.WorkerServer(
        tf.data.experimental.service.WorkerConfig(dispatcher_address=dispatcher_address), start=True
    )
    print("Dispatcher target is ", dispatcher.target)
    return dispatcher, worker, dispatcher.target


def apply_transformations(ds_train):
    ds_train = ds_train.map(normalize_img, num_parallel_calls=tf.data.experimental.AUTOTUNE)
    ds_train = ds_train.cache()
    ds_train = ds_train.shuffle(60000)
    ds_train = ds_train.batch(128)
    ds_train = ds_train.prefetch(tf.data.experimental.AUTOTUNE)
    return ds_train


(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
x_train = x_train / np.float32(255)
y_train = y_train.astype(np.int64)
ds_train = tf.data.Dataset.from_tensor_slices((x_train, y_train))


def normalize_img(image, label):
    """Normalizes images: `uint8` -> `float32`."""
    return tf.cast(image, tf.float32) / 255.0, label


ds_train = apply_transformations(ds_train)
# Create dataset however you were before using the tf.data service.
dataset = ds_train
if FLAGS.local:
    dispatcher, worker, service = local_service()
else:
    dispatcher_address = "localhost"
    dispatcher_port = "50050"
    service = "grpc://{}:{}".format(dispatcher_address, dispatcher_port)
if FLAGS.distribute:
    processing_mode = "distributed_epoch"
else:
    processing_mode = "parallel_epochs"

# This will register the dataset with the tf.data service cluster so that
# tf.data workers can run the dataset to produce elements. The dataset returned
# from applying `distribute` will fetch elements produced by tf.data workers.
dataset = dataset.apply(
    tf.data.experimental.service.distribute(processing_mode=processing_mode, service=service)
)

for (x1, y1), (x2, y2) in zip(dataset, ds_train):
    np.allclose(x1, x2)
    np.allclose(y1, y2)

print("verified mnist dataset locally vs over service")

# script to run 
python -m pip install --upgrade pip
python -m pip install tensorflow==2.18.0
python -m pip install 'protobuf<4'
screen -d -m python start_dataservice.py
python3 test_dataset_service.py --local=False
```

### Relevant log output

```shell
2025-01-14 21:56:19.778399: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:477] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
E0000 00:00:1736891779.795141    9168 cuda_dnn.cc:8310] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
E0000 00:00:1736891779.800177    9168 cuda_blas.cc:1418] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2025-01-14 21:56:19.815971: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: SSE3 SSE4.1 SSE4.2 AVX AVX2 AVX512F FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
I0000 00:00:1736891783.518634    9168 gpu_device.cc:2022] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 889 MB memory:  -> device: 0, name: NVIDIA A100-SXM4-40GB, pci bus id: 0000:10:1c.0, compute capability: 8.0
I0000 00:00:1736891783.520395    9168 gpu_device.cc:2022] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 37945 MB memory:  -> device: 1, name: NVIDIA A100-SXM4-40GB, pci bus id: 0000:10:1d.0, compute capability: 8.0
I0000 00:00:1736891783.522012    9168 gpu_device.cc:2022] Created device /job:localhost/replica:0/task:0/device:GPU:2 with 37945 MB memory:  -> device: 2, name: NVIDIA A100-SXM4-40GB, pci bus id: 0000:20:1c.0, compute capability: 8.0
I0000 00:00:1736891783.523626    9168 gpu_device.cc:2022] Created device /job:localhost/replica:0/task:0/device:GPU:3 with 37945 MB memory:  -> device: 3, name: NVIDIA A100-SXM4-40GB, pci bus id: 0000:20:1d.0, compute capability: 8.0
I0000 00:00:1736891783.525222    9168 gpu_device.cc:2022] Created device /job:localhost/replica:0/task:0/device:GPU:4 with 37945 MB memory:  -> device: 4, name: NVIDIA A100-SXM4-40GB, pci bus id: 0000:90:1c.0, compute capability: 8.0
I0000 00:00:1736891783.526807    9168 gpu_device.cc:2022] Created device /job:localhost/replica:0/task:0/device:GPU:5 with 37945 MB memory:  -> device: 5, name: NVIDIA A100-SXM4-40GB, pci bus id: 0000:90:1d.0, compute capability: 8.0
I0000 00:00:1736891783.528377    9168 gpu_device.cc:2022] Created device /job:localhost/replica:0/task:0/device:GPU:6 with 37945 MB memory:  -> device: 6, name: NVIDIA A100-SXM4-40GB, pci bus id: 0000:a0:1c.0, compute capability: 8.0
I0000 00:00:1736891783.529933    9168 gpu_device.cc:2022] Created device /job:localhost/replica:0/task:0/device:GPU:7 with 37945 MB memory:  -> device: 7, name: NVIDIA A100-SXM4-40GB, pci bus id: 0000:a0:1d.0, compute capability: 8.0
/test/bin/testDataservice: line 5:  9168 Segmentation fault      (core dumped) python ${BIN_DIR}/test_dataset_service.py --local=False
```