[CI] use multiple cache from (#232)

Signed-off-by: Junpu Fan <junpu@amazon.com>

Author: junpuf

buildkite/bootstrap.sh

@@ -73,56 +73,42 @@ resolve_ecr_cache_vars() {
     #  -  BUILDKITE_BRANCH 
     #  -  BUILDKITE_PULL_REQUEST
     #  -  BUILDKITE_PULL_REQUEST_BASE_BRANCH
+    # Export environment variables: 
+    #  -  CACHE_FROM: primary cache source
+    #  -  CACHE_FROM_BASE_BRANCH: secondary cache source
+    #  -  CACHE_FROM_MAIN: fallback cache source
+    #  -  CACHE_TO: cache destination
+    # Note: CACHE_FROM, CACHE_FROM_BASE_BRANCH, CACHE_FROM_MAIN could be the same.
+    #     This is intended behavior to allow BuildKit to merge all possible cache source 
+    #     to maximize cache hit potential, see https://docs.docker.com/build/cache/backends/#multiple-caches
 
     # Define ECR repository URLs for test and main cache
     local TEST_CACHE_ECR="936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache"
     local MAIN_CACHE_ECR="936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache"
 
-    # login to ECR to check ceche availability
-    aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com
-    
     if [[ "$BUILDKITE_PULL_REQUEST" == "false" ]]; then
-        # For non-PR builds
         if [[ "$BUILDKITE_BRANCH" == "main" ]]; then
-            # For main branch: use main cache for both source and destination
-            CACHE_TO="${MAIN_CACHE_ECR}:latest"
-            CACHE_FROM="${MAIN_CACHE_ECR}:latest"
+            local cache="${MAIN_CACHE_ECR}:latest"
         else
-            # For other branches: use branch-specific cache if it exists
             local clean_branch=$(clean_docker_tag "$BUILDKITE_BRANCH")
-            CACHE_TO="${TEST_CACHE_ECR}:${clean_branch}"
-            if docker manifest inspect "${TEST_CACHE_ECR}:${clean_branch}" &>/dev/null; then
-                # Use branch-specific cache if available
-                CACHE_FROM="${TEST_CACHE_ECR}:${clean_branch}"
-            else
-                # Fallback to main cache if branch cache doesn't exist
-                CACHE_FROM="${MAIN_CACHE_ECR}:latest"
-            fi
+            local cache="${TEST_CACHE_ECR}:${clean_branch}"
         fi
+        CACHE_TO="$cache"
+        CACHE_FROM="$cache"
+        CACHE_FROM_BASE_BRANCH="$cache"
     else
-        # For PR builds
-        # Always set cache destination to PR-specific tag
         CACHE_TO="${TEST_CACHE_ECR}:pr-${BUILDKITE_PULL_REQUEST}"
-        
-        if docker manifest inspect "${TEST_CACHE_ECR}:pr-${BUILDKITE_PULL_REQUEST}" &>/dev/null; then
-            # Use PR-specific cache if it exists
-            CACHE_FROM="${TEST_CACHE_ECR}:pr-${BUILDKITE_PULL_REQUEST}"
-        elif [[ "$BUILDKITE_PULL_REQUEST_BASE_BRANCH" != "main" ]]; then
-            # For PRs targeting non-main branches, try using base branch cache
-            local clean_base=$(clean_docker_tag "$BUILDKITE_PULL_REQUEST_BASE_BRANCH")
-            if docker manifest inspect "${TEST_CACHE_ECR}:${clean_base}" &>/dev/null; then
-                CACHE_FROM="${TEST_CACHE_ECR}:${clean_base}"
-            else
-                # Fallback to main cache if base branch cache doesn't exist
-                CACHE_FROM="${MAIN_CACHE_ECR}:latest"
-            fi
+        CACHE_FROM="${TEST_CACHE_ECR}:pr-${BUILDKITE_PULL_REQUEST}"
+        if [[ "$BUILDKITE_PULL_REQUEST_BASE_BRANCH" == "main" ]]; then
+            CACHE_FROM_BASE_BRANCH="${MAIN_CACHE_ECR}:latest"
         else
-            # Fallback to main cache for PRs targeting main branch
-            CACHE_FROM="${MAIN_CACHE_ECR}:latest"
+            local clean_base=$(clean_docker_tag "$BUILDKITE_PULL_REQUEST_BASE_BRANCH")
+            CACHE_FROM_BASE_BRANCH="${TEST_CACHE_ECR}:${clean_base}"
         fi
     fi
-    # Export variables
-    export CACHE_FROM CACHE_TO
+    
+    CACHE_FROM_MAIN="${MAIN_CACHE_ECR}:latest"
+    export CACHE_FROM CACHE_FROM_BASE_BRANCH CACHE_FROM_MAIN CACHE_TO
 }
 
 upload_pipeline() {
@@ -157,11 +143,13 @@ upload_pipeline() {
 
     # Resolve CACHE_FROM and CACHE_TO for ECR Registry Caching
     resolve_ecr_cache_vars
-    if [[ -z "${CACHE_FROM:-}" ]] || [[ -z "${CACHE_TO:-}" ]]; then
-        echo "Error: CACHE_FROM or CACHE_TO not set after resolve_ecr_cache_vars"
+    if [[ -z "${CACHE_FROM:-}" ]] || [[ -z "${CACHE_FROM_BASE_BRANCH:-}" ]] || [[ -z "${CACHE_FROM_MAIN:-}" ]] || [[ -z "${CACHE_TO:-}" ]]; then
+        echo "Error: CACHE_FROM, CACHE_FROM_BASE_BRANCH, CACHE_FROM_MAIN, or CACHE_TO not set after resolve_ecr_cache_vars"
         exit 1
     else
         echo "Resolved CACHE_FROM: ${CACHE_FROM}"
+        echo "Resolved CACHE_FROM_BASE_BRANCH: ${CACHE_FROM_BASE_BRANCH}"
+        echo "Resolved CACHE_FROM_MAIN: ${CACHE_FROM_MAIN}"
         echo "Resolved CACHE_TO: ${CACHE_TO}"
     fi
 
@@ -181,6 +169,8 @@ upload_pipeline() {
             -D cov_enabled="$COV_ENABLED" \
             -D vllm_ci_branch="$VLLM_CI_BRANCH" \
             -D cache_from="$CACHE_FROM" \
+            -D cache_from_base_branch="$CACHE_FROM_BASE_BRANCH" \
+            -D cache_from_main="$CACHE_FROM_MAIN" \
             -D cache_to="$CACHE_TO" \
             | sed '/^[[:space:]]*$/d' \
             > pipeline.yaml

buildkite/test-template-ci.j2

@@ -280,8 +280,10 @@ steps:
         {% if vllm_use_precompiled == "1" %}
         --build-arg VLLM_MERGE_BASE_COMMIT={{ vllm_merge_base_commit }}
         {% endif %}
+        --cache-to type=registry,ref={{ cache_to }},mode=max,compression=zstd
         --cache-from type=registry,ref={{ cache_from }},mode=max
-        --cache-to type=registry,ref={{ cache_to }},mode=max
+        --cache-from type=registry,ref={{ cache_from_base_branch }},mode=max
+        --cache-from type=registry,ref={{ cache_from_main }},mode=max
         --tag {{ docker_image }}
         {% if branch == "main" %}
         --tag {{ docker_image_latest }}