p

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

scripts/pipeline_generator/pipeline_generator.py

@@ -3,11 +3,11 @@
 from typing import List, Dict, Union
 import os
 
-from plugin import (
+from .plugin import (
     get_kubernetes_plugin_config,
     get_docker_plugin_config,
 )
-from utils import (
+from .utils import (
     AgentQueue,
     AMD_REPO,
     A100_GPU,
@@ -21,7 +21,7 @@
     get_full_test_command,
     get_multi_node_test_command,
 )
-from step import (
+from .step import (
     TestStep, 
     BuildkiteStep, 
     BuildkiteBlockStep, 
@@ -54,7 +54,7 @@ def step_should_run(self, step: TestStep) -> bool:
     def process_step(self, step: TestStep) -> List[Union[BuildkiteStep, BuildkiteBlockStep]]:
         """Process test step and return corresponding BuildkiteStep."""
         steps = []
-        current_step = self._create_buildkite_step(step)
+        current_step = self.create_buildkite_step(step)
 
         if step.num_nodes > 1:
             self._configure_multi_node_step(current_step, step)
@@ -86,6 +86,12 @@ def generate_build_step(self) -> BuildkiteStep:
             commands=build_commands,
             depends_on=None,
         )
+
+    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 get_external_hardware_tests(self, test_steps: List[TestStep]) -> List[Union[BuildkiteStep, BuildkiteBlockStep]]:
         """Process the external hardware tests from the yaml file and convert to Buildkite steps."""
@@ -95,28 +101,28 @@ def get_external_hardware_tests(self, test_steps: List[TestStep]) -> List[Union[
 
     def get_plugin_config(self, step: TestStep) -> Dict:
         """Returns the plugin configuration for the step."""
-        test_commands = [step.command] if step.command else step.commands
+        test_step_commands = [step.command] if step.command else step.commands
         test_bash_command = [
             "bash",
             "-c",
-            get_full_test_command(test_commands, step.working_dir)
+            get_full_test_command(test_step_commands, step.working_dir)
         ]
-        docker_image_path = f"{VLLM_ECR_REPO}:{self.commit}"
+        test_bash_command[-1] = f"'{test_bash_command[-1]}'"
+        container_image = f"{VLLM_ECR_REPO}:{self.commit}"
 
         if step.gpu == A100_GPU:
-            test_bash_command[-1] = f"'{test_bash_command[-1]}'"
             return get_kubernetes_plugin_config(
-                docker_image_path,
+                container_image,
                 test_bash_command,
                 step.num_gpus
             )
         return get_docker_plugin_config(
-            docker_image_path,
+            container_image,
             test_bash_command,
             step.no_gpu
         )
 
-    def _create_buildkite_step(self, step: TestStep) -> BuildkiteStep:
+    def create_buildkite_step(self, step: TestStep) -> BuildkiteStep:
         return BuildkiteStep(
             label=step.label, 
             key=get_step_key(step.label), 
@@ -206,10 +212,8 @@ def main(run_all: str = "-1", list_file_diff: str = None):
         *pipeline_generator.get_external_hardware_tests(test_steps)
     ]
 
-    buildkite_steps_dict = {"steps": [step.dict(exclude_none=True) for step in buildkite_steps]}
+    pipeline_generator.write_buildkite_steps(buildkite_steps, PIPELINE_FILE_PATH)
 
-    with open(PIPELINE_FILE_PATH, "w") as f:
-        yaml.dump(buildkite_steps_dict, f, sort_keys=False)
 
 if __name__ == "__main__":
     main()

scripts/pipeline_generator/plugin.py

@@ -1,7 +1,7 @@
 from pydantic import BaseModel, Field
 from typing import List, Dict, Any, Optional
 
-from utils import HF_HOME
+from .utils import HF_HOME
 
 DOCKER_PLUGIN_NAME = "docker#v5.2.0"
 KUBERNETES_PLUGIN_NAME = "kubernetes"
@@ -103,7 +103,7 @@ def get_kubernetes_plugin_config(container_image: str, test_bash_command: List[s
 def get_docker_plugin_config(docker_image_path: str, test_bash_command: List[str], no_gpu: bool) -> Dict:
     docker_plugin_config = DockerPluginConfig(
         image=docker_image_path,
-        command=test_bash_command
+        command=[" ".join(test_bash_command)]
     )
     if no_gpu:
         docker_plugin_config.gpus = None

scripts/pipeline_generator/step.py

@@ -1,7 +1,7 @@
 from pydantic import BaseModel, Field
 from typing import List, Dict, Any, Optional
 
-from utils import AgentQueue
+from .utils import AgentQueue
 
 BUILD_STEP_KEY = "build"
 

scripts/pipeline_generator/utils.py

@@ -40,7 +40,15 @@ def get_full_test_command(test_commands: List[str], step_working_dir: str) -> st
     """Convert test commands into one-line command with the right directory."""
     working_dir = step_working_dir or DEFAULT_WORKING_DIR
     test_commands_str = ";\n".join(test_commands)
-    return f"cd {working_dir};\n{test_commands_str}"
+    # Always add these commands before running the tests
+    commands = [
+        "(command nvidia-smi || true)",
+        "export VLLM_LOGGING_LEVEL=DEBUG",
+        "export VLLM_ALLOW_DEPRECATED_BEAM_SEARCH=1",
+        f"cd {working_dir}",
+        test_commands_str
+    ]
+    return ";\n".join(commands)
 
 
 def get_multi_node_test_command(

scripts/tests/pipeline_generator/test_pipeline_generator.py

@@ -0,0 +1,70 @@
+import pytest
+
+from scripts.pipeline_generator.pipeline_generator import PipelineGenerator
+from scripts.pipeline_generator.step import TestStep, BuildkiteStep, BuildkiteBlockStep
+
+TEST_COMMIT = "123456789abcdef123456789abcdef123456789a"
+TEST_FILE_PATH = "scripts/tests/pipeline_generator/tests.yaml"
+
+def get_test_pipeline_generator():
+    pipeline_generator = PipelineGenerator(run_all=False, list_file_diff=[])
+    pipeline_generator.commit = TEST_COMMIT
+    return pipeline_generator
+
+def test_read_test_steps():
+    pipeline_generator = get_test_pipeline_generator()
+    steps = pipeline_generator.read_test_steps(TEST_FILE_PATH)
+    assert len(steps) == 4
+    for i in range(4):
+        assert steps[i].label == f"Test {i}"
+    assert steps[0].source_file_dependencies == ["dir1/", "dir2/file1"]
+    assert steps[0].commands == ["pytest -v -s a", "pytest -v -s b.py"]
+    assert steps[1].working_dir == "/tests"
+    assert steps[2].num_gpus == 2
+    assert steps[2].num_nodes == 2
+    assert steps[3].gpu == "a100"
+    assert steps[3].optional == True
+
+@pytest.mark.parametrize(
+    ("test_step", "expected_plugin_config"),
+    [
+        (
+            TestStep(
+                label="Test 0",
+                source_file_dependencies=["dir1/", "dir2/file1"],
+                commands=["test command 1", "test command 2"]
+            ),
+            {
+                "plugin": "docker"
+            }
+        ),
+        (
+            TestStep(
+                label="Test 1",
+                commands=["test command 1", "test command 2"]
+                gpu="a100"
+            ),
+            {
+                "plugin": "kubernetes"
+            }
+        )
+    ]
+)
+@mock.patch("scripts.pipeline_generator.pipeline_generator.get_docker_plugin_config")
+@mock.patch("scripts.pipeline_generator.pipeline_generator.get_kubernetes_plugin_config")
+@mock.patch("scripts.pipeline_generator.utils.get_full_test_command")
+def test_get_plugin_config(mock_get_full_test_command, mock_get_kubernetes_plugin_config, mock_get_docker_plugin_config, test_step, expected_plugin_config):
+    pipeline_generator = get_test_pipeline_generator()
+    mock_get_full_test_command.return_value = "test command 1;\ntest command 2"
+    mock_get_docker_plugin_config.return_value = {"plugin": "docker"}
+    mock_get_kubernetes_plugin_config.return_value = {"plugin": "kubernetes"}
+    container_image_path = f"{VLLM_ECR_REPO}:{TEST_COMMIT}"
+
+    plugin_config = pipeline_generator.get_plugin_config(test_step)
+    assert plugin_config == expected_plugin_config
+    if test_step.gpu == "a100":
+        assert mock_get_kubernetes_plugin_config.called_once_with(container_image_path, )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

scripts/tests/pipeline_generator/test_utils.py

@@ -27,9 +27,9 @@ def test_get_agent_queue(no_gpu: bool, gpu_type: str, num_gpus: int, expected_re
 @pytest.mark.parametrize(
     ("test_commands", "step_working_dir", "expected_result"),
     [
-        (["echo 'hello'"], None, "cd /vllm-workspace/tests;\necho 'hello'"),
-        (["echo 'hello'"], "/vllm-workspace/tests", "cd /vllm-workspace/tests;\necho 'hello'"),
-        (["echo 'hello1'", "echo 'hello2'"], None, "cd /vllm-workspace/tests;\necho 'hello1';\necho 'hello2'"),
+        (["echo 'hello'"], None, "(command nvidia-smi || true);\nexport VLLM_LOGGING_LEVEL=DEBUG;\nexport VLLM_ALLOW_DEPRECATED_BEAM_SEARCH=1;\ncd /vllm-workspace/tests;\necho 'hello'"),
+        (["echo 'hello'"], "/vllm-workspace/tests", "(command nvidia-smi || true);\nexport VLLM_LOGGING_LEVEL=DEBUG;\nexport VLLM_ALLOW_DEPRECATED_BEAM_SEARCH=1;\ncd /vllm-workspace/tests;\necho 'hello'"),
+        (["echo 'hello1'", "echo 'hello2'"], None, "(command nvidia-smi || true);\nexport VLLM_LOGGING_LEVEL=DEBUG;\nexport VLLM_ALLOW_DEPRECATED_BEAM_SEARCH=1;\ncd /vllm-workspace/tests;\necho 'hello1';\necho 'hello2'"),
     ],
 )
 def test_get_full_test_command(test_commands: List[str], step_working_dir: str, expected_result: str):

scripts/tests/pipeline_generator/tests.yaml

@@ -0,0 +1,25 @@
+steps:
+- label: Test 0
+  source_file_dependencies:
+  - dir1/
+  - dir2/file1
+  commands:
+  - pytest -v -s a
+  - pytest -v -s b.py
+- label: Test 1
+  working_dir: "/tests"
+  commands:
+  - pytest -v -s d
+- label: Test 2
+  num_gpus: 2
+  num_nodes: 2
+  commands:
+  - pytest -v -s e && pytest -v -s f
+  - pytest -v -s g
+- label: Test 3
+  working_dir: "/tests"
+  gpu: a100
+  num_gpus: 4
+  optional: true
+  commands:
+  - pytest -v -s d