Some fixes on model server (#362)

* Some fixes on model server

* Remove prompt_prefilling message

* Fix logging

* Fix poetry issues

* Improve logging and update the support for text truncation

* Fix tests

* Fix tests

* Fix tests

* Fix modelserver tests

* Update modelserver tests

model_server/src/cli.py

@@ -1,22 +1,17 @@
 import importlib
-import logging
-from os import path
 import os
-from signal import SIGKILL
 import sys
 import subprocess
 import argparse
+import signal
 import tempfile
 import time
-
 import requests
 
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-)
+import src.commons.utils as utils
 
-logger = logging.getLogger(__name__)
+
+logger = utils.get_model_server_logger()
 
 
 def get_version():
@@ -42,76 +37,9 @@ def wait_for_health_check(url, timeout=300):
     return False
 
 
-def parse_args():
-    parser = argparse.ArgumentParser(description="Manage the Uvicorn server.")
-    parser.add_argument(
-        "action",
-        choices=["start", "stop", "restart"],
-        default="start",
-        nargs="?",
-        help="Action to perform on the server (default: start).",
-    )
-    parser.add_argument(
-        "--port",
-        type=int,
-        default=51000,
-        help="Port number for the server (default: 51000).",
-    )
-
-    parser.add_argument(
-        "--foreground",
-        default=False,
-        action="store_true",
-        help="Run the server in the foreground (default: False).",
-    )
-
-    return parser.parse_args()
-
-
 def get_pid_file():
     temp_dir = tempfile.gettempdir()
-    return path.join(temp_dir, "model_server.pid")
-
-
-def stop_server():
-    """Stop the Uvicorn server."""
-    pid_file = get_pid_file()
-    if os.path.exists(pid_file):
-        logger.info(f"PID file found, shutting down the server.")
-        # read pid from file
-        with open(pid_file, "r") as f:
-            pid = int(f.read())
-            logger.info(f"Killing model server {pid}")
-            try:
-                os.kill(pid, SIGKILL)
-            except ProcessLookupError:
-                logger.info(f"Process {pid} not found")
-        os.remove(pid_file)
-    else:
-        logger.info("No PID file found, server is not running.")
-
-
-def restart_server(port=51000, foreground=False):
-    """Restart the Uvicorn server."""
-    stop_server()
-    start_server(port, foreground)
-
-
-def run_server():
-    """Start, stop, or restart the Uvicorn server based on command-line arguments."""
-
-    args = parse_args()
-    action = args.action
-
-    if action == "start":
-        start_server(args.port, args.foreground)
-    elif action == "stop":
-        stop_server()
-    elif action == "restart":
-        restart_server(args.port, args.foreground)
-    else:
-        logger.info(f"Unknown action: {action}")
-        sys.exit(1)
+    return os.path.join(temp_dir, "model_server.pid")
 
 
 def ensure_killed(process):
@@ -131,7 +59,7 @@ def ensure_killed(process):
 def start_server(port=51000, foreground=False):
     """Start the Uvicorn server."""
 
-    logging.info("model server version: %s", get_version())
+    logger.info("model server version: %s", get_version())
 
     stop_server()
 
@@ -196,6 +124,57 @@ def start_server(port=51000, foreground=False):
             ensure_killed(process)
 
 
+def stop_server():
+    """Stop the Uvicorn server."""
+
+    pid_file = get_pid_file()
+    if os.path.exists(pid_file):
+        logger.info("PID file found, shutting down the server.")
+        # read pid from file
+        with open(pid_file, "r") as f:
+            pid = int(f.read())
+            logger.info(f"Killing model server {pid}")
+            try:
+                os.kill(pid, signal.SIGKILL)
+            except ProcessLookupError:
+                logger.info(f"Process {pid} not found")
+        os.remove(pid_file)
+    else:
+        logger.info("No PID file found, server is not running.")
+
+
+def restart_server(port=51000, foreground=False):
+    """Restart the Uvicorn server."""
+    stop_server()
+    start_server(port, foreground)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Manage the Uvicorn server.")
+    parser.add_argument(
+        "action",
+        choices=["start", "stop", "restart"],
+        default="start",
+        nargs="?",
+        help="Action to perform on the server (default: start).",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=51000,
+        help="Port number for the server (default: 51000).",
+    )
+
+    parser.add_argument(
+        "--foreground",
+        default=False,
+        action="store_true",
+        help="Run the server in the foreground (default: False).",
+    )
+
+    return parser.parse_args()
+
+
 def main():
     """
     Start, stop, or restart the Uvicorn server based on command-line arguments.
@@ -204,11 +183,14 @@ def main():
     args = parse_args()
 
     if args.action == "start":
+        logger.info("[CLI] - Starting server")
         start_server(args.port, args.foreground)
     elif args.action == "stop":
+        logger.info("[CLI] - Stopping server")
         stop_server()
     elif args.action == "restart":
+        logger.info("[CLI] - Restarting server")
         restart_server(args.port)
     else:
-        logger.error(f"Unknown action: {args.action}")
+        logger.error(f"[CLI] - Unknown action: {args.action}")
         sys.exit(1)