From 067a34fb211a33e8707bc7e737fe024c0c0c05ed Mon Sep 17 00:00:00 2001 From: jamesbeedy Date: Mon, 25 Nov 2024 15:44:47 +0000 Subject: [PATCH] enhancement: support acct_gather.conf These changes add a data model and editor for acct_gather.conf. --- slurmutils/editors/acct_gatherconfig.py | 84 ++++++++++++++++++++ slurmutils/models/__init__.py | 1 + slurmutils/models/acct_gather.py | 49 ++++++++++++ slurmutils/models/cgroup.py | 2 +- slurmutils/models/option.py | 29 +++++++ tests/unit/editors/constants.py | 10 +++ tests/unit/editors/test_acct_gatherconfig.py | 60 ++++++++++++++ 7 files changed, 234 insertions(+), 1 deletion(-) create mode 100644 slurmutils/editors/acct_gatherconfig.py create mode 100644 slurmutils/models/acct_gather.py create mode 100644 tests/unit/editors/test_acct_gatherconfig.py diff --git a/slurmutils/editors/acct_gatherconfig.py b/slurmutils/editors/acct_gatherconfig.py new file mode 100644 index 0000000..1246943 --- /dev/null +++ b/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 . + +"""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) diff --git a/slurmutils/models/__init__.py b/slurmutils/models/__init__.py index 1bd6776..141a0be 100644 --- a/slurmutils/models/__init__.py +++ b/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 diff --git a/slurmutils/models/acct_gather.py b/slurmutils/models/acct_gather.py new file mode 100644 index 0000000..9b35d74 --- /dev/null +++ b/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 . + +"""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))) diff --git a/slurmutils/models/cgroup.py b/slurmutils/models/cgroup.py index 6062918..b5e7da1 100644 --- a/slurmutils/models/cgroup.py +++ b/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): diff --git a/slurmutils/models/option.py b/slurmutils/models/option.py index c979c1d..46ce033 100644 --- a/slurmutils/models/option.py +++ b/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.""" diff --git a/tests/unit/editors/constants.py b/tests/unit/editors/constants.py index 2081302..08bb6cb 100644 --- a/tests/unit/editors/constants.py +++ b/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 +""" diff --git a/tests/unit/editors/test_acct_gatherconfig.py b/tests/unit/editors/test_acct_gatherconfig.py new file mode 100644 index 0000000..5515c36 --- /dev/null +++ b/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 . + +"""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")