local-llm-server/daemon.py

import argparse
import os
import sys
import time
from pathlib import Path

from redis import Redis

from llm_server.cluster.cluster_config import cluster_config
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.v1.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__":
    parser = argparse.ArgumentParser(
        description='Daemon microservice.')
    parser.add_argument('--no-reset', action='store_true', help="Don't clear the Redis server databases.")
    args = parser.parse_args()

    if not args.no_reset:
        Redis().flushall()
        print('Flushed Redis.')

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

    create_db()

    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)