tf.data.experimental.service.DispatchServer segfaults after normal start/shutdown sequence

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

source

TensorFlow version

tf 2.21.0

Custom code

Yes

OS platform and distribution

Ubuntu 24.04.3 LTS x86_64

Mobile device

No response

Python version

3.10.20

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

4 x NVIDIA GeForce RTX 2080 Ti, 11264 MiB each

Current behavior?

Current behavior

I re-ran this reproducer locally and confirmed the failure.

tf.data.experimental.service.DispatchServer goes through an apparently normal lifecycle:

• the dispatcher starts successfully
• TensorFlow logs that the server is running
• TensorFlow logs that the server shuts down successfully

but the process then ends with a segmentation fault:

Segmentation fault (core dumped)

The process exits with code 139.

This is particularly important because this is a hard segfault after an apparently successful start/shutdown sequence, not a normal exception and not an out-of-memory condition.

Expected behavior

Starting and shutting down tf.data.experimental.service.DispatchServer should not crash the interpreter. If the API usage is invalid, TensorFlow should return a normal Python/TensorFlow exception rather than causing a segmentation fault.

Minimal reproducer

import tensorflow as tf
import numpy as np

def generate_data(num_samples=1000, num_features=10):
    x = np.random.rand(num_samples, num_features).astype(np.float32)
    y = np.random.randint(0, 2, size=(num_samples,)).astype(np.int32)
    return x, y

def create_dataset(x, y):
    dataset = tf.data.Dataset.from_tensor_slices((x, y))
    dataset = dataset.shuffle(buffer_size=1000).batch(32)
    return dataset

def start_dispatch_server(start=True):
    return tf.data.experimental.service.DispatchServer(start=start)

x, y = generate_data()
dataset = create_dataset(x, y)

dispatch_server = start_dispatch_server(start=True)
dispatch_server = start_dispatch_server(start=False)

print(f"Dispatch server started at: {dispatch_server.target}")

### Standalone code to reproduce the issue

```shell
Environment requirements

Ubuntu 24.04.3 LTS x86_64
Python 3.10
TensorFlow 2.21.0
NumPy 2.2.6
Reproduced in local conda environment: tf_test

conda create -n tf_test python=3.10 -y
conda activate tf_test
pip install tensorflow==2.21.0 numpy==2.2.6

conda run --no-capture-output -n tf_test python repro.py

Relevant log output

(tf_test) user@user-Standard-PC-Q35-ICH9-2009:/home/lhj$ python /home/lhj/callChainBuild/src_tensorflow/gen_program/validation_results_rerun/crashes/4300.py
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1774354528.998577 3824177 port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
I0000 00:00:1774354529.146547 3824177 cpu_feature_guard.cc:227] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1774354530.615380 3824177 port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
I0000 00:00:1774354532.205013 3824177 gpu_device.cc:2043] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 9606 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 2080 Ti, pci bus id: 0000:01:00.0, compute capability: 7.5
I0000 00:00:1774354532.206587 3824177 gpu_device.cc:2043] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 9606 MB memory:  -> device: 1, name: NVIDIA GeForce RTX 2080 Ti, pci bus id: 0000:02:00.0, compute capability: 7.5
I0000 00:00:1774354532.208159 3824177 gpu_device.cc:2043] Created device /job:localhost/replica:0/task:0/device:GPU:2 with 9606 MB memory:  -> device: 2, name: NVIDIA GeForce RTX 2080 Ti, pci bus id: 0000:03:00.0, compute capability: 7.5
I0000 00:00:1774354532.209530 3824177 gpu_device.cc:2043] Created device /job:localhost/replica:0/task:0/device:GPU:3 with 9606 MB memory:  -> device: 3, name: NVIDIA GeForce RTX 2080 Ti, pci bus id: 0000:04:00.0, compute capability: 7.5
I0000 00:00:1774354532.437703 3824177 dispatcher_impl.cc:230] Started tf.data service dispatcher in non-fault-tolerant mode with config: goo.gle/debugonly   
protocol: "grpc"
job_gc_check_interval_ms: 600000
job_gc_timeout_ms: 300000
client_timeout_ms: 300000
worker_timeout_ms: 600000
worker_max_concurrent_snapshots: 3

It will not recover its state on restart.
I0000 00:00:1774354532.437809 3824177 server_lib.cc:88] Started tf.data DispatchServer running at 0.0.0.0:44709
I0000 00:00:1774354532.438009 3824177 server_lib.cc:100] Shut down DispatchServer server running at port 44709
Dispatch server started at: grpc://localhost:0
Segmentation fault (core dumped)