Implement old session cleanup

Author: adrgs

pyproject.toml

@@ -11,6 +11,8 @@ dependencies = [
     ### Invariant extra dependencies
     "presidio-analyzer>=2.2.354",
     "spacy>=3.7.5",
+    "fastapi-utils>=0.7.0",
+    "typing-inspect>=0.9.0",
 ]
 readme = "README.md"
 requires-python = ">= 3.10"

requirements-dev.lock

@@ -67,9 +67,12 @@ exceptiongroup==1.2.1
 face==22.0.0
     # via glom
 fastapi==0.111.0
+    # via fastapi-utils
     # via invariant-server
 fastapi-cli==0.0.4
     # via fastapi
+fastapi-utils==0.7.0
+    # via invariant-server
 filelock==3.15.4
     # via tldextract
 glom==22.1.0
@@ -118,6 +121,8 @@ murmurhash==1.0.10
     # via preshed
     # via spacy
     # via thinc
+mypy-extensions==1.0.0
+    # via typing-inspect
 numpy==1.26.4
     # via blis
     # via spacy
@@ -143,9 +148,12 @@ preshed==3.0.9
     # via thinc
 presidio-analyzer==2.2.354
     # via invariant-server
+psutil==5.9.8
+    # via fastapi-utils
 pydantic==2.8.2
     # via confection
     # via fastapi
+    # via fastapi-utils
     # via invariant
     # via pydantic-settings
     # via spacy
@@ -240,6 +248,9 @@ typing-extensions==4.12.2
     # via semgrep
     # via sqlalchemy
     # via typer
+    # via typing-inspect
+typing-inspect==0.9.0
+    # via invariant-server
 ujson==5.10.0
     # via fastapi
 urllib3==2.2.2

requirements.lock

@@ -67,9 +67,12 @@ exceptiongroup==1.2.1
 face==22.0.0
     # via glom
 fastapi==0.111.0
+    # via fastapi-utils
     # via invariant-server
 fastapi-cli==0.0.4
     # via fastapi
+fastapi-utils==0.7.0
+    # via invariant-server
 filelock==3.15.4
     # via tldextract
 glom==22.1.0
@@ -116,6 +119,8 @@ murmurhash==1.0.10
     # via preshed
     # via spacy
     # via thinc
+mypy-extensions==1.0.0
+    # via typing-inspect
 numpy==1.26.4
     # via blis
     # via spacy
@@ -138,9 +143,12 @@ preshed==3.0.9
     # via thinc
 presidio-analyzer==2.2.354
     # via invariant-server
+psutil==5.9.8
+    # via fastapi-utils
 pydantic==2.8.2
     # via confection
     # via fastapi
+    # via fastapi-utils
     # via invariant
     # via pydantic-settings
     # via spacy
@@ -233,6 +241,9 @@ typing-extensions==4.12.2
     # via semgrep
     # via sqlalchemy
     # via typer
+    # via typing-inspect
+typing-inspect==0.9.0
+    # via invariant-server
 ujson==5.10.0
     # via fastapi
 urllib3==2.2.2

server/config.py

@@ -3,6 +3,7 @@
 
 class Settings(BaseSettings):
     production: bool = False
+    idle_timeout: int = 10 * 60  # 10 minutes of inactivity before stopping the process
 
 
 settings = Settings()

server/ipc/controller.py

@@ -3,6 +3,7 @@
 import json
 import os
 from server.config import settings
+import time
 
 
 class IpcController:
@@ -25,6 +26,7 @@ def _execute_function(self, session_id, func_name, *args, **kwargs):
         rpc_info = self.rpc_map[session_id]
         message = json.dumps({"func_name": func_name, "args": args, "kwargs": kwargs})
         try:
+            rpc_info["timestamp"] = time.time()
             rpc_info["stdin"].write(message + "\n")
             rpc_info["stdin"].flush()
             result = rpc_info["stdout"].readline().strip()
@@ -67,7 +69,11 @@ def start_process(self, session_id):
                 text=True,
             )
         self.process_map[session_id] = process
-        self.rpc_map[session_id] = {"stdin": process.stdin, "stdout": process.stdout}
+        self.rpc_map[session_id] = {
+            "stdin": process.stdin,
+            "stdout": process.stdout,
+            "timestamp": time.time(),
+        }
 
     def call_function(self, session_id, func_name, *args, **kwargs):
         if session_id not in self.rpc_map:
@@ -80,9 +86,19 @@ def stop_process(self, session_id):
             self.rpc_map[session_id]["stdin"].close()
             self.rpc_map[session_id]["stdout"].close()
             process = self.process_map.pop(session_id, None)
+            del self.rpc_map[session_id]
             if process:
                 process.terminate()
 
+    def cleanup(self):
+        for session_id in list(self.rpc_map.keys()):
+            if (
+                session_id in self.rpc_map
+                and time.time() - self.rpc_map[session_id]["timestamp"]
+                > settings.idle_timeout
+            ):
+                self.stop_process(session_id)
+
 
 class IpcControllerSingleton:
     def __init__(self):

server/main.py

@@ -1,7 +1,9 @@
-from fastapi import FastAPI
+from fastapi import Depends, FastAPI
+from server.ipc.controller import IpcController, get_ipc_controller
 from server.routers import session, policy, monitor
 from server.database import engine
 from server.models import Base
+from fastapi_utils.tasks import repeat_every
 
 Base.metadata.create_all(bind=engine)
 
@@ -20,3 +22,10 @@ def index():
 app.include_router(session.router, prefix="/session", tags=["session"])
 app.include_router(policy.router, prefix="/policy", tags=["policy"])
 app.include_router(monitor.router, prefix="/monitor", tags=["monitor"])
+
+
+@app.on_event("startup")
+@repeat_every(seconds=60)  # 1 minute
+def cleanup_old_ipc() -> None:
+    ipc_controller = get_ipc_controller()
+    ipc_controller.cleanup()