feat(cgroup): add editor for cgroup.conf file

slurmutils/editors/__init__.py

@@ -14,5 +14,6 @@
 
 """Editors for Slurm workload manager configuration files."""
 
+from . import cgroupconfig as cgroupconfig
 from . import slurmconfig as slurmconfig
 from . import slurmdbdconfig as slurmdbdconfig

slurmutils/editors/cgroupconfig.py

@@ -0,0 +1,105 @@
+# 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 cgroup.conf files."""
+
+__all__ = ["dump", "dumps", "load", "loads", "edit"]
+
+import logging
+import os
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Union
+
+from ..models import CgroupConfig
+from ..models.option import CgroupConfigOptionSet
+from .editor import (
+    clean,
+    dumper,
+    loader,
+    marshall_content,
+    parse_line,
+)
+
+_logger = logging.getLogger("slurmutils")
+
+
+@loader
+def load(file: Union[str, os.PathLike]) -> CgroupConfig:
+    """Load `cgroup.conf` data model from cgroup.conf file."""
+    return loads(Path(file).read_text())
+
+
+def loads(content: str) -> CgroupConfig:
+    """Load `cgroup.conf` data model from string."""
+    return _parse(content)
+
+
+@dumper
+def dump(config: CgroupConfig, file: Union[str, os.PathLike]) -> None:
+    """Dump `cgroup.conf` data model into cgroup.conf file."""
+    Path(file).write_text(dumps(config))
+
+
+def dumps(config: CgroupConfig) -> str:
+    """Dump `cgroup.conf` data model into a string."""
+    return _marshall(config)
+
+
+@contextmanager
+def edit(file: Union[str, os.PathLike]) -> CgroupConfig:
+    """Edit a cgroup.conf file.
+
+    Args:
+        file: Path to cgroup.conf file to edit. If cgroup.conf does
+            not exist at the specified file path, it will be created.
+    """
+    if not os.path.exists(file):
+        _logger.warning("file %s not found. creating new empty cgroup.conf configuration", file)
+        config = CgroupConfig()
+    else:
+        config = load(file)
+
+    yield config
+    dump(config, file)
+
+
+def _parse(content: str) -> CgroupConfig:
+    """Parse contents of `cgroup.conf`.
+
+    Args:
+        content: Contents of `cgroup.conf`.
+    """
+    data = {}
+    lines = content.splitlines()
+    for index, line in enumerate(lines):
+        config = clean(line)
+        if config is None:
+            _logger.debug("ignoring line %s at index %s in cgroup.conf", line, index)
+            continue
+
+        data.update(parse_line(CgroupConfigOptionSet, config))
+
+    return CgroupConfig.from_dict(data)
+
+
+def _marshall(config: CgroupConfig) -> str:
+    """Marshall `cgroup.conf` data model back into cgroup.conf format.
+
+    Args:
+        config: `cgroup.conf` data model to marshall.
+    """
+    result = []
+    result.extend(marshall_content(CgroupConfigOptionSet, config.dict()))
+    return "\n".join(result)

slurmutils/editors/editor.py

@@ -49,7 +49,7 @@ def parse_line(options, line: str) -> Dict[str, Any]:
             raise EditorError(
                 (
                     f"unable to parse configuration option {k}. "
-                    + f"valid configuration options are {[option.name for option in options]}"
+                    + f"valid configuration options are {list(options.keys())}"
                 )
             )
 

slurmutils/models/__init__.py

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

slurmutils/models/cgroup.py

@@ -0,0 +1,29 @@
+# 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 `cgroup.conf` configuration file."""
+
+from .model import BaseModel, format_key, generate_descriptors
+from .option import CgroupConfigOptionSet
+
+
+class CgroupConfig(BaseModel):
+    """`cgroup.conf` data model."""
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(CgroupConfigOptionSet, **kwargs)
+
+
+for opt in CgroupConfigOptionSet.keys():
+    setattr(CgroupConfig, format_key(opt), property(*generate_descriptors(opt)))

slurmutils/models/option.py

@@ -15,6 +15,7 @@
 """Configuration options for Slurm data models."""
 
 __all__ = [
+    "CgroupConfigOptionSet",
     "SlurmdbdConfigOptionSet",
     "SlurmConfigOptionSet",
     "NodeOptionSet",
@@ -47,6 +48,29 @@ def keys(cls) -> Iterable[str]:
             yield field.name
 
 
+@dataclass(frozen=True)
+class CgroupConfigOptionSet(_OptionSet):
+    """`cgroup.conf` configuration options."""
+
+    CgroupMountpoint: Callback = Callback()
+    CgroupPlugin: Callback = Callback()
+    SystemdTimeout: Callback = Callback()
+    IgnoreSystemd: Callback = Callback()
+    IgnoreSystemdOnFailure: Callback = Callback()
+    EnableControllers: Callback = Callback()
+    AllowedRAMSpace: Callback = Callback()
+    AllowedSwapSpace: Callback = Callback()
+    ConstrainCores: Callback = Callback()
+    ConstrainDevices: Callback = Callback()
+    ConstrainRAMSpace: Callback = Callback()
+    ConstrainSwapSpace: Callback = Callback()
+    MaxRAMPercent: Callback = Callback()
+    MaxSwapPercent: Callback = Callback()
+    MemorySwappiness: Callback = Callback()
+    MinRAMSpace: Callback = Callback()
+    SignalChildrenProcesses: Callback = Callback()
+
+
 @dataclass(frozen=True)
 class SlurmdbdConfigOptionSet(_OptionSet):
     """`slurmdbd.conf` configuration options."""

tests/unit/editors/test_cgroupconfig.py

@@ -0,0 +1,69 @@
+#!/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 cgroup.conf editor."""
+
+import unittest
+from pathlib import Path
+
+from slurmutils.editors import cgroupconfig
+
+EXAMPLE_CGROUP_CONF = """#
+# `cgroup.conf` file generated at 2024-09-18 15:10:44.652017 by slurmutils.
+#
+ConstrainCores=yes
+ConstrainDevices=yes
+ConstrainRAMSpace=yes
+ConstrainSwapSpace=yes
+"""
+
+
+class TestCgroupConfigEditor(unittest.TestCase):
+    """Unit tests for cgroup.conf file editor."""
+
+    def setUp(self) -> None:
+        Path("cgroup.conf").write_text(EXAMPLE_CGROUP_CONF)
+
+    def test_loads(self) -> None:
+        """Test `loads` method of the cgroupconfig module."""
+        config = cgroupconfig.loads(EXAMPLE_CGROUP_CONF)
+        self.assertEqual(config.constrain_cores, "yes")
+        self.assertEqual(config.constrain_devices, "yes")
+        self.assertEqual(config.constrain_ram_space, "yes")
+        self.assertEqual(config.constrain_swap_space, "yes")
+
+    def test_dumps(self) -> None:
+        """Test `dumps` method of the cgroupconfig module."""
+        config = cgroupconfig.loads(EXAMPLE_CGROUP_CONF)
+        # The new config and old config should not be equal since the
+        # timestamps in the header will be different.
+        self.assertNotEqual(cgroupconfig.dumps(config), EXAMPLE_CGROUP_CONF)
+
+    def test_edit(self) -> None:
+        """Test `edit` context manager from the cgroupconfig module."""
+        with cgroupconfig.edit("cgroup.conf") as config:
+            config.constrain_cores = "no"
+            config.constrain_devices = "no"
+            config.constrain_ram_space = "no"
+            config.constrain_swap_space = "no"
+
+        config = cgroupconfig.load("cgroup.conf")
+        self.assertEqual(config.constrain_cores, "no")
+        self.assertEqual(config.constrain_devices, "no")
+        self.assertEqual(config.constrain_ram_space, "no")
+        self.assertEqual(config.constrain_swap_space, "no")
+
+    def tearDown(self) -> None:
+        Path("cgroup.conf").unlink()