Merge pull request #2 from bgunnar5/feature/redis-connection-manager

add RedisConnectionManager

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.
+# <[email protected]>
+#
+# 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,
+        )

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: