Docker Ci update

.github/workflows/docker-ci.yaml

@@ -19,14 +19,25 @@ jobs:
       - uses: actions/checkout@v3
         with:
           submodules: recursive
+          repository: ${{ github.event.pull_request.head.repo.full_name }}
+          fetch-depth: 0 
 
       - name: Test build_image.sh script with custom tagging and gpu flag
         working-directory: docker
         run: ./test_build_image_tagging.sh ${{ matrix.python-version }}
 
+      - name: Check if fork
+        run: |
+          CURR_HEAD=${{ github.event.pull_request.head.repo.full_name }}
+          BASE_HEAD=${{ github.event.pull_request.base.repo.full_name }}
+          if [ "${CURR_HEAD}" != "${BASE_HEAD}" ]; then
+            echo "FORK=${CURR_HEAD}" >> $GITHUB_ENV
+          fi
+
       - name: Determine the branch name
         id: branch-name
         run: |
+          git fetch --all
           if [[ "${GITHUB_REF}" == refs/heads/* ]]; then
             GITHUB_BRANCH=${GITHUB_REF#refs/heads/}
           elif [[ "${GITHUB_REF}" == refs/pull/*/merge ]]; then
@@ -39,7 +50,11 @@ jobs:
         working-directory: docker
         run: |
           IMAGE_TAG=test-image-${{ matrix.python-version }}
-          ./build_image.sh -py "${{ matrix.python-version }}" -t "${IMAGE_TAG}" -b ${{ steps.branch-name.outputs.GITHUB_BRANCH }} -s
+          if [[ -z "${{ env.FORK }}" ]]; then
+            ./build_image.sh -py "${{ matrix.python-version }}" -t "${IMAGE_TAG}" -b "${{ steps.branch-name.outputs.GITHUB_BRANCH }}" -s
+          else
+            ./build_image.sh -py "${{ matrix.python-version }}" -t "${IMAGE_TAG}" -b "${{ steps.branch-name.outputs.GITHUB_BRANCH }}" -s -f "${{ env.FORK }}"
+          fi
           echo "IMAGE_TAG=${IMAGE_TAG}" >> $GITHUB_OUTPUT
 
       - name: Container Healthcheck

docker/Dockerfile

@@ -34,6 +34,7 @@ ARG BUILD_NIGHTLY
 ARG BUILD_FROM_SRC
 ARG LOCAL_CHANGES
 ARG BRANCH_NAME
+ARG FORK_NAME
 ENV PYTHONUNBUFFERED TRUE
 
 RUN --mount=type=cache,id=apt-dev,target=/var/cache/apt \
@@ -71,12 +72,17 @@ COPY ./ serve
 
 RUN \
     if echo "$LOCAL_CHANGES" | grep -q "false"; then \
-        rm -rf serve;\
-        git clone --recursive https://github.com/pytorch/serve.git -b $BRANCH_NAME; \
+        if [ -n "$FORK_NAME" ]; then \
+            rm -rf serve; \
+            git clone --recursive https://github.com/${FORK_NAME}.git -b $BRANCH_NAME; \
+        else \
+            rm -rf serve; \
+            git clone --recursive https://github.com/pytorch/serve.git -b $BRANCH_NAME; \
+        fi \
     fi
 
 
-WORKDIR "serve"
+WORKDIR $(basename $FORK_NAME)
 
 RUN cp docker/dockerd-entrypoint.sh /usr/local/bin/dockerd-entrypoint.sh
 

docker/build_image.sh

@@ -4,6 +4,7 @@ set -o errexit -o nounset -o pipefail
 
 MACHINE=cpu
 BRANCH_NAME="master"
+FORK_NAME=""
 DOCKER_TAG="pytorch/torchserve:latest-cpu"
 BUILD_TYPE="production"
 BASE_IMAGE="ubuntu:20.04"
@@ -37,6 +38,7 @@ do
           echo "-n, --nightly specify to build with TorchServe nightly"
           echo "-s, --source specify to build with TorchServe from source"
           echo "-r, --remote specify to use local TorchServe"
+          echo "-f, --fork specify a fork to use"
           exit 0
           ;;
         -b|--branch_name)
@@ -110,6 +112,18 @@ do
           LOCAL_CHANGES=false
           shift
           ;;
+        -f|--fork)
+          if test $
+          then
+            FORK_NAME="$2"
+            LOCAL_CHANGES=false
+            shift
+          else
+            echo "Error! fork_name not provided"
+            exit 1
+          fi
+          shift
+          ;;
         # With default ubuntu version 20.04
         -cv|--cudaversion)
           CUDA_VERSION="$2"
@@ -202,15 +216,15 @@ fi
 
 if [ "${BUILD_TYPE}" == "production" ]
 then
-  DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg USE_CUDA_VERSION="${CUDA_VERSION}"  --build-arg PYTHON_VERSION="${PYTHON_VERSION}" --build-arg BUILD_NIGHTLY="${BUILD_NIGHTLY}" --build-arg BRANCH_NAME="${BRANCH_NAME}" --build-arg BUILD_FROM_SRC="${BUILD_FROM_SRC}" --build-arg LOCAL_CHANGES="${LOCAL_CHANGES}" -t "${DOCKER_TAG}" --target production-image  ../
+  DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg USE_CUDA_VERSION="${CUDA_VERSION}"  --build-arg PYTHON_VERSION="${PYTHON_VERSION}" --build-arg BUILD_NIGHTLY="${BUILD_NIGHTLY}" --build-arg BRANCH_NAME="${BRANCH_NAME}" --build-arg BUILD_FROM_SRC="${BUILD_FROM_SRC}" --build-arg LOCAL_CHANGES="${LOCAL_CHANGES}" --build-arg FORK_NAME="${FORK_NAME}" -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 USE_CUDA_VERSION="${CUDA_VERSION}"  --build-arg PYTHON_VERSION="${PYTHON_VERSION}" --build-arg BUILD_NIGHTLY="${BUILD_NIGHTLY}" --build-arg BRANCH_NAME="${BRANCH_NAME}" --build-arg BUILD_FROM_SRC="${BUILD_FROM_SRC}" --build-arg LOCAL_CHANGES="${LOCAL_CHANGES}" -t "${DOCKER_TAG}" --target ci-image  ../
+  DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg USE_CUDA_VERSION="${CUDA_VERSION}"  --build-arg PYTHON_VERSION="${PYTHON_VERSION}" --build-arg BUILD_NIGHTLY="${BUILD_NIGHTLY}" --build-arg BRANCH_NAME="${BRANCH_NAME}" --build-arg BUILD_FROM_SRC="${BUILD_FROM_SRC}" --build-arg LOCAL_CHANGES="${LOCAL_CHANGES}" --build-arg FORK_NAME="${FORK_NAME}" -t "${DOCKER_TAG}" --target ci-image  ../
 else
   if [ "${BUILD_CPP}" == "true" ]
   then
-    DOCKER_BUILDKIT=1 docker build --file Dockerfile.cpp --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg USE_CUDA_VERSION="${CUDA_VERSION}" --build-arg PYTHON_VERSION="${PYTHON_VERSION}" --build-arg BRANCH_NAME="${BRANCH_NAME}" -t "${DOCKER_TAG}" --target cpp-dev-image .
+    DOCKER_BUILDKIT=1 docker build --file Dockerfile.cpp --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg USE_CUDA_VERSION="${CUDA_VERSION}" --build-arg PYTHON_VERSION="${PYTHON_VERSION}" --build-arg BRANCH_NAME="${BRANCH_NAME}" --build-arg FORK_NAME="${FORK_NAME}" -t "${DOCKER_TAG}" --target cpp-dev-image .
   else
-    DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg USE_CUDA_VERSION="${CUDA_VERSION}"  --build-arg PYTHON_VERSION="${PYTHON_VERSION}" --build-arg BUILD_NIGHTLY="${BUILD_NIGHTLY}" --build-arg BRANCH_NAME="${BRANCH_NAME}" --build-arg BUILD_FROM_SRC="${BUILD_FROM_SRC}" --build-arg LOCAL_CHANGES="${LOCAL_CHANGES}" --build-arg BUILD_WITH_IPEX="${BUILD_WITH_IPEX}"  -t "${DOCKER_TAG}" --target dev-image  ../
+    DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg USE_CUDA_VERSION="${CUDA_VERSION}"  --build-arg PYTHON_VERSION="${PYTHON_VERSION}" --build-arg BUILD_NIGHTLY="${BUILD_NIGHTLY}" --build-arg BRANCH_NAME="${BRANCH_NAME}" --build-arg BUILD_FROM_SRC="${BUILD_FROM_SRC}" --build-arg LOCAL_CHANGES="${LOCAL_CHANGES}" --build-arg BUILD_WITH_IPEX="${BUILD_WITH_IPEX}" --build-arg FORK_NAME="${FORK_NAME}" -t "${DOCKER_TAG}" --target dev-image  ../
   fi
 fi