[1/n][pipeline-gen] Add default test commands (#37)

* p Signed-off-by: kevin <[email protected]> * P Signed-off-by: kevin <[email protected]> --------- Signed-off-by: kevin <[email protected]>
vllm-project · Sep 30, 2024 · d4efddc · d4efddc
1 parent 11cb8cd
commit d4efddc
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 5 deletions.
diff --git a/scripts/pipeline_generator/utils.py b/scripts/pipeline_generator/utils.py
@@ -15,6 +15,12 @@
 PIPELINE_FILE_PATH = ".buildkite/pipeline.yaml"
 MULTI_NODE_TEST_SCRIPT = ".buildkite/run-multi-node-test.sh"
 
+TEST_DEFAULT_COMMANDS = [
+    "(command nvidia-smi || true)", # Sanity check for Nvidia GPU setup
+    "export VLLM_LOGGING_LEVEL=DEBUG",
+    "export VLLM_ALLOW_DEPRECATED_BEAM_SEARCH=1",
+]
+
 STEPS_TO_BLOCK = []
 
 
@@ -33,14 +39,19 @@ def get_agent_queue(no_gpu: Optional[bool], gpu_type: Optional[str], num_gpus: O
         return AgentQueue.AWS_SMALL_CPU
     if gpu_type == A100_GPU:
         return AgentQueue.A100
-    return AgentQueue.AWS_1xL4 if num_gpus == 1 else AgentQueue.AWS_4xL4
+    return AgentQueue.AWS_1xL4 if not num_gpus or num_gpus == 1 else AgentQueue.AWS_4xL4
 
 
 def get_full_test_command(test_commands: List[str], step_working_dir: str) -> str:
     """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}"
+    full_test_commands = [
+        *TEST_DEFAULT_COMMANDS,
+        f"cd {working_dir}",
+        test_commands_str
+    ]
+    return ";\n".join(full_test_commands)
 
 
 def get_multi_node_test_command(

diff --git a/scripts/tests/pipeline_generator/test_utils.py b/scripts/tests/pipeline_generator/test_utils.py
@@ -8,8 +8,10 @@
     get_multi_node_test_command,
     AgentQueue,
     MULTI_NODE_TEST_SCRIPT,
+    TEST_DEFAULT_COMMANDS,
 )
 
+TEST_DEFAULT_COMMANDS_STR = ";\n".join(TEST_DEFAULT_COMMANDS)
 
 @pytest.mark.parametrize(
     ("no_gpu", "gpu_type", "num_gpus", "expected_result"),
@@ -27,9 +29,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, f"{TEST_DEFAULT_COMMANDS_STR};\ncd /vllm-workspace/tests;\necho 'hello'"),
+        (["echo 'hello'"], "/vllm-workspace/tests", f"{TEST_DEFAULT_COMMANDS_STR};\ncd /vllm-workspace/tests;\necho 'hello'"),
+        (["echo 'hello1'", "echo 'hello2'"], "/sample_tests", f"{TEST_DEFAULT_COMMANDS_STR};\ncd /sample_tests;\necho 'hello1';\necho 'hello2'"),
     ],
 )
 def test_get_full_test_command(test_commands: List[str], step_working_dir: str, expected_result: str):