update use of CUDA images (#21062)

Author: Borda

.azure/gpu-benchmarks.yml

@@ -46,7 +46,7 @@ jobs:
     variables:
       DEVICES: $( python -c 'print("$(Agent.Name)".split("_")[-1])' )
     container:
-      image: "pytorchlightning/pytorch_lightning:base-cuda-py3.12-torch2.5-cuda12.1.0"
+      image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.8"
       options: "--gpus=all --shm-size=32g"
     strategy:
       matrix:

.azure/gpu-tests-fabric.yml

@@ -57,16 +57,16 @@ jobs:
     strategy:
       matrix:
         "Fabric | oldest":
-          image: "pytorchlightning/pytorch_lightning:base-cuda-py3.10-torch2.1-cuda12.1.1"
+          image: "pytorchlightning/pytorch_lightning:base-cuda12.1.1-py3.10-torch2.1"
           PACKAGE_NAME: "fabric"
         "Fabric | latest":
-          image: "pytorchlightning/pytorch_lightning:base-cuda-py3.12-torch2.8-cuda12.6.3"
+          image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.8"
           PACKAGE_NAME: "fabric"
         #"Fabric | future":
-        #  image: "pytorchlightning/pytorch_lightning:base-cuda-py3.12-torch2.7-cuda12.6.3"
+        #  image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.7"
         #  PACKAGE_NAME: "fabric"
         "Lightning | latest":
-          image: "pytorchlightning/pytorch_lightning:base-cuda-py3.12-torch2.8-cuda12.6.3"
+          image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.8"
           PACKAGE_NAME: "lightning"
     workspace:
       clean: all

.azure/gpu-tests-pytorch.yml

@@ -50,16 +50,16 @@ jobs:
     strategy:
       matrix:
         "PyTorch | oldest":
-          image: "pytorchlightning/pytorch_lightning:base-cuda-py3.10-torch2.1-cuda12.1.1"
+          image: "pytorchlightning/pytorch_lightning:base-cuda12.1.1-py3.10-torch2.1"
           PACKAGE_NAME: "pytorch"
         "PyTorch | latest":
-          image: "pytorchlightning/pytorch_lightning:base-cuda-py3.12-torch2.8-cuda12.6.3"
+          image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.8"
           PACKAGE_NAME: "pytorch"
         #"PyTorch | future":
-        #  image: "pytorchlightning/pytorch_lightning:base-cuda-py3.12-torch2.7-cuda12.6.3"
+        #  image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.7"
         #  PACKAGE_NAME: "pytorch"
         "Lightning | latest":
-          image: "pytorchlightning/pytorch_lightning:base-cuda-py3.12-torch2.8-cuda12.6.3"
+          image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.8"
           PACKAGE_NAME: "lightning"
     pool: lit-rtx-3090
     variables:

.lightning/workflows/fabric.yml

@@ -10,13 +10,13 @@ parametrize:
   matrix: {}
   include:
     # note that this is setting also all oldest requirements which is linked to Torch == 2.0
-    - image: "pytorchlightning/pytorch_lightning:base-cuda-py3.10-torch2.1-cuda12.1.1"
+    - image: "pytorchlightning/pytorch_lightning:base-cuda12.1.1-py3.10-torch2.1"
       PACKAGE_NAME: "fabric"
-    - image: "pytorchlightning/pytorch_lightning:base-cuda-py3.12-torch2.7-cuda12.6.3"
+    - image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.7"
       PACKAGE_NAME: "fabric"
-    # - image: "pytorchlightning/pytorch_lightning:base-cuda-py3.12-torch2.7-cuda12.6.3"
+    # - image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.7"
     #   PACKAGE_NAME: "fabric"
-    - image: "pytorchlightning/pytorch_lightning:base-cuda-py3.12-torch2.7-cuda12.6.3"
+    - image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.7"
       PACKAGE_NAME: "lightning"
   exclude: []
 

.lightning/workflows/pytorch.yml

@@ -10,13 +10,13 @@ parametrize:
   matrix: {}
   include:
     # note that this is setting also all oldest requirements which is linked to Torch == 2.0
-    - image: "pytorchlightning/pytorch_lightning:base-cuda-py3.10-torch2.1-cuda12.1.1"
+    - image: "pytorchlightning/pytorch_lightning:base-cuda12.1.1-py3.10-torch2.1"
       PACKAGE_NAME: "pytorch"
-    - image: "pytorchlightning/pytorch_lightning:base-cuda-py3.12-torch2.7-cuda12.6.3"
+    - image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.7"
       PACKAGE_NAME: "pytorch"
-    # - image: "pytorchlightning/pytorch_lightning:base-cuda-py3.12-torch2.7-cuda12.6.3"
+    # - image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.7"
     #   PACKAGE_NAME: "pytorch"
-    - image: "pytorchlightning/pytorch_lightning:base-cuda-py3.12-torch2.7-cuda12.6.3"
+    - image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.7"
       PACKAGE_NAME: "lightning"
   exclude: []
 

dockers/README.md

@@ -11,7 +11,13 @@ git clone https://github.com/Lightning-AI/lightning.git
 docker image build -t pytorch-lightning:latest -f dockers/base-cuda/Dockerfile .
 
 # build with specific arguments
-docker image build -t pytorch-lightning:base-cuda-py3.9-torch1.13-cuda11.7.1 -f dockers/base-cuda/Dockerfile --build-arg PYTHON_VERSION=3.9 --build-arg PYTORCH_VERSION=1.13 --build-arg CUDA_VERSION=11.7.1 .
+docker image build \
+  -t pytorch-lightning:base-cuda12.6.3-py3.10-torch2.8 \
+  -f dockers/base-cuda/Dockerfile \
+  --build-arg PYTHON_VERSION=3.10 \
+  --build-arg PYTORCH_VERSION=2.8 \
+  --build-arg CUDA_VERSION=12.6.3 \
+  .
 ```
 
 To run your docker use
@@ -45,18 +51,18 @@ sudo systemctl restart docker
 and later run the docker image with `--gpus all`. For example,
 
 ```
-docker run --rm -it --gpus all pytorchlightning/pytorch_lightning:base-cuda-py3.9-torch1.13-cuda11.7.1
+docker run --rm -it --gpus all pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.10-torch2.8
 ```
 
 ## Run Jupyter server
 
 1. Build the docker image:
    ```bash
-   docker image build -t pytorch-lightning:v1.6.5 -f dockers/nvidia/Dockerfile --build-arg LIGHTNING_VERSION=1.6.5 .
+   docker image build -t pytorch-lightning:v2.5.1 -f dockers/nvidia/Dockerfile --build-arg LIGHTNING_VERSION=2.5.1 .
    ```
 1. start the server and map ports:
    ```bash
-   docker run --rm -it --gpus=all -p 8888:8888 pytorch-lightning:v1.6.5
+   docker run --rm -it --gpus=all -p 8888:8888 pytorch-lightning:v2.5.1
    ```
 1. Connect in local browser:
    - copy the generated path e.g. `http://hostname:8888/?token=0719fa7e1729778b0cec363541a608d5003e26d4910983c6`

dockers/release/Dockerfile

@@ -16,7 +16,7 @@ ARG PYTHON_VERSION=3.10
 ARG PYTORCH_VERSION=2.8
 ARG CUDA_VERSION=12.6.3
 
-FROM pytorchlightning/pytorch_lightning:base-cuda-py${PYTHON_VERSION}-torch${PYTORCH_VERSION}-cuda${CUDA_VERSION}
+FROM pytorchlightning/pytorch_lightning:base-cuda${CUDA_VERSION}-py${PYTHON_VERSION}-torch${PYTORCH_VERSION}
 
 LABEL maintainer="Lightning-AI <https://github.com/Lightning-AI>"