enhancement: support acct_gather.conf

These changes add a data model and editor for acct_gather.conf.

slurmutils/editors/acct_gatherconfig.py

@@ -0,0 +1,84 @@
+# 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 acct_gather.conf file."""
+
+__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 AcctGatherConfig
+from .editor import dumper, loader, set_file_permissions
+
+_logger = logging.getLogger("slurmutils")
+
+
+@loader
+def load(file: Union[str, os.PathLike]) -> AcctGatherConfig:
+    """Load `acct_gather.conf` data model from acct_gather.conf file."""
+    return loads(Path(file).read_text())
+
+
+def loads(content: str) -> AcctGatherConfig:
+    """Load `acct_gather.conf` data model from string."""
+    return AcctGatherConfig.from_str(content)
+
+
+@dumper
+def dump(
+    config: AcctGatherConfig,
+    file: Union[str, os.PathLike],
+    mode: int = 0o644,
+    user: Optional[Union[str, int]] = None,
+    group: Optional[Union[str, int]] = None,
+) -> None:
+    """Dump `acct_gather.conf` data model into acct_gather.conf file."""
+    Path(file).write_text(dumps(config))
+    set_file_permissions(file, mode, user, group)
+
+
+def dumps(config: AcctGatherConfig) -> str:
+    """Dump `acct_gather.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,
+) -> AcctGatherConfig:
+    """Edit a acct_gather.conf file.
+
+    Args:
+        file: acct_gather.conf file to edit. An empty config will be created if it does not exist.
+        mode: Access mode to apply to the acct_gather.conf file. (Default: rw-r--r--)
+        user: User to set as owner of the acct_gather.conf file. (Default: $USER)
+        group: Group to set as owner of the acct_gather.conf file. (Default: None)
+    """
+    if not os.path.exists(file):
+        _logger.warning(
+            "file %s not found. creating new empty acct_gather.conf configuration", file
+        )
+        config = AcctGatherConfig()
+    else:
+        config = load(file)
+
+    yield config
+    dump(config, file, mode, user, group)

slurmutils/models/__init__.py

@@ -14,6 +14,7 @@
 
 """Data models for common Slurm objects."""
 
+from .acct_gather import AcctGatherConfig as AcctGatherConfig
 from .cgroup import CgroupConfig as CgroupConfig
 from .slurm import DownNodes as DownNodes
 from .slurm import FrontendNode as FrontendNode

slurmutils/models/acct_gather.py

@@ -0,0 +1,49 @@
+# 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 `acct_gather.conf` configuration file."""
+
+from .model import BaseModel, clean, format_key, generate_descriptors, marshall_content, parse_line
+from .option import AcctGatherConfigOptionSet
+
+
+class AcctGatherConfig(BaseModel):
+    """`acct_gather.conf` data model."""
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(AcctGatherConfigOptionSet, **kwargs)
+
+    @classmethod
+    def from_str(cls, content: str) -> "AcctGatherConfig":
+        """Construct AcctGatherConfig data model from acct_gather.conf format."""
+        data = {}
+        lines = content.splitlines()
+        for index, line in enumerate(lines):
+            config = clean(line)
+            if config is None:
+                continue
+
+            data.update(parse_line(AcctGatherConfigOptionSet, config))
+
+        return AcctGatherConfig.from_dict(data)
+
+    def __str__(self) -> str:
+        """Return AcctGatherConfig data model in acct_gather.conf format."""
+        result = []
+        result.extend(marshall_content(AcctGatherConfigOptionSet, self.dict()))
+        return "\n".join(result)
+
+
+for opt in AcctGatherConfigOptionSet.keys():
+    setattr(AcctGatherConfig, format_key(opt), property(*generate_descriptors(opt)))

slurmutils/models/cgroup.py

@@ -26,7 +26,7 @@ def __init__(self, **kwargs) -> None:
 
     @classmethod
     def from_str(cls, content: str) -> "CgroupConfig":
-        """Construct SlurmdbdConfig data model from slurmdbd.conf format."""
+        """Construct CgroupConfig data model from cgroup.conf format."""
         data = {}
         lines = content.splitlines()
         for index, line in enumerate(lines):

slurmutils/models/option.py

@@ -15,6 +15,7 @@
 """Configuration options for Slurm data models."""
 
 __all__ = [
+    "AcctGatherConfigOptionSet",
     "CgroupConfigOptionSet",
     "SlurmdbdConfigOptionSet",
     "SlurmConfigOptionSet",
@@ -48,6 +49,34 @@ def keys(cls) -> Iterable[str]:
             yield field.name
 
 
+@dataclass(frozen=True)
+class AcctGatherConfigOptionSet(_OptionSet):
+    """`acct_gather.conf` configuration options."""
+
+    AcctGatherEnergyType: Callback = Callback()
+    EnergyIPMIFrequency: Callback = Callback()
+    EnergyIPMICalcAdjustment: Callback = Callback()
+    EnergyIPMIPowerSensors: Callback = Callback()
+    EnergyIPMIUsername: Callback = Callback()
+    EnergyIPMIPassword: Callback = Callback()
+    EnergyIPMIFrequency: Callback = Callback()
+    EnergyIPMITimeout: Callback = Callback()
+    AcctGatherFilesystemType: Callback = Callback()
+    AcctGatherProfileType: Callback = Callback()
+    ProfileHDF5Dir: Callback = Callback()
+    ProfileHDF5Default: Callback = Callback()
+    ProfileInfluxDBDatabase: Callback = Callback()
+    ProfileInfluxDBDefault: Callback = Callback()
+    ProfileInfluxDBHost: Callback = Callback()
+    ProfileInfluxDBPass: Callback = Callback()
+    PProfileInfluxDBRTPolicy: Callback = Callback()
+    ProfileInfluxDBUser: Callback = Callback()
+    ProfileInfluxDBTimeout: Callback = Callback()
+    AcctGatherInterconnectType: Callback = Callback()
+    InfinibandOFEDPort: Callback = Callback()
+    SysfsInterfaces: Callback = Callback()
+
+
 @dataclass(frozen=True)
 class CgroupConfigOptionSet(_OptionSet):
     """`cgroup.conf` configuration options."""

tests/unit/editors/constants.py

@@ -106,3 +106,13 @@
 ConstrainRAMSpace=yes
 ConstrainSwapSpace=yes
 """
+
+EXAMPLE_ACCT_GATHER_CONFIG = """#
+# `acct_gather.conf` file generated at 2024-09-18 15:10:44.652017 by slurmutils.
+#
+AcctGatherEnergyType=acct_gather_energy/gpu
+AcctGatherFilesystemType=acct_gather_filesystem/lustre
+AcctGatherProfileType=acct_gather_profile/hdf5
+ProfileHDF5Dir=/mydir
+ProfileHDF5Default=ALL
+"""

tests/unit/editors/test_acct_gatherconfig.py

@@ -0,0 +1,60 @@
+#!/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/>.
+
+"""Unit tests for the acct_gather.conf editor."""
+
+
+from constants import EXAMPLE_ACCT_GATHER_CONFIG
+from pyfakefs.fake_filesystem_unittest import TestCase
+
+from slurmutils.editors import acct_gatherconfig
+
+
+class TestAcctGatherConfigEditor(TestCase):
+    """Unit tests for acct_gather.conf file editor."""
+
+    def setUp(self) -> None:
+        self.setUpPyfakefs()
+        self.fs.create_file("/etc/slurm/acct_gather.conf", contents=EXAMPLE_ACCT_GATHER_CONFIG)
+
+    def test_loads(self) -> None:
+        """Test `loads` method of the acct_gatherconfig module."""
+        config = acct_gatherconfig.loads(EXAMPLE_ACCT_GATHER_CONFIG)
+        self.assertEqual(config.acct_gather_energy_type, "acct_gather_energy/gpu")
+        self.assertEqual(config.acct_gather_filesystem_type, "acct_gather_filesystem/lustre")
+        self.assertEqual(config.acct_gather_profile_type, "acct_gather_profile/hdf5")
+        self.assertEqual(config.profile_hdf5_dir, "/mydir")
+        self.assertEqual(config.profile_hdf5_default, "ALL")
+
+        config = acct_gatherconfig.loads(EXAMPLE_ACCT_GATHER_CONFIG)
+        # The new config and old config should not be equal since the
+        # timestamps in the header will be different.
+        self.assertNotEqual(acct_gatherconfig.dumps(config), EXAMPLE_ACCT_GATHER_CONFIG)
+
+    def test_edit(self) -> None:
+        """Test `edit` context manager from the acct_gatherconfig module."""
+        with acct_gatherconfig.edit("/etc/slurm/acct_gather.conf") as config:
+            config.acct_gather_energy_type = "acct_gather_energy/ipmi"
+            config.acct_gather_filesystem_type = "acct_gather_filesystem/other"
+            config.acct_gather_profile_type = "acct_gather_profile/influxdb"
+            config.profile_hdf5_dir = "/mydir1234"
+            config.profile_hdf5_default = "NONE"
+
+        config = acct_gatherconfig.load("/etc/slurm/acct_gather.conf")
+        self.assertEqual(config.acct_gather_energy_type, "acct_gather_energy/ipmi")
+        self.assertEqual(config.acct_gather_filesystem_type, "acct_gather_filesystem/other")
+        self.assertEqual(config.acct_gather_profile_type, "acct_gather_profile/influxdb")
+        self.assertEqual(config.profile_hdf5_dir, "/mydir1234")
+        self.assertEqual(config.profile_hdf5_default, "NONE")