update NGC docker (Lightning-AI#13136)

* update docker
* Apply suggestions from code review

Co-authored-by: Akihiro Nitta <[email protected]>
Co-authored-by: Carlos Mocholí <[email protected]>

Author: 3 people

dockers/nvidia/Dockerfile

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-ARG PYTORCH_VERSION=21.11
+ARG PYTORCH_VERSION=22.04
 
 # https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes
 FROM nvcr.io/nvidia/pytorch:${PYTORCH_VERSION}-py3
@@ -37,20 +37,19 @@ RUN \
         cd .. ; \
     fi && \
 # save the examples
-    mv pytorch-lightning/_notebooks notebooks && \
+    mv pytorch-lightning/_notebooks/.notebooks/ notebooks && \
     mv pytorch-lightning/pl_examples . && \
 
 # Installations \
     pip install -q fire && \
-    python ./pytorch-lightning/.actions/assistant.py requirements_prune_pkgs horovod --req_files ./pytorch-lightning/requirements/extra.txt  && \
-    pip install "Pillow>=8.2, !=8.3.0" "cryptography>=3.4" "py>=1.10" --no-cache-dir --upgrade-strategy only-if-needed && \
-    pip install -r ./pytorch-lightning/requirements/extra.txt --no-cache-dir --upgrade-strategy only-if-needed && \
-    pip install -r ./pytorch-lightning/requirements/examples.txt --no-cache-dir --upgrade-strategy only-if-needed && \
-    pip install ./pytorch-lightning --no-cache-dir && \
+    pip install "Pillow>=8.2, !=8.3.0" "cryptography>=3.4" "py>=1.10" --no-cache-dir && \
+    pip install ./pytorch-lightning["extra","loggers","strategies"] --no-cache-dir && \
+    pip install -r ./pytorch-lightning/requirements/examples.txt --no-cache-dir && \
     rm -rf pytorch-lightning && \
-    pip install jupyterlab[all] -U && \
     pip list
 
+RUN pip install jupyterlab[all] -U
+
 RUN pip install lightning-grid -U && \
     pip install "py>=1.10" "protobuf>=3.15.6" --upgrade-strategy only-if-needed
 

dockers/release/Dockerfile

@@ -36,7 +36,7 @@ RUN \
         mv pytorch-lightning-*/ pytorch-lightning ; \
         rm *.zip ; \
     fi && \
-    pip install ./pytorch-lightning["extra"] --no-cache-dir && \
+    pip install ./pytorch-lightning["extra","loggers","strategies"] --no-cache-dir && \
     rm -rf pytorch-lightning
 
 RUN python --version && \