[6/n][pipeline-gen] Update buildkite step format and add validators (#46

) * p Signed-off-by: kevin <[email protected]> * p Signed-off-by: kevin <[email protected]> --------- Signed-off-by: kevin <[email protected]>
vllm-project · Oct 11, 2024 · 717659c · 717659c
1 parent 1de5202
commit 717659c
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/scripts/pipeline_generator/step.py b/scripts/pipeline_generator/step.py
@@ -58,7 +58,7 @@ def validate_multi_node(self) -> Self:
 class BuildkiteStep(BaseModel):
     """This class represents a step in Buildkite format."""
     label: str
-    agents: Dict[str, AgentQueue] = {"queue": AgentQueue.AWS_CPU}
+    agents: Dict[str, str] = {"queue": AgentQueue.AWS_CPU.value}
     commands: List[str]
     key: Optional[str] = None
     plugins: Optional[List[Dict]] = None
@@ -68,12 +68,18 @@ class BuildkiteStep(BaseModel):
     env: Optional[Dict[str, str]] = None
     retry: Optional[Dict[str, Any]] = None
 
+    @model_validator(mode="after")
+    def validate_agent_queue(self) -> Self:
+        queue = self.agents.get("queue")
+        if not AgentQueue(queue):
+            raise ValueError(f"Invalid agent queue: {queue}")
+
 
 class BuildkiteBlockStep(BaseModel):
     """This class represents a block step in Buildkite format."""
     block: str
-    depends_on: Optional[str] = BUILD_STEP_KEY
     key: str
+    depends_on: Optional[str] = BUILD_STEP_KEY
 
 
 def get_step_key(step_label: str) -> str: