Merge pull request #37 from NucciTheBoss/nuccitheboss/feat/gres-editor

README.md

@@ -16,6 +16,7 @@ slurmutils package include:
 #### `from slurmutils.editors import ...`
 
 * `acctgatherconfig`: An editor for _acct_gather.conf_ configuration files.
+* `gresconfig`: An editor for _gres.conf_ configuration files.
 * `cgroupconfig`: An editor for _cgroup.conf_ configuration files.
 * `slurmconfig`: An editor for _slurm.conf_ configuration files.
 * `slurmdbdconfig`: An editor for _slurmdbd.conf_ configuration files.
@@ -84,6 +85,33 @@ with cgroupconfig.edit("/etc/slurm/cgroup.conf") as config:
     config.constrain_swap_space = "yes"
 ```
 
+##### `gresconfig`
+
+###### Edit a pre-existing _gres.conf_ configuration file
+
+```python
+from slurmutils.editors import gresconfig
+from slurmutils.models import GRESName, GRESNode
+
+with gresconfig.edit("/etc/slurm/gres.conf") as config:
+    name = GRESName(
+            Name="gpu",
+            Type="epyc",
+            File="/dev/amd4",
+            Cores=["0", "1"],
+        )
+    node = GRESNode(
+        NodeName="juju-abc654-[1-20]",
+        Name="gpu",
+        Type="epyc",
+        File="/dev/amd[0-3]",
+        Count="12G",
+    )
+    config.auto_detect = "rsmi"
+    config.names.append(name.dict())
+    config.nodes.updaten(node.dict())
+```
+
 ##### `slurmconfig`
 
 ###### Edit a pre-existing _slurm.conf_ configuration file

pyproject.toml

@@ -18,7 +18,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "slurmutils"
-version = "0.9.0"
+version = "0.10.0"
 description = "Utilities and APIs for interfacing with the Slurm workload manager."
 repository = "https://github.com/charmed-hpc/slurmutils"
 authors = ["Jason C. Nucciarone <[email protected]>"]

slurmutils/editors/__init__.py

@@ -16,5 +16,6 @@
 
 from . import acctgatherconfig as acctgatherconfig
 from . import cgroupconfig as cgroupconfig
+from . import gresconfig as gresconfig
 from . import slurmconfig as slurmconfig
 from . import slurmdbdconfig as slurmdbdconfig

slurmutils/editors/gresconfig.py

@@ -0,0 +1,82 @@
+# Copyright 2024 Canonical Ltd.
+#
+# This program is free software; you can redistribute it and/or
+# modify it under the terms of the GNU Lesser General Public
+# License version 3 as published by the Free Software Foundation.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+"""Edit gres.conf files."""
+
+__all__ = ["dump", "dumps", "load", "loads", "edit"]
+
+import logging
+import os
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Optional, Union
+
+from ..models import GRESConfig
+from .editor import dumper, loader, set_file_permissions
+
+_logger = logging.getLogger("slurmutils")
+
+
+@loader
+def load(file: Union[str, os.PathLike]) -> GRESConfig:
+    """Load `gres.conf` data model from gres.conf file."""
+    return loads(Path(file).read_text())
+
+
+def loads(content: str) -> GRESConfig:
+    """Load `gres.conf` data model from string."""
+    return GRESConfig.from_str(content)
+
+
+@dumper
+def dump(
+    config: GRESConfig,
+    file: Union[str, os.PathLike],
+    mode: int = 0o644,
+    user: Optional[Union[str, int]] = None,
+    group: Optional[Union[str, int]] = None,
+) -> None:
+    """Dump `gres.conf` data model into gres.conf file."""
+    Path(file).write_text(dumps(config))
+    set_file_permissions(file, mode, user, group)
+
+
+def dumps(config: GRESConfig) -> str:
+    """Dump `gres.conf` data model into a string."""
+    return str(config)
+
+
+@contextmanager
+def edit(
+    file: Union[str, os.PathLike],
+    mode: int = 0o644,
+    user: Optional[Union[str, int]] = None,
+    group: Optional[Union[str, int]] = None,
+) -> GRESConfig:
+    """Edit a gres.conf file.
+
+    Args:
+        file: gres.conf file to edit. An empty config will be created if it does not exist.
+        mode: Access mode to apply to the gres.conf file. (Default: rw-r--r--)
+        user: User to set as owner of the gres.conf file. (Default: $USER)
+        group: Group to set as owner of the gres.conf file. (Default: None)
+    """
+    if not os.path.exists(file):
+        _logger.warning("file %s not found. creating new empty gres.conf configuration", file)
+        config = GRESConfig()
+    else:
+        config = load(file)
+
+    yield config
+    dump(config, file, mode, user, group)

slurmutils/models/__init__.py

@@ -16,6 +16,9 @@
 
 from .acctgather import AcctGatherConfig as AcctGatherConfig
 from .cgroup import CgroupConfig as CgroupConfig
+from .gres import GRESConfig as GRESConfig
+from .gres import GRESName as GRESName
+from .gres import GRESNode as GRESNode
 from .slurm import DownNodes as DownNodes
 from .slurm import FrontendNode as FrontendNode
 from .slurm import Node as Node