[7/n][pipeline-gen] Method to write Buildkite steps to yaml file (#45)

* p Signed-off-by: kevin <[email protected]> * p Signed-off-by: kevin <[email protected]> * p Signed-off-by: kevin <[email protected]> * p Signed-off-by: kevin <[email protected]> --------- Signed-off-by: kevin <[email protected]>
vllm-project · Oct 11, 2024 · fa69836 · fa69836
1 parent 717659c
commit fa69836
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 5 deletions.
diff --git a/scripts/pipeline_generator/pipeline_generator.py b/scripts/pipeline_generator/pipeline_generator.py
@@ -1,12 +1,13 @@
 import click
 import os
 import re
+from typing import List, Optional, Union
 import yaml
-from typing import List, Optional
 from pydantic import BaseModel, field_validator
 
-from .step import TestStep
+from .step import BuildkiteStep, BuildkiteBlockStep, TestStep
 from .utils import VLLM_ECR_URL, VLLM_ECR_REPO
+
 class PipelineGeneratorConfig:
     def __init__(
         self,
@@ -21,7 +22,7 @@ def __init__(
         self.container_registry = container_registry
         self.container_registry_repo = container_registry_repo
         self.commit = commit
-    
+
     @property
     def container_image(self):
         return f"{self.container_registry}/{self.container_registry_repo}:{self.commit}"
@@ -42,12 +43,19 @@ def __init__(
         config.validate()
         self.config = config
 
+
 def read_test_steps(file_path: str) -> List[TestStep]:
     """Read test steps from test pipeline yaml and parse them into TestStep objects."""
     with open(file_path, "r") as f:
         content = yaml.safe_load(f)
     return [TestStep(**step) for step in content["steps"]]
 
+def write_buildkite_steps(steps: List[Union[BuildkiteStep, BuildkiteBlockStep]], file_path: str) -> None:
+    """Write the buildkite steps to the Buildkite pipeline yaml file."""
+    buildkite_steps_dict = {"steps": [step.dict(exclude_none=True) for step in steps]}
+    with open(file_path, "w") as f:
+        yaml.dump(buildkite_steps_dict, f, sort_keys=False)
+
 @click.command()
 @click.option("--test_path", type=str, required=True, help="Path to the test pipeline yaml file")
 @click.option("--run_all", type=str, help="If set to 1, run all tests")

diff --git a/scripts/tests/pipeline_generator/test_files/expected_pipeline.yaml b/scripts/tests/pipeline_generator/test_files/expected_pipeline.yaml
@@ -0,0 +1,23 @@
+steps:
+- label: Test 1
+  agents:
+    queue: cpu_queue
+  commands:
+  - echo "Test1.1"
+  - echo "Test1.2"
+  depends_on: build
+- label: Test 2
+  agents:
+    queue: gpu_1_queue
+  commands:
+  - command3
+  depends_on: build
+- block: Run Test 3
+  key: block-test-3
+  depends_on: build
+- label: Test 3
+  agents:
+    queue: cpu_queue
+  commands:
+  - command4
+  depends_on: block-test-3
diff --git a/scripts/tests/pipeline_generator/test_pipeline_generator.py b/scripts/tests/pipeline_generator/test_pipeline_generator.py
@@ -2,9 +2,11 @@
 import sys
 import os
 import tempfile
+import yaml
 
-from scripts.pipeline_generator.pipeline_generator import PipelineGeneratorConfig, PipelineGenerator, read_test_steps
-from scripts.pipeline_generator.step import DEFAULT_TEST_WORKING_DIR
+from scripts.pipeline_generator.pipeline_generator import PipelineGeneratorConfig, PipelineGenerator, read_test_steps, write_buildkite_steps
+from scripts.pipeline_generator.step import BuildkiteStep, BuildkiteBlockStep, DEFAULT_TEST_WORKING_DIR
+from scripts.pipeline_generator.utils import AgentQueue
 
 TEST_COMMIT = "abcdef0123456789abcdef0123456789abcdef01"
 TEST_CONTAINER_REGISTRY = "container.registry"
@@ -60,5 +62,27 @@ def test_read_test_steps():
     assert test_steps[3].num_gpus == 4
 
 
+def test_write_buildkite_steps():
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    expected_output_path = os.path.join(current_dir, "test_files/expected_pipeline.yaml")
+    with open(expected_output_path, "r") as f:
+        expected_output = yaml.safe_load(f)
+
+    steps = [
+        BuildkiteStep(label="Test 1", commands=['echo "Test1.1"', 'echo "Test1.2"']),
+        BuildkiteStep(label="Test 2", commands=["command3"], agents = {"queue": AgentQueue.AWS_1xL4.value}),
+        BuildkiteBlockStep(block="Run Test 3", key="block-test-3"),
+        BuildkiteStep(label="Test 3", commands=["command4"], depends_on="block-test-3"),
+    ]
+    with tempfile.TemporaryDirectory() as temp_dir:
+        output_file_path = os.path.join(temp_dir, "output.yaml")
+        write_buildkite_steps(steps, output_file_path)
+        with open(output_file_path, "r") as f:
+            content = f.read()
+        with open(expected_output_path, "r") as f:
+            expected_content = f.read()
+        assert content == expected_content
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-v", __file__]))