[8/n][pipeline-gen] Add method to generate build step

scripts/pipeline_generator/pipeline_generator.py

@@ -6,7 +6,8 @@
 from pydantic import BaseModel, field_validator
 
 from .step import BuildkiteStep, BuildkiteBlockStep, TestStep
-from .utils import VLLM_ECR_URL, VLLM_ECR_REPO
+from .utils import VLLM_ECR_URL, VLLM_ECR_REPO, AgentQueue
+from .pipeline_generator_helper import get_build_commands
 
 class PipelineGeneratorConfig:
     def __init__(
@@ -43,6 +44,24 @@ def __init__(
         config.validate()
         self.config = config
 
+    def generate_build_step(self) -> BuildkiteStep:
+        """Build the Docker image and push it to container registry."""
+        build_commands = get_build_commands(self.config.container_registry, self.config.commit, self.config.container_image)
+
+        return BuildkiteStep(
+            label=":docker: build image",
+            key="build",
+            agents={"queue": AgentQueue.AWS_CPU.value},
+            env={"DOCKER_BUILDKIT": "1"},
+            retry={
+                "automatic": [
+                    {"exit_status": -1, "limit": 2},
+                    {"exit_status": -10, "limit": 2}
+                ]
+            },
+            commands=build_commands,
+            depends_on=None,
+        )
 
 def read_test_steps(file_path: str) -> List[TestStep]:
     """Read test steps from test pipeline yaml and parse them into TestStep objects."""

scripts/pipeline_generator/pipeline_generator_helper.py

@@ -0,0 +1,26 @@
+from typing import List
+
+def get_build_commands(container_registry: str, buildkite_commit: str, container_image: str) -> List[str]:
+    ecr_login_command = (
+        "aws ecr-public get-login-password --region us-east-1 | "
+        f"docker login --username AWS --password-stdin {container_registry}"
+    )
+    image_check_command = f"""#!/bin/bash
+if [[ -z $(docker manifest inspect {container_image}) ]]; then
+echo "Image not found, proceeding with build..."
+else
+echo "Image found"
+exit 0
+fi
+"""
+    docker_build_command = (
+        f"docker build "
+        f"--build-arg max_jobs=64 "
+        f"--build-arg buildkite_commit={buildkite_commit} "
+        f"--build-arg USE_SCCACHE=1 "
+        f"--tag {container_image} "
+        f"--target test "
+        f"--progress plain ."
+    )
+    docker_push_command = f"docker push {container_image}"
+    return [ecr_login_command, image_check_command, docker_build_command, docker_push_command]