Pipeline generator steps (#32)

* 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]> * 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]> * p Signed-off-by: kevin <[email protected]> * p Signed-off-by: kevin <[email protected]> --------- Signed-off-by: kevin <[email protected]>
vllm-project · Sep 23, 2024 · 24f4deb · 24f4deb
1 parent 8272857
commit 24f4deb
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 0 deletions.
diff --git a/scripts/pipeline_generator/step.py b/scripts/pipeline_generator/step.py
@@ -0,0 +1,28 @@
+from pydantic import BaseModel, Field
+from typing import List, Dict, Any, Optional
+
+from .utils import AgentQueue
+
+BUILD_STEP_KEY = "build"
+
+class BuildkiteBlockStep(BaseModel):
+    """This class represents a block step in Buildkite format."""
+    block: str
+    depends_on: Optional[str] = BUILD_STEP_KEY
+    key: str
+
+
+def get_step_key(step_label: str) -> str:
+    step_key = ""
+    skip_chars = "()% "
+    for char in step_label.lower():
+        if char in ", " and step_key[-1] != "-":
+            step_key += "-"
+        elif char not in skip_chars:
+            step_key += char
+
+    return step_key
+
+
+def get_block_step(step_label: str) -> BuildkiteBlockStep:
+    return BuildkiteBlockStep(block=f"Run {step_label}", key=f"block-{get_step_key(step_label)}")
diff --git a/scripts/tests/pipeline_generator/test_step.py b/scripts/tests/pipeline_generator/test_step.py
@@ -0,0 +1,35 @@
+import pytest
+import sys
+
+
+from scripts.pipeline_generator.step import get_step_key, get_block_step, BuildkiteBlockStep
+
+
+@pytest.mark.parametrize(
+    ("step_label", "expected_result"),
+    [
+        ("Test Step", "test-step"),
+        ("Test Step 2", "test-step-2"),
+        ("Test (Step)", "test-step"),
+        ("Test A, B, C", "test-a-b-c"),
+    ],
+)
+def test_get_step_key(step_label: str, expected_result: str):
+    assert get_step_key(step_label) == expected_result
+
+
+@pytest.mark.parametrize(
+    ("step_label", "expected_result"),
+    [
+        ("Test Step", BuildkiteBlockStep(block="Run Test Step", key="block-test-step")),
+        ("Test Step 2", BuildkiteBlockStep(block="Run Test Step 2", key="block-test-step-2")),
+        ("Test (Step)", BuildkiteBlockStep(block="Run Test (Step)", key="block-test-step")),
+        ("Test A, B, C", BuildkiteBlockStep(block="Run Test A, B, C", key="block-test-a-b-c")),
+    ],
+)
+def test_get_block_step(step_label: str, expected_result: BuildkiteBlockStep):
+    assert get_block_step(step_label) == expected_result
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))