Added expiration timeout and handling for models

Signed-off-by: Michael Engel <[email protected]>

Author: engelmi

ramalama/daemon/daemon.py

@@ -4,6 +4,7 @@
 import signal
 import socketserver
 import threading
+from datetime import datetime, timedelta
 
 from ramalama.daemon.handler.ramalama import RamalamaHandler
 from ramalama.daemon.logging import configure_logger, logger
@@ -17,35 +18,66 @@ def __init__(self, server: "RamalamaServer") -> None:
     def handle_kill(self, signum, frame):
         self.server.shutdown()
 
+    def handle_alarm(self, signum, frame):
+        # check for expiration of all models, stopping them if necessary and
+        # stop shutdown server if no models are running
+        self.server.check_model_expiration()
+        if not self.server.model_runner.managed_models:
+            self.server.shutdown()
+            return
+
+        # register alarm again for next check
+        signal.alarm(self.server.idle_check_interval.seconds)
+
     def __enter__(self):
         signal.signal(signal.SIGINT, self.handle_kill)
         signal.signal(signal.SIGTERM, self.handle_kill)
 
+        # set initial idle check to 300s == 5min to prevent service from stopping
+        # right afer being started
+        signal.signal(signal.SIGALRM, self.handle_alarm)
+        signal.alarm(300)
+
     def __exit__(self, type, value, traceback):
         pass
 
 
 class RamalamaServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
 
-    def __init__(self, host: str, port: int, model_store_path: str, bind_and_activate=True):
+    def __init__(
+        self, host: str, port: int, model_store_path: str, idle_check_interval: timedelta, bind_and_activate=True
+    ):
         # Do not pass a RequestHandlerClass here, we will create a custom handler in finish_request
         super().__init__((host, port), None, bind_and_activate)
 
         self.model_store_path: str = model_store_path
         self.model_runner: ModelRunner = ModelRunner()
+        self.idle_check_interval = idle_check_interval
 
         self.allow_reuse_address = True
 
     def finish_request(self, request, client_address):
         RamalamaHandler(self.model_store_path, self.model_runner, request, client_address, self)
 
+    def check_model_expiration(self):
+        curr_time = datetime.now()
+        for name, m in self.model_runner.managed_models.items():
+            if m.expiration_date > curr_time:
+                continue
+
+            try:
+                logger.error(f"Stopping expired model '{name}'...")
+                self.model_runner.stop_model(m.id)
+            except Exception as e:
+                logger.error(f"Failed to stop expired model '{name}': {e}")
+
     def shutdown(self):
         logger.info("Shutting down ramalama daemon...")
 
         for name, managed_model in self.model_runner.managed_models.items():
             try:
                 logger.info(f"Stopping model runner {name}...")
-                managed_model.stop()
+                self.model_runner.stop_model(managed_model.id)
             except Exception as e:
                 logger.error(f"Error stopping model runner {name}: {e}")
 
@@ -64,7 +96,7 @@ def parse_args():
 def run(host: str = "0.0.0.0", port: int = 8080, model_store_path: str = "/models"):
     configure_logger("DEBUG")
     logger.info(f"Starting Ramalama daemon on {host}:{port}...")
-    with RamalamaServer(host, port, model_store_path) as httpd:
+    with RamalamaServer(host, port, model_store_path, timedelta(seconds=10)) as httpd:
         with ShutdownHandler(httpd):
             server_thread = threading.Thread(target=httpd.serve_forever, daemon=True)
             server_thread.start()

ramalama/daemon/handler/base.py

@@ -1,7 +1,6 @@
 import http.server
 import json
 from abc import ABC, abstractmethod
-from datetime import datetime, timedelta
 
 from ramalama.daemon.dto.model import RunningModelResponse, running_model_list_to_dict
 from ramalama.daemon.service.model_runner import ModelRunner
@@ -40,7 +39,6 @@ def _handle_get_running_models(self, handler: http.server.SimpleHTTPRequestHandl
             full_model_name = (
                 f"{m.model.model_type}://{m.model.model_organization}/{m.model.model_name}:{m.model.model_tag}"
             )
-            expiration = datetime.now() + timedelta(minutes=5)
             models.append(
                 RunningModelResponse(
                     id=m.id,
@@ -49,7 +47,7 @@ def _handle_get_running_models(self, handler: http.server.SimpleHTTPRequestHandl
                     tag=m.model.model_tag,
                     source=m.model.type,
                     model=full_model_name,
-                    expires_at=expiration.strftime("%Y-%m-%dT%H:%M:%SZ"),
+                    expires_at=m.expiration_date.strftime("%Y-%m-%dT%H:%M:%SZ"),
                     size_vram=0,
                     digest=m.id.replace("sha-", ""),
                     cmd=" ".join(m.run_cmd),

ramalama/daemon/handler/daemon.py

@@ -1,6 +1,6 @@
 import http.server
 import json
-from datetime import datetime
+from datetime import datetime, timedelta
 
 from ramalama.arg_types import StoreArgs
 from ramalama.common import generate_sha256
@@ -123,7 +123,7 @@ def _handle_post_serve(self, handler: http.server.SimpleHTTPRequestHandler):
 
         logger.info(f"Starting model runner for {serve_request.model_name} with command: {cmd}")
         id = ModelRunner.generate_model_id(model.model_name, model.model_tag, model.model_organization)
-        model = ManagedModel(id, model, cmd, port)
+        model = ManagedModel(id, model, cmd, port, timedelta(seconds=30))
         self.model_runner.add_model(model)
         self.model_runner.start_model(id)
 

ramalama/daemon/handler/proxy.py

@@ -74,6 +74,8 @@ def _forward_request(self, handler: http.server.SimpleHTTPRequestHandler, is_ref
             return
 
         managed_model = self.model_runner.managed_models[model_id]
+        managed_model.update_expiration_date()
+
         target_url = f"http://0.0.0.0:{managed_model.port}{path}"
         method = handler.command
         headers = handler.headers

ramalama/daemon/service/model_runner.py

@@ -1,4 +1,5 @@
 import subprocess
+from datetime import datetime, timedelta
 from typing import Optional
 
 from ramalama.common import generate_sha256
@@ -7,16 +8,28 @@
 
 class ManagedModel:
 
-    def __init__(self, id: str, model: CLASS_MODEL_TYPES, run_cmd: list[str], port: int):
+    def __init__(
+        self,
+        id: str,
+        model: CLASS_MODEL_TYPES,
+        run_cmd: list[str],
+        port: int,
+        expires_after: timedelta = timedelta(minutes=5),
+    ):
         self.id = id
         self.model = model
         self.run_cmd: list[str] = run_cmd
         self.port: str = port
+
+        self.expires_after = expires_after
+        self.expiration_date: Optional[datetime] = None
+
         self.process: Optional[subprocess.Popen] = None
 
     def start(self):
         if self.process is not None:
             raise RuntimeError(f"Model {self.id} is already running.")
+        self.update_expiration_date()
         self.process = subprocess.Popen(self.run_cmd)
 
     def stop(self):
@@ -25,6 +38,9 @@ def stop(self):
             self.process.wait()
             self.process = None
 
+    def update_expiration_date(self):
+        self.expiration_date = datetime.now() + self.expires_after
+
 
 class ModelRunner:
 
@@ -47,7 +63,7 @@ def next_available_port(self) -> int:
 
     @staticmethod
     def generate_model_id(model_name: str, model_tag: str, model_organization: str) -> str:
-        return generate_sha256(f"{model_name}-{model_tag}-{model_organization}")
+        return generate_sha256(f"{model_name}-{model_tag}-{model_organization}", with_sha_prefix=False)
 
     def add_model(self, model: ManagedModel):
         if model.id in self._models: