Repository: tensorflow
Issue: Current tensorflow[and-cuda] installed by pip pulls ptxas which causes Jupyter kernel restart
Labels: stat:awaiting tensorflower,type:build/install,subtype: ubuntu/linux,TF 2.12

### Issue type

Build/Install

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

No

### Source

binary

### TensorFlow version

2.12.1

### Custom code

Yes

### OS platform and distribution

Red Hat Enterprise Linux release 8.5 (Ootpa) x86_64

### Mobile device

_No response_

### Python version

3.12

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

System: 12.3; Python env: nvidia-cuda-runtime-cu12==12.3.101 nvidia-cudnn-cu12==8.9.7.29

### GPU model and memory

Nvidia a10 24GB

### Current behavior?

### Test case
- The first linked standalone code implements a CNN for the CIFAR-10 dataset as a minimal example and test case.
- The second linked code is a kernel.json which implements the environment variables for the production environment.  

### Reproduction of environment 
I set up a conda environment as follows:
`$ conda create -n tf-2.16 python==3.12`
We then source the environment:
`$ source /mmfs1/apps/pyenvs/anaconda3-2022.05/bin/activate tf-2.16`
I then install Tensorflow as per [documentation](https://www.tensorflow.org/install/pip) and add jupyter to this environment:
`$ pip install tensorflow[and-cuda] jupyterlab`
A kernel is produced: 
`$ python -m ipykernel install --user --name=MyEnvName`

### Terminal behavior 
Within the terminal, normal execution of the test code works with some normally expected warnings.
```
# export environment variables for system CUDA and cudnn (nvidia upstream)
$ module load cuda/12.3
$ module load cudnn/8.9
$ source /mmfs1/apps/pyenvs/anaconda3-2022.05/bin/activate tf-2.16
$ python tf_gpu.py
```
When the code calls `model.fit()` we get the following STDERR:
```
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1717178625.148066  121055 service.cc:146] XLA service 0x155468007cb0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
I0000 00:00:1717178625.148196  121055 service.cc:154]   StreamExecutor device (0): NVIDIA A10, Compute Capability 8.6
2024-05-31 13:03:45.227617: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:268] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.
2024-05-31 13:03:45.621453: I external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:531] Loaded cuDNN version 8907
2024-05-31 13:03:48.511964: I external/local_xla/xla/stream_executor/cuda/cuda_asm_compiler.cc:393] ptxas warning : Registers are spilled to local memory in function 'gemm_fusion_dot_1290', 4 bytes spill stores, 4 bytes spill loads

I0000 00:00:1717178635.476383  121055 device_compiler.h:188] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.
```
Despite warnings, the training and test work begins and main exits normally. 

### Unexpected Jupyter behavior
Within jupyter, the same code run within blocks fails, with **Kernel Restarting: The kernel for .ipynb appears to have died. It will restart automatically** which is unhelpful. The STDERR from Jupyter is logged, and this gives
```
2024-05-31 11:15:06.061841: F external/local_xla/xla/service/gpu/nvptx_compiler.cc:619] ptxas returned an error during compilation of ptx to sass: 'INTERNAL: ptxas 12.3.103 has a bug that we think can affect XLA. Please use a different version.'  If the error message indicates that a file could not be written, please verify that sufficient filesystem space is provided.
```
So, the ptxas shipped with the pip distribution of `tensorflow[and-cuda]` as per current documentation pulls a bad version of `nvidia-cuda-nvcc-cu12`. 
 
### Fix which confirms the bug
I fixed this on production by bumping the offending version:
`$ pip install nvidia-cuda-nvcc-cu12==12.5.40`
Now, `fit()` finishes successfully, and a new warning is logged:
```
2024-05-31 13:19:06.243730: W external/local_xla/xla/service/gpu/nvptx_compiler.cc:764] The NVIDIA driver's CUDA version is 12.3 which is older than the ptxas CUDA version (12.5.40). Because the driver is older than the ptxas version, XLA is disabling parallel compilation, which may slow down compilation. You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.
```
The version of `nvidia-cuda-nvcc-cu12` which current stable pip installed during `pip install tensorflow[and-cuda]` does not work inside of Jupyter if not linked against some other system CUDA distribution (which is not always possible inside of a heterogeneous compute environment). 
A better fix would be to bump the nvidia pip packages so that the offending version of ptxas are avoided completely. 

### nvidia-smi output
```
$ nvidia-smi
Fri May 31 13:25:06 2024
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 545.23.08              Driver Version: 545.23.08    CUDA Version: 12.3     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  NVIDIA A10                     Off | 00000000:03:00.0 Off |                    0 |
|  0%   29C    P8               8W / 150W |      4MiB / 23028MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   1  NVIDIA A10                     Off | 00000000:27:00.0 Off |                    0 |
|  0%   30C    P8               9W / 150W |      4MiB / 23028MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   2  NVIDIA A10                     Off | 00000000:83:00.0 Off |                    0 |
|  0%   27C    P8               8W / 150W |      7MiB / 23028MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   3  NVIDIA A10                     Off | 00000000:A3:00.0 Off |                    0 |
|  0%   29C    P8              15W / 150W |      4MiB / 23028MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+

+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|  No running processes found                                                           |
+---------------------------------------------------------------------------------------+
```

### Standalone code to reproduce the issue

```shell
https://gist.github.com/StephenSzwiec/d4f6f587441f54b9b6c5fed5efd8002d
https://gist.github.com/StephenSzwiec/8107202f0bbdc52deb3c9ebcd5c7f828
```


### Relevant log output

```shell
2024-05-31 11:15:06.061841: F external/local_xla/xla/service/gpu/nvptx_compiler.cc:619] ptxas returned an error during compilation of ptx to sass: 'INTERNAL: ptxas 12.3.103 has a bug that we think can affect XLA. Please use a different version.'  If the error message indicates that a file could not be written, please verify that sufficient filesystem space is provided.
```