p

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

scripts/pipeline_generator/pipeline_generator.py

@@ -11,7 +11,6 @@
 from .utils import (
     AgentQueue,
     AMD_REPO,
-    A100_GPU,
     TEST_PATH,
     EXTERNAL_HARDWARE_TEST_PATH,
     PIPELINE_FILE_PATH,
@@ -36,19 +35,61 @@
     get_build_commands,
 )
 
-class PipelineGeneratorConfig(BaseModel):
-    run_all: bool
-    list_file_diff: List[str]
-    container_registry: str
-    container_registry_repo: str
-    commit: str
-    test_path: str
-    external_hardware_test_path: str
-    pipeline_file_path: str
-
+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) -> str:
+    def container_image(self):
         return f"{self.container_registry}/{self.container_registry_repo}:{self.commit}"
+
+    def validate(self):
+        """Validate the configuration."""
+        # Check if commit is a valid Git commit hash
+        pattern = r"^[0-9a-f]{40}$"
+        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")
+
+def read_test_steps(self, 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(
+        self,
+        buildkite_steps: List[Union[BuildkiteStep, BuildkiteBlockStep]],
+        output_file_path: str
+        ) -> None:
+    """Output the buildkite steps to the Buildkite pipeline yaml file."""
+    buildkite_steps_dict = {"steps": [step.dict(exclude_none=True) for step in buildkite_steps]}
+    with open(output_file_path, "w") as f:
+        yaml.dump(buildkite_steps_dict, f, sort_keys=False)
 
 class PipelineGenerator:
     def __init__(
@@ -76,12 +117,6 @@ def generate_build_step(self) -> BuildkiteStep:
             depends_on=None,
         )
 
-    def read_test_steps(self, file_path: str) -> List[TestStep]:
-        """Read test steps from test pipeline yaml and parse them into Step objects."""
-        with open(file_path, "r") as f:
-            content = yaml.safe_load(f)
-        return [TestStep(**step) for step in content["steps"]]
-
     def convert_test_step_to_buildkite_steps(self, step: TestStep) -> List[Union[BuildkiteStep, BuildkiteBlockStep]]:
         """Process test step and return corresponding BuildkiteStep."""
         steps = []
@@ -142,15 +177,6 @@ def _mirror_amd_test_steps(self, test_steps: List[TestStep]) -> List[BuildkiteSt
                 mirrored_buildkite_steps.append(mirrored_buildkite_step)
         return mirrored_buildkite_steps
 
-    def write_buildkite_steps(
-            self,
-            buildkite_steps: List[Union[BuildkiteStep, BuildkiteBlockStep]],
-            output_file_path: str
-            ) -> None:
-        """Output the buildkite steps to the Buildkite pipeline yaml file."""
-        buildkite_steps_dict = {"steps": [step.dict(exclude_none=True) for step in buildkite_steps]}
-        with open(output_file_path, "w") as f:
-            yaml.dump(buildkite_steps_dict, f, sort_keys=False)
 
     def generate(self):
         test_steps = self.read_test_steps(self.config.test_path)
@@ -185,3 +211,11 @@ def main(run_all: str = "-1", list_file_diff: str = None):
 
 if __name__ == "__main__":
     main()
+import os
+import re
+from typing import List, Optional
+
+from pydantic import BaseModel, field_validator
+
+
+

scripts/pipeline_generator/pipeline_generator_helper.py

@@ -1,6 +1,6 @@
 from typing import List, Dict
 from .plugin import get_kubernetes_plugin_config, get_docker_plugin_config
-from .utils import get_agent_queue, get_full_test_command, get_multi_node_test_command, A100_GPU
+from .utils import get_agent_queue, get_full_test_command, get_multi_node_test_command, GPUType
 from .step import BuildkiteStep, TestStep, get_step_key
 
 def step_should_run(step: TestStep, run_all: bool, list_file_diff: List[str]) -> bool:
@@ -21,7 +21,7 @@ def get_plugin_config(step: TestStep, container_image: str) -> Dict:
         "-c",
         get_full_test_command(test_step_commands, step.working_dir)
     ]
-    if step.gpu == A100_GPU:
+    if step.gpu == GPUType.A100:
         return get_kubernetes_plugin_config(
             container_image,
             test_bash_command,

scripts/pipeline_generator/step.py

@@ -28,11 +28,11 @@ class TestStep(BaseModel):
     @classmethod
     def validate_and_convert_command(cls, values) -> Any:
         """
-        Validate that either 'command' or 'commands' is defined.
+        Validate that either 'command' or 'commands' or `plugins` is defined.
         If 'command' is defined, convert it to 'commands'.
         """
-        if not values.get("command") and not values.get("commands"):
-            raise ValueError("Either 'command' or 'commands' must be defined.")
+        if not values.get("command") and not values.get("commands") and not values.get("plugins"):
+            raise ValueError("Either 'command' or 'commands' or `plugins` must be defined.")
         if values.get("command") and values.get("commands"):
             raise ValueError("Only one of 'command' or 'commands' can be defined.")
         if values.get("command"):
@@ -59,7 +59,7 @@ class BuildkiteStep(BaseModel):
     """This class represents a step in Buildkite format."""
     label: str
     agents: Dict[str, AgentQueue] = {"queue": AgentQueue.AWS_CPU}
-    commands: List[str]
+    commands: Optional[List[str]] = None
     key: Optional[str] = None
     plugins: Optional[List[Dict]] = None
     parallelism: Optional[int] = None

scripts/tests/pipeline_generator/test_pipeline_generator.py

@@ -359,5 +359,68 @@ def test_generate_build_step():
     assert build_step == expected_build_step
 
 
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))
+import pytest
+import sys
+import os
+import tempfile
+
+from scripts.pipeline_generator.pipeline_generator import PipelineGeneratorConfig, PipelineGenerator
+
+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")
+
+    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"
+
+
+@pytest.mark.parametrize(
+    "commit",
+    [
+        "abcdefghijklmnopqrstuvwxyz1234567890abcd", # Invalid, not in a-f 0-9
+        "1234567890abcdef", # Invalid, not 40 characters
+    ]
+)
+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)
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-v", __file__]))