Repository: tensorflow
Issue: Regression: TF 2.18 crashes with cudaSetDevice failing due to GPU being busy
Labels: stat:awaiting tensorflower,type:build/install,comp:gpu,TF 2.18

### Issue type

Bug

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

No

### Source

binary

### TensorFlow version

tensorflow[and-cuda] 2.18

### Custom code

No

### OS platform and distribution

Rocky 9.1

### Mobile device

_No response_

### Python version

3.12.7

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

get_build_info gives cuda 12.5.1, cudnn version 9, Driver version 525.78.01, system CUDA version 12.0

### GPU model and memory

RTX 3090

### Current behavior?

Tensorflow[and-cuda] 2.17 installed cuda_version 12.3 and cudnn 8, and the code below works with the normal error spew about being unable to register cuda factories. With 2.18, I get a crash about the gpu is busy.

This is on a desktop workstation with a window manager, so TF cannot have the entire GPU, but I'm not sure if that's the cause of the problem.

These are pretty normal conda environments, with `pip install "tensorflow[and-cuda]"` and then 

```
export XLA_FLAGS="--xla_gpu_cuda_data_dir=/path/to/my/conda/environment"


NVIDIA_DIR=$(dirname $(dirname $(python -c "import nvidia.cudnn; print(nvidia.cudnn.__file__)")))

pathAccum=""

for dir in $NVIDIA_DIR/*; do
    if [ -d "$dir/lib" ]; then
        pathAccum="$dir/lib:$pathAccum"
    fi
done

export LD_LIBRARY_PATH="${pathAccum}${LD_LIBRARY_PATH}"
```

inside the activate.d for the conda environment. 
I have tried with and without this activate script. It is necessary for 2.17, and I get the same error when I remove it for 2.18. I'm not sure if this is due to "hermetic cuda", since it's not at all clear what that means to an end user.

Interestingly, `tf.config.list_physical_devices("GPU")` correctly identifies the GPU, it returns `PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')`. 

### Standalone code to reproduce the issue

```shell
#!/usr/bin/env python3
# TensorFlow and tf.keras
from tensorflow import keras

model = keras.Sequential([
    keras.layers.Flatten(input_shape=(28, 28)),
    keras.layers.Dense(128, activation='relu'),
    keras.layers.Dense(10)
])
```


### Relevant log output

```shell
TF 2.17 (no bug):

2024-10-25 14:07:40.179745: I tensorflow/core/util/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`.
2024-10-25 14:07:40.187066: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2024-10-25 14:07:40.194965: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2024-10-25 14:07:40.197313: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2024-10-25 14:07:40.203777: 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 AVX_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-10-25 14:07:40.624911: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
/scratch/miniconda/install/envs/bpreveal-teak/lib/python3.12/site-packages/keras/src/layers/reshaping/flatten.py:37: UserWarning: Do not pass an `input_shape`/`input_dim` argument to a layer. When using Sequential models, prefer using an `Input(shape)` object as the first layer in the model instead.
  super().__init__(**kwargs)
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1729883261.043484 2360745 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1729883261.057001 2360745 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1729883261.057122 2360745 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1729883261.058905 2360745 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1729883261.059226 2360745 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1729883261.059269 2360745 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1729883261.521319 2360745 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1729883261.521431 2360745 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1729883261.521503 2360745 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-10-25 14:07:41.521550: I tensorflow/core/common_runtime/gpu/gpu_device.cc:2021] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 18871 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 3090 Ti, pci bus id: 0000:01:00.0, compute capability: 8.6



TF 2.18 (with bug):

2024-10-25 14:02:53.202498: I tensorflow/core/util/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`.
2024-10-25 14:02:53.209011: 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:1729882973.216967 2360012 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:1729882973.219326 2360012 cuda_blas.cc:1418] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2024-10-25 14:02:53.227106: 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 AVX_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
/scratch/miniconda/install/envs/bpreveal-testing/lib/python3.12/site-packages/keras/src/layers/reshaping/flatten.py:37: UserWarning: Do not pass an `input_shape`/`input_dim` argument to a layer. When using Sequential models, prefer using an `Input(shape)` object as the first layer in the model instead.
  super().__init__(**kwargs)
Traceback (most recent call last):
  File "/home/cm2363/scratch/./tfTest.py", line 5, in <module>
    model = keras.Sequential([
            ^^^^^^^^^^^^^^^^^^
  File "/scratch/miniconda/install/envs/bpreveal-testing/lib/python3.12/site-packages/keras/src/models/sequential.py", line 76, in __init__
    self._maybe_rebuild()
  File "/scratch/miniconda/install/envs/bpreveal-testing/lib/python3.12/site-packages/keras/src/models/sequential.py", line 141, in _maybe_rebuild
    self.build(input_shape)
  File "/scratch/miniconda/install/envs/bpreveal-testing/lib/python3.12/site-packages/keras/src/layers/layer.py", line 226, in build_wrapper
    original_build_method(*args, **kwargs)
  File "/scratch/miniconda/install/envs/bpreveal-testing/lib/python3.12/site-packages/keras/src/models/sequential.py", line 187, in build
    x = layer(x)
        ^^^^^^^^
  File "/scratch/miniconda/install/envs/bpreveal-testing/lib/python3.12/site-packages/keras/src/utils/traceback_utils.py", line 122, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/scratch/miniconda/install/envs/bpreveal-testing/lib/python3.12/site-packages/keras/src/backend/tensorflow/core.py", line 125, in convert_to_tensor
    return tf.convert_to_tensor(x, dtype=dtype)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
tensorflow.python.framework.errors_impl.InternalError: cudaSetDevice() on GPU:0 failed. Status: CUDA-capable device(s) is/are busy or unavailable
```