From b73936b3661ecac138b9cc9c6ceb04d5fab2adbf Mon Sep 17 00:00:00 2001
From: "Jason C. Nucciarone" <nuccitheboss@ubuntu.com>
Date: Wed, 11 Dec 2024 15:00:10 -0500
Subject: [PATCH] feat(slurm_ops): add gres.conf editor to `SlurmctldManager`

Other changes:

* Reordered configuration manager classes to be in alphabetical order.
* Bump the API version of slurmutils being pulled to 0.10.

Signed-off-by: Jason C. Nucciarone <nuccitheboss@ubuntu.com>
---
 lib/charms/hpc_libs/v0/slurm_ops.py | 81 +++++++++++++++++++++--------
 1 file changed, 58 insertions(+), 23 deletions(-)

diff --git a/lib/charms/hpc_libs/v0/slurm_ops.py b/lib/charms/hpc_libs/v0/slurm_ops.py
index 5652db6..c156ed2 100644
--- a/lib/charms/hpc_libs/v0/slurm_ops.py
+++ b/lib/charms/hpc_libs/v0/slurm_ops.py
@@ -77,8 +77,20 @@ def _on_install(self, _) -> None:
 import yaml
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import rsa
-from slurmutils.editors import acctgatherconfig, cgroupconfig, slurmconfig, slurmdbdconfig
-from slurmutils.models import AcctGatherConfig, CgroupConfig, SlurmConfig, SlurmdbdConfig
+from slurmutils.editors import (
+    acctgatherconfig,
+    cgroupconfig,
+    gresconfig,
+    slurmconfig,
+    slurmdbdconfig,
+)
+from slurmutils.models import (
+    AcctGatherConfig,
+    CgroupConfig,
+    GRESConfig,
+    SlurmConfig,
+    SlurmdbdConfig,
+)
 
 try:
     import charms.operator_libs_linux.v0.apt as apt
@@ -104,7 +116,7 @@ def _on_install(self, _) -> None:
     "cryptography~=44.0.0",
     "pyyaml>=6.0.2",
     "python-dotenv~=1.0.1",
-    "slurmutils~=0.9.0",
+    "slurmutils~=0.10.0",
     "distro~=1.9.0",
 ]
 
@@ -245,26 +257,6 @@ def edit(self):
         """Edit the current configuration file."""
 
 
-class _SlurmConfigManager(_ConfigManager):
-    """Control the `slurm.conf` configuration file."""
-
-    def load(self) -> SlurmConfig:
-        """Load the current `slurm.conf` configuration file."""
-        return slurmconfig.load(self._config_path)
-
-    def dump(self, config: SlurmConfig) -> None:
-        """Dump new configuration into `slurm.conf` configuration file."""
-        slurmconfig.dump(config, self._config_path, mode=0o644, user=self._user, group=self._group)
-
-    @contextmanager
-    def edit(self) -> SlurmConfig:
-        """Edit the current `slurm.conf` configuration file."""
-        with slurmconfig.edit(
-            self._config_path, mode=0o644, user=self._user, group=self._group
-        ) as config:
-            yield config
-
-
 class _AcctGatherConfigManager(_ConfigManager):
     """Manage the `acct_gather.conf` configuration file."""
 
@@ -309,6 +301,46 @@ def edit(self) -> CgroupConfig:
             yield config
 
 
+class _GRESConfigManager(_ConfigManager):
+    """Manage the `gres.conf` configuration file."""
+
+    def load(self) -> GRESConfig:
+        """Load the current `gres.conf` configuration files."""
+        return gresconfig.load(self._config_path)
+
+    def dump(self, config: GRESConfig) -> None:
+        """Dump new configuration into `gres.conf` configuration file."""
+        gresconfig.dump(config, self._config_path, mode=0o644, user=self._user, group=self._group)
+
+    @contextmanager
+    def edit(self) -> GRESConfig:
+        """Edit the current `gres.conf` configuration file."""
+        with gresconfig.edit(
+            self._config_path, mode=0o644, user=self._user, group=self._group
+        ) as config:
+            yield config
+
+
+class _SlurmConfigManager(_ConfigManager):
+    """Control the `slurm.conf` configuration file."""
+
+    def load(self) -> SlurmConfig:
+        """Load the current `slurm.conf` configuration file."""
+        return slurmconfig.load(self._config_path)
+
+    def dump(self, config: SlurmConfig) -> None:
+        """Dump new configuration into `slurm.conf` configuration file."""
+        slurmconfig.dump(config, self._config_path, mode=0o644, user=self._user, group=self._group)
+
+    @contextmanager
+    def edit(self) -> SlurmConfig:
+        """Edit the current `slurm.conf` configuration file."""
+        with slurmconfig.edit(
+            self._config_path, mode=0o644, user=self._user, group=self._group
+        ) as config:
+            yield config
+
+
 class _SlurmdbdConfigManager(_ConfigManager):
     """Control the `slurmdbd.conf` configuration file."""
 
@@ -951,6 +983,9 @@ def __init__(self, *args, **kwargs) -> None:
         self.cgroup = _CgroupConfigManager(
             self._ops_manager.etc_path / "cgroup.conf", self.user, self.group
         )
+        self.gres = _GRESConfigManager(
+            self._ops_manager.etc_path / "gres.conf", self.user, self.group
+        )
 
 
 class SlurmdManager(_SlurmManagerBase):