p

Signed-off-by: kevin <[email protected]>

scripts/pipeline_generator/pipeline_generator.py

@@ -1,9 +1,12 @@
 import os
 import re
-from typing import List, Optional
+from typing import List, Optional, Union
+import yaml
+from enum import Enum
 
 from pydantic import BaseModel, field_validator
 
+from .step import BuildkiteStep, BuildkiteBlockStep
 
 class PipelineGeneratorConfig:
     def __init__(
@@ -53,3 +56,10 @@ def __init__(
         ):
         config.validate()
         self.config = config
+
+
+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)

scripts/pipeline_generator/step.py

@@ -58,8 +58,8 @@ def validate_multi_node(self) -> Self:
 class BuildkiteStep(BaseModel):
     """This class represents a step in Buildkite format."""
     label: str
-    agents: Dict[str, AgentQueue] = {"queue": AgentQueue.AWS_CPU}
-    commands: List[str]
+    agents: Dict[str, str] = {"queue": AgentQueue.AWS_CPU.value}
+    commands: Optional[List[str]] = None
     key: Optional[str] = None
     plugins: Optional[List[Dict]] = None
     parallelism: Optional[int] = None
@@ -68,12 +68,17 @@ class BuildkiteStep(BaseModel):
     env: Optional[Dict[str, str]] = None
     retry: Optional[Dict[str, Any]] = None
 
+    @model_validator(mode="after")
+    def validate_agent_queue(self) -> Self:
+        queue = self.agents.get("queue")
+        if not AgentQueue(queue):
+            raise ValueError(f"Invalid agent queue: {queue}")
 
 class BuildkiteBlockStep(BaseModel):
     """This class represents a block step in Buildkite format."""
     block: str
-    depends_on: Optional[str] = BUILD_STEP_KEY
     key: str
+    depends_on: Optional[str] = BUILD_STEP_KEY
 
 
 def get_step_key(step_label: str) -> str:

scripts/tests/pipeline_generator/test_files/expected_pipeline.yaml

@@ -0,0 +1,26 @@
+steps:
+  - label: Test 1
+    depends_on: build
+    agents:
+      queue: cpu_queue
+    commands:
+      - echo "Test1.1"
+      - echo "Test1.2"
+
+  - label: Test 2
+    depends_on: build
+    agents:
+      queue: gpu_1_queue
+    commands:
+      - command3
+
+  - block: "Run Test 3"
+    depends_on: build
+    key: "block-test-3"
+
+  - label: Test 3
+    agents:
+      queue: cpu_queue
+    depends_on: block-test-3
+    commands:
+      - command4

scripts/tests/pipeline_generator/test_pipeline_generator.py

@@ -2,8 +2,11 @@
 import sys
 import os
 import tempfile
+import yaml
 
-from scripts.pipeline_generator.pipeline_generator import PipelineGeneratorConfig, PipelineGenerator
+from scripts.pipeline_generator.pipeline_generator import PipelineGeneratorConfig, PipelineGenerator, write_buildkite_steps
+from scripts.pipeline_generator.step import BuildkiteStep, BuildkiteBlockStep
+from scripts.pipeline_generator.utils import AgentQueue
 
 TEST_COMMIT = "abcdef0123456789abcdef0123456789abcdef01"
 TEST_FILE_PATH = "tests.yaml"
@@ -59,5 +62,26 @@ def test_get_pipeline_generator_fail_nonexistent_test_file():
         with pytest.raises(FileNotFoundError, match="Test file"):
             _ = PipelineGenerator(config)
 
+
+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:
+            output = yaml.safe_load(f)
+        assert output == expected_output
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-v", __file__]))