TAO 5.0 Release - PyTorch

Author: Arun-George-Zachariah

README.md

@@ -82,7 +82,7 @@ Inorder to maintain a uniform development enviroment across all users, TAO Toolk
 ```sh
 usage: tao_pt [-h] [--gpus GPUS] [--volume VOLUME] [--env ENV]
               [--mounts_file MOUNTS_FILE] [--shm_size SHM_SIZE]
-              [--run_as_user] [--tag TAG] [--ulimit ULIMIT] [--port PORT]
+              [--run_as_user] [--ulimit ULIMIT] [--port PORT]
 
 Tool to run the pytorch container.
 
@@ -94,7 +94,6 @@ optional arguments:
   --mounts_file MOUNTS_FILE Path to the mounts file.
   --shm_size SHM_SIZE       Shared memory size for docker
   --run_as_user             Flag to run as user
-  --tag TAG                 The tag value for the local dev docker.
   --ulimit ULIMIT           Docker ulimits for the host machine.
   --port PORT               Port mapping (e.g. 8889:8889).
 
@@ -123,11 +122,10 @@ cd $NV_TAO_PYTORCH_TOP/docker
 ```
 
 #### <a name='Testthenewlybuiltbasedocker'></a>Test the newly built base docker
-
-The build script tags the newly built base docker with the username of the account in the user's local machine. Therefore, the developers may tests their new docker by using the `tao_pt` command with the `--tag` option.
+Developers may tests their new docker by using the `tao_pt` command.
 
 ```sh
-tao_pt --tag $USER -- script args
+tao_pt -- script args
 ```
 
 #### <a name='Updatethenewdocker'></a>Update the new docker

docker/build.sh

@@ -3,11 +3,10 @@
 set -eo pipefail
 cd "$( dirname "${BASH_SOURCE[0]}" )"
 
-registry="nvcr.io"
-repository="nvidia/tao/tao-toolkit"
-
-tag="5.0.0-pyt-base"
-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"
@@ -55,11 +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 \
+    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"

docker/manifest.json

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

runner/tao_pt.py

@@ -28,7 +28,7 @@
 
 DOCKER_REGISTRY = docker_config["registry"]
 DOCKER_REPOSITORY = docker_config["repository"]
-DOCKER_DIGEST = docker_config["digest"]
+DOCKER_TAG = docker_config["tag"]
 DOCKER_COMMAND = "docker"
 HOME_PATH = os.path.expanduser("~")
 MOUNTS_PATH = os.path.join(HOME_PATH, ".tao_mounts.json")
@@ -57,14 +57,14 @@ def format_mounts(mount_points):
 
 def check_image_exists(docker_image):
     """Check if the image exists locally."""
-    check_command = '{} images | grep "\\<{}\\>" >/dev/null 2>&1'.format(DOCKER_COMMAND, docker_image)
+    check_command = '{} images | grep "\\<{}\\>" | grep "{}" >/dev/null 2>&1'.format(DOCKER_COMMAND, docker_image, DOCKER_TAG)
     rc = subprocess.call(check_command, stdout=sys.stderr, shell=True)
     return rc == 0
 
 
 def pull_base_container(docker_image):
     """Pull the default base container."""
-    pull_command = "{} pull {}@{}".format(DOCKER_COMMAND, docker_image, DOCKER_DIGEST)
+    pull_command = "{} pull {}:{}".format(DOCKER_COMMAND, docker_image, DOCKER_TAG)
     rc = subprocess.call(pull_command, stdout=sys.stderr, shell=True)
     return rc == 0
 
@@ -110,16 +110,23 @@ def get_docker_gpus_prefix(gpus):
     return gpu_string
 
 
+def create_base_docker():
+    """Function to create the base docker."""
+    create_command = "bash {}/docker/build.sh --build".format(ROOT_DIR)
+    try:
+        subprocess.run(create_command, stdout=sys.stderr, shell=True, check=True)
+    except subprocess.CalledProcessError as e:
+        raise RuntimeError(f"Container build failed with error {e}")
+
+
 def instantiate_dev_docker(gpus, mount_file,
                            mount_cli_list,
                            env_var_list,
-                           tag, command, ulimit=None,
+                           command, ulimit=None,
                            shm_size="16G", run_as_user=False,
                            port_mapping=None):
     """Instiate the docker container."""
-    docker_image = "{}/{}@{}".format(DOCKER_REGISTRY, DOCKER_REPOSITORY, DOCKER_DIGEST)
-    if tag is not None:
-        docker_image = "{}/{}:{}".format(DOCKER_REGISTRY, DOCKER_REPOSITORY, tag)
+    docker_image = "{}/{}:{}".format(DOCKER_REGISTRY, DOCKER_REPOSITORY, DOCKER_TAG)
 
     # Invoking the nvidia docker.
     gpu_string = get_docker_gpus_prefix(gpus)
@@ -195,7 +202,6 @@ def parse_cli_args(args=None):
     parser.add_argument("--shm_size", help="Shared memory size for docker", default="16G", type=str)
     parser.add_argument("--run_as_user", help="Flag to run as user", action="store_true", default=False)
 
-    parser.add_argument("--tag", help="The tag value for the local dev docker.", default=None, type=str)
     parser.add_argument("--ulimit", action='append', help="Docker ulimits for the host machine." )
     parser.add_argument(
         "--port",
@@ -221,17 +227,17 @@ def main(cl_args=None):
     # parse command line args.
     args = parse_cli_args(tao_pt_args)
     docker_image = "{}/{}".format(DOCKER_REGISTRY, DOCKER_REPOSITORY)
-    if args["tag"] is not None:
-        docker_image = "{}:{}".format(docker_image, args["tag"])
+
     if not check_image_exists(docker_image):
-        assert pull_base_container(docker_image), "The base container doesn't exist locally and " "the pull failed."
+        if not pull_base_container(docker_image):
+            print("The base container doesn't exist locally and the pull failed. Hence creating the base container")
+            create_base_docker()
     try:
         instantiate_dev_docker(
             args["gpus"], args["mounts_file"],
             args["volume"], args["env"],
-            args["tag"], command_args,
-            args["ulimit"], args["shm_size"],
-            args["run_as_user"],
+            command_args, args["ulimit"],
+            args["shm_size"], args["run_as_user"],
             args['port']
         )
     except subprocess.CalledProcessError:

scripts/envsetup.sh

@@ -51,6 +51,11 @@ function _check_tlt_pytorch_requirements(){
         return 1
     fi
 
+    if ! command -v jq >/dev/null; then
+        echo -e "\033[1;31mERROR:\033[0m jq not found"
+        return 1
+    fi
+
     # Check if docker was installed.
     if ! command -v docker >/dev/null; then
         warnings+=("docker not found")