diff --git a/scripts/pipeline_generator/pipeline_generator.py b/scripts/pipeline_generator/pipeline_generator.py index b1053b8..0319d7a 100644 --- a/scripts/pipeline_generator/pipeline_generator.py +++ b/scripts/pipeline_generator/pipeline_generator.py @@ -1,10 +1,12 @@ +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, @@ -12,9 +14,6 @@ def __init__( 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 @@ -22,9 +21,6 @@ def __init__( 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"), + ) diff --git a/scripts/tests/pipeline_generator/test_files/test-pipeline.yaml b/scripts/tests/pipeline_generator/test_files/test-pipeline.yaml new file mode 100644 index 0000000..404ea00 --- /dev/null +++ b/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" diff --git a/scripts/tests/pipeline_generator/test_pipeline_generator.py b/scripts/tests/pipeline_generator/test_pipeline_generator.py index 145d11e..5d0696f 100644 --- a/scripts/tests/pipeline_generator/test_pipeline_generator.py +++ b/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__]))