From e5b4fd00ed66ce950069e623ac6206188c6c2bab Mon Sep 17 00:00:00 2001
From: Adil Hafeez <adil.hafeez@gmail.com>
Date: Tue, 10 Dec 2024 19:02:20 -0800
Subject: [PATCH] fix model_server test bootup

---
 model_server/src/commons/utils.py          | 56 ++++++++++++++++++++++
 model_server/tests/test_cli_stop_server.py |  2 +-
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/model_server/src/commons/utils.py b/model_server/src/commons/utils.py
index 99fe03af..0e9d9e10 100644
--- a/model_server/src/commons/utils.py
+++ b/model_server/src/commons/utils.py
@@ -167,6 +167,62 @@ def find_processes_by_port(port=51000):
         return []
 
 
+def kill_process(port=51000, wait=True, timeout=10):
+    """Stop the running Uvicorn server."""
+    try:
+        # Run the function to check and install lsof if necessary
+        # 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 model server process with PID {pid}")
+            subprocess.run(f"kill {pid}", shell=True)
+
+            if wait:
+                # Step 4: Wait for the process to be killed by checking if it's still running
+                start_time = time.time()
+
+                while True:
+                    check_process = subprocess.run(
+                        f"ps -p {pid}", shell=True, capture_output=True, text=True
+                    )
+                    if check_process.returncode != 0:
+                        print(f"Process {pid} has been killed.")
+                        break
+
+                    elapsed_time = time.time() - start_time
+                    if elapsed_time > timeout:
+                        print(
+                            f"Process {pid} did not terminate within {timeout} seconds."
+                        )
+                        print(f"Attempting to force kill process {pid}...")
+                        subprocess.run(f"kill -9 {pid}", shell=True)  # SIGKILL
+                        break
+
+                    print(
+                        f"Waiting for process {pid} to be killed... ({elapsed_time:.2f} seconds)"
+                    )
+                    time.sleep(0.5)
+
+    except Exception as e:
+        print(f"Error occurred: {e}")
+
+
 def kill_processes(port_processes, wait=True, timeout=10):
     """Kill processes on a specific port."""
 
diff --git a/model_server/tests/test_cli_stop_server.py b/model_server/tests/test_cli_stop_server.py
index d9d380d5..4c7183ac 100644
--- a/model_server/tests/test_cli_stop_server.py
+++ b/model_server/tests/test_cli_stop_server.py
@@ -1,7 +1,7 @@
 import unittest
 
 from unittest.mock import patch, MagicMock
-from src.core.cli import kill_process
+from src.commons.utils import kill_process
 
 
 class TestStopServer(unittest.TestCase):