From 1a4d416f97ead9725e0ca636d9b0f61141e370dd Mon Sep 17 00:00:00 2001 From: Brian Gunnarson Date: Thu, 22 Aug 2024 14:49:02 -0700 Subject: [PATCH] move managers to their own folder and fix ssl problems --- merlin/managers/__init__.py | 0 merlin/{study => managers}/celerymanager.py | 97 +++++++++++---------- merlin/managers/redis_connection.py | 81 +++++++++++++++++ merlin/study/celeryadapter.py | 2 +- merlin/study/celerymanageradapter.py | 2 +- 5 files changed, 132 insertions(+), 50 deletions(-) create mode 100644 merlin/managers/__init__.py rename merlin/{study => managers}/celerymanager.py (81%) create mode 100644 merlin/managers/redis_connection.py diff --git a/merlin/managers/__init__.py b/merlin/managers/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/merlin/study/celerymanager.py b/merlin/managers/celerymanager.py similarity index 81% rename from merlin/study/celerymanager.py rename to merlin/managers/celerymanager.py index ddefab02c..99914545c 100644 --- a/merlin/study/celerymanager.py +++ b/merlin/managers/celerymanager.py @@ -35,6 +35,7 @@ import psutil import redis +from merlin.managers.redis_connection import RedisConnectionManager LOG = logging.getLogger(__name__) @@ -55,54 +56,54 @@ class WorkerStatus: } -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 import results_backend - from merlin.config.configfile import CONFIG - - conn_string = results_backend.get_connection_string() - base, _ = conn_string.rsplit("/", 1) - new_db_num = CONFIG.results_backend.db_num + self.db_num - conn_string = f"{base}/{new_db_num}" - LOG.debug(f"MANAGER: Connecting to redis at db_num: {new_db_num}") - return redis.from_url(conn_string, decode_responses=True) - # 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, - # ) +# 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 import results_backend +# from merlin.config.configfile import CONFIG + +# conn_string = results_backend.get_connection_string() +# base, _ = conn_string.rsplit("/", 1) +# new_db_num = CONFIG.results_backend.db_num + self.db_num +# conn_string = f"{base}/{new_db_num}" +# LOG.debug(f"MANAGER: Connecting to redis at db_num: {new_db_num}") +# return redis.from_url(conn_string, decode_responses=True) +# # 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, +# # ) class CeleryManager: diff --git a/merlin/managers/redis_connection.py b/merlin/managers/redis_connection.py new file mode 100644 index 000000000..8749fcb64 --- /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/celeryadapter.py b/merlin/study/celeryadapter.py index e392b1795..651cd7a4b 100644 --- a/merlin/study/celeryadapter.py +++ b/merlin/study/celeryadapter.py @@ -46,8 +46,8 @@ from merlin.common.dumper import dump_handler from merlin.config import Config +from merlin.managers.celerymanager import CeleryManager from merlin.study.batch import batch_check_parallel, batch_worker_launch -from merlin.study.celerymanager import CeleryManager from merlin.study.celerymanageradapter import add_monitor_workers, remove_monitor_workers from merlin.utils import apply_list_of_regex, check_machines, get_procs, get_yaml_var, is_running diff --git a/merlin/study/celerymanageradapter.py b/merlin/study/celerymanageradapter.py index d195eb966..31072d23e 100644 --- a/merlin/study/celerymanageradapter.py +++ b/merlin/study/celerymanageradapter.py @@ -32,7 +32,7 @@ import psutil -from merlin.study.celerymanager import WORKER_INFO, CeleryManager, WorkerStatus +from merlin.managers.celerymanager import WORKER_INFO, CeleryManager, WorkerStatus LOG = logging.getLogger(__name__)