diff --git a/daemon.py b/daemon.py index 35c1d59..305c202 100644 --- a/daemon.py +++ b/daemon.py @@ -1,3 +1,4 @@ +import argparse import os import sys import time @@ -20,8 +21,14 @@ else: config_path = Path(script_path, 'config', 'config.yml') if __name__ == "__main__": - Redis().flushall() - print('Flushed Redis.') + 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: diff --git a/llm_server/cluster/redis_cycle.py b/llm_server/cluster/redis_cycle.py index 7cff2c4..266241d 100644 --- a/llm_server/cluster/redis_cycle.py +++ b/llm_server/cluster/redis_cycle.py @@ -9,11 +9,15 @@ def redis_cycle(list_name): :param list_name: :return: """ - to_move = redis_cycler_db.rpop(list_name) + pipeline = redis_cycler_db.pipeline() + pipeline.lpop(list_name) + to_move = pipeline.execute()[0] if not to_move: return [] - redis_cycler_db.lpush(list_name, to_move) - new_list = redis_cycler_db.lrange(list_name, 0, -1) + pipeline.rpush(list_name, to_move) + pipeline.lrange(list_name, 0, -1) + results = pipeline.execute() + new_list = results[-1] return [x.decode('utf-8') for x in new_list]