adding graviton docker image

pytorch · Sep 12, 2024 · 500db43 · 500db43
1 parent d6ea6e7
commit 500db43
Show file tree

Hide file tree

Showing 3 changed files with 99 additions and 43 deletions.
diff --git a/.github/workflows/docker-nightly-build.yml b/.github/workflows/docker-nightly-build.yml
@@ -1,10 +1,11 @@
 name: Push Docker Nightly
 
 on:
-  # run every day at 1:15pm
+  # Run every day at 1:15pm
   schedule:
     - cron: "15 13 * * *"
   workflow_dispatch:
+
 jobs:
   nightly:
     runs-on: [self-hosted, ci-gpu]
@@ -37,7 +38,6 @@ jobs:
         run: |
           cd kubernetes/kserve
           python docker_nightly.py --cleanup
-
       - name: Open issue on failure
         if: ${{ failure() && github.event_name  == 'schedule' }}
         uses: dacbd/create-issue-action@v1
@@ -46,3 +46,31 @@ jobs:
           title: Nightly Docker build failed
           body: Commit ${{ github.sha }} daily scheduled [CI run](https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}) failed, please check why
           assignees: ""
+
+  graviton_job:
+    runs-on: [graviton-test]
+    steps:
+      - name: Clean up previous run on Graviton
+        run: |
+          echo "Cleaning up previous Graviton"
+          ls -la ./
+          sudo rm -rf ./* || true
+          sudo rm -rf ./.??* || true
+          ls -la ./
+      - name: Setup Python 3.9 on Graviton
+        uses: actions/setup-python@v5
+        with:
+          python-version: 3.9
+          architecture: aarch64
+      - name: Checkout TorchServe
+        uses: actions/checkout@v3
+        with:
+          submodules: recursive
+      - name: Login to Docker
+        env:
+          DOCKER_PASSWORD: ${{secrets.DOCKER_PASSWORD}}
+        run: docker login --username pytorchbot --password "$DOCKER_PASSWORD"
+      - name: Push Docker Nightly
+        run: |
+          cd docker
+          python docker_nightly.py --cleanup --arm
diff --git a/docker/build_image.sh b/docker/build_image.sh
@@ -18,6 +18,7 @@ BUILD_NIGHTLY=false
 BUILD_FROM_SRC=false
 LOCAL_CHANGES=true
 PYTHON_VERSION=3.9
+ARCH="linux/amd64"
 
 for arg in "$@"
 do
@@ -39,6 +40,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 "-a, --arch to arm64"
           exit 0
           ;;
         -b|--branch_name)
@@ -101,6 +103,10 @@ do
           BUILD_CPP=true
           shift
           ;;
+        -a|--arch)
+          ARCH="linux/arm64"
+          shift
+          ;;
         -n|--nightly)
           BUILD_NIGHTLY=true
           shift
@@ -218,7 +224,7 @@ 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 REPO_URL="${REPO_URL}" --build-arg BUILD_FROM_SRC="${BUILD_FROM_SRC}"\
-  --build-arg LOCAL_CHANGES="${LOCAL_CHANGES}" -t "${DOCKER_TAG}" --target production-image  ../
+  --build-arg LOCAL_CHANGES="${LOCAL_CHANGES}" -t "${DOCKER_TAG}" --platform "${ARCH}" --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}"\

diff --git a/docker/docker_nightly.py b/docker/docker_nightly.py
@@ -28,60 +28,82 @@
         action="store_true",
         help="Delete all built docker images",
     )
+    parser.add_argument("--arm", action="store_true", help="Sets architecture to arm")
     args = parser.parse_args()
     dry_run = args.dry_run
     organization = args.organization
 
     project = "torchserve-nightly"
     cpu_version = f"{project}:cpu-{get_nightly_version()}"
+    graviton_version = f"{project}:cpu-graviton-{get_nightly_version()}"
     gpu_version = f"{project}:gpu-{get_nightly_version()}"
     cpp_dev_cpu_version = f"{project}:cpp-dev-cpu-{get_nightly_version()}"
     cpp_dev_gpu_version = f"{project}:cpp-dev-gpu-{get_nightly_version()}"
 
     # Build Nightly images and append the date in the name
-    try_and_handle(f"./build_image.sh -n -t {organization}/{cpu_version}", dry_run)
-    try_and_handle(
-        f"./build_image.sh -g -cv cu121 -n -t {organization}/{gpu_version}",
-        dry_run,
-    )
-    try_and_handle(
-        f"./build_image.sh -bt dev -cpp -t {organization}/{cpp_dev_cpu_version}",
-        dry_run,
-    )
-    try_and_handle(
-        f"./build_image.sh -bt dev -g -cv cu121 -cpp -t {organization}/{cpp_dev_gpu_version}",
-        dry_run,
-    )
+    if args.arch:
+        try_and_handle(
+            f"./build_image.sh -n -a -t {organization}/{graviton_version}", dry_run
+        )
 
-    # Push Nightly images to official PyTorch Dockerhub account
-    try_and_handle(f"docker push {organization}/{cpu_version}", dry_run)
-    try_and_handle(f"docker push {organization}/{gpu_version}", dry_run)
-    try_and_handle(f"docker push {organization}/{cpp_dev_cpu_version}", dry_run)
-    try_and_handle(f"docker push {organization}/{cpp_dev_gpu_version}", dry_run)
+        try_and_handle(f"docker push {organization}/{graviton_version}", dry_run)
 
-    # Tag nightly images with latest
-    try_and_handle(
-        f"docker tag {organization}/{cpu_version} {organization}/{project}:latest-cpu",
-        dry_run,
-    )
-    try_and_handle(
-        f"docker tag {organization}/{gpu_version} {organization}/{project}:latest-gpu",
-        dry_run,
-    )
-    try_and_handle(
-        f"docker tag {organization}/{cpp_dev_cpu_version} {organization}/{project}:latest-cpp-dev-cpu",
-        dry_run,
-    )
-    try_and_handle(
-        f"docker tag {organization}/{cpp_dev_gpu_version} {organization}/{project}:latest-cpp-dev-gpu",
-        dry_run,
-    )
+        try_and_handle(
+            f"docker tag {organization}/{cpu_version} {graviton_version}/{project}:latest-cpu-graviton",
+            dry_run,
+        )
+
+        try_and_handle(
+            f"docker push {organization}/{project}:latest-cpu-graviton", dry_run
+        )
+    else:
+        try_and_handle(f"./build_image.sh -n -t {organization}/{cpu_version}", dry_run)
+        try_and_handle(
+            f"./build_image.sh -g -cv cu121 -n -t {organization}/{gpu_version}",
+            dry_run,
+        )
+        try_and_handle(
+            f"./build_image.sh -bt dev -cpp -t {organization}/{cpp_dev_cpu_version}",
+            dry_run,
+        )
+        try_and_handle(
+            f"./build_image.sh -bt dev -g -cv cu121 -cpp -t {organization}/{cpp_dev_gpu_version}",
+            dry_run,
+        )
+
+        # Push Nightly images to official PyTorch Dockerhub account
+        try_and_handle(f"docker push {organization}/{cpu_version}", dry_run)
+        try_and_handle(f"docker push {organization}/{gpu_version}", dry_run)
+        try_and_handle(f"docker push {organization}/{cpp_dev_cpu_version}", dry_run)
+        try_and_handle(f"docker push {organization}/{cpp_dev_gpu_version}", dry_run)
+
+        # Tag nightly images with latest
+        try_and_handle(
+            f"docker tag {organization}/{cpu_version} {organization}/{project}:latest-cpu",
+            dry_run,
+        )
+        try_and_handle(
+            f"docker tag {organization}/{gpu_version} {organization}/{project}:latest-gpu",
+            dry_run,
+        )
+        try_and_handle(
+            f"docker tag {organization}/{cpp_dev_cpu_version} {organization}/{project}:latest-cpp-dev-cpu",
+            dry_run,
+        )
+        try_and_handle(
+            f"docker tag {organization}/{cpp_dev_gpu_version} {organization}/{project}:latest-cpp-dev-gpu",
+            dry_run,
+        )
 
-    # Push images with latest tag
-    try_and_handle(f"docker push {organization}/{project}:latest-cpu", dry_run)
-    try_and_handle(f"docker push {organization}/{project}:latest-gpu", dry_run)
-    try_and_handle(f"docker push {organization}/{project}:latest-cpp-dev-cpu", dry_run)
-    try_and_handle(f"docker push {organization}/{project}:latest-cpp-dev-gpu", dry_run)
+        # Push images with latest tag
+        try_and_handle(f"docker push {organization}/{project}:latest-cpu", dry_run)
+        try_and_handle(f"docker push {organization}/{project}:latest-gpu", dry_run)
+        try_and_handle(
+            f"docker push {organization}/{project}:latest-cpp-dev-cpu", dry_run
+        )
+        try_and_handle(
+            f"docker push {organization}/{project}:latest-cpp-dev-gpu", dry_run
+        )
 
     # Cleanup built images
     if args.cleanup: