From 3e1e393e70c0b44bd71159285196c0f3549a4243 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jos=C3=A9=20Juli=C3=A1n=20Espina?= Date: Mon, 9 Sep 2024 15:08:22 -0600 Subject: [PATCH] Address comments --- lib/charms/hpc_libs/v0/slurm_ops.py | 73 +++-------------------------- tests/unit/test_slurm_ops.py | 5 -- 2 files changed, 6 insertions(+), 72 deletions(-) diff --git a/lib/charms/hpc_libs/v0/slurm_ops.py b/lib/charms/hpc_libs/v0/slurm_ops.py index 60b085a..b1c1c3c 100644 --- a/lib/charms/hpc_libs/v0/slurm_ops.py +++ b/lib/charms/hpc_libs/v0/slurm_ops.py @@ -129,27 +129,6 @@ def _snap(*args) -> str: return _call("snap", *args) -def format_key(key: str) -> str: - """Format Slurm configuration keys from SlurmCASe into kebab case. - - Args: - key: Slurm configuration key to convert to kebab case. - - Notes: - Slurm configuration syntax does not follow proper PascalCasing - format, so we cannot put keys directly through a kebab case converter - to get the desired format. Some additional processing is needed for - certain keys before the key can properly kebabized. - - For example, without additional preprocessing, the key `CPUs` will - become `cp-us` if put through a kebabizer with being preformatted to `Cpus`. - """ - if "CPUs" in key: - key = key.replace("CPUs", "Cpus") - key = _acronym.sub(r"-", key) - return _kebabize.sub(r"-", key).lower() - - class SlurmOpsError(Exception): """Exception raised when a slurm operation failed.""" @@ -233,12 +212,14 @@ def generate(self) -> None: class _EnvManager: - """Control configuration of environment variables used in Slurm components.""" + """Control configuration of environment variables used in Slurm components. + + Every configuration value is automatically uppercased and prefixed with the service name. + """ - def __init__(self, file: Union[str, os.PathLike], prefix: str, keys: [str]) -> None: + def __init__(self, file: Union[str, os.PathLike], prefix: str) -> None: self._file: Path = Path(file) self._service = prefix - self._keys: [str] = keys def _config_to_env_var(self, key: str) -> str: """Get the environment variable corresponding to the configuration `key`.""" @@ -246,29 +227,15 @@ def _config_to_env_var(self, key: str) -> str: def get(self, key: str) -> Optional[str]: """Get specific environment variable for service.""" - if key not in self._keys: - raise SlurmOpsError(f"invalid configuration key `{key}` for service `{self._service}`") - return dotenv.get_key(self._file, self._config_to_env_var(key)) def set(self, config: Mapping[str, Any]) -> None: """Set environment variable for service.""" - # Check to avoid modifying the .env if the input keys are invalid. - for key in config.keys(): - if key not in self._keys: - raise SlurmOpsError( - f"invalid configuration key `{key}` for service `{self._service}`" - ) - for key, value in config.items(): dotenv.set_key(self._file, self._config_to_env_var(key), str(value)) def unset(self, key: str) -> None: """Unset environment variable for service.""" - # Check to avoid modifying the .env if the input keys are invalid. - if key not in self._keys: - raise SlurmOpsError(f"invalid configuration key `{key}` for service `{self._service}`") - dotenv.unset_key(self._file, self._config_to_env_var(key)) @@ -307,25 +274,6 @@ class MungeManager: def __init__(self, ops_manager: SlurmOpsManager) -> None: self.service = ops_manager.service_manager_for(ServiceType.MUNGED) self.key = ops_manager.munge_key_manager() - self._env_manager = ops_manager._env_manager_for(ServiceType.MUNGED) - - @property - def max_thread_count(self) -> Optional[int]: - """Get the max number of threads that munged can use.""" - if not (mtc := self._env_manager.get("max-thread-count")): - return None - - return int(mtc) - - @max_thread_count.setter - def max_thread_count(self, count: int) -> None: - """Set the max number of threads that munged can use.""" - self._env_manager.set({"max-thread-count": count}) - - @max_thread_count.deleter - def max_thread_count(self, count: int) -> None: - """Unset the max number of threads that munged can use.""" - self._env_manager.unset("max-thread-count") class PrometheusExporterManager: @@ -486,15 +434,6 @@ def generate(self) -> None: class SnapManager(SlurmOpsManager): """Slurm ops manager that uses Snap as its package manager.""" - _ENV_KEYS = { - ServiceType.SLURMCTLD: [], - ServiceType.SLURMDBD: [], - ServiceType.MUNGED: ["max-thread-count"], - ServiceType.SLURMD: ["config-server"], - ServiceType.PROMETHEUS_EXPORTER: [], - ServiceType.SLURMRESTD: ["max-connections", "max-thread-count"], - } - def install(self) -> None: """Install Slurm using the `slurm` snap.""" # FIXME: Pin slurm to the stable channel @@ -519,7 +458,7 @@ def service_manager_for(self, type: ServiceType) -> ServiceManager: def _env_manager_for(self, type: ServiceType) -> _EnvManager: """Return the `_EnvManager` for the specified `ServiceType`.""" return _EnvManager( - file="/var/snap/slurm/common/.env", prefix=type.value, keys=self._ENV_KEYS[type] + file="/var/snap/slurm/common/.env", prefix=type.value ) def munge_key_manager(self) -> MungeKeyManager: diff --git a/tests/unit/test_slurm_ops.py b/tests/unit/test_slurm_ops.py index f9b2ba6..e47e7db 100644 --- a/tests/unit/test_slurm_ops.py +++ b/tests/unit/test_slurm_ops.py @@ -71,11 +71,6 @@ @patch("charms.hpc_libs.v0.slurm_ops.subprocess.check_output") class TestSlurmOps(TestCase): - def test_format_key(self, _) -> None: - """Test that `kebabize` properly formats slurm keys.""" - self.assertEqual(slurm.format_key("CPUs"), "cpus") - self.assertEqual(slurm.format_key("AccountingStorageHost"), "accounting-storage-host") - def test_error_message(self, *_) -> None: """Test that `SlurmOpsError` stores the correct message.""" message = "error message!"