revert to main vllm server script

Author: chibu

src/automation/vllm/server.py

@@ -14,34 +14,25 @@ def start_vllm_server(
     vllm_args, 
     model_id, 
     target, 
-    server_wait_time,
-    gpu_count,
+    server_wait_time, 
 ):
     task = Task.current_task()
 
-    print("Inside start vllm server")
-
     executable_path = os.path.dirname(sys.executable)
     vllm_path = os.path.join(executable_path, "vllm")
 
-    available_gpus = list(range(torch.cuda.device_count()))
-    selected_gpus = available_gpus[:gpu_count]
-
-    subprocess_env = os.environ.copy()
-    subprocess_env["CUDA_VISIBLE_DEVICES"] = ",".join(str(i) for i in selected_gpus)
+    num_gpus = torch.cuda.device_count()
 
     parsed_target = urlparse(target)
-    print(f"vllm path is: {vllm_path}")
 
     server_command = [
         f"{vllm_path}", "serve", 
         model_id,
         "--host", parsed_target.hostname, 
         "--port", str(parsed_target.port),
-        "--tensor-parallel-size", str(gpu_count),
+        "--tensor-parallel-size", str(num_gpus)
     ]
 
-    print(server_command)
     subprocess_env = os.environ.copy()
 
     for k, v in vllm_args.items():
@@ -52,20 +43,17 @@ def start_vllm_server(
                 server_command.append(f"--{k}")
             else:
                 server_command.extend([f"--{k}", str(v)])
-
+                
 
     server_log_file_name = f"{SERVER_LOG_PREFIX}_{task.id}.txt"
     server_log_file = open(server_log_file_name, "w")
-    print("Server command:", " ".join(server_command))
-    print(f"VLLM logs are located at: {server_log_file} in {os.getcwd()}")
     server_process = subprocess.Popen(server_command, stdout=server_log_file, stderr=server_log_file, shell=False, env=subprocess_env)
 
     delay = 5
     server_initialized = False
     for _ in range(server_wait_time // delay):
         try:
             response = requests.get(target + "/models")
-            print(f"response: {response}")
             if response.status_code == 200:
                 print("Server initialized")
                 server_initialized = True