improved localvllm logic

Author: Michael Harrison

deploy_and_run.sh

@@ -0,0 +1,77 @@
+#!/bin/bash
+
+export PYTHONPATH="$(pwd):$PYTHONPATH"
+model_name="microsoft/phi-4"
+exp_config="IFEval_PIPELINE"
+current_datetime=$(date +"%Y-%m-%d-%H:%M:%S")
+log_dir="logs/deploy_vllm_and_run_eval/$current_datetime"
+mkdir -p $log_dir
+
+# vLLM args
+num_servers=4
+tensor_parallel_size=1 
+pipeline_parallel_size=1
+base_port=8000
+gpus_per_port=$((tensor_parallel_size * pipeline_parallel_size))
+
+# Add any additional args accepted by vLLM serve here
+VLLM_ARGS="\
+    --tensor-parallel-size=${tensor_parallel_size} \
+    --pipeline-parallel-size=${pipeline_parallel_size} \
+    --gpu-memory-utilization=0.9 \
+"
+
+# Start servers
+echo "Spinning up servers..."
+for (( i = 0; i < $num_servers; i++ )) do
+    port=$((base_port + i))
+    first_gpu=$((i * gpus_per_port))
+    last_gpu=$((first_gpu + gpus_per_port - 1))
+    devices=$(seq -s, $first_gpu $last_gpu)
+    CUDA_VISIBLE_DEVICES=${devices} vllm serve ${model_name} "$@" --port ${port} ${VLLM_ARGS} >> $log_dir/${port}.log 2>&1 &
+done
+
+# Wait for servers to come online
+while true; do
+
+    servers_online=0
+    for (( i = 0; i < $num_servers; i++ )) do
+        port=$((base_port + i))
+        url="http://0.0.0.0:${port}/health"
+        response=$(curl -s -o /dev/null -w "%{http_code}" "$url")
+
+        if [ "$response" -eq 200 ]; then
+            servers_online=$((servers_online + 1))
+        fi
+    done
+
+    if [ $servers_online -eq $num_servers ]; then
+        echo "All servers are online."
+        break
+    else
+        echo "Waiting for $((num_servers - servers_online)) more servers to come online..."
+    fi
+
+    sleep 10
+done
+
+# Call Eureka to initiate evals
+ports=$(seq -s ' ' $base_port $((base_port + num_servers - 1)))
+EUREKA_ARGS="\
+    --model_config=${model_name} \
+    --exp_config=${exp_config} \
+    --local_vllm \
+    --ports ${ports} \
+"
+echo "Starting evals..."
+python main.py ${EUREKA_ARGS} >> $log_dir/out.log 2>&1
+
+# Shut down servers
+echo "Shutting down vLLM servers..."
+for (( i = 0; i < $num_servers; i++ )) do
+    port=$((base_port + i))
+    logfile="$log_dir/${port}.log"
+    pid=$(grep "Started server process" $logfile | grep -o '[0-9]\+')
+    echo "Shutting down server on port ${port} (PID ${pid})"
+    kill -INT $pid
+done

deploy_vllm_and_run_eval.sh

@@ -2,12 +2,14 @@
 
 import json
 import logging
+import random
+import re
 import requests
 import time
 import urllib.request
 from abc import ABC, abstractmethod
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
-import random
 
 import anthropic
 import tiktoken
@@ -1157,7 +1159,7 @@ def model_template_fn(self, text_prompt, system_message=None):
         raise NotImplementedError
 
 
-class LocalVLLMDeploymentHandler:
+class _LocalVLLMDeploymentHandler:
     """This class is used to handle the deployment of vLLM servers."""
 
     # Chose against dataclass here so we have the option to accept kwargs
@@ -1241,22 +1243,35 @@ def get_healthy_ports(self) -> list[str]:
         return healthy_ports
 
     def deploy_servers(self):
+        """Deploy vLLM servers in background threads using the specified parameters."""
+
         logging.info(f"No vLLM servers are running. Starting {self.num_servers} new servers at {self.ports}.")
-        import os, subprocess, sys, datetime
-        
-        env = os.environ.copy()
-        env['NUM_SERVERS'] = str(self.num_servers)
-        env['CURRENT_PYTHON_EXEC'] = sys.executable
-        env['GPU_SKIP'] = str(self.pipeline_parallel_size * self.tensor_parallel_size)
+        import os, datetime
 
+        gpus_per_port = self.pipeline_parallel_size * self.tensor_parallel_size
         date = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S.%f")
         log_dir = os.path.join("logs", "local_vllm_deployment_logs", f"{date}")
         os.makedirs(log_dir)
-        env['LOCAL_VLLM_LOG_DIR'] = log_dir
+
+        executor = ThreadPoolExecutor(max_workers=self.num_servers)
+        futures = [executor.submit(lambda index: self.deploy_server(index, gpus_per_port, log_dir), i) for i in range(self.num_servers)]
+
+    def deploy_server(self, index: int, gpus_per_port: int, log_dir: str):
+        """Deploy a single vLLM server using gpus_per_port many gpus starting at index*gpus_per_port."""
+        
+        import os, subprocess
+
+        port = 8000 + index
+        first_gpu = index * gpus_per_port
+        last_gpu = first_gpu + gpus_per_port - 1
+        devices = ",".join(str(gpu_num) for gpu_num in range(first_gpu, last_gpu + 1))
+        log_file = os.path.join(log_dir, f"{port}.log")
 
         command = [
-            os.path.dirname(os.path.abspath(__file__)) + "/vllm_deployment_script.sh",
-            "--model", self.model_name,
+            "CUDA_VISIBLE_DEVICES=" + devices,
+            "vllm serve",
+            self.model_name,
+            "--port", str(port),
             "--tensor_parallel_size", str(self.tensor_parallel_size),
             "--pipeline_parallel_size", str(self.pipeline_parallel_size),
             "--dtype", self.dtype,
@@ -1269,17 +1284,11 @@ def deploy_servers(self):
             command.append(self.quantization)
         if self.trust_remote_code:
             command.append("--trust_remote_code")
+        #command.append(">> " + log_file + " 2>&1 &")
+        command = " ".join(command)
         logging.info(f"Running command: {command}")
-        response = subprocess.run(command, text=True, env=env)
-        return response
-    
-    @classmethod
-    def shutdown_servers(cls):
-        # Consider whether this is appropriate since it will probably kill all vLLM servers.
-        import subprocess
-        logging.info(f"Shutting down vLLM servers.")
-        command = [f'pgrep -f "vllm.entrypoints.openai.api_server --model" | xargs kill -INT']
-        subprocess.run(command, shell=True)
+        with open(log_file, 'w') as log_writer:
+            subprocess.run(command, shell=True, stdout=log_writer, stderr=log_writer)
 
 
 @dataclass
@@ -1301,7 +1310,7 @@ class LocalVLLMModel(Model, OpenAICommonRequestResponseMixIn):
 
     # Deployment handler
     ports: list = None
-    handler: LocalVLLMDeploymentHandler = None
+    handler: _LocalVLLMDeploymentHandler = None
 
     # Inference parameters
     temperature: float = 0.01
@@ -1312,7 +1321,7 @@ class LocalVLLMModel(Model, OpenAICommonRequestResponseMixIn):
     def __post_init__(self):
         if not self.model_name:
             raise ValueError("LocalVLLM model_name must be specified.")
-        self.handler = LocalVLLMDeploymentHandler(
+        self.handler = _LocalVLLMDeploymentHandler(
             model_name=self.model_name,
             num_servers=self.num_servers,
             trust_remote_code=self.trust_remote_code,
@@ -1351,7 +1360,7 @@ def generate(self, text_prompt, query_images=None, system_message=None):
         response_dict = {}
 
         if text_prompt:
-            # Format request for OpenAI API using create_request from OpenAIRequestResponseMixIn
+            # Format request for OpenAI API using create_request from OpenAICommonRequestResponseMixIn
             request = self.create_request(text_prompt, query_images, system_message)
             try:
                 response_dict.update(self._generate(request))

eureka_ml_insights/models/models.py

@@ -81,12 +81,3 @@
     logging.info(f"Saving experiment logs in {pipeline_config.log_dir}.")
     pipeline = Pipeline(pipeline_config.component_configs, pipeline_config.log_dir)
     pipeline.run()
-
-    # Shut down vllm servers.
-    if args.local_vllm:
-        try:
-            from eureka_ml_insights.models.models import LocalVLLMDeploymentHandler
-            LocalVLLMDeploymentHandler.shutdown_servers()
-        except:
-            logging.warning("Failed to shut down local vllm servers.")
-