TAO 5.5 Release - PyTorch

Author: achafale

docker/Dockerfile

@@ -46,11 +46,6 @@ RUN pip uninstall -y jupyterlab \
   && rm requirements-pip.txt \
   && mim install "mmengine==0.10.4" "mmpretrain>=1.0.0rc8"
 
-  COPY tao-core tao-core
-  RUN cd tao-core && bash release/python/build_wheel.sh && \
-      find dist/ -name "nvidia_tao_core*.whl" -type f | xargs -n 1 pip install && \
-      rm -rf *
-
 ENV FORCE_CUDA="1"
 RUN mkdir mmcv_src && \
   cd mmcv_src && \

docker/build.sh

@@ -3,12 +3,10 @@
 set -eo pipefail
 cd "$( dirname "${BASH_SOURCE[0]}" )"
 
-registry="nvcr.io"
-repository="nvstaging/tao/tao_pytorch_base_image"
-base_image="nvcr.io/nvidia/pytorch:24.04-py3"
-
-tag="$USER-$(date +%Y%m%d%H%M)"
-local_tag="$USER"
+# Read parameters from manifest.json
+registry=`jq -r '.registry' $NV_TAO_PYTORCH_TOP/docker/manifest.json`
+repository=`jq -r '.repository' $NV_TAO_PYTORCH_TOP/docker/manifest.json`
+tag=`jq -r '.tag' $NV_TAO_PYTORCH_TOP/docker/manifest.json`
 
 # Build parameters.
 BUILD_DOCKER="0"
@@ -56,12 +54,10 @@ if [ $BUILD_DOCKER = "1" ]; then
     else
         NO_CACHE=""
     fi
-    DOCKER_BUILDKIT=1 docker build --pull -f $NV_TAO_PYTORCH_TOP/docker/Dockerfile -t $registry/$repository:$local_tag $NO_CACHE \
-        --network=host --build-arg PYTORCH_BASE_IMAGE=${base_image} \
-        $NV_TAO_PYTORCH_TOP/.
+    DOCKER_BUILDKIT=1 docker build --pull -f $NV_TAO_PYTORCH_TOP/docker/Dockerfile -t $registry/$repository:$tag $NO_CACHE \
+        --network=host $NV_TAO_PYTORCH_TOP/.
     if [ $PUSH_DOCKER = "1" ]; then
         echo "Pusing docker ..."
-        docker tag $registry/$repository:$local_tag $registry/$repository:$tag
         docker push $registry/$repository:$tag
         digest=$(docker inspect --format='{{index .RepoDigests 0}}' $registry/$repository:$tag)
         echo -e "\033[1;33mUpdate the digest in the manifest.json file to:\033[0m"
@@ -72,4 +68,4 @@ if [ $BUILD_DOCKER = "1" ]; then
 # Exit by printing usage.
 else
     echo "Usage: ./build.sh [--build] [--push] [--force] [--default]"
-fi
+fi

docker/manifest.json

@@ -1,5 +1,6 @@
 {
     "registry": "nvcr.io",
     "repository": "nvidia/tao/tao-toolkit",
-    "digest": "sha256:ff920080025e2dd12d5740c3d722f0630f1f84630ae6f55e56f872d4914fca03"
+    "digest": "sha256:ff920080025e2dd12d5740c3d722f0630f1f84630ae6f55e56f872d4914fca03",
+    "tag": "5.5.0-pyt-base"
 }

release/docker/Dockerfile

@@ -29,11 +29,6 @@ COPY release/docker/entrypoint.d/* /opt/nvidia/entrypoint.d/
 # Copying tlt-converter and building it in the docker.
 ENV TRT_LIB_PATH="/usr/lib/x86_64-linux-gnu"
 ENV TRT_INCLUDE_PATH="/usr/include/x86_64-linux-gnu"
-COPY tao-converter /opt/nvidia/tools
-WORKDIR /opt/nvidia/tools
-RUN make \
-    && rm -rf Makefile *.cc *.h *.o
-RUN ln -s tao-converter converter
 
 ENV NVIDIA_PRODUCT_NAME "TAO Toolkit"
 ENV TAO_TOOLKIT_VERSION="5.2.0"

release/docker/deploy.sh

@@ -4,18 +4,16 @@ set -eo pipefail
 # cd "$( dirname "${BASH_SOURCE[0]}" )"
 
 registry="nvcr.io"
-pytorch_version="2.1.0"
-tao_version="5.2.0"
-repository="nvstaging/tao/tao-toolkit-pyt"
-build_id="01"
-tag="v${tao_version}-pyt${pytorch_version}-py3-${build_id}"
+tao_version="5.0.0"
+repository="nvidia/tao/tao-toolkit"
+tag="${tao_version}-pyt-base"
 
 # Required for the tao-converter since it is now a submodule.
 git submodule update --init --recursive
 
 # Build parameters.
 BUILD_DOCKER="0"
-BUILD_WHEEL="0"
+BUILD_WHEELS="0"
 PUSH_DOCKER="0"
 FORCE="0"
 
@@ -83,7 +81,6 @@ if [ $BUILD_DOCKER = "1" ]; then
           mkdir -p $wheel_dir
         fi
         echo "Building source code wheel ..."
-        # tao_pt --env 'TORCH_CUDA_ARCH_LIST="5.3 6.0 6.1 7.0 7.5 8.0 8.6 9.0"' -- bash /tao-pt/release/docker/build_wheel.sh
         tao_pt --env 'TORCH_CUDA_ARCH_LIST="5.3 6.0 6.1 7.0 7.5 8.0 8.6 9.0"' -- python setup.py bdist_wheel
     else
         echo "Skipping wheel builds ..."
@@ -108,12 +105,12 @@ if [ $BUILD_DOCKER = "1" ]; then
     fi
 elif [ $RUN_DOCKER ="1" ]; then
     echo "Running docker interactively..."
-    docker run --gpus all -v $HOME/tlt-experiments:/workspace/tlt-experiments \
+    docker run --gpus all -v $HOME/tlt-experiments:/workspace/tlt-experiments  \
                           --network=host \
                           --shm-size=30g \
                           --ulimit memlock=-1 \
                           --ulimit stack=67108864 \
                           --rm -it $registry/$repository:$tag /bin/bash
 else
     echo "Usage: ./deploy.sh [--build] [--wheel] [--run] [--push] [--default]"
-fi
+fi