75 lines
2.3 KiB
Python
75 lines
2.3 KiB
Python
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)
|