Lint and formatting

Author: rzabarazesh

buildkite/pipeline_generator/ci/ci_pipeline.py

@@ -3,11 +3,7 @@
 from typing import Any, Dict, List, Union
 
 from ..core.amd_tests import generate_amd_group
-from ..core.docker_builds import (
-    generate_cpu_build_step,
-    generate_cu118_build_steps,
-    generate_main_build_step,
-)
+from ..core.docker_builds import generate_cpu_build_step, generate_cu118_build_steps, generate_main_build_step
 from ..core.hardware_tests import generate_all_hardware_tests
 from ..core.manual_trigger_rules import should_block_ci_test
 from ..core.test_step_converter import convert_test_step_to_buildkite_step

buildkite/pipeline_generator/core/__init__.py

@@ -1,2 +1 @@
 """Core pipeline generation modules - unified CI and Fastcheck logic."""
-

buildkite/pipeline_generator/core/amd_tests.py

@@ -40,19 +40,19 @@ def get_amd_queue_ci(label: str, num_gpus: Optional[int] = None) -> str:
 def generate_amd_group(test_steps: List[TestStep], config: PipelineGeneratorConfig) -> Dict[str, Any]:
     """
     Generate the AMD tests group.
-    
+
     CI mode: All matching tests, no blocks, soft_fail=false, uses AMD MI325 queues
     Fastcheck mode: Only Basic Correctness Test, has block, soft_fail=true, uses MI300_1 queue
     """
     amd_steps = []
 
     # Add AMD build step (now returns dict directly)
     amd_build_dict = generate_amd_build_step(config)
-    
+
     # Fastcheck needs depends_on: null explicitly
     if config.pipeline_mode == PipelineMode.FASTCHECK:
         amd_build_dict["depends_on"] = None
-    
+
     amd_steps.append(amd_build_dict)
 
     # Add AMD mirror tests
@@ -65,7 +65,7 @@ def generate_amd_group(test_steps: List[TestStep], config: PipelineGeneratorConf
         if config.pipeline_mode == PipelineMode.FASTCHECK:
             if test_step.label != TestLabels.BASIC_CORRECTNESS_TEST:
                 continue
-            
+
             # Fastcheck adds a block for Basic Correctness Test
             block_key = f"block-amd-{get_step_key(test_step.label)}"
             amd_steps.append(
@@ -105,4 +105,3 @@ def generate_amd_group(test_steps: List[TestStep], config: PipelineGeneratorConf
         amd_steps.append(amd_step_dict)
 
     return {"group": "AMD Tests", "depends_on": None, "steps": amd_steps}
-

buildkite/pipeline_generator/core/docker_builds.py

@@ -1,6 +1,6 @@
 """Unified Docker build step generation for both CI and Fastcheck modes."""
 
-from typing import Any, Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 from ..data_models.buildkite_step import BuildkiteBlockStep, BuildkiteStep
 from ..docker_build_configs import (
@@ -18,10 +18,10 @@
 def generate_main_build_step(config: PipelineGeneratorConfig) -> Union[BuildkiteStep, Dict]:
     """
     Build the main Docker CUDA image.
-    
+
     CI: Uses postmerge/test repo based on branch
     Fastcheck: Always uses test-repo with premerge queue
-    
+
     Returns dict directly to preserve all fields (especially 'key')
     """
     if config.pipeline_mode == PipelineMode.FASTCHECK:
@@ -52,12 +52,12 @@ def generate_cu118_build_steps(
 ) -> List[Union[BuildkiteStep, BuildkiteBlockStep]]:
     """
     Build the CUDA 11.8 Docker image.
-    
+
     CI only - returns empty list for Fastcheck mode.
     """
     if config.pipeline_mode == PipelineMode.FASTCHECK:
         return []  # Fastcheck doesn't build cu118
-    
+
     # CI mode
     queue = AgentQueue.CPU_QUEUE_POSTMERGE_US_EAST_1 if config.branch == "main" else AgentQueue.CPU_QUEUE_PREMERGE_US_EAST_1
 
@@ -80,13 +80,13 @@ def generate_cu118_build_steps(
 def generate_cpu_build_step(config: PipelineGeneratorConfig) -> Optional[Dict]:
     """
     Build the CPU Docker image.
-    
+
     CI only - returns None for Fastcheck mode.
     Returns dict directly to preserve all fields.
     """
     if config.pipeline_mode == PipelineMode.FASTCHECK:
         return None  # Fastcheck doesn't build CPU image
-    
+
     # CI mode
     queue = AgentQueue.CPU_QUEUE_POSTMERGE_US_EAST_1 if config.branch == "main" else AgentQueue.CPU_QUEUE_PREMERGE_US_EAST_1
 
@@ -99,7 +99,7 @@ def generate_cpu_build_step(config: PipelineGeneratorConfig) -> Optional[Dict]:
 def generate_torch_nightly_build_step(config: PipelineGeneratorConfig, depends_on: Optional[str]) -> Dict:
     """
     Build the torch nightly Docker image.
-    
+
     CI only - should not be called in Fastcheck mode.
     Returns dict directly to preserve all fields.
     """
@@ -124,7 +124,7 @@ def generate_torch_nightly_build_step(config: PipelineGeneratorConfig, depends_o
 def generate_amd_build_step(config: PipelineGeneratorConfig) -> Dict:
     """
     Build the AMD Docker image.
-    
+
     Different configurations for CI vs Fastcheck.
     Returns dict directly to preserve all fields.
     """
@@ -137,4 +137,3 @@ def generate_amd_build_step(config: PipelineGeneratorConfig) -> Dict:
 
     # Return dict directly
     return amd_config.to_buildkite_step()
-

buildkite/pipeline_generator/core/docker_plugins.py

@@ -74,7 +74,7 @@ def build_full_docker_command(test_step: TestStep, config: PipelineGeneratorConf
         docker_command = build_docker_command_fastcheck(test_step, config)
     else:  # CI mode
         docker_command = build_docker_command_ci(test_step, config)
-    
+
     working_dir = test_step.working_dir or DEFAULT_WORKING_DIR
     return f"{ShellCommands.CHECK_NVIDIA_GPU} && {ShellCommands.SETUP_DEPRECATED_BEAM_SEARCH} && cd {working_dir} && {docker_command}"
 
@@ -116,12 +116,12 @@ def build_environment(test_step: TestStep, config: PipelineGeneratorConfig) -> D
 def build_docker_plugin(test_step: TestStep, container_image: str, config: PipelineGeneratorConfig) -> Dict:
     """Build standard Docker plugin configuration."""
     full_command = build_full_docker_command(test_step, config)
-    
+
     # CI mode adds trailing space; Fastcheck doesn't
     if config.pipeline_mode == PipelineMode.CI:
         if full_command and not full_command.endswith(" "):
             full_command += " "
-    
+
     bash_flags = "-xce" if config.fail_fast else "-xc"
 
     # Build environment configuration (mode-aware)
@@ -132,9 +132,9 @@ def build_docker_plugin(test_step: TestStep, container_image: str, config: Pipel
 
     # Determine if mount_buildkite_agent is needed
     if config.pipeline_mode == PipelineMode.FASTCHECK:
-        mount_agent = test_step.label == TestLabels.BENCHMARKS or test_step.mount_buildkite_agent
+        mount_agent = test_step.label == TestLabels.BENCHMARKS or bool(test_step.mount_buildkite_agent)
     else:  # CI mode
-        mount_agent = test_step.label == TestLabels.BENCHMARKS or test_step.mount_buildkite_agent or config.cov_enabled
+        mount_agent = test_step.label == TestLabels.BENCHMARKS or bool(test_step.mount_buildkite_agent) or config.cov_enabled
 
     docker_config = StandardDockerConfig(
         image=container_image,
@@ -152,12 +152,12 @@ def build_docker_plugin(test_step: TestStep, container_image: str, config: Pipel
 def build_special_gpu_plugin(test_step: TestStep, container_image: str, config: PipelineGeneratorConfig) -> Dict:
     """Build Docker plugin for special GPUs (H200, B200)."""
     full_command = build_full_docker_command(test_step, config)
-    
+
     # CI mode adds trailing space
     if config.pipeline_mode == PipelineMode.CI:
         if full_command and not full_command.endswith(" "):
             full_command += " "
-    
+
     bash_flags = "-xce" if config.fail_fast else "-xc"
 
     gpu_type = test_step.gpu.value if test_step.gpu else "h200"
@@ -191,7 +191,7 @@ def build_special_gpu_plugin(test_step: TestStep, container_image: str, config:
 def build_fastcheck_a100_kubernetes_plugin(test_step: TestStep, container_image: str) -> Dict:
     """Build Kubernetes plugin for A100 in fastcheck mode (fastcheck-specific)."""
     # Build command without coverage
-    commands = flatten_commands(test_step.commands)
+    commands = flatten_commands(test_step.commands or [])
     commands = normalize_commands(commands)
     docker_command = " && ".join(commands)
     working_dir = test_step.working_dir or DEFAULT_WORKING_DIR
@@ -282,7 +282,7 @@ def build_plugin_for_test_step(test_step: TestStep, container_image: str, config
             return build_fastcheck_a100_kubernetes_plugin(test_step, container_image)
         # All other GPUs (H100, H200, B200, etc.) use standard Docker in fastcheck
         return build_docker_plugin(test_step, container_image, config)
-    
+
     # CI mode: special handling for multiple GPU types
     # CI mode uses Kubernetes for H100 and A100
     if test_step.gpu in [GPUType.H100, GPUType.A100]:
@@ -294,4 +294,3 @@ def build_plugin_for_test_step(test_step: TestStep, container_image: str, config
 
     # Standard Docker for all others
     return build_docker_plugin(test_step, container_image, config)
-

buildkite/pipeline_generator/core/hardware_tests.py

@@ -18,13 +18,13 @@
     get_tpu_notification_config,
     get_tpu_v0_notification_config,
 )
-from ..utils.constants import AgentQueue, BlockLabels, HardwareLabels, PipelineMode, Scripts
+from ..utils.constants import AgentQueue, BlockLabels, HardwareLabels, Scripts
 
 
 def generate_all_hardware_tests(branch: str, nightly: bool) -> List[Dict[str, Any]]:
     """
     Generate all hardware-specific test steps using data-driven configuration.
-    
+
     Used by CI mode. Returns CI-specific hardware tests without blocking steps.
     """
     steps = []
@@ -68,7 +68,9 @@ def generate_all_hardware_tests(branch: str, nightly: bool) -> List[Dict[str, An
     if branch == "main":
         tpu_depends_on = get_tpu_notification_config()
         # Build notification command with proper indentation
-        tpu_notif_command = """if [[ $$(buildkite-agent step get "outcome" --step "run-tpu-v1-test") != "passed" || $$(buildkite-agent step get "outcome" --step "run-tpu-v1-test-part2") != "passed" ]]; then
+        tpu_notif_command = (
+            """if [[ $$(buildkite-agent step get "outcome" --step "run-tpu-v1-test") != "passed" || """
+            """$$(buildkite-agent step get "outcome" --step "run-tpu-v1-test-part2") != "passed" ]]; then
    cat <<- YAML | buildkite-agent pipeline upload
    steps:
      - label: "Notify owners about failing test"
@@ -81,6 +83,7 @@ def generate_all_hardware_tests(branch: str, nightly: bool) -> List[Dict[str, An
                - "vllm#tpu-ci-notifications"
 YAML
 fi"""
+        )
         steps.append(
             {
                 "label": "TPU V1 Test Notification",
@@ -205,4 +208,3 @@ def add_neuron_test_fastcheck(steps: List) -> None:
     }
     steps.append(neuron_block)
     steps.append(neuron_test)
-

buildkite/pipeline_generator/core/manual_trigger_rules.py

@@ -93,11 +93,10 @@ def should_block_torch_nightly_test(test_step: TestStep, config: PipelineGenerat
 def should_block_test(test_step: TestStep, config: PipelineGeneratorConfig) -> bool:
     """
     Unified function to determine if a test needs a manual trigger block.
-    
+
     Routes to mode-specific logic based on config.pipeline_mode.
     """
     if config.pipeline_mode == PipelineMode.FASTCHECK:
         return should_block_fastcheck_test(test_step, config)
     else:  # CI mode
         return should_block_ci_test(test_step, config)
-

buildkite/pipeline_generator/core/test_filtering.py

@@ -4,13 +4,12 @@
 
 from ..data_models.test_step import TestStep
 from ..pipeline_config import PipelineGeneratorConfig
-from ..utils.constants import PipelineMode
 
 
 def should_run_step(test_step: TestStep, config: PipelineGeneratorConfig) -> bool:
     """
     Determine if a step should run based on configuration and file changes.
-    
+
     Used by CI mode for intelligent test selection.
     """
     # Always run if run_all or nightly is enabled
@@ -176,4 +175,3 @@ def extract_pytest_markers(commands) -> str:
             return f" -m {marker}"
 
     return ""
-

buildkite/pipeline_generator/core/test_step_converter.py

@@ -6,16 +6,7 @@
 from ..data_models.test_step import TestStep
 from ..pipeline_config import PipelineGeneratorConfig
 from ..utils.agent_queues import get_agent_queue
-from ..utils.constants import (
-    DEFAULT_WORKING_DIR,
-    AgentQueue,
-    GPUType,
-    PipelineMode,
-    PriorityValues,
-    RetryConfig,
-    Scripts,
-    TestLabels,
-)
+from ..utils.constants import DEFAULT_WORKING_DIR, AgentQueue, GPUType, PipelineMode, PriorityValues, RetryConfig, Scripts, TestLabels
 from .docker_plugins import build_plugin_for_test_step
 
 
@@ -30,8 +21,8 @@ def convert_multi_node_test_step_ci(test_step: TestStep, container_image: str, c
         node_commands = test_step.commands  # type: ignore
     else:
         # Fallback: use the same commands for all nodes
-        simple_commands: List[str] = test_step.commands if test_step.commands else []
-        node_commands = [simple_commands] * (test_step.num_nodes or 2)
+        simple_commands = test_step.commands if test_step.commands else []
+        node_commands = [simple_commands] * (test_step.num_nodes or 2)  # type: ignore
 
     # Build the multi-node command
     quoted_node_commands = []
@@ -78,8 +69,8 @@ def convert_multi_node_test_step_fastcheck(test_step: TestStep, container_image:
     if test_step.commands and len(test_step.commands) > 0 and isinstance(test_step.commands[0], list):
         node_commands = test_step.commands  # type: ignore
     else:
-        simple_commands: List[str] = test_step.commands if test_step.commands else []
-        node_commands = [simple_commands] * (test_step.num_nodes or 2)
+        simple_commands = test_step.commands if test_step.commands else []
+        node_commands = [simple_commands] * (test_step.num_nodes or 2)  # type: ignore
 
     # Build the multi-node command
     quoted_node_commands = []
@@ -123,7 +114,7 @@ def get_agent_queue_fastcheck(test_step: TestStep) -> str:
 def convert_test_step_to_buildkite_step(test_step: TestStep, container_image: str, config: PipelineGeneratorConfig) -> BuildkiteStep:
     """
     Unified function to convert TestStep into BuildkiteStep.
-    
+
     Routes to mode-specific logic based on config.pipeline_mode.
     """
     # Check if this is a multi-node test
@@ -143,10 +134,10 @@ def convert_test_step_to_buildkite_step(test_step: TestStep, container_image: st
             agent_queue = AgentQueue.A100_QUEUE
         else:
             agent_queue = get_agent_queue_fastcheck(test_step)
-        
+
         # Fastcheck uses retry_limit = 5 (except for A100 which is handled below)
         retry_limit = 5
-        
+
         # A100 in fastcheck has special handling
         if test_step.gpu == GPUType.A100:
             return BuildkiteStep(
@@ -192,4 +183,3 @@ def convert_test_step_to_buildkite_step(test_step: TestStep, container_image: st
     )
 
     return buildkite_step
-

buildkite/pipeline_generator/data_models/docker_config.py

@@ -214,6 +214,15 @@ def to_plugin_dict(self) -> Dict[str, Any]:
                         },
                         {"name": "NCCL_CUMEM_HOST_ENABLE", "value": "0"},
                         {"name": "HF_HOME", "value": HF_HOME},
+                        {
+                            "name": "HF_TOKEN",
+                            "valueFrom": {
+                                "secretKeyRef": {
+                                    "name": KubernetesConstants.HF_TOKEN_SECRET_NAME,
+                                    "key": KubernetesConstants.HF_TOKEN_SECRET_KEY,
+                                }
+                            },
+                        },
                     ],
                 }
             ],
@@ -235,18 +244,6 @@ def to_plugin_dict(self) -> Dict[str, Any]:
 
         if self.priority_class:
             pod_spec["priorityClassName"] = self.priority_class
-            # Add HF_TOKEN secret for A100
-            pod_spec["containers"][0]["env"].append(
-                {
-                    "name": "HF_TOKEN",
-                    "valueFrom": {
-                        "secretKeyRef": {
-                            "name": KubernetesConstants.HF_TOKEN_SECRET_NAME,
-                            "key": KubernetesConstants.HF_TOKEN_SECRET_KEY,
-                        }
-                    },
-                }
-            )
 
         return {PluginNames.KUBERNETES: {"podSpec": pod_spec}}