Denial of Service in tf.raw_ops.Unstage

[Zoeeeeey]

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

No

Source

source

TensorFlow version

tf 2.16.1

Custom code

Yes

OS platform and distribution

Ubuntu 20.04

Mobile device

No response

Python version

3.9

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

12.0

GPU model and memory

No response

Current behavior?

Hi,
When using tf.raw_ops.Unstage, a Denial of Service was encountered under normal invocation, with all passed parameters meeting the requirements of the API documentation.
After a long wait, there was still no response.

Standalone code to reproduce the issue

import tensorflow as tf

tensor = tf.constant([1.0, 2.0, 3.0])
stage_op = tf.raw_ops.Stage(values=[tensor])
print(tensor.dtype)  # <dtype: 'float32'>

unstage_op = tf.raw_ops.Unstage(dtypes=[tensor.dtype])  # Hang

with tf.compat.v1.Session() as sess:
    print("here")
    sess.run(stage_op)
    result = sess.run(unstage_op)

print(result)  # Expected output: [1. 2. 3.]

Relevant log output

2024-04-22 02:23:12.166249: I tensorflow/core/util/port.cc:113] 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`.
2024-04-22 02:23:12.166825: I external/local_tsl/tsl/cuda/cudart_stub.cc:32] Could not find cuda drivers on your machine, GPU will not be used.
2024-04-22 02:23:12.169923: I external/local_tsl/tsl/cuda/cudart_stub.cc:32] Could not find cuda drivers on your machine, GPU will not be used.
2024-04-22 02:23:12.215234: 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: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-04-22 02:23:13.066132: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
2024-04-22 02:23:13.789282: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2251] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...
<dtype: 'float32'>

[SuryanarayanaY]

Hi @Zoeeeeey ,

I have tested the code and observed that the execution hangs at Unstage Op though the inputs seems valid. Attached gist for reference. Needs to look into the issue.

Thanks!