Add label to ts image (aws#3155)

* Add label to ts image * reformat * PT 1.13 x86 test * graviton pt 2.0 test * x86 pt 1.13 test * neuron pt 1.13 test * neuronx pt 1.13 test * x86 pt 1.12 test * graviton pt 1.12 test * neuron pt 1.12 test * x86 pt 1.11 test * revert toml * Build image label using image_build.py * reformat * fix typo. * x86 pt 2.0 test 2 * graviton pt 2.0 test 2 * x86 pt 1.13 test 2 * neuron pt 1.13 test 2 * Pass ts version to dockerfile. * x86 pt 2.0 test 3 * fix image_build.py * x86 pt 2.0 test 3 * Add ts version in buildspec * graviton pt 2.0 test 3 * x86 pt 1.13 test 3 * neuron pt 1.13 test 3 * neuronx pt 1.13 test 3 * neuron pt 1.13 test 4 * neuronx pt 1.13 test 4 * revert toml * remove Toolkit version from dockerfile. --------- Co-authored-by: Boyu Chen <[email protected]>
chen3933 · Jul 13, 2023 · 49953ec · 49953ec
1 parent bf9e4b8
commit 49953ec
Show file tree

Hide file tree

Showing 18 changed files with 128 additions and 22 deletions.
diff --git a/pytorch/inference/buildspec-1-13-neuron.yml b/pytorch/inference/buildspec-1-13-neuron.yml
@@ -40,6 +40,8 @@ images:
     tag_python_version: &TAG_PYTHON_VERSION py38
     neuron_sdk_version: &NEURON_SDK_VERSION sdk2.10.0
     os_version: &OS_VERSION ubuntu20.04
+    torch_serve_version: &TORCHSERVE_VERSION 0.7.1
+    tool_kit_version: &SM_TOOLKIT_VERSION 2.0.14
     tag: !join [ *VERSION, "-", *DEVICE_TYPE, "-", *TAG_PYTHON_VERSION, "-", *NEURON_SDK_VERSION, "-", *OS_VERSION ]
     docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /, *NEURON_SDK_VERSION, /Dockerfile., *DEVICE_TYPE ]
     context:

diff --git a/pytorch/inference/buildspec-1-13-neuronx.yml b/pytorch/inference/buildspec-1-13-neuronx.yml
@@ -40,6 +40,8 @@ images:
     tag_python_version: &TAG_PYTHON_VERSION py38
     neuron_sdk_version: &NEURON_SDK_VERSION sdk2.10.0
     os_version: &OS_VERSION ubuntu20.04
+    torch_serve_version: &TORCHSERVE_VERSION 0.7.1
+    tool_kit_version: &SM_TOOLKIT_VERSION 2.0.14
     tag: !join [ *VERSION, "-", *DEVICE_TYPE, "-", *TAG_PYTHON_VERSION, "-", *NEURON_SDK_VERSION, "-", *OS_VERSION ]
     docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /, *NEURON_SDK_VERSION, /Dockerfile., *DEVICE_TYPE ]
     context:

diff --git a/pytorch/inference/buildspec-1-13.yml b/pytorch/inference/buildspec-1-13.yml
@@ -36,6 +36,8 @@ images:
     python_version: &DOCKER_PYTHON_VERSION py3
     tag_python_version: &TAG_PYTHON_VERSION py39
     os_version: &OS_VERSION ubuntu20.04
+    torch_serve_version: &TORCHSERVE_VERSION 0.7.1
+    tool_kit_version: &SM_TOOLKIT_VERSION 2.0.14
     tag: !join [ *VERSION, "-", *DEVICE_TYPE, "-", *TAG_PYTHON_VERSION, "-", *OS_VERSION, "-ec2" ]
     docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /Dockerfile., *DEVICE_TYPE ]
     target: ec2
@@ -50,6 +52,8 @@ images:
     tag_python_version: &TAG_PYTHON_VERSION py39
     cuda_version: &CUDA_VERSION cu117
     os_version: &OS_VERSION ubuntu20.04
+    torch_serve_version: &TORCHSERVE_VERSION 0.7.1
+    tool_kit_version: &SM_TOOLKIT_VERSION 2.0.14
     tag: !join [ *VERSION, "-", *DEVICE_TYPE, "-", *TAG_PYTHON_VERSION, "-", *CUDA_VERSION, "-", *OS_VERSION, "-ec2" ]
     docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /, *CUDA_VERSION, /Dockerfile.,
                          *DEVICE_TYPE ]
@@ -64,6 +68,8 @@ images:
     python_version: &DOCKER_PYTHON_VERSION py3
     tag_python_version: &TAG_PYTHON_VERSION py39
     os_version: &OS_VERSION ubuntu20.04
+    torch_serve_version: &TORCHSERVE_VERSION 0.7.1
+    tool_kit_version: &SM_TOOLKIT_VERSION 2.0.14
     tag: !join [ *VERSION, "-", *DEVICE_TYPE, "-", *TAG_PYTHON_VERSION, "-", *OS_VERSION, "-sagemaker" ]
     docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /Dockerfile., *DEVICE_TYPE ]
     target: sagemaker
@@ -78,6 +84,8 @@ images:
     tag_python_version: &TAG_PYTHON_VERSION py39
     cuda_version: &CUDA_VERSION cu117
     os_version: &OS_VERSION ubuntu20.04
+    torch_serve_version: &TORCHSERVE_VERSION 0.7.1
+    tool_kit_version: &SM_TOOLKIT_VERSION 2.0.14
     tag: !join [ *VERSION, "-", *DEVICE_TYPE, "-", *TAG_PYTHON_VERSION, "-", *CUDA_VERSION, "-", *OS_VERSION, "-sagemaker" ]
     docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /, *CUDA_VERSION, /Dockerfile.,
                          *DEVICE_TYPE ]

diff --git a/pytorch/inference/buildspec-eia.yml b/pytorch/inference/buildspec-eia.yml
@@ -33,6 +33,8 @@
       python_version: &DOCKER_PYTHON_VERSION py38
       tag_python_version: &TAG_PYTHON_VERSION py38
       os_version: &OS_VERSION ubuntu20.04
+      torch_serve_version: &TORCHSERVE_VERSION 0.8.0
+      tool_kit_version: &SM_TOOLKIT_VERSION 2.0.14
       tag: !join [ *VERSION, "-", "cpu", "-", *TAG_PYTHON_VERSION, "-", *OS_VERSION ]
       docker_file: !join [ docker/, *VERSION, /, *DOCKER_PYTHON_VERSION, /Dockerfile., eia ]
       context:

diff --git a/pytorch/inference/buildspec-graviton.yml b/pytorch/inference/buildspec-graviton.yml
@@ -34,6 +34,8 @@ images:
     image_size_baseline: 10000
     device_type: &DEVICE_TYPE cpu
     os_version: &OS_VERSION ubuntu20.04
+    torch_serve_version: &TORCHSERVE_VERSION 0.8.0
+    tool_kit_version: &SM_TOOLKIT_VERSION 2.0.14
     python_version: &DOCKER_PYTHON_VERSION py3
     tag_python_version: &TAG_PYTHON_VERSION py310
     tag: !join [ *VERSION, "-", *DEVICE_TYPE, "-", *TAG_PYTHON_VERSION, "-", *OS_VERSION , "-ec2"]
@@ -47,6 +49,8 @@ images:
     image_size_baseline: 10000
     device_type: &DEVICE_TYPE cpu
     os_version: &OS_VERSION ubuntu20.04
+    torch_serve_version: &TORCHSERVE_VERSION 0.8.0
+    tool_kit_version: &SM_TOOLKIT_VERSION 2.0.14
     python_version: &DOCKER_PYTHON_VERSION py3
     tag_python_version: &TAG_PYTHON_VERSION py310
     tag: !join [ *VERSION, "-", *DEVICE_TYPE, "-", *TAG_PYTHON_VERSION, "-", *OS_VERSION , "-sagemaker"]

diff --git a/pytorch/inference/buildspec-neuron.yml b/pytorch/inference/buildspec-neuron.yml
@@ -40,6 +40,8 @@ images:
     tag_python_version: &TAG_PYTHON_VERSION py38
     neuron_sdk_version: &NEURON_SDK_VERSION sdk2.10.0
     os_version: &OS_VERSION ubuntu20.04
+    torch_serve_version: &TORCHSERVE_VERSION 0.7.1
+    tool_kit_version: &SM_TOOLKIT_VERSION 2.0.14
     tag: !join [ *VERSION, "-", *DEVICE_TYPE, "-", *TAG_PYTHON_VERSION, "-", *NEURON_SDK_VERSION, "-", *OS_VERSION ]
     docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /, *NEURON_SDK_VERSION, /Dockerfile., *DEVICE_TYPE ]
     context:

diff --git a/pytorch/inference/buildspec-neuronx.yml b/pytorch/inference/buildspec-neuronx.yml
@@ -40,6 +40,8 @@ images:
     tag_python_version: &TAG_PYTHON_VERSION py38
     neuron_sdk_version: &NEURON_SDK_VERSION sdk2.10.0
     os_version: &OS_VERSION ubuntu20.04
+    torch_serve_version: &TORCHSERVE_VERSION 0.7.1
+    tool_kit_version: &SM_TOOLKIT_VERSION 2.0.14
     tag: !join [ *VERSION, "-", *DEVICE_TYPE, "-", *TAG_PYTHON_VERSION, "-", *NEURON_SDK_VERSION, "-", *OS_VERSION ]
     docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /, *NEURON_SDK_VERSION, /Dockerfile., *DEVICE_TYPE ]
     context:

diff --git a/pytorch/inference/buildspec.yml b/pytorch/inference/buildspec.yml
@@ -36,6 +36,8 @@ images:
     python_version: &DOCKER_PYTHON_VERSION py3
     tag_python_version: &TAG_PYTHON_VERSION py310
     os_version: &OS_VERSION ubuntu20.04
+    torch_serve_version: &TORCHSERVE_VERSION 0.8.0
+    tool_kit_version: &SM_TOOLKIT_VERSION 2.0.14
     tag: !join [ *VERSION, "-", *DEVICE_TYPE, "-", *TAG_PYTHON_VERSION, "-", *OS_VERSION, "-ec2" ]
     docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /Dockerfile., *DEVICE_TYPE ]
     target: ec2
@@ -50,6 +52,8 @@ images:
     tag_python_version: &TAG_PYTHON_VERSION py310
     cuda_version: &CUDA_VERSION cu118
     os_version: &OS_VERSION ubuntu20.04
+    torch_serve_version: &TORCHSERVE_VERSION 0.8.0
+    tool_kit_version: &SM_TOOLKIT_VERSION 2.0.14
     tag: !join [ *VERSION, "-", *DEVICE_TYPE, "-", *TAG_PYTHON_VERSION, "-", *CUDA_VERSION, "-", *OS_VERSION, "-ec2" ]
     docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /, *CUDA_VERSION, /Dockerfile.,
                          *DEVICE_TYPE ]
@@ -64,6 +68,8 @@ images:
     python_version: &DOCKER_PYTHON_VERSION py3
     tag_python_version: &TAG_PYTHON_VERSION py310
     os_version: &OS_VERSION ubuntu20.04
+    torch_serve_version: &TORCHSERVE_VERSION 0.8.0
+    tool_kit_version: &SM_TOOLKIT_VERSION 2.0.14
     tag: !join [ *VERSION, "-", *DEVICE_TYPE, "-", *TAG_PYTHON_VERSION, "-", *OS_VERSION, "-sagemaker" ]
     docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /Dockerfile., *DEVICE_TYPE ]
     target: sagemaker
@@ -78,6 +84,8 @@ images:
     tag_python_version: &TAG_PYTHON_VERSION py310
     cuda_version: &CUDA_VERSION cu118
     os_version: &OS_VERSION ubuntu20.04
+    torch_serve_version: &TORCHSERVE_VERSION 0.8.0
+    tool_kit_version: &SM_TOOLKIT_VERSION 2.0.14
     tag: !join [ *VERSION, "-", *DEVICE_TYPE, "-", *TAG_PYTHON_VERSION, "-", *CUDA_VERSION, "-", *OS_VERSION, "-sagemaker" ]
     docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /, *CUDA_VERSION, /Dockerfile.,
                          *DEVICE_TYPE ]

diff --git a/pytorch/inference/docker/1.13/py3/Dockerfile.cpu b/pytorch/inference/docker/1.13/py3/Dockerfile.cpu
@@ -8,8 +8,8 @@ ARG TORCH_URL=https://aws-pytorch-unified-cicd-binaries.s3.us-west-2.amazonaws.c
 ARG TORCHVISION_VERSION=0.14.1+cpu
 ARG TORCHAUDIO_VERSION=0.13.1+cpu
 ARG TORCHDATA_VERSION=0.5.1
-ARG TORCHSERVE_VERSION=0.7.1
-ARG SM_TOOLKIT_VERSION=2.0.14
+ARG TORCHSERVE_VERSION
+ARG SM_TOOLKIT_VERSION
 
 ########################################################
 #  _____ ____ ____    ___
@@ -222,6 +222,7 @@ LABEL com.amazonaws.sagemaker.capabilities.multi-models=true
 
 ARG PYTHON
 ARG SM_TOOLKIT_VERSION
+ARG TORCHSERVE_VERSION
 
 ENV SAGEMAKER_SERVING_MODULE sagemaker_pytorch_serving_container.serving:main
 

diff --git a/pytorch/inference/docker/1.13/py3/cu117/Dockerfile.gpu b/pytorch/inference/docker/1.13/py3/cu117/Dockerfile.gpu
@@ -12,8 +12,8 @@ ARG TORCH_URL=https://aws-pytorch-unified-cicd-binaries.s3.us-west-2.amazonaws.c
 ARG TORCHVISION_VERSION=0.14.1+cu117
 ARG TORCHAUDIO_VERSION=0.13.1+cu117
 ARG TORCHDATA_VERSION=0.5.1
-ARG TORCHSERVE_VERSION=0.7.1
-ARG SM_TOOLKIT_VERSION=2.0.14
+ARG TORCHSERVE_VERSION
+ARG SM_TOOLKIT_VERSION
 
 ########################################################
 #  _____ ____ ____    ___
@@ -288,6 +288,7 @@ LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true
 
 ARG PYTHON
 ARG SM_TOOLKIT_VERSION
+ARG TORCHSERVE_VERSION
 
 ENV SAGEMAKER_SERVING_MODULE sagemaker_pytorch_serving_container.serving:main
 

diff --git a/pytorch/inference/docker/1.13/py3/sdk2.10.0/Dockerfile.neuron b/pytorch/inference/docker/1.13/py3/sdk2.10.0/Dockerfile.neuron
@@ -6,7 +6,8 @@ LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true
 
 ARG PYTHON=python3.8
 ARG PYTHON_VERSION=3.8.16
-ARG TS_VERSION=0.7.1
+ARG TORCHSERVE_VERSION
+ARG SM_TOOLKIT_VERSION
 ARG MAMBA_VERSION=23.1.0-1
 
 # Neuron SDK components version numbers
@@ -109,9 +110,9 @@ RUN pip install --no-cache-dir -U \
 RUN pip install neuron-cc[tensorflow]==$NEURON_CC_VERSION --extra-index-url https://pip.repos.neuron.amazonaws.com \
  && pip install torch-neuron==$NEURON_FRAMEWORK_VERSION --extra-index-url https://pip.repos.neuron.amazonaws.com \
  && pip install "protobuf<4" \
- && pip install torchserve==$TS_VERSION \
+ && pip install torchserve==$TORCHSERVE_VERSION \
  && pip install --no-deps --no-cache-dir -U torchvision==0.14.* \
- && pip install torch-model-archiver==$TS_VERSION
+ && pip install torch-model-archiver==$TORCHSERVE_VERSION
 
 RUN useradd -m model-server \
  && mkdir -p /home/model-server/tmp /opt/ml/model \
@@ -130,7 +131,7 @@ ADD https://raw.githubusercontent.com/aws/deep-learning-containers/master/src/de
 
 RUN chmod +x /usr/local/bin/deep_learning_container.py
 
-RUN pip install --no-cache-dir "sagemaker-pytorch-inference==2.0.14"
+RUN pip install --no-cache-dir "sagemaker-pytorch-inference==${SM_TOOLKIT_VERSION}"
 
 RUN HOME_DIR=/root \
  && curl -o ${HOME_DIR}/oss_compliance.zip https://aws-dlinfra-utilities.s3.amazonaws.com/oss_compliance.zip \

diff --git a/pytorch/inference/docker/1.13/py3/sdk2.10.0/Dockerfile.neuronx b/pytorch/inference/docker/1.13/py3/sdk2.10.0/Dockerfile.neuronx
@@ -6,7 +6,8 @@ LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true
 
 ARG PYTHON=python3.8
 ARG PYTHON_VERSION=3.8.16
-ARG TS_VERSION=0.7.1
+ARG TORCHSERVE_VERSION
+ARG SM_TOOLKIT_VERSION
 ARG MAMBA_VERSION=23.1.0-1
 
 # Neuron SDK components version numbers
@@ -111,9 +112,9 @@ RUN pip install --no-cache-dir -U \
 RUN pip install neuronx-cc==$NEURONX_CC_VERSION --extra-index-url https://pip.repos.neuron.amazonaws.com \
  && pip install torch-neuronx==$NEURONX_FRAMEWORK_VERSION --extra-index-url https://pip.repos.neuron.amazonaws.com \
  && pip install "protobuf<4" \
- && pip install torchserve==$TS_VERSION \
+ && pip install torchserve==$TORCHSERVE_VERSION \
  && pip install --no-deps --no-cache-dir -U torchvision==0.14.* \
- && pip install torch-model-archiver==$TS_VERSION
+ && pip install torch-model-archiver==$TORCHSERVE_VERSION
 
 RUN useradd -m model-server \
  && mkdir -p /home/model-server/tmp /opt/ml/model \
@@ -132,7 +133,7 @@ ADD https://raw.githubusercontent.com/aws/deep-learning-containers/master/src/de
 
 RUN chmod +x /usr/local/bin/deep_learning_container.py
 
-RUN pip install --no-cache-dir "sagemaker-pytorch-inference==2.0.14"
+RUN pip install --no-cache-dir "sagemaker-pytorch-inference==${SM_TOOLKIT_VERSION}"
 
 # patch default_pytorch_inference_handler.py to import torch_neuronx
 RUN DEST_DIR=$(python -c "import os.path, sagemaker_pytorch_serving_container; print(os.path.dirname(sagemaker_pytorch_serving_container.__file__))") \

diff --git a/pytorch/inference/docker/1.13/py3/sdk2.9.0/Dockerfile.neuronx b/pytorch/inference/docker/1.13/py3/sdk2.9.0/Dockerfile.neuronx
@@ -6,7 +6,8 @@ LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true
 
 ARG PYTHON=python3.8
 ARG PYTHON_VERSION=3.8.16
-ARG TS_VERSION=0.7.1
+ARG TORCHSERVE_VERSION
+ARG SM_TOOLKIT_VERSION
 ARG MAMBA_VERSION=23.1.0-1
 
 # Neuron SDK components version numbers
@@ -111,9 +112,9 @@ RUN pip install --no-cache-dir -U \
 RUN pip install neuronx-cc==$NEURONX_CC_VERSION --extra-index-url https://pip.repos.neuron.amazonaws.com \
  && pip install torch-neuronx==$NEURONX_FRAMEWORK_VERSION --extra-index-url https://pip.repos.neuron.amazonaws.com \
  && pip install "protobuf<4" \
- && pip install torchserve==$TS_VERSION \
+ && pip install torchserve==$TORCHSERVE_VERSION \
  && pip install --no-deps --no-cache-dir -U torchvision==0.14.* \
- && pip install torch-model-archiver==$TS_VERSION
+ && pip install torch-model-archiver==$TORCHSERVE_VERSION
 
 RUN useradd -m model-server \
  && mkdir -p /home/model-server/tmp /opt/ml/model \
@@ -132,7 +133,7 @@ ADD https://raw.githubusercontent.com/aws/deep-learning-containers/master/src/de
 
 RUN chmod +x /usr/local/bin/deep_learning_container.py
 
-RUN pip install --no-cache-dir "sagemaker-pytorch-inference==2.0.14"
+RUN pip install --no-cache-dir "sagemaker-pytorch-inference==${SM_TOOLKIT_VERSION}"
 
 RUN HOME_DIR=/root \
  && curl -o ${HOME_DIR}/oss_compliance.zip https://aws-dlinfra-utilities.s3.amazonaws.com/oss_compliance.zip \

diff --git a/pytorch/inference/docker/2.0/py3/Dockerfile.cpu b/pytorch/inference/docker/2.0/py3/Dockerfile.cpu
@@ -9,8 +9,8 @@ ARG TORCHVISION_URL=https://framework-binaries.s3.us-west-2.amazonaws.com/pytorc
 ARG TORCHAUDIO_URL=https://framework-binaries.s3.us-west-2.amazonaws.com/pytorch/v2.0.1/cpu/2023-05-23-20-21/torchaudio-2.0.2%2Bcpu-cp310-cp310-linux_x86_64.whl
 ARG TORCHTEXT_URL=https://framework-binaries.s3.us-west-2.amazonaws.com/pytorch/v2.0.1/cpu/2023-05-23-20-21/torchtext-0.15.2%2Bcpu-cp310-cp310-linux_x86_64.whl
 ARG TORCHDATA_URL=https://framework-binaries.s3.us-west-2.amazonaws.com/pytorch/v2.0.1/cpu/2023-05-23-20-21/torchdata-0.6.1%2Bcpu-cp310-cp310-linux_x86_64.whl
-ARG TORCHSERVE_VERSION=0.8.0
-ARG SM_TOOLKIT_VERSION=2.0.14
+ARG TORCHSERVE_VERSION
+ARG SM_TOOLKIT_VERSION
 
 ########################################################
 #  _____ ____ ____    ___
@@ -221,6 +221,7 @@ LABEL com.amazonaws.sagemaker.capabilities.multi-models=true
 
 ARG PYTHON
 ARG SM_TOOLKIT_VERSION
+ARG TORCHSERVE_VERSION
 
 ENV SAGEMAKER_SERVING_MODULE sagemaker_pytorch_serving_container.serving:main
 

diff --git a/pytorch/inference/docker/2.0/py3/Dockerfile.graviton.cpu b/pytorch/inference/docker/2.0/py3/Dockerfile.graviton.cpu
@@ -9,8 +9,8 @@ ARG TORCHVISION_URL=https://framework-binaries.s3.us-west-2.amazonaws.com/pytorc
 ARG TORCHAUDIO_URL=https://framework-binaries.s3.us-west-2.amazonaws.com/pytorch/v2.0.1/graviton/2023-05-23-21-36/torchaudio-2.0.2%2Bcpu-cp310-cp310-linux_aarch64.whl
 ARG TORCHTEXT_URL=https://framework-binaries.s3.us-west-2.amazonaws.com/pytorch/v2.0.1/graviton/2023-05-23-21-36/torchtext-0.15.2%2Bcpu-cp310-cp310-linux_aarch64.whl
 ARG TORCHDATA_URL=https://framework-binaries.s3.us-west-2.amazonaws.com/pytorch/v2.0.1/graviton/2023-05-23-21-36/torchdata-0.6.1%2Bcpu-cp310-cp310-linux_aarch64.whl
-ARG TORCHSERVE_VERSION=0.8.0
-ARG SM_TOOLKIT_VERSION=2.0.14
+ARG TORCHSERVE_VERSION
+ARG SM_TOOLKIT_VERSION
 
 ########################################################
 #  _____ ____ ____    ___
@@ -224,6 +224,7 @@ LABEL com.amazonaws.sagemaker.capabilities.multi-models=true
 
 ARG PYTHON
 ARG SM_TOOLKIT_VERSION
+ARG TORCHSERVE_VERSION
 
 ENV SAGEMAKER_SERVING_MODULE sagemaker_pytorch_serving_container.serving:main
 

diff --git a/pytorch/inference/docker/2.0/py3/cu118/Dockerfile.gpu b/pytorch/inference/docker/2.0/py3/cu118/Dockerfile.gpu
@@ -15,8 +15,8 @@ ARG TORCHVISION_URL=https://framework-binaries.s3.us-west-2.amazonaws.com/pytorc
 ARG TORCHAUDIO_URL=https://framework-binaries.s3.us-west-2.amazonaws.com/pytorch/v2.0.1/cu11.8.0/2023-05-24-03-41/torchaudio-2.0.2%2Bcu118-cp310-cp310-linux_x86_64.whl
 ARG TORCHTEXT_URL=https://framework-binaries.s3.us-west-2.amazonaws.com/pytorch/v2.0.1/cu11.8.0/2023-05-24-03-41/torchtext-0.15.2%2Bcu118-cp310-cp310-linux_x86_64.whl
 ARG TORCHDATA_URL=https://framework-binaries.s3.us-west-2.amazonaws.com/pytorch/v2.0.1/cu11.8.0/2023-05-24-03-41/torchdata-0.6.1%2Bcu118-cp310-cp310-linux_x86_64.whl
-ARG TORCHSERVE_VERSION=0.8.0
-ARG SM_TOOLKIT_VERSION=2.0.14
+ARG TORCHSERVE_VERSION
+ARG SM_TOOLKIT_VERSION
 
 ########################################################
 #  _____ ____ ____    ___
@@ -278,6 +278,7 @@ LABEL com.amazonaws.sagemaker.capabilities.multi-models=true
 
 ARG PYTHON
 ARG SM_TOOLKIT_VERSION
+ARG TORCHSERVE_VERSION
 
 ENV SAGEMAKER_SERVING_MODULE sagemaker_pytorch_serving_container.serving:main
 ENV TS_DISABLE_SYSTEM_METRICS=true

diff --git a/src/image_builder.py b/src/image_builder.py
@@ -177,6 +177,13 @@ def image_builder(buildspec, image_types=[], device_types=[]):
                     f"HuggingFace buildspec.yml must contain 'datasets_version' field for each image"
                 )
 
+        torchserve_version = image_config.get("torch_serve_version")
+        inference_toolkit_version = image_config.get("tool_kit_version")
+        if torchserve_version:
+            extra_build_args["TORCHSERVE_VERSION"] = torchserve_version
+        if inference_toolkit_version:
+            extra_build_args["SM_TOOLKIT_VERSION"] = inference_toolkit_version
+
         ARTIFACTS.update(
             {
                 "dockerfile": {
@@ -242,6 +249,10 @@ def image_builder(buildspec, image_types=[], device_types=[]):
                 labels[
                     f"com.amazonaws.ml.engines.{cx_type}.dlc.lib.transformers.{label_transformers_version}"
                 ] = "true"
+            if torchserve_version and inference_toolkit_version:
+                labels[
+                    f"com.amazonaws.ml.engines.{cx_type}.dlc.inference-toolkit.{inference_toolkit_version}.torchserve.{torchserve_version}"
+                ] = "true"
 
         """
         Override parameters from parent in child.