From 286eb58c87e91ddf4cd3dd6c7cec85ade46e1b72 Mon Sep 17 00:00:00 2001
From: cotran <cotran2@utexas.edu>
Date: Wed, 23 Oct 2024 14:39:17 -0700
Subject: [PATCH] fix black

---
 model_server/app/cli.py | 63 ++++++++++++++++++++++++++++++++---------
 1 file changed, 50 insertions(+), 13 deletions(-)

diff --git a/model_server/app/cli.py b/model_server/app/cli.py
index 6a3f81ea..adb49c1b 100644
--- a/model_server/app/cli.py
+++ b/model_server/app/cli.py
@@ -19,7 +19,7 @@ log.setLevel(logging.INFO)
 PID_FILE = os.path.join(tempfile.gettempdir(), "model_server.pid")
 
 
-def run_server():
+def run_server(port=51000):
     """Start, stop, or restart the Uvicorn server based on command-line arguments."""
     if len(sys.argv) > 1:
         action = sys.argv[1]
@@ -27,17 +27,17 @@ def run_server():
         action = "start"
 
     if action == "start":
-        start_server()
+        start_server(port)
     elif action == "stop":
-        stop_server()
+        stop_server(port)
     elif action == "restart":
-        restart_server()
+        restart_server(port)
     else:
         log.info(f"Unknown action: {action}")
         sys.exit(1)
 
 
-def start_server():
+def start_server(port=51000):
     """Start the Uvicorn server and save the process ID."""
     if os.path.exists(PID_FILE):
         log.info("Server is already running. Use 'model_server restart' to restart it.")
@@ -55,7 +55,7 @@ def start_server():
             "--host",
             "0.0.0.0",
             "--port",
-            "51000",
+            f"{port}",
         ],
         start_new_session=True,
         bufsize=1,
@@ -64,7 +64,7 @@ def start_server():
         stderr=subprocess.PIPE,  # Suppress standard error. There is a logger that model_server prints to
     )
 
-    if wait_for_health_check("http://0.0.0.0:51000/healthz"):
+    if wait_for_health_check(f"http://0.0.0.0:{port}/healthz"):
         # Write the process ID to the PID file
         with open(PID_FILE, "w") as f:
             f.write(str(process.pid))
@@ -89,7 +89,35 @@ def wait_for_health_check(url, timeout=180):
     return False
 
 
-def stop_server():
+def kill_process_on_port(port=51000):
+    try:
+        # Step 1: Run lsof command to get the process using the port
+        lsof_command = f"lsof -n | grep {port} | grep -i LISTEN"
+        result = subprocess.run(
+            lsof_command, shell=True, capture_output=True, text=True
+        )
+
+        if result.returncode != 0:
+            print(f"No process found listening on port {port}.")
+            return
+
+        # Step 2: Parse the process IDs from the output
+        process_ids = [line.split()[1] for line in result.stdout.splitlines()]
+
+        if not process_ids:
+            print(f"No process found listening on port {port}.")
+            return
+
+        # Step 3: Kill each process using its PID
+        for pid in process_ids:
+            print(f"Killing process with PID {pid}")
+            subprocess.run(f"kill {pid}", shell=True)
+
+    except Exception as e:
+        print(f"Error occurred: {e}")
+
+
+def stop_server(port=51000):
     """Stop the running Uvicorn server."""
     log.info("Stopping model server")
     if not os.path.exists(PID_FILE):
@@ -105,12 +133,21 @@ def stop_server():
         process = psutil.Process(pid)
 
         # Gracefully terminate the process
-        process.terminate()  # Sends SIGTERM by default
-        process.wait(timeout=10)  # Wait for up to 10 seconds for the process to exit
+        for child in process.children(recursive=True):
+            child.terminate()
+        process.terminate()
+
+        process.wait(timeout=20)  # Wait for up to 20 seconds for the process to exit
+
+        if process.is_running():
+            log.info(f"Process with PID {pid} is still running. Forcing shutdown.")
+            process.kill()  # Forcefully kill the process
 
         log.info(f"Model server with PID {pid} stopped.")
         os.remove(PID_FILE)
 
+        kill_process_on_port(port)
+
     except psutil.NoSuchProcess:
         log.info(f"Model server with PID {pid} not found. Cleaning up PID file.")
         os.remove(PID_FILE)
@@ -122,7 +159,7 @@ def stop_server():
         os.remove(PID_FILE)
 
 
-def restart_server():
+def restart_server(port=51000):
     """Restart the Uvicorn server."""
-    stop_server()
-    start_server()
+    stop_server(port)
+    start_server(port)