multiple cache from in pipeline-gen and fix/add tests

Signed-off-by: Junpu Fan <[email protected]>

Author: junpuf

buildkite/pipeline_generator/buildkite_step.py

@@ -1,11 +1,11 @@
 from pydantic import BaseModel
 from typing import Dict, List, Optional, Any, Union
-from step import Step
-from utils_lib.docker_utils import get_image, get_ecr_cache_registry
-from global_config import get_global_config
-from plugin.k8s_plugin import get_k8s_plugin
-from plugin.docker_plugin import get_docker_plugin
-from constants import GPUType, AgentQueue
+from .step import Step
+from .utils_lib.docker_utils import get_image, resolve_ecr_cache_vars
+from .global_config import get_global_config
+from .plugin.k8s_plugin import get_k8s_plugin
+from .plugin.docker_plugin import get_docker_plugin
+from .constants import GPUType, AgentQueue
 
 
 class BuildkiteCommandStep(BaseModel):
@@ -87,7 +87,7 @@ def _get_variables_to_inject() -> Dict[str, str]:
     if global_config["name"] != "vllm_ci":
         return {}
 
-    cache_from_tag, cache_to_tag = get_ecr_cache_registry()
+    cache_from, cache_from_base_branch, cache_from_main, cache_to = resolve_ecr_cache_vars()
     return {
         "$REGISTRY": global_config["registries"],
         "$REPO": global_config["repositories"]["main"]
@@ -97,8 +97,10 @@ def _get_variables_to_inject() -> Dict[str, str]:
         "$BRANCH": global_config["branch"],
         "$VLLM_USE_PRECOMPILED": "1" if global_config["use_precompiled"] else "0",
         "$VLLM_MERGE_BASE_COMMIT": global_config["merge_base_commit"],
-        "$CACHE_FROM": cache_from_tag,
-        "$CACHE_TO": cache_to_tag,
+        "$CACHE_FROM": cache_from,
+        "$CACHE_FROM_BASE_BRANCH": cache_from_base_branch,
+        "$CACHE_FROM_MAIN": cache_from_main,
+        "$CACHE_TO": cache_to,
     }
 
 

buildkite/pipeline_generator/global_config.py

@@ -3,7 +3,7 @@
 import os
 import re
 import requests
-from utils_lib.git_utils import get_merge_base_commit, get_list_file_diff, get_pr_labels
+from .utils_lib.git_utils import get_merge_base_commit, get_list_file_diff, get_pr_labels
 
 
 class GlobalConfig(TypedDict):

buildkite/pipeline_generator/pipeline_generator.py

@@ -2,9 +2,9 @@
 import yaml
 import subprocess
 import sys
-from step import read_steps_from_job_dir, group_steps
-from buildkite_step import convert_group_step_to_buildkite_step
-from global_config import init_global_config, get_global_config
+from .step import read_steps_from_job_dir, group_steps
+from .buildkite_step import convert_group_step_to_buildkite_step
+from .global_config import init_global_config, get_global_config
 
 
 class PipelineGenerator:

buildkite/pipeline_generator/plugin/docker_plugin.py

@@ -1,5 +1,5 @@
-from step import Step
-from constants import GPUType
+from ..step import Step
+from ..constants import GPUType
 import copy
 
 docker_plugin_template = {

buildkite/pipeline_generator/plugin/k8s_plugin.py

@@ -1,6 +1,6 @@
 import copy
-from step import Step
-from constants import GPUType
+from ..step import Step
+from ..constants import GPUType
 
 HF_HOME = "/root/.cache/huggingface"
 

buildkite/pipeline_generator/step.py

@@ -3,7 +3,7 @@
 from pydantic import model_validator
 from typing_extensions import Self
 from collections import defaultdict
-from global_config import get_global_config
+from .global_config import get_global_config
 import os
 import yaml
 
@@ -91,3 +91,8 @@ def group_steps(steps: List[Step]) -> Dict[str, List[Step]]:
     for group, steps in grouped_steps.items():
         sorted_grouped_steps[group] = sorted(steps, key=lambda x: x.label)
     return sorted_grouped_steps
+
+
+def sort_steps(steps: List[Step]) -> List[Step]:
+    """Sort steps by group alphabetically, then by label within each group."""
+    return sorted(steps, key=lambda x: (x.group, x.label))

buildkite/pipeline_generator/utils_lib/docker_utils.py

@@ -2,7 +2,7 @@
 import os
 import re
 from typing import Tuple
-from global_config import get_global_config
+from ..global_config import get_global_config
 
 
 def get_image(cpu: bool = False) -> str:
@@ -21,80 +21,63 @@ def get_image(cpu: bool = False) -> str:
     return image
 
 
-def _clean_docker_tag(tag: str) -> str:
-    # Only allows alphanumeric, dashes and underscores for Docker tags, and replaces others with '-'
-    return re.sub(r"[^a-zA-Z0-9_.-]", "-", tag or "")
+def clean_docker_tag(tag: str) -> str:
+    """
+    Function to replace invalid characters in Docker image tags and truncate to 128 chars
+    Valid characters: a-z, A-Z, 0-9, _, ., -
+    """
+    # Replace invalid characters with underscore and truncate to 128 chars
+    cleaned = re.sub(r"[^a-zA-Z0-9._-]", "_", tag or "")
+    return cleaned[:128]
 
 
-def _docker_manifest_exists(image_tag: str) -> bool:
-    try:
-        subprocess.run(
-            ["docker", "manifest", "inspect", image_tag],
-            stdout=subprocess.DEVNULL,
-            stderr=subprocess.DEVNULL,
-            check=True,
-        )
-        return True
-    except subprocess.CalledProcessError:
-        return False
 
-
-def get_ecr_cache_registry() -> Tuple[str, str]:
+def resolve_ecr_cache_vars() -> Tuple[str, str, str, str]:
+    """
+    Resolve ECR cache-from, cache-to using buildkite environment variables:
+     -  BUILDKITE_BRANCH 
+     -  BUILDKITE_PULL_REQUEST
+     -  BUILDKITE_PULL_REQUEST_BASE_BRANCH
+    Return tuple of: 
+     -  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.
+    """
     global_config = get_global_config()
     branch = global_config["branch"]
+    pull_request = global_config["pull_request"]
+    
+    # Define ECR repository URLs for test and main cache
     test_cache_ecr = "936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache"
-    postmerge_cache_ecr = (
-        "936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache"
-    )
-    cache_from_tag, cache_to_tag = None, None
-    # Authenticate Docker to AWS ECR
-    login_cmd = ["aws", "ecr", "get-login-password", "--region", "us-east-1"]
-    try:
-        proc = subprocess.Popen(login_cmd, stdout=subprocess.PIPE)
-        subprocess.run(
-            [
-                "docker",
-                "login",
-                "--username",
-                "AWS",
-                "--password-stdin",
-                "936637512419.dkr.ecr.us-east-1.amazonaws.com",
-            ],
-            stdin=proc.stdout,
-            check=True,
-        )
-        proc.stdout.close()
-        proc.wait()
-    except Exception as e:
-        raise RuntimeError(f"Failed to authenticate with AWS ECR: {e}")
-
-    if global_config["pull_request"]:  # PR build
-        cache_to_tag = f"{test_cache_ecr}:pr-{global_config['pull_request']}"
-        if _docker_manifest_exists(cache_to_tag):  # use PR cache if exists
-            cache_from_tag = cache_to_tag
-        elif (
-            os.getenv("BUILDKITE_PULL_REQUEST_BASE_BRANCH") != "main"
-        ):  # use base branch cache if exists
-            clean_base = _clean_docker_tag(
-                os.getenv("BUILDKITE_PULL_REQUEST_BASE_BRANCH")
-            )
-            if _docker_manifest_exists(f"{test_cache_ecr}:{clean_base}"):
-                cache_from_tag = f"{test_cache_ecr}:{clean_base}"
-            else:  # fall back to postmerge cache ecr if base branch cache does not exist
-                cache_from_tag = f"{postmerge_cache_ecr}:latest"
+    main_cache_ecr = "936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache"
+    
+    if not pull_request or pull_request == "false":
+        # Not a PR
+        if branch == "main":
+            cache = f"{main_cache_ecr}:latest"
         else:
-            cache_from_tag = f"{postmerge_cache_ecr}:latest"
-    else:  # non-PR build
-        if branch == "main":  # postmerge
-            cache_to_tag = f"{postmerge_cache_ecr}:latest"
-            cache_from_tag = f"{postmerge_cache_ecr}:latest"
+            clean_branch = clean_docker_tag(branch)
+            cache = f"{test_cache_ecr}:{clean_branch}"
+        
+        cache_to = cache
+        cache_from = cache
+        cache_from_base_branch = cache
+    else:
+        # PR build
+        cache_to = f"{test_cache_ecr}:pr-{pull_request}"
+        cache_from = f"{test_cache_ecr}:pr-{pull_request}"
+        
+        base_branch = os.getenv("BUILDKITE_PULL_REQUEST_BASE_BRANCH", "main")
+        if base_branch == "main":
+            cache_from_base_branch = f"{main_cache_ecr}:latest"
         else:
-            clean_branch = _clean_docker_tag(branch)
-            cache_to_tag = f"{test_cache_ecr}:{clean_branch}"
-            if _docker_manifest_exists(f"{test_cache_ecr}:{clean_branch}"):
-                cache_from_tag = f"{test_cache_ecr}:{clean_branch}"
-            else:
-                cache_from_tag = f"{postmerge_cache_ecr}:latest"
-    if not cache_from_tag or not cache_to_tag:
-        raise RuntimeError("Failed to get ECR cache tags")
-    return cache_from_tag, cache_to_tag
+            clean_base = clean_docker_tag(base_branch)
+            cache_from_base_branch = f"{test_cache_ecr}:{clean_base}"
+    
+    cache_from_main = f"{main_cache_ecr}:latest"
+    
+    return cache_from, cache_from_base_branch, cache_from_main, cache_to