[10/n][pipeline-gen] Add helper method to convert test step into Buil…

…dkite step (#50) * p Signed-off-by: kevin <[email protected]> * p Signed-off-by: kevin <[email protected]> * p Signed-off-by: kevin <[email protected]> * p Signed-off-by: kevin <[email protected]> --------- Signed-off-by: kevin <[email protected]>
vllm-project · Oct 29, 2024 · 49ee880 · 49ee880
1 parent 64fde4c
commit 49ee880
Show file tree

Hide file tree

Showing 3 changed files with 124 additions and 2 deletions.
diff --git a/scripts/pipeline_generator/pipeline_generator_helper.py b/scripts/pipeline_generator/pipeline_generator_helper.py
@@ -1,4 +1,50 @@
-from typing import List
+from typing import Dict, List, Optional
+
+from .utils import GPUType, get_agent_queue, get_full_test_command, get_multi_node_test_command
+from .step import TestStep, BuildkiteStep, get_step_key
+from .plugin import get_docker_plugin_config, get_kubernetes_plugin_config
+
+def get_plugin_config(
+        container_image: str,
+        no_gpu: Optional[bool] = None,
+        gpu_type: Optional[GPUType] = None,
+        num_gpus: Optional[int] = None
+    ) -> Dict:
+    """Returns the plugin configuration for the Buildkite step."""
+    if gpu_type and gpu_type == GPUType.A100 and num_gpus:
+        return get_kubernetes_plugin_config(
+            container_image,
+            num_gpus
+        )
+    return get_docker_plugin_config(
+        container_image,
+        no_gpu or False,
+    )
+
+
+def convert_test_step_to_buildkite_step(step: TestStep, container_image: str) -> BuildkiteStep:
+    """Convert TestStep into BuildkiteStep."""
+    buildkite_step = BuildkiteStep(
+        label=step.label,
+        key=get_step_key(step.label),
+        commands=step.commands,
+        parallelism=step.parallelism,
+        soft_fail=step.soft_fail,
+        plugins=[get_plugin_config(container_image, step.no_gpu, step.gpu, step.num_gpus)],
+        agents={"queue": get_agent_queue(step.no_gpu, step.gpu, step.num_gpus).value}
+    )
+    # If test is multi-node, configure step to run with custom script
+    if step.num_nodes and step.num_nodes > 1:
+        buildkite_step.commands = [get_multi_node_test_command(
+                step.commands,
+                step.working_dir,
+                step.num_nodes,
+                step.num_gpus,
+                container_image
+            )
+        ]
+        buildkite_step.plugins = None
+    return buildkite_step
 
 def get_build_commands(container_registry: str, buildkite_commit: str, container_image: str) -> List[str]:
     ecr_login_command = (

diff --git a/scripts/pipeline_generator/step.py b/scripts/pipeline_generator/step.py
@@ -39,7 +39,7 @@ def validate_and_convert_command(cls, values) -> Any:
             values["commands"] = [values["command"]]
             del values["command"]
         return values
-    
+
     @model_validator(mode="after")
     def validate_gpu(self) -> Self:
         if self.gpu and self.no_gpu:

diff --git a/scripts/tests/pipeline_generator/test_pipeline_generator_helper.py b/scripts/tests/pipeline_generator/test_pipeline_generator_helper.py
@@ -0,0 +1,76 @@
+import pytest
+import sys
+from unittest import mock
+
+from scripts.pipeline_generator.pipeline_generator_helper import get_plugin_config, convert_test_step_to_buildkite_step
+from scripts.pipeline_generator.utils import GPUType
+from scripts.pipeline_generator.step import TestStep, BuildkiteStep
+
+@mock.patch("scripts.pipeline_generator.pipeline_generator_helper.get_kubernetes_plugin_config")
+@mock.patch("scripts.pipeline_generator.pipeline_generator_helper.get_docker_plugin_config")
+def test_get_plugin_config(mock_get_docker_plugin_config, mock_get_kubernetes_plugin_config):
+    mock_get_docker_plugin_config.return_value = {"docker": "plugin"}
+    mock_get_kubernetes_plugin_config.return_value = {"kubernetes": "plugin"}
+    plugin = get_plugin_config("image:latest")
+    assert plugin == {"docker": "plugin"}
+    assert mock_get_docker_plugin_config.call_count == 1
+    assert mock_get_docker_plugin_config.call_args[0] == ("image:latest", False)
+
+@mock.patch("scripts.pipeline_generator.pipeline_generator_helper.get_kubernetes_plugin_config")
+@mock.patch("scripts.pipeline_generator.pipeline_generator_helper.get_docker_plugin_config")
+def test_get_plugin_config_kubernetes(mock_get_docker_plugin_config, mock_get_kubernetes_plugin_config):
+    mock_get_docker_plugin_config.return_value = {"docker": "plugin"}
+    mock_get_kubernetes_plugin_config.return_value = {"kubernetes": "plugin"}
+    plugin = get_plugin_config(container_image="image:latest", gpu_type=GPUType.A100, num_gpus=4)
+    assert plugin == {"kubernetes": "plugin"}
+    assert mock_get_kubernetes_plugin_config.call_count == 1
+    assert mock_get_kubernetes_plugin_config.call_args[0] == ("image:latest", 4)
+
+@pytest.mark.parametrize(
+    ("test_step", "expected_buildkite_step"),
+    [
+        # Regular test with plugin
+        (
+            TestStep(
+                label="First test",
+                commands=["echo A", "echo B", "echo C"],
+                num_gpus=1,
+            ),
+            BuildkiteStep(
+                label="First test",
+                key="first-test",
+                commands=["echo A", "echo B", "echo C"],
+                plugins=[{"plugin": "config"}],
+                agents={"queue": "gpu_1_queue"}
+            )
+        ),
+        # Multi node test without plugin and custom command for multi-node
+        (
+            TestStep(
+                label="Second test",
+                commands=["echo D", "echo E"],
+                num_nodes=2,
+                num_gpus=2,
+            ),
+            BuildkiteStep(
+                label="Second test",
+                key="second-test",
+                commands=["multi-node-command"],
+                plugins=None,
+                agents={"queue": "gpu_4_queue"}
+            )
+        )
+    ]
+)
+@mock.patch("scripts.pipeline_generator.pipeline_generator_helper.get_multi_node_test_command")
+@mock.patch("scripts.pipeline_generator.pipeline_generator_helper.get_plugin_config")
+def test_convert_test_step(mock_get_plugin_config, mock_get_multi_node_test_command, test_step, expected_buildkite_step):
+    mock_get_plugin_config.return_value = {"plugin": "config"}
+    mock_get_multi_node_test_command.return_value = "multi-node-command"
+
+    buildkite_step = convert_test_step_to_buildkite_step(test_step, "image:latest")
+    assert buildkite_step == expected_buildkite_step
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))