build/ci: enable bazel remote caching

zml · Sep 12, 2024 · d6a4fe7 · d6a4fe7
1 parent 41b1c82
commit d6a4fe7
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/.github/runs-on.yml b/.github/runs-on.yml
@@ -0,0 +1,11 @@
+images:
+  ubuntu24-amd64:
+    platform: "linux"
+    arch: "amd64"
+    ami: "ami-04a92520784b93e73"
+    preinstall: |
+      #!/bin/bash
+      wget https://github.com/jbdalido/bazel-remote/releases/download/test/bazel-remote -O /usr/local/bin/bzlcache
+      chmod +x /usr/local/bin/bzlcache
+      echo "Run bazel cache"
+      systemd-run --property=Restart=always /usr/local/bin/bzlcache --max_size=1000 --grpc_address 127.0.0.1:15501 --dir=/tmp/cache --s3.region eu-west-3 --s3.bucket runs-on-s3bucketcache-jycgicewbssh --s3.auth_method iam_role --s3.prefix cache --s3.endpoint=s3.eu-west-3.amazonaws.com --storage_mode uncompressed
diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
@@ -74,21 +74,21 @@ jobs:
             config: "--config=cuda"
             artifact: libpjrt_c_api_gpu_plugin.so
             renamed_artifact: libpjrt_cuda.so
-            runs_on: ['runs-on','runner=32cpu-linux-x64']
+            runs_on: ['runs-on','runner=32cpu-linux-x64', 'image=ubuntu24-amd64']
             platform: linux-amd64
             bazel_target: //xla/pjrt/c:pjrt_c_api_gpu_plugin
           - target: rocm
             config: "--config=rocm"
             artifact: libpjrt_c_api_gpu_plugin.so
             renamed_artifact: libpjrt_rocm.so
-            runs_on: ['runs-on','runner=32cpu-linux-x64']
+            runs_on: ['runs-on','runner=32cpu-linux-x64', 'image=ubuntu24-amd64']
             platform: linux-amd64
             bazel_target: //xla/pjrt/c:pjrt_c_api_gpu_plugin
           - target: cpu
             config: ""
             artifact: libpjrt_c_api_cpu_plugin.so
             renamed_artifact: libpjrt_cpu.so
-            runs_on: ['runs-on','runner=32cpu-linux-x64']
+            runs_on: ['runs-on','runner=32cpu-linux-x64', 'image=ubuntu24-amd64']
             platform: linux-amd64
             bazel_target: //xla/pjrt/c:pjrt_c_api_cpu_plugin
 
@@ -114,10 +114,12 @@ jobs:
           docker run \
             --name=pjrt-${{ matrix.pjrt.target }} \
             -w=/xla \
+            --net=host \
             -v=/tmp/cache:/tmp/cache \
             distroless/${{ matrix.pjrt.target }}_builder:latest \
               bazelisk \
               build \
+              --remote_cache=grpc://127.0.0.1:15501 \
               --disk_cache=/tmp/cache \
               ${{ matrix.pjrt.config }} --repo_env=HERMETIC_PYTHON_VERSION=3.11 \
               ${{ matrix.pjrt.bazel_target }}