Added arg to build docker with torchserve nightly

.github/workflows/regression_tests_docker.yml

@@ -35,12 +35,12 @@ jobs:
         if: contains(matrix.hardware, 'ubuntu')
         run: |
           cd docker
-          ./build_image.sh -bt ci -b $GITHUB_REF_NAME -t pytorch/torchserve:ci
+          ./build_image.sh -bt ci -n -b $GITHUB_REF_NAME -t pytorch/torchserve:ci
       - name: Build GPU Docker Image
         if: false == contains(matrix.hardware, 'ubuntu')
         run: |
           cd docker
-          ./build_image.sh -g -cv cu117 -bt ci -b $GITHUB_REF_NAME -t pytorch/torchserve:ci
+          ./build_image.sh -g -cv cu117 -bt ci -n -b $GITHUB_REF_NAME -t pytorch/torchserve:ci
       - name: Torchserve GPU Regression Tests
         if: false == contains(matrix.hardware, 'ubuntu')
         run: |

docker/Dockerfile

@@ -30,6 +30,7 @@ ARG PYTHON_VERSION=3.9
 FROM ${BASE_IMAGE} AS compile-image
 ARG BASE_IMAGE=ubuntu:rolling
 ARG PYTHON_VERSION
+ARG BUILD_NIGHTLY
 ENV PYTHONUNBUFFERED TRUE
 
 RUN --mount=type=cache,id=apt-dev,target=/var/cache/apt \
@@ -82,7 +83,12 @@ RUN \
     fi
 
 # Make sure latest version of torchserve is uploaded before running this
-RUN python -m pip install --no-cache-dir torchserve torch-model-archiver torch-workflow-archiver
+RUN \
+    if echo "$BUILD_NIGHTLY" | grep -q "false"; then \
+        python -m pip install --no-cache-dir torchserve torch-model-archiver torch-workflow-archiver;\
+    else \
+        python -m pip install --no-cache-dir torchserve-nightly torch-model-archiver-nightly torch-workflow-archiver-nightly;\
+    fi
 
 # Final image for production
 FROM ${BASE_IMAGE} AS production-image

docker/README.md

@@ -38,6 +38,7 @@ Use `build_image.sh` script to build the docker images. The script builds the `p
 |-t, --tag|Tag name for image. If not specified, script uses torchserve default tag names.|
 |-cv, --cudaversion| Specify to cuda version to use. Supported values `cu92`, `cu101`, `cu102`, `cu111`, `cu113`, `cu116`, `cu117`, `cu118`. Default `cu117`|
 |-ipex, --build-with-ipex| Specify to build with intel_extension_for_pytorch. If not specified, script builds without intel_extension_for_pytorch.|
+|-n, --nightly| Specify to build with TorchServe nightly.|
 |--codebuild| Set if you need [AWS CodeBuild](https://aws.amazon.com/codebuild/)|
 |-py, --pythonversion| Specify the python version to use. Supported values `3.8`, `3.9`, `3.10`. Default `3.9`|
 

docker/build_image.sh

@@ -11,6 +11,7 @@ USE_CUSTOM_TAG=false
 CUDA_VERSION=""
 USE_LOCAL_SERVE_FOLDER=false
 BUILD_WITH_IPEX=false
+BUILD_NIGHTLY=false
 PYTHON_VERSION=3.9
 
 for arg in "$@"
@@ -27,6 +28,7 @@ do
           echo "-lf, --use-local-serve-folder specify this option for the benchmark image if the current 'serve' folder should be used during automated benchmarks"
           echo "-ipex, --build-with-ipex specify to build with intel_extension_for_pytorch"
           echo "-py, --pythonversion specify to python version to use: Possible values: 3.8 3.9 3.10"
+          echo "-n, --nightly specify to build with TorchServe nightly"
           exit 0
           ;;
         -b|--branch_name)
@@ -66,6 +68,10 @@ do
           BUILD_WITH_IPEX=true
           shift
           ;;
+        -n|--nightly)
+          BUILD_NIGHTLY=true
+          shift
+          ;;
         -py|--pythonversion)
           PYTHON_VERSION="$2"
           if [[ $PYTHON_VERSION = 3.8 || $PYTHON_VERSION = 3.9 || $PYTHON_VERSION = 3.10 ]]; then
@@ -137,10 +143,10 @@ fi
 
 if [ "${BUILD_TYPE}" == "production" ]
 then
-  DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg CUDA_VERSION="${CUDA_VERSION}"  --build-arg PYTHON_VERSION="${PYTHON_VERSION}" -t "${DOCKER_TAG}" --target production-image  .
+  DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg CUDA_VERSION="${CUDA_VERSION}"  --build-arg PYTHON_VERSION="${PYTHON_VERSION}" --build-arg BUILD_NIGHTLY="${BUILD_NIGHTLY}" -t "${DOCKER_TAG}" --target production-image  .
 elif [ "${BUILD_TYPE}" == "ci" ]
 then
-  DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg CUDA_VERSION="${CUDA_VERSION}"  --build-arg PYTHON_VERSION="${PYTHON_VERSION}" --build-arg BRANCH_NAME="${BRANCH_NAME}"  -t "${DOCKER_TAG}" --target ci-image  .
+  DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg CUDA_VERSION="${CUDA_VERSION}"  --build-arg PYTHON_VERSION="${PYTHON_VERSION}" --build-arg BUILD_NIGHTLY="${BUILD_NIGHTLY}" --build-arg BRANCH_NAME="${BRANCH_NAME}"  -t "${DOCKER_TAG}" --target ci-image  .
 elif [ "${BUILD_TYPE}" == "benchmark" ]
 then
   DOCKER_BUILDKIT=1 docker build --pull --no-cache --file Dockerfile.benchmark --build-arg USE_LOCAL_SERVE_FOLDER=$USE_LOCAL_SERVE_FOLDER --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg BRANCH_NAME="${BRANCH_NAME}" --build-arg CUDA_VERSION="${CUDA_VERSION}" --build-arg MACHINE_TYPE="${MACHINE}" --build-arg PYTHON_VERSION="${PYTHON_VERSION}" -t "${DOCKER_TAG}" .

docker/docker_nightly.py

@@ -32,9 +32,9 @@
     gpu_version = f"{project}:gpu-{get_nightly_version()}"
 
     # Build Nightly images and append the date in the name
-    try_and_handle(f"./build_image.sh -t {organization}/{cpu_version}", dry_run)
+    try_and_handle(f"./build_image.sh -n -t {organization}/{cpu_version}", dry_run)
     try_and_handle(
-        f"./build_image.sh -g -cv cu117 -t {organization}/{gpu_version}",
+        f"./build_image.sh -g -cv cu117 -n -t {organization}/{gpu_version}",
         dry_run,
     )