diff --git a/merlin/managers/redis_connection.py b/merlin/managers/redis_connection.py new file mode 100644 index 00000000..8749fcb6 --- /dev/null +++ b/merlin/managers/redis_connection.py @@ -0,0 +1,81 @@ +############################################################################### +# Copyright (c) 2023, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory +# Written by the Merlin dev team, listed in the CONTRIBUTORS file. +# +# +# LLNL-CODE-797170 +# All rights reserved. +# This file is part of Merlin, Version: 1.12.2b1. +# +# For details, see https://github.com/LLNL/merlin. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +############################################################################### +""" +This module stores a manager for redis connections. +""" +import logging +import redis + +LOG = logging.getLogger(__name__) + + +class RedisConnectionManager: + """ + A context manager for handling redis connections. + This will ensure safe opening and closing of Redis connections. + """ + + def __init__(self, db_num: int): + self.db_num = db_num + self.connection = None + + def __enter__(self): + self.connection = self.get_redis_connection() + return self.connection + + def __exit__(self, exc_type, exc_val, exc_tb): + if self.connection: + LOG.debug(f"MANAGER: Closing connection at db_num: {self.db_num}") + self.connection.close() + + def get_redis_connection(self) -> redis.Redis: + """ + Generic redis connection function to get the results backend redis server with a given db number increment. + + :return: Redis connection object that can be used to access values for the manager. + """ + from merlin.config.results_backend import get_backend_password + from merlin.config.configfile import CONFIG + + password_file = CONFIG.results_backend.password + try: + password = get_backend_password(password_file) + except IOError: + password = CONFIG.results_backend.password + return redis.Redis( + host=CONFIG.results_backend.server, + port=CONFIG.results_backend.port, + db=CONFIG.results_backend.db_num + self.db_num, # Increment db_num to avoid conflicts + username=CONFIG.results_backend.username, + password=password, + decode_responses=True, + ssl=True, + ssl_cert_reqs=CONFIG.results_backend.cert_reqs, + ) diff --git a/merlin/study/vineadapter.py b/merlin/study/vineadapter.py index 9eda8100..2d67d3ac 100644 --- a/merlin/study/vineadapter.py +++ b/merlin/study/vineadapter.py @@ -46,6 +46,7 @@ from merlin.common.dumper import dump_handler from merlin.config import Config +from merlin.managers.redis_connection import RedisConnectionManager from merlin.spec.specification import MerlinSpec from merlin.study.batch import batch_check_parallel, batch_worker_launch from merlin.utils import apply_list_of_regex, check_machines, get_procs, get_yaml_var, is_running @@ -54,6 +55,17 @@ LOG = logging.getLogger(__name__) +def get_vine_redis_connection() -> RedisConnectionManager: + """ + Helper function to obtain the connection to the redis database + that's storing the process ID's for the vine factory and vine stem + processes. + + :returns: A RedisConnectionManager context manager that's connected to redis + """ + return RedisConnectionManager(1) + + def run_taskvine(study, run_mode=None): """ Run the given MerlinStudy object. If the run mode is set to "local" @@ -92,10 +104,10 @@ def start_taskvine_workers( command = ["vine_factory"] command.append("-T") - command.append(f"{spec.batch["type"]}") + command.append(f"{spec.batch['type']}") command.append("-M") - command.append(f"{spec.merlin["resources"]["workers"][worker]["manager"]}") + command.append(f"{spec.merlin['resources']['workers'][worker]['manager']}") cores = spec.merlin["resources"]["workers"][worker]["cores"] if cores: