clean up debugs

Author: chibu

src/automation/tasks/scripts/guidellm_script.py

@@ -6,7 +6,7 @@
 from pyhocon import ConfigFactory
 from automation.configs import DEFAULT_GUIDELLM_SCENARIO
 
-def main():
+def main(configurations=None):
     task = Task.current_task()
 
     args = task.get_parameters_as_dict(cast=True)
@@ -53,15 +53,12 @@ def clean_hocon_value(v):
     # Resolve model_id
     model_id = resolve_model_id(args["Args"]["model"], clearml_model, force_download)
 
-    gpu_count = int(guidellm_args.get("gpu_count", 1)) 
-
     # Start vLLM server
     server_process, server_initialized, server_log = start_vllm_server(
         vllm_args,
         model_id,
         guidellm_args["target"],
         args["Args"]["server_wait_time"],
-        gpu_count,
     )
 
     if not server_initialized:
@@ -94,7 +91,6 @@ def clean_hocon_value(v):
     else:
         filepath = Path(os.path.join(".", "src", "automation", "standards", "benchmarking", f"{DEFAULT_GUIDELLM_SCENARIO}.json"))
         current_scenario = GenerativeTextScenario.from_file(filepath, dict(guidellm_args))
-    print(current_scenario.model_fields)
 
     # Ensure output_path is set and consistent
     output_path = Path(guidellm_args.get("output_path", "guidellm-output.json"))

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