remove uv changes

Author: Jyothirmaikottu

buildspec.yml

@@ -10,7 +10,7 @@ phases:
       - start-dockerd
       - echo Logging in to Amazon ECR...
       - $(aws ecr get-login --no-include-email --region $AWS_DEFAULT_REGION)
-      - uv pip install -r src/requirements.txt
+      - pip install -r src/requirements.txt
       - bash src/setup.sh $FRAMEWORK
       - python src/parse_partner_developers.py
       - python src/send_status.py --status 2

pytorch/training/docker/2.8/py3/cu129/Dockerfile.gpu

@@ -53,8 +53,6 @@ ENV PYTHONUNBUFFERED=1
 ENV PYTHONIOENCODING=UTF-8
 ENV LANG=C.UTF-8
 ENV LC_ALL=C.UTF-8
-ENV UV_VERSION=0.8.22 
-ENV UV_SYSTEM_PYTHON=1
 
 ENV TORCH_NVCC_FLAGS="-Xfatbin -compress-all"
 
@@ -82,15 +80,8 @@ ENV LD_LIBRARY_PATH="/usr/local/lib:/opt/amazon/ofi-nccl/lib/x86_64-linux-gnu:/o
 # Python Path
 ENV PATH="/usr/local/bin:${PATH}"
 
-
-ENV UV_CACHE_DIR=/root/.cache/uv \
-    UV_SYSTEM_PYTHON=1
-
-RUN curl -LsSf https://astral.sh/uv/install.sh | sh \
-    && ln -s /root/.local/bin/uv /usr/local/bin/uv 
-
 # Install common conda packages
-RUN uv pip install --no-cache-dir \
+RUN pip install --no-cache-dir \
     cython \
     cryptography \
     pyOpenSSL \
@@ -123,11 +114,11 @@ RUN uv pip install --no-cache-dir \
     tornado>=6.5.1
 
 # Install PyTorch
-RUN uv pip install --no-cache-dir -U torch==${PYTORCH_VERSION} \
+RUN pip install --no-cache-dir -U torch==${PYTORCH_VERSION} \
     torchvision==${TORCHVISION_VERSION} \
     torchaudio==${TORCHAUDIO_VERSION} \
     --index-url https://download.pytorch.org/whl/cu129 \
-    && uv pip install --no-cache-dir -U torchtnt==${TORCHTNT_VERSION} \
+    && pip install --no-cache-dir -U torchtnt==${TORCHTNT_VERSION} \
     torchdata==${TORCHDATA_VERSION} \
     triton \
     s3torchconnector \
@@ -140,18 +131,18 @@ RUN uv pip install --no-cache-dir -U torch==${PYTORCH_VERSION} \
     thinc==8.3.4 \
     blis \
     numpy \
- && uv pip uninstall dataclasses
+ && pip uninstall -y dataclasses
 
 # Install flash attn and NVIDIA transformer engine.
 # Optionally set NVTE_FRAMEWORK to avoid bringing in additional frameworks during TE install
 ENV NVTE_FRAMEWORK=pytorch
 
 RUN curl -LO https://github.com/Dao-AILab/flash-attention/releases/download/v${FLASH_ATTN_VERSION}/flash_attn-${FLASH_ATTN_VERSION}+cu12torch2.8cxx11abiTRUE-cp312-cp312-linux_x86_64.whl \
-    && uv pip install flash_attn-${FLASH_ATTN_VERSION}+cu12torch2.8cxx11abiTRUE-cp312-cp312-linux_x86_64.whl --no-build-isolation \
+    && pip install flash_attn-${FLASH_ATTN_VERSION}+cu12torch2.8cxx11abiTRUE-cp312-cp312-linux_x86_64.whl --no-build-isolation \
     && rm flash_attn-${FLASH_ATTN_VERSION}+cu12torch2.8cxx11abiTRUE-cp312-cp312-linux_x86_64.whl
 
 # Install TE using instructions from https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/installation.html
-RUN uv pip install --no-cache-dir git+https://github.com/NVIDIA/TransformerEngine.git@release_v${TE_VERSION} --no-build-isolation
+RUN pip install --no-cache-dir git+https://github.com/NVIDIA/TransformerEngine.git@release_v${TE_VERSION} --no-build-isolation
 
 RUN curl -o /license.txt https://aws-dlc-licenses.s3.amazonaws.com/pytorch-2.8/license.txt
 
@@ -254,15 +245,15 @@ ARG PYTHON
 WORKDIR /
 
 # Install SM packages
-RUN uv pip install --no-cache-dir -U \
+RUN pip install --no-cache-dir -U \
     smclarify \
     "sagemaker>=2" \
     sagemaker-experiments \
     sagemaker-pytorch-training \
     sagemaker-training
 
 # Install extra packages
-RUN uv pip install --no-cache-dir -U \
+RUN pip install --no-cache-dir -U \
     bokeh \
     imageio \
     numba \

tensorflow/training/docker/2.18/py3/Dockerfile.cpu

@@ -264,7 +264,7 @@ RUN ${PIP} install --no-cache-dir -U \
     "tensorflow-datasets==4.9.7"
 
 RUN $PYTHON -m pip install --no-cache-dir -U \
-    numba \
+    numba==0.61.0 \
     bokeh \
     imageio \
     opencv-python \
@@ -279,8 +279,9 @@ RUN $PYTHON -m pip install --no-cache-dir -U \
 RUN $PYTHON -m pip install --no-cache-dir -U \
     sagemaker-experiments==0.1.45
 
+#pinning old version because of protobuf dependency with tensorflow-metadata
 RUN $PYTHON -m pip install --no-cache-dir -U \
-    sagemaker-training
+    sagemaker-training==4.8.4
 
     RUN $PYTHON -m pip install --no-cache-dir -U \
     sagemaker-tensorflow-training==20.4.1