From a5265abe0c295022775a0cd146259b4870857f71 Mon Sep 17 00:00:00 2001
From: Dominic Sloan-Murphy <dominic.sloanmurphy@canonical.com>
Date: Fri, 20 Dec 2024 17:08:46 +0000
Subject: [PATCH] HACK: update `slurm_ops` and modify to use latest
 `slurmutils`.

Done temporarily to get GPU support working. Undo this once
`slurm_ops` is properly updated.
---
 external/lib/charms/hpc_libs/v0/slurm_ops.py | 83 ++++++++++++++------
 1 file changed, 59 insertions(+), 24 deletions(-)

diff --git a/external/lib/charms/hpc_libs/v0/slurm_ops.py b/external/lib/charms/hpc_libs/v0/slurm_ops.py
index 5652db6..ee85541 100644
--- a/external/lib/charms/hpc_libs/v0/slurm_ops.py
+++ b/external/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
@@ -97,14 +109,14 @@ def _on_install(self, _) -> None:
 
 # Increment this PATCH version before using `charmcraft publish-lib` or reset
 # to 0 if you are raising the major API version
-LIBPATCH = 10
+LIBPATCH = 11
 
 # Charm library dependencies to fetch during `charmcraft pack`.
 PYDEPS = [
     "cryptography~=44.0.0",
     "pyyaml>=6.0.2",
     "python-dotenv~=1.0.1",
-    "slurmutils~=0.9.0",
+    "slurmutils~=0.11.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):