docker

Author: Borda

dockers/release/Dockerfile

@@ -40,6 +40,12 @@ RUN \
     # otherwise there is collision with folder name and pkg name on Pypi
     cd pytorch-lightning && \
     pip install setuptools==75.6.0 && \
+    pip install -U "lightning-utilities[cli]" && \
+    # drop deepspeed since it is not supported by our minimal Torch requirements \
+    if [[ "$PYTORCH_VERSION" =~ ^(2\.1|2\.2|2\.3|2\.4)$ ]]; then \
+        python -m lightning_utilities.cli requirements prune-pkgs --packages deepspeed --req_files requirements/fabric/strategies.txt ; \
+        python -m lightning_utilities.cli requirements prune-pkgs --packages deepspeed --req_files requirements/pytorch/strategies.txt ; \
+    fi && \
     PACKAGE_NAME=lightning pip install '.[extra,loggers,strategies]' --no-cache-dir && \
     PACKAGE_NAME=pytorch pip install '.[extra,loggers,strategies]' --no-cache-dir && \
     cd .. && \