local-llm-server/daemon.py

import argparse
import logging
import os
import sys
import time

from redis import Redis

from llm_server.cluster.cluster_config import cluster_config
from llm_server.config.global_config import GlobalConfig
from llm_server.config.load import load_config
from llm_server.custom_redis import redis
from llm_server.database.conn import Database
from llm_server.database.create import create_db
from llm_server.database.database import get_number_of_rows
from llm_server.helpers import resolve_path
from llm_server.logging import create_logger, logging_info, init_logging
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 = resolve_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.")
    parser.add_argument('-d', '--debug', action='store_true', help='Enable debug logging.')
    args = parser.parse_args()

    if args.debug:
        logging_info.level = logging.DEBUG

    init_logging()
    logger = create_logger('daemon')
    logger.debug('Debug logging enabled.')

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

    success, msg = load_config(config_path)
    if not success:
        logger.info(f'Failed to load config: {msg}')
        sys.exit(1)

    Database.initialise(**GlobalConfig.get().postgresql.dict())
    create_db()

    cluster_config.clear()
    cluster_config.load()

    logger.info('Loading backend stats...')
    generate_stats(regen=True)

    if GlobalConfig.get().load_num_prompts:
        redis.set('proompts', get_number_of_rows('messages'))

    start_background()

    # Give some time for the background threads to get themselves ready to go.
    time.sleep(2)

    redis.set('daemon_started', 1)
    logger.info('== Daemon Setup Complete ==')

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