model_manager.py

import threading
import time
import weakref

class ModelTimeoutManager:
    def __init__(self, timeout=60):
        self.timeout = timeout
        self.models = {}
        self.lock = threading.Lock()
        self.check_thread = threading.Thread(target=self._check_timeout)
        self.check_thread.daemon = True
        self.check_thread.start()

    def set_model(self, model_name, model):
        with self.lock:
            self.models[model_name] = {
                "model": weakref.ref(model),
                "last_access_time": time.time()
            }

    def get_model(self, model_name):
        with self.lock:
            if model_name in self.models:
                self.models[model_name]["last_access_time"] = time.time()
                return self.models[model_name]["model"]()
            return None

    def _check_timeout(self):
        while True:
            time.sleep(10)
            with self.lock:
                current_time = time.time()
                for model_name in list(self.models.keys()):
                    if current_time - self.models[model_name]["last_access_time"] > self.timeout:
                        model = self.models[model_name]["model"]()
                        if model is not None:
                            del model
                        del self.models[model_name]
                        print(f"Model '{model_name}' released due to timeout")