diff --git a/scripts/pipeline_generator/pipeline_generator.py b/scripts/pipeline_generator/pipeline_generator.py index 0319d7a..b97d8de 100644 --- 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 new file mode 100644 index 0000000..d36f1dd --- /dev/null +++ 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 index 5d0696f..d5d3b38 100644 --- 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__]))