fix: move setuptools install to EC2/SageMaker stages for pkg_resources

Author: Bhanu Teja Goshikonda

pytorch/training/docker/2.10/py3/Dockerfile.cpu

@@ -229,9 +229,6 @@ RUN echo 'source /usr/local/bin/bash_telemetry.sh' >> /etc/bash.bashrc
 # Removing the cache as it is needed for security verification
 RUN rm -rf /root/.cache | true
 
-# Install setuptools for pkg_resources required by piplicenses (Python 3.12+ compatibility)
-RUN pip install --no-cache-dir setuptools
-
 ########################################################
 #  _____ ____ ____    ___
 # | ____/ ___|___ \  |_ _|_ __ ___   __ _  __ _  ___
@@ -251,6 +248,9 @@ FROM common AS ec2
 
 WORKDIR /
 
+# Ensure setuptools is available for OSS compliance (pkg_resources needed by piplicenses)
+RUN pip install --no-cache-dir setuptools
+
 COPY setup_oss_compliance.sh setup_oss_compliance.sh
 RUN bash setup_oss_compliance.sh ${PYTHON} && rm setup_oss_compliance.sh
 
@@ -311,6 +311,9 @@ COPY start_with_right_hostname.sh /usr/local/bin/start_with_right_hostname.sh
 
 RUN chmod +x /usr/local/bin/start_with_right_hostname.sh
 
+# Ensure setuptools is available for OSS compliance (pkg_resources needed by piplicenses)
+RUN pip install --no-cache-dir setuptools
+
 COPY setup_oss_compliance.sh setup_oss_compliance.sh
 RUN bash setup_oss_compliance.sh ${PYTHON} && rm setup_oss_compliance.sh
 

pytorch/training/docker/2.10/py3/cu130/Dockerfile.gpu

@@ -184,9 +184,6 @@ RUN apt-get update \
 # Removing the cache as it is needed for security verification
 RUN rm -rf /root/.cache | true
 
-# Install setuptools for pkg_resources required by piplicenses (Python 3.12+ compatibility)
-RUN pip install --no-cache-dir setuptools
-
 ########################################################
 #  _____ ____ ____    ___
 # | ____/ ___|___ \  |_ _|_ __ ___   __ _  __ _  ___
@@ -211,6 +208,9 @@ WORKDIR /
 COPY dockerd_entrypoint.sh /usr/local/bin/dockerd_entrypoint.sh
 RUN chmod +x /usr/local/bin/dockerd_entrypoint.sh
 
+# Ensure setuptools is available for OSS compliance (pkg_resources needed by piplicenses)
+RUN pip install --no-cache-dir setuptools
+
 COPY setup_oss_compliance.sh setup_oss_compliance.sh
 RUN bash setup_oss_compliance.sh ${PYTHON} && rm setup_oss_compliance.sh
 
@@ -272,6 +272,9 @@ RUN pip install --no-cache-dir -U \
     seaborn \
     cloudpickle
 
+# Ensure setuptools is available for OSS compliance (pkg_resources needed by piplicenses)
+RUN pip install --no-cache-dir setuptools
+
 COPY setup_oss_compliance.sh setup_oss_compliance.sh
 RUN bash setup_oss_compliance.sh ${PYTHON} && rm setup_oss_compliance.sh