[5/n][pipeline-gen] Method to read test steps from Yaml file

scripts/pipeline_generator/pipeline_generator.py

@@ -1,30 +1,26 @@
+import click
 import os
 import re
+import yaml
 from typing import List, Optional
-
 from pydantic import BaseModel, field_validator
 
-
+from .step import TestStep
+from .utils import VLLM_ECR_URL, VLLM_ECR_REPO
 class PipelineGeneratorConfig:
     def __init__(
         self,
         container_registry: str,
         container_registry_repo: str,
         commit: str,
         list_file_diff: List[str],
-        test_path: str,  # List of tests
-        external_hardware_test_path: str,  # List of external hardware tests
-        pipeline_file_path: str,  # Path to the output pipeline file
         run_all: bool = False,
     ):
         self.run_all = run_all
         self.list_file_diff = list_file_diff
         self.container_registry = container_registry
         self.container_registry_repo = container_registry_repo
         self.commit = commit
-        self.test_path = test_path
-        self.external_hardware_test_path = external_hardware_test_path
-        self.pipeline_file_path = pipeline_file_path
 
     @property
     def container_image(self):
@@ -37,14 +33,6 @@ def validate(self):
         if not re.match(pattern, self.commit):
             raise ValueError(f"Commit {self.commit} is not a valid Git commit hash")
 
-        # Check if test_path exists
-        if not os.path.isfile(self.test_path):
-            raise FileNotFoundError(f"Test file {self.test_path} not found")
-
-        # Check if external_hardware_test_path exists
-        if not os.path.isfile(self.external_hardware_test_path):
-            raise FileNotFoundError(f"External hardware test file {self.external_hardware_test_path} not found")
-
 
 class PipelineGenerator:
     def __init__(
@@ -53,3 +41,24 @@ 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"]]
+
+@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")
+@click.option("--list_file_diff", type=str, help="List of files in the diff between current branch and main")
+def main(test_path: str, external_hardware_test_path: str, run_all: str, list_file_diff: str):
+    test_steps = read_test_steps(test_path)
+
+    pipeline_generator_config = PipelineGeneratorConfig(
+        run_all=run_all == "1",
+        list_file_diff=list_file_diff,
+        container_registry=VLLM_ECR_URL,
+        container_registry_repo=VLLM_ECR_REPO,
+        commit=os.getenv("BUILDKITE_COMMIT"),
+    )

scripts/tests/pipeline_generator/test_files/test-pipeline.yaml

@@ -0,0 +1,23 @@
+steps:
+- label: Test 1
+  command: echo "Test 1"
+
+- label: Test 2
+  working_dir: "/tests2/"
+  no_gpu: true
+  command: echo "Test 2"
+
+- label: Test 3
+  commands:
+  - echo "Test 3"
+  - echo "Test 3.1"
+  source_file_dependencies:
+  - "file1"
+  - "src/file2"
+
+- label: Test 4
+  num_nodes: 2
+  num_gpus: 4
+  commands:
+  - echo "Test 4.1"
+  - echo "Test 4.2"

scripts/tests/pipeline_generator/test_pipeline_generator.py

@@ -3,38 +3,27 @@
 import os
 import tempfile
 
-from scripts.pipeline_generator.pipeline_generator import PipelineGeneratorConfig, PipelineGenerator
+from scripts.pipeline_generator.pipeline_generator import PipelineGeneratorConfig, PipelineGenerator, read_test_steps
+from scripts.pipeline_generator.step import DEFAULT_TEST_WORKING_DIR
 
 TEST_COMMIT = "abcdef0123456789abcdef0123456789abcdef01"
-TEST_FILE_PATH = "tests.yaml"
-EXTERNAL_HARDWARE_TEST_FILE_PATH = "external_hardware_tests.yaml"
-PIPELINE_OUTPUT_FILE_PATH = "pipeline.yaml"
 TEST_CONTAINER_REGISTRY = "container.registry"
 TEST_CONTAINER_REGISTRY_REPO = "test"
 
 
-def _get_pipeline_generator_config(test_dir: str):
-    with open(os.path.join(test_dir, TEST_FILE_PATH), "w") as f:
-        f.write("test-content")
-    with open(os.path.join(test_dir, EXTERNAL_HARDWARE_TEST_FILE_PATH), "w") as f:
-        f.write("external-hardware-test-content")
-
+def _get_pipeline_generator_config():
     return PipelineGeneratorConfig(
         container_registry=TEST_CONTAINER_REGISTRY,
         container_registry_repo=TEST_CONTAINER_REGISTRY_REPO,
         commit=TEST_COMMIT,
         list_file_diff=[],
-        test_path=os.path.join(test_dir, TEST_FILE_PATH),
-        external_hardware_test_path=os.path.join(test_dir, EXTERNAL_HARDWARE_TEST_FILE_PATH),
-        pipeline_file_path=os.path.join(test_dir, PIPELINE_OUTPUT_FILE_PATH)
     )
 
 
 def test_pipeline_generator_config_get_container_image():
-    with tempfile.TemporaryDirectory() as temp_dir:
-        config = _get_pipeline_generator_config(temp_dir)
-        config.validate()
-        assert config.container_image == "container.registry/test:abcdef0123456789abcdef0123456789abcdef01"
+    config = _get_pipeline_generator_config()
+    config.validate()
+    assert config.container_image == "container.registry/test:abcdef0123456789abcdef0123456789abcdef01"
 
 
 @pytest.mark.parametrize(
@@ -45,19 +34,31 @@ def test_pipeline_generator_config_get_container_image():
     ]
 )
 def test_get_pipeline_generator_config_invalid_commit(commit):
-    with tempfile.TemporaryDirectory() as temp_dir:
-        config = _get_pipeline_generator_config(temp_dir)
-        config.commit = commit
-        with pytest.raises(ValueError, match="not a valid Git commit hash"):
-            config.validate()
-
-
-def test_get_pipeline_generator_fail_nonexistent_test_file():
-    with tempfile.TemporaryDirectory() as temp_dir:
-        config = _get_pipeline_generator_config(temp_dir)
-        config.test_path = "non-existent-file"
-        with pytest.raises(FileNotFoundError, match="Test file"):
-            _ = PipelineGenerator(config)
+    config = _get_pipeline_generator_config()
+    config.commit = commit
+    with pytest.raises(ValueError, match="not a valid Git commit hash"):
+        config.validate()
+
+
+def test_read_test_steps():
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    test_path = os.path.join(current_dir, "test_files/test-pipeline.yaml")
+    test_steps = read_test_steps(test_path)
+    assert len(test_steps) == 4
+    assert test_steps[0].commands == ['echo "Test 1"']
+    assert test_steps[0].command is None
+    assert test_steps[0].working_dir == DEFAULT_TEST_WORKING_DIR
+
+    assert test_steps[1].working_dir == "/tests2/"
+    assert test_steps[1].no_gpu is True
+
+    assert test_steps[2].commands == ['echo "Test 3"', 'echo "Test 3.1"']
+    assert test_steps[2].source_file_dependencies == ["file1", "src/file2"]
+
+    assert test_steps[3].commands == ['echo "Test 4.1"', 'echo "Test 4.2"']
+    assert test_steps[3].num_nodes == 2
+    assert test_steps[3].num_gpus == 4
+
 
 if __name__ == "__main__":
     sys.exit(pytest.main(["-v", __file__]))