Merge pull request #165 from danielferr85/main

Improvements in Expiration Date for models in long inference sessions

Author: NotPunchnox

README.md

@@ -1,6 +1,6 @@
 # RKLLama: LLM Server and Client for Rockchip 3588/3576
 
-### [Version: 0.0.74](#New-Version)
+### [Version: 0.0.75](#New-Version)
 
 Video demo ( version 0.0.1 ):
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rkllama"
-version = "0.0.74"
+version = "0.0.75"
 authors = [
     { name="NotPunchnox", email="punchnoxpro@gmail.com" },
     { name="TomJacobsUK", email="tom@tomjacobs.co.uk" },

src/rkllama/api/server_utils.py

@@ -342,6 +342,10 @@ def generate():
             while not thread_finished or not final_sent:
                 if parent_pipe.poll(timeout):  # Timeout in seconds
                     token = parent_pipe.recv()
+
+                    # Updating expiration date for the model during token generation to prevent expiration
+                    variables.worker_manager_rkllm.update_expiration_date_for_model(model_name)
+
                 else:
                     # Abort the current inference
                     variables.worker_manager_rkllm.workers[model_name].abort_flag.value = True
@@ -704,6 +708,10 @@ def generate():
             while not thread_finished or not final_sent:
                 if parent_pipe.poll(timeout):  # Timeout in seconds
                     token = parent_pipe.recv()
+
+                    # Updating expiration date for the model during token generation to prevent expiration
+                    variables.worker_manager_rkllm.update_expiration_date_for_model(model_name)
+
                 else:
                     # Abort the current inference
                     variables.worker_manager_rkllm.workers[model_name].abort_flag.value = True

src/rkllama/api/worker.py

@@ -1062,18 +1062,37 @@ def send_task(self, model_name, task):
             # Send the TASK to the model with the task queue and the pipe to send the response
             self.workers[model_name].task_queue.put((child_conn,) + task)
 
-            # CLose the not needed connection
-            # child_conn.close()
-            
             # Update the worker model info with the invocation
-            self.workers[model_name].worker_model_info.last_call = datetime.now()
-            self.workers[model_name].worker_model_info.expires_at = datetime.now() + timedelta(minutes=int(rkllama.config.get("model", "max_minutes_loaded_in_memory")),)
+            self.update_last_call_for_model(model_name)
 
             # Return the request parent pipe
             return parent_conn
 
         return None
 
+    
+    def update_last_call_for_model(self, model_name):
+        """
+        Update the metadata of a worker with the last time called
+        Args:
+        model_name (str): Model to update
+        """
+        # Update the last call of the model
+        self.workers[model_name].worker_model_info.last_call = datetime.now()
+
+        # Update the expiration of the model
+        self.update_expiration_date_for_model(model_name)
+
+
+    def update_expiration_date_for_model(self, model_name):
+        """
+        Update the metadata of a worker with the expiration date for the model
+        Args:
+        model_name (str): Model to update
+        """
+        # Update the expiration of the model
+        self.workers[model_name].worker_model_info.expires_at = datetime.now() + timedelta(minutes=int(rkllama.config.get("model", "max_minutes_loaded_in_memory")),)
+
 
     def stop_worker(self, model_name, timeout=30):
         """

src/rkllama/server/server.py

@@ -1299,7 +1299,7 @@ def embeddings_ollama():
 def ollama_version():
     """Return a dummy version to be compatible with Ollama clients"""
     return jsonify({
-        "version": "0.0.74"
+        "version": "0.0.75"
     }), 200
 
 
@@ -1587,6 +1587,9 @@ def forward_request_to_llama_cpp_worker(is_openai_request,request):
     logger.debug(f"Routing request to llama.cpp whith this URL: {proxy_route_url} with this data:\n{data}")
 
     try:
+        # Updating last call info for the model to prevent expiration
+        variables.worker_manager_rkllm.update_last_call_for_model(model_name)
+
         # Set the header for the llama-server call
         headers = {
             "Authorization": f"Bearer NOT_IN_USE",
@@ -1623,6 +1626,10 @@ def generate():
 
                         # Loop over the chunks returned by llama.cpp
                         for line in response.iter_lines():
+                            
+                            # Updating expiration date for the model during token generation to prevent expiration
+                            variables.worker_manager_rkllm.update_expiration_date_for_model(model_name)
+
                             # Decode the bytes line 
                             line = line.decode("utf-8")