test multinode

Author: Jyothirmaikottu

test/vllm/ec2/infra/setup_ec2.py

@@ -444,7 +444,7 @@ def wait_for_instances(instance_ids):
     time.sleep(30)
 
     if resources.get("sg_fsx"):
-        max_attempts = 5
+        max_attempts = 10
         for attempt in range(max_attempts):
             try:
                 ec2_cli.delete_security_group(GroupId=resources["sg_fsx"])

test/vllm/ec2/test_artifacts/test_ec2.py

@@ -118,6 +118,7 @@ def test_vllm_benchmark_on_multi_node(head_connection, worker_connection, image_
 
         for conn in [head_connection, worker_connection]:
             setup_docker_image(conn, image_uri)
+            setup_env(conn)
 
         head_connection.put(
             "vllm/ec2/utils/head_node_setup.sh", "/home/ec2-user/head_node_setup.sh"
@@ -160,10 +161,9 @@ def test_vllm_benchmark_on_multi_node(head_connection, worker_connection, image_
         print("Model serving started successfully")
 
         # Run benchmark
-        setup_env(worker_connection)
         print("Running benchmark...")
         benchmark_cmd = "source vllm_env/bin/activate" + create_benchmark_command()
-        benchmark_result = worker_connection.run(benchmark_cmd, timeout=7200)
+        benchmark_result = head_connection.run(benchmark_cmd, timeout=7200)
         print(f"Benchmark completed: {benchmark_result.stdout}")
 
         return benchmark_result