local-llm-server/daemon.py

import os
import sys
import time
from pathlib import Path

from llm_server.cluster.cluster_config import cluster_config
from llm_server.cluster.redis_cycle import redis_cycler_db
from llm_server.cluster.stores import redis_running_models
from llm_server.config.load import load_config, parse_backends
from llm_server.custom_redis import redis
from llm_server.database.create import create_db
from llm_server.routes.queue import priority_queue
from llm_server.routes.v2.generate_stats import generate_stats
from llm_server.workers.threader import start_background

script_path = os.path.dirname(os.path.realpath(__file__))
config_path_environ = os.getenv("CONFIG_PATH")
if config_path_environ:
    config_path = config_path_environ
else:
    config_path = Path(script_path, 'config', 'config.yml')

if __name__ == "__main__":
    flushed_keys = redis.flush()
    print('Flushed', len(flushed_keys), 'keys from Redis.')

    redis_cycler_db.flushall()
    redis_running_models.flush()

    success, config, msg = load_config(config_path)
    if not success:
        print('Failed to load config:', msg)
        sys.exit(1)

    create_db()

    priority_queue.flush()
    cluster_config.clear()
    cluster_config.load(parse_backends(config))

    print('Loading backend stats...')
    generate_stats()

    start_background()

    redis.set('daemon_started', 1)
    print('== Daemon Setup Complete ==\n')

    try:
        while True:
            time.sleep(3600)
    except KeyboardInterrupt:
        redis.set('daemon_started', 0)