diff --git a/README.md b/README.md
index 56b89b7..3c12db7 100644
--- a/README.md
+++ b/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
diff --git a/pyproject.toml b/pyproject.toml
index e777aab..4f81dc0 100644
--- a/pyproject.toml
+++ b/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 <nuccitheboss@ubuntu.com>"]
diff --git a/slurmutils/editors/__init__.py b/slurmutils/editors/__init__.py
index 9a4a0f1..40ba550 100644
--- a/slurmutils/editors/__init__.py
+++ b/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
diff --git a/slurmutils/editors/gresconfig.py b/slurmutils/editors/gresconfig.py
new file mode 100644
index 0000000..62b20db
--- /dev/null
+++ b/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)
diff --git a/slurmutils/models/__init__.py b/slurmutils/models/__init__.py
index 67dbae2..f43c0e9 100644
--- a/slurmutils/models/__init__.py
+++ b/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
diff --git a/slurmutils/models/gres.py b/slurmutils/models/gres.py
new file mode 100644
index 0000000..142f5f3
--- /dev/null
+++ b/slurmutils/models/gres.py
@@ -0,0 +1,330 @@
+# 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/>.
+
+"""Data models for `gres.conf` configuration file."""
+
+__all__ = ["GRESConfig", "GRESName", "GRESNode"]
+
+import copy
+from typing import Any
+
+from .model import BaseModel, clean, marshall_content, parse_line
+from .option import GRESConfigOptionSet, GRESNameOptionSet, GRESNodeOptionSet
+
+
+class GRESName(BaseModel):
+ """`gres.conf` name data model."""
+
+ def __init__(self, *, Name, **kwargs) -> None: # noqa N803
+ super().__init__(GRESNameOptionSet, Name=Name, **kwargs)
+
+ @classmethod
+ def from_str(cls, content: str) -> "GRESName":
+ """Construct `GRESName` data model from a gres.conf configuration line."""
+ return cls(**parse_line(GRESNameOptionSet, content))
+
+ def __str__(self) -> str:
+ """Return `GRESName` data model as a gres.conf configuration line."""
+ return " ".join(marshall_content(GRESNameOptionSet, self.data))
+
+ @property
+ def auto_detect(self) -> str | None:
+ """Hardware detection mechanism to enable for automatic GRES configuration.
+
+ Warnings:
+ * Setting this option will override the configured global automatic
+ hardware detection mechanism for this generic resource.
+ """
+ return self.data.get("AutoDetect", None)
+
+ @auto_detect.setter
+ def auto_detect(self, value: str) -> None:
+ self.data["AutoDetect"] = value
+
+ @auto_detect.deleter
+ def auto_detect(self) -> None:
+ try:
+ del self.data["AutoDetect"]
+ except KeyError:
+ pass
+
+ @property
+ def count(self) -> str | None:
+ """Number of resources of this name/type available on the node."""
+ return self.data.get("Count", None)
+
+ @count.setter
+ def count(self, value: str) -> None:
+ self.data["Count"] = value
+
+ @count.deleter
+ def count(self) -> None:
+ try:
+ del self.data["Count"]
+ except KeyError:
+ pass
+
+ @property
+ def cores(self) -> list[str] | None:
+ """Core index numbers for the specific cores which can use this resource."""
+ return self.data.get("Cores", None)
+
+ @cores.setter
+ def cores(self, value: list[str]) -> None:
+ self.data["Cores"] = value
+
+ @cores.deleter
+ def cores(self) -> None:
+ try:
+ del self.data["Cores"]
+ except KeyError:
+ pass
+
+ @property
+ def file(self) -> str | None:
+ """Fully qualified pathname of the device files associated with a resource."""
+ return self.data.get("File", None)
+
+ @file.setter
+ def file(self, value: str) -> None:
+ self.data["File"] = value
+
+ @file.deleter
+ def file(self) -> None:
+ try:
+ del self.data["File"]
+ except KeyError:
+ pass
+
+ @property
+ def flags(self) -> list[str] | None:
+ """Flags to change the configured behavior of the generic resource."""
+ return self.data.get("Flags", None)
+
+ @flags.setter
+ def flags(self, value: list[str]) -> None:
+ self.data["Flags"] = value
+
+ @flags.deleter
+ def flags(self) -> None:
+ try:
+ del self.data["Flags"]
+ except KeyError:
+ pass
+
+ @property
+ def links(self) -> list[str] | None:
+ """Numbers identifying the number of connections between this device and other devices."""
+ return self.data.get("Links", None)
+
+ @links.setter
+ def links(self, value: list[str]) -> None:
+ self.data["Links"] = value
+
+ @links.deleter
+ def links(self) -> None:
+ try:
+ del self.data["Links"]
+ except KeyError:
+ pass
+
+ @property
+ def multiple_files(self) -> str | None:
+ """Fully qualified pathname of the device files associated with a resource.
+
+ Warnings:
+ * Uses `files` instead if not using GPUs with multi-instance functionality.
+ * `files` and `multiple_files` cannot be used together.
+ """
+ return self.data.get("MultipleFiles", None)
+
+ @multiple_files.setter
+ def multiple_files(self, value: str) -> None:
+ self.data["MultipleFiles"] = value
+
+ @multiple_files.deleter
+ def multiple_files(self) -> None:
+ try:
+ del self.data["MultipleFiles"]
+ except KeyError:
+ pass
+
+ @property
+ def name(self) -> str:
+ """Name of generic resource."""
+ return self.data.get("Name")
+
+ @name.setter
+ def name(self, value: str) -> None:
+ self.data["Name"] = value
+
+ @property
+ def type(self) -> str | None:
+ """Arbitrary string identifying the type of the generic resource."""
+ return self.data.get("Type", None)
+
+ @type.setter
+ def type(self, value: str) -> None:
+ self.data["Type"] = value
+
+ @type.deleter
+ def type(self) -> None:
+ try:
+ del self.data["Type"]
+ except KeyError:
+ pass
+
+
+class GRESNode(GRESName):
+ """`gres.conf` node data model."""
+
+ def __init__(self, *, NodeName: str, **kwargs): # noqa N803
+ self.__node_name = NodeName
+ # Want to share `GRESName` descriptors, but not constructor.
+ BaseModel.__init__(self, GRESNodeOptionSet, **kwargs)
+
+ @classmethod
+ def from_dict(cls, data: dict[str, Any]) -> "GRESNode":
+ """Construct `GRESNode` data model from dictionary object."""
+ node_name = list(data.keys())[0]
+ return cls(NodeName=node_name, **data[node_name])
+
+ @classmethod
+ def from_str(cls, content: str) -> "GRESNode":
+ """Construct `GRESNode` data model from a gres.conf configuration line."""
+ return cls(**parse_line(GRESNodeOptionSet, content))
+
+ def dict(self) -> dict[str, Any]:
+ """Return `GRESNode` data model as a dictionary object."""
+ return copy.deepcopy({self.__node_name: self.data})
+
+ def __str__(self) -> str:
+ """Return `GRESNode` data model as a gres.conf configuration line."""
+ line = [f"NodeName={self.__node_name}"]
+ line.extend(marshall_content(GRESNodeOptionSet, self.data))
+ return " ".join(line)
+
+ @property
+ def node_name(self) -> str:
+ """Node(s) the generic resource configuration will be applied to.
+
+ Value `NodeName` specification can use a Slurm hostlist specification.
+ """
+ return self.__node_name
+
+ @node_name.setter
+ def node_name(self, value: str) -> None:
+ self.__node_name = value
+
+
+class GRESConfig(BaseModel):
+ """`gres.conf` data model."""
+
+ def __init__(
+ self,
+ *,
+ Names: list[str] | None = None, # noqa N803
+ Nodes: dict[str, Any] | None = None, # noqa N803
+ **kwargs,
+ ) -> None:
+ super().__init__(GRESConfigOptionSet, **kwargs)
+ self.data["Names"] = Names or []
+ self.data["Nodes"] = Nodes or {}
+
+ @classmethod
+ def from_str(cls, content: str) -> "GRESConfig":
+ """Construct `gres.conf` data model from a gres.conf configuration file."""
+ data = {}
+ lines = content.splitlines()
+ for line in lines:
+ config = clean(line)
+ if config is None:
+ continue
+
+ if config.startswith("Name"):
+ data["Names"] = data.get("Names", []) + [GRESName.from_str(config).dict()]
+ elif config.startswith("NodeName"):
+ nodes = data.get("Nodes", {})
+ nodes.update(GRESNode.from_str(config).dict())
+ data["Nodes"] = nodes
+ else:
+ data.update(parse_line(GRESConfigOptionSet, config))
+
+ return GRESConfig.from_dict(data)
+
+ def __str__(self) -> str:
+ """Return `gres.conf` data model in gres.conf format."""
+ data = self.dict()
+ global_auto_detect = data.pop("AutoDetect", None)
+ names = data.pop("Names", [])
+ nodes = data.pop("Nodes", {})
+
+ content = []
+ if global_auto_detect:
+ content.append(f"AutoDetect={global_auto_detect}")
+ if names:
+ content.extend([str(GRESName(**name)) for name in names])
+ if nodes:
+ content.extend([str(GRESNode(NodeName=k, **v)) for k, v in nodes.items()])
+
+ return "\n".join(content) + "\n"
+
+ @property
+ def auto_detect(self) -> str | None:
+ """Get global `AutoDetect` configuration in `gres.conf`.
+
+ Warnings:
+ * Setting this option will configure the automatic hardware detection mechanism
+ globally within `gres.conf`. Inline `AutoDetect` can be set used on
+ `GRESNode` and`GRESName` to override the global automatic hardware
+ detection mechanism for specific nodes or resource names.
+ """
+ return self.data.get("AutoDetect", None)
+
+ @auto_detect.setter
+ def auto_detect(self, value: str) -> None:
+ self.data["AutoDetect"] = value
+
+ @auto_detect.deleter
+ def auto_detect(self) -> None:
+ try:
+ del self.data["AutoDetect"]
+ except KeyError:
+ pass
+
+ @property
+ def names(self) -> list[dict[str, Any]] | None:
+ """List of configured generic resources."""
+ return self.data.get("Names", None)
+
+ @names.setter
+ def names(self, value: list[dict[str, Any]]) -> None:
+ self.data["Names"] = value
+
+ @names.deleter
+ def names(self) -> None:
+ self.data["Names"] = []
+
+ @property
+ def nodes(self) -> dict[str, dict[str, Any]]:
+ """Map of nodes with configured generic resources."""
+ return self.data["Nodes"]
+
+ @nodes.setter
+ def nodes(self, value: dict[str, GRESNode]) -> None:
+ self.data["Nodes"] = value
+
+ @nodes.deleter
+ def nodes(self) -> None:
+ self.data["Nodes"] = {}
diff --git a/slurmutils/models/option.py b/slurmutils/models/option.py
index 7ca6611..a01ac86 100644
--- a/slurmutils/models/option.py
+++ b/slurmutils/models/option.py
@@ -17,6 +17,9 @@
__all__ = [
"AcctGatherConfigOptionSet",
"CgroupConfigOptionSet",
+ "GRESConfigOptionSet",
+ "GRESNameOptionSet",
+ "GRESNodeOptionSet",
"SlurmdbdConfigOptionSet",
"SlurmConfigOptionSet",
"NodeOptionSet",
@@ -97,6 +100,34 @@ class CgroupConfigOptionSet(_OptionSet):
SignalChildrenProcesses: Callback = Callback()
+@dataclass(frozen=True)
+class GRESConfigOptionSet(_OptionSet):
+ """`gres.conf` configuration options."""
+
+ AutoDetect: Callback = Callback()
+
+
+@dataclass(frozen=True)
+class GRESNameOptionSet(GRESConfigOptionSet):
+ """`gres.conf` generic configuration options."""
+
+ Count: Callback = Callback()
+ Cores: Callback = CommaSeparatorCallback
+ File: Callback = Callback()
+ Flags: Callback = CommaSeparatorCallback
+ Links: Callback = CommaSeparatorCallback
+ MultipleFiles: Callback = Callback()
+ Name: Callback = Callback()
+ Type: Callback = Callback()
+
+
+@dataclass(frozen=True)
+class GRESNodeOptionSet(GRESNameOptionSet):
+ """`gres.conf` node configuration options."""
+
+ NodeName: Callback = Callback()
+
+
@dataclass(frozen=True)
class SlurmdbdConfigOptionSet(_OptionSet):
"""`slurmdbd.conf` configuration options."""
diff --git a/tests/unit/editors/constants.py b/tests/unit/editors/constants.py
index c534fb6..caa3403 100644
--- a/tests/unit/editors/constants.py
+++ b/tests/unit/editors/constants.py
@@ -12,6 +12,23 @@
# 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/>.
+EXAMPLE_GRES_CONFIG = """#
+# `gres.conf` file generated at 2024-12-10 14:17:35.161642 by slurmutils.
+#
+AutoDetect=nvml
+Name=gpu Type=gp100 File=/dev/nvidia0 Cores=0,1
+Name=gpu Type=gp100 File=/dev/nvidia1 Cores=0,1
+Name=gpu Type=p6000 File=/dev/nvidia2 Cores=2,3
+Name=gpu Type=p6000 File=/dev/nvidia3 Cores=2,3
+Name=mps Count=200 File=/dev/nvidia0
+Name=mps Count=200 File=/dev/nvidia1
+Name=mps Count=100 File=/dev/nvidia2
+Name=mps Count=100 File=/dev/nvidia3
+Name=bandwidth Type=lustre Count=4G Flags=CountOnly
+
+NodeName=juju-c9c6f-[1-10] Name=gpu Type=rtx File=/dev/nvidia[0-3] Count=8G
+"""
+
EXAMPLE_SLURM_CONFIG = """#
# `slurm.conf` file generated at 2024-01-30 17:18:36.171652 by slurmutils.
#
diff --git a/tests/unit/editors/test_gresconfig.py b/tests/unit/editors/test_gresconfig.py
new file mode 100644
index 0000000..f9513d4
--- /dev/null
+++ b/tests/unit/editors/test_gresconfig.py
@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+# 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/>.
+
+from constants import EXAMPLE_GRES_CONFIG
+from pyfakefs.fake_filesystem_unittest import TestCase
+
+from slurmutils.editors import gresconfig
+from slurmutils.models import GRESName, GRESNode
+
+
+class TestGRESConfigEditor(TestCase):
+ """Unit tests for gres.conf configuration file editor."""
+
+ def setUp(self) -> None:
+ self.setUpPyfakefs()
+ self.fs.create_file("/etc/slurm/gres.conf", contents=EXAMPLE_GRES_CONFIG)
+
+ def test_loads(self) -> None:
+ """Test `loads` function from the `gresconfig` editor module."""
+ config = gresconfig.loads(EXAMPLE_GRES_CONFIG)
+ self.assertEqual(config.auto_detect, "nvml")
+ self.assertListEqual(
+ config.names,
+ [
+ {"Name": "gpu", "Type": "gp100", "File": "/dev/nvidia0", "Cores": ["0", "1"]},
+ {"Name": "gpu", "Type": "gp100", "File": "/dev/nvidia1", "Cores": ["0", "1"]},
+ {"Name": "gpu", "Type": "p6000", "File": "/dev/nvidia2", "Cores": ["2", "3"]},
+ {"Name": "gpu", "Type": "p6000", "File": "/dev/nvidia3", "Cores": ["2", "3"]},
+ {"Name": "mps", "Count": "200", "File": "/dev/nvidia0"},
+ {"Name": "mps", "Count": "200", "File": "/dev/nvidia1"},
+ {"Name": "mps", "Count": "100", "File": "/dev/nvidia2"},
+ {"Name": "mps", "Count": "100", "File": "/dev/nvidia3"},
+ {"Name": "bandwidth", "Type": "lustre", "Count": "4G", "Flags": ["CountOnly"]},
+ ],
+ )
+ self.assertDictEqual(
+ config.nodes,
+ {
+ "juju-c9c6f-[1-10]": {
+ "Name": "gpu",
+ "Type": "rtx",
+ "File": "/dev/nvidia[0-3]",
+ "Count": "8G",
+ }
+ },
+ )
+
+ def test_dumps(self) -> None:
+ """Test `dumps` function from the `gresconfig` editor module."""
+ config = gresconfig.loads(EXAMPLE_GRES_CONFIG)
+ # New `gres.conf` will be different since the comments have been
+ # stripped out by the editor.
+ self.assertNotEqual(gresconfig.dumps(config), EXAMPLE_GRES_CONFIG)
+
+ def test_edit(self) -> None:
+ """Test `edit` context manager from the `gresconfig` editor module."""
+ 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",
+ )
+
+ # Set new values with each accessor.
+ with gresconfig.edit("/etc/slurm/gres.conf") as config:
+ config.auto_detect = "rsmi"
+ config.names = [name.dict()]
+ config.nodes = node.dict()
+
+ config = gresconfig.load("/etc/slurm/gres.conf")
+ self.assertEqual(config.auto_detect, "rsmi")
+ self.assertListEqual(config.names, [name.dict()])
+ self.assertDictEqual(config.nodes, node.dict())
+
+ # Delete all configured values from GRES configuration.
+ with gresconfig.edit("/etc/slurm/gres.conf") as config:
+ del config.auto_detect
+ del config.names
+ del config.nodes
+
+ config = gresconfig.load("/etc/slurm/gres.conf")
+ self.assertIsNone(config.auto_detect)
+ self.assertListEqual(config.names, [])
+ self.assertDictEqual(config.nodes, {})
diff --git a/tests/unit/models/test_gres.py b/tests/unit/models/test_gres.py
new file mode 100644
index 0000000..3cf72ab
--- /dev/null
+++ b/tests/unit/models/test_gres.py
@@ -0,0 +1,156 @@
+#!/usr/bin/env python3
+# 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/>.
+
+from unittest import TestCase
+
+from slurmutils.models import GRESConfig, GRESName, GRESNode
+
+
+class TestGRESConfig(TestCase):
+ """Unit tests for `gres.conf` data model."""
+
+ def setUp(self) -> None:
+ self.config = GRESConfig()
+ self.names = GRESName.from_dict(
+ {"Name": "gpu", "Type": "gp100", "File": "/dev/nvidia0", "Cores": ["0", "1"]}
+ )
+ self.nodes = GRESNode.from_dict(
+ {
+ "juju-c9c6f-[1-10]": {
+ "Name": "gpu",
+ "Type": "rtx",
+ "File": "/dev/nvidia[0-3]",
+ "Count": "8G",
+ }
+ }
+ )
+
+ def test_auto_detect(self) -> None:
+ """Test global `AutoDetect` descriptor."""
+ del self.config.auto_detect
+ self.config.auto_detect = "rsmi"
+ self.assertEqual(self.config.auto_detect, "rsmi")
+ del self.config.auto_detect
+ self.assertIsNone(self.config.auto_detect)
+
+ def test_names(self) -> None:
+ """Test `Names` descriptor."""
+ self.config.names = [self.names.dict()]
+ self.assertListEqual(self.config.names, [self.names.dict()])
+ del self.config.names
+ self.assertListEqual(self.config.names, [])
+
+ def test_nodes(self) -> None:
+ """Test `Nodes` descriptor."""
+ self.config.nodes = self.nodes.dict()
+ self.assertDictEqual(self.config.nodes, self.nodes.dict())
+ del self.config.nodes
+ self.assertDictEqual(self.config.nodes, {})
+
+
+class TestGRESName(TestCase):
+ """Unit tests for `GRESName` data model."""
+
+ def setUp(self) -> None:
+ self.config = GRESName(Name="gpu")
+
+ def test_auto_detect(self) -> None:
+ """Test in-line `AutoDetect` descriptor."""
+ del self.config.auto_detect
+ self.config.auto_detect = "rsmi"
+ self.assertEqual(self.config.auto_detect, "rsmi")
+ del self.config.auto_detect
+ self.assertIsNone(self.config.auto_detect)
+
+ def test_count(self) -> None:
+ """Test `Count` descriptor."""
+ del self.config.count
+ self.config.count = "10G"
+ self.assertEqual(self.config.count, "10G")
+ del self.config.count
+ self.assertIsNone(self.config.count)
+
+ def test_cores(self) -> None:
+ """Test `Cores` descriptor."""
+ del self.config.cores
+ self.config.cores = ["0", "1"]
+ self.assertListEqual(self.config.cores, ["0", "1"])
+ del self.config.cores
+ self.assertIsNone(self.config.cores)
+
+ def test_file(self) -> None:
+ """Test `File` descriptor."""
+ del self.config.file
+ self.config.file = "/dev/amd[0-4]"
+ self.assertEqual(self.config.file, "/dev/amd[0-4]")
+ del self.config.file
+ self.assertIsNone(self.config.file)
+
+ def test_flags(self) -> None:
+ """Test `Flags` descriptor."""
+ del self.config.flags
+ self.config.flags = ["CountOnly", "amd_gpu_env"]
+ self.assertListEqual(self.config.flags, ["CountOnly", "amd_gpu_env"])
+ del self.config.flags
+ self.assertIsNone(self.config.flags)
+
+ def test_links(self) -> None:
+ """Test `Links` descriptor."""
+ del self.config.links
+ self.config.links = ["-1", "16", "16", "16"]
+ self.assertListEqual(self.config.links, ["-1", "16", "16", "16"])
+ del self.config.links
+ self.assertIsNone(self.config.links)
+
+ def test_multiple_files(self) -> None:
+ """Test `MultipleFiles` descriptor."""
+ del self.config.multiple_files
+ self.config.multiple_files = "/dev/amd[0-4]"
+ self.assertEqual(self.config.multiple_files, "/dev/amd[0-4]")
+ del self.config.multiple_files
+ self.assertIsNone(self.config.multiple_files)
+
+ def test_name(self) -> None:
+ """Test `Name` descriptor."""
+ self.assertEqual(self.config.name, "gpu")
+ self.config.name = "shard"
+ self.assertEqual(self.config.name, "shard")
+ # Ensure that `Name` cannot be deleted.
+ with self.assertRaises(AttributeError):
+ del self.config.name # noqa
+
+ def test_type(self) -> None:
+ """Test `Type` descriptor."""
+ del self.config.type
+ self.config.type = "epyc"
+ self.assertEqual(self.config.type, "epyc")
+ del self.config.type
+ self.assertIsNone(self.config.type)
+
+
+class TestGRESNode(TestCase):
+ """Unit tests for `GRESNode` data model."""
+
+ def setUp(self) -> None:
+ self.config = GRESNode(NodeName="juju-c9c6f-[1-10]")
+
+ def test_node_name(self) -> None:
+ """Test `NodeName` descriptor."""
+ self.assertEqual(self.config.node_name, "juju-c9c6f-[1-10]")
+ self.config.node_name = "juju-c9c6f-[1-5]"
+ self.assertEqual(self.config.node_name, "juju-c9c6f-[1-5]")
+ # Ensure that `NodeName` cannot be deleted.
+ with self.assertRaises(AttributeError):
+ del self.config.node_name # noqa