diff --git a/slurmutils/editors/cgroupconfig.py b/slurmutils/editors/cgroupconfig.py index fb53fbf..965d94a 100644 --- a/slurmutils/editors/cgroupconfig.py +++ b/slurmutils/editors/cgroupconfig.py @@ -23,14 +23,7 @@ from typing import Union from ..models import CgroupConfig -from ..models.option import CgroupConfigOptionSet -from .editor import ( - clean, - dumper, - loader, - marshall_content, - parse_line, -) +from .editor import dumper, loader _logger = logging.getLogger("slurmutils") @@ -43,7 +36,7 @@ def load(file: Union[str, os.PathLike]) -> CgroupConfig: def loads(content: str) -> CgroupConfig: """Load `cgroup.conf` data model from string.""" - return _parse(content) + return CgroupConfig.from_str(content) @dumper @@ -54,7 +47,7 @@ def dump(config: CgroupConfig, file: Union[str, os.PathLike]) -> None: def dumps(config: CgroupConfig) -> str: """Dump `cgroup.conf` data model into a string.""" - return _marshall(config) + return str(config) @contextmanager @@ -73,33 +66,3 @@ def edit(file: Union[str, os.PathLike]) -> CgroupConfig: 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) diff --git a/slurmutils/editors/editor.py b/slurmutils/editors/editor.py index 6544a04..becde3b 100644 --- a/slurmutils/editors/editor.py +++ b/slurmutils/editors/editor.py @@ -15,73 +15,12 @@ """Base methods for Slurm workload manager configuration file editors.""" import logging -import shlex from functools import wraps from os import path -from typing import Any, Dict, List, Optional - -from ..exceptions import EditorError _logger = logging.getLogger("slurmutils") -def clean(line: str) -> Optional[str]: - """Clean line before further processing. - - Returns: - Line with inline comments removed. `None` if line is a comment. - """ - return cleaned if (cleaned := line.split("#", maxsplit=1)[0]) != "" else None - - -def parse_line(options, line: str) -> Dict[str, Any]: - """Parse configuration line. - - Args: - options: Available options for line. - line: Configuration line to parse. - """ - data = {} - opts = shlex.split(line) # Use `shlex.split(...)` to preserve quotation strings. - for opt in opts: - k, v = opt.split("=", maxsplit=1) - if not hasattr(options, k): - raise EditorError( - ( - f"unable to parse configuration option {k}. " - + f"valid configuration options are {list(options.keys())}" - ) - ) - - parse = getattr(options, k).parser - data[k] = parse(v) if parse else v - - return data - - -def marshall_content(options, line: Dict[str, Any]) -> List[str]: - """Marshall data model content back into configuration line. - - Args: - options: Available options for line. - line: Data model to marshall into line. - """ - result = [] - for k, v in line.items(): - if not hasattr(options, k): - raise EditorError( - ( - f"unable to marshall configuration option {k}. " - + f"valid configuration options are {[option.name for option in options]}" - ) - ) - - marshall = getattr(options, k).marshaller - result.append(f"{k}={marshall(v) if marshall else v}") - - return result - - def loader(func): """Wrap function that loads configuration data from file.""" diff --git a/slurmutils/editors/slurmconfig.py b/slurmutils/editors/slurmconfig.py index 072f72c..1f0aac3 100644 --- a/slurmutils/editors/slurmconfig.py +++ b/slurmutils/editors/slurmconfig.py @@ -22,15 +22,8 @@ from pathlib import Path from typing import Union -from ..models import DownNodes, FrontendNode, Node, NodeSet, Partition, SlurmConfig -from ..models.option import SlurmConfigOptionSet -from .editor import ( - clean, - dumper, - loader, - marshall_content, - parse_line, -) +from ..models import SlurmConfig +from .editor import dumper, loader _logger = logging.getLogger("slurmutils") @@ -43,7 +36,7 @@ def load(file: Union[str, os.PathLike]) -> SlurmConfig: def loads(content: str) -> SlurmConfig: """Load `slurm.conf` data model from string.""" - return _parse(content) + return SlurmConfig.from_str(content) @dumper @@ -54,7 +47,7 @@ def dump(config: SlurmConfig, file: Union[str, os.PathLike]) -> None: def dumps(config: SlurmConfig) -> str: """Dump `slurm.conf` data model into a string.""" - return _marshall(config) + return str(config) @contextmanager @@ -73,94 +66,3 @@ def edit(file: Union[str, os.PathLike]) -> SlurmConfig: yield config dump(config, file) - - -def _parse(content: str) -> SlurmConfig: - """Parse contents of `slurm.conf`. - - Args: - content: Contents of `slurm.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 slurm.conf", line, index) - continue - - if config.startswith("Include"): - _, v = config.split(maxsplit=1) - data["Include"] = data.get("Include", []) + [v] - elif config.startswith("SlurmctldHost"): - _, v = config.split("=", maxsplit=1) - data["SlurmctldHost"] = data.get("SlurmctldHost", []) + [v] - elif config.startswith("NodeName"): - nodes = data.get("Nodes", {}) - nodes.update(Node.from_str(config).dict()) - data["Nodes"] = nodes - elif config.startswith("DownNodes"): - data["DownNodes"] = data.get("DownNodes", []) + [DownNodes.from_str(config).dict()] - elif config.startswith("FrontendNode"): - frontend_nodes = data.get("FrontendNodes", {}) - frontend_nodes.update(FrontendNode.from_str(config).dict()) - data["FrontendNodes"] = frontend_nodes - elif config.startswith("NodeSet"): - node_sets = data.get("NodeSets", {}) - node_sets.update(NodeSet.from_str(config).dict()) - data["NodeSets"] = node_sets - elif config.startswith("PartitionName"): - partitions = data.get("Partitions", {}) - partitions.update(Partition.from_str(config).dict()) - data["Partitions"] = partitions - else: - data.update(parse_line(SlurmConfigOptionSet, config)) - - return SlurmConfig.from_dict(data) - - -def _marshall(config: SlurmConfig) -> str: - """Marshall `slurm.conf` data model back into slurm.conf format. - - Args: - config: `slurm.conf` data model to marshall. - """ - result = [] - data = config.dict() - include = data.pop("Include", None) - slurmctld_host = data.pop("SlurmctldHost", None) - nodes = data.pop("Nodes", {}) - down_nodes = data.pop("DownNodes", []) - frontend_nodes = data.pop("FrontendNodes", {}) - node_sets = data.pop("NodeSets", {}) - partitions = data.pop("Partitions", {}) - - if include: - result.extend([f"Include {i}" for i in include]) - - if slurmctld_host: - result.extend([f"SlurmctldHost={host}" for host in slurmctld_host]) - - result.extend(marshall_content(SlurmConfigOptionSet, data)) - - if nodes: - for k, v in nodes.items(): - result.append(str(Node(NodeName=k, **v))) - - if down_nodes: - for entry in down_nodes: - result.append(str(DownNodes(**entry))) - - if frontend_nodes: - for k, v in frontend_nodes.items(): - result.append(str(FrontendNode(FrontendName=k, **v))) - - if node_sets: - for k, v in node_sets.items(): - result.append(str(NodeSet(NodeSet=k, **v))) - - if partitions: - for k, v in partitions.items(): - result.append(str(Partition(PartitionName=k, **v))) - - return "\n".join(result) diff --git a/slurmutils/editors/slurmdbdconfig.py b/slurmutils/editors/slurmdbdconfig.py index 7aab0cb..70f9a04 100644 --- a/slurmutils/editors/slurmdbdconfig.py +++ b/slurmutils/editors/slurmdbdconfig.py @@ -24,14 +24,7 @@ from slurmutils.models import SlurmdbdConfig -from ..models.option import SlurmdbdConfigOptionSet -from .editor import ( - clean, - dumper, - loader, - marshall_content, - parse_line, -) +from .editor import dumper, loader _logger = logging.getLogger("slurmutils") @@ -44,7 +37,7 @@ def load(file: Union[str, os.PathLike]) -> SlurmdbdConfig: def loads(content: str) -> SlurmdbdConfig: """Load `slurmdbd.conf` data model from string.""" - return _parse(content) + return SlurmdbdConfig.from_str(content) @dumper @@ -55,7 +48,7 @@ def dump(config: SlurmdbdConfig, file: Union[str, os.PathLike]) -> None: def dumps(config: SlurmdbdConfig) -> str: """Dump `slurmdbd.conf` data model into a string.""" - return _marshall(config) + return str(config) @contextmanager @@ -74,33 +67,3 @@ def edit(file: Union[str, os.PathLike]) -> SlurmdbdConfig: yield config dump(config, file) - - -def _parse(content: str) -> SlurmdbdConfig: - """Parse contents of `slurmdbd.conf`. - - Args: - content: Contents of `slurmdbd.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 slurmdbd.conf", line, index) - continue - - data.update(parse_line(SlurmdbdConfigOptionSet, config)) - - return SlurmdbdConfig.from_dict(data) - - -def _marshall(config: SlurmdbdConfig) -> str: - """Marshall `slurmdbd.conf` data model back into slurmdbd.conf format. - - Args: - config: `slurmdbd.conf` data model to marshall. - """ - result = [] - result.extend(marshall_content(SlurmdbdConfigOptionSet, config.dict())) - return "\n".join(result) diff --git a/slurmutils/models/cgroup.py b/slurmutils/models/cgroup.py index ce82b9c..6062918 100644 --- a/slurmutils/models/cgroup.py +++ b/slurmutils/models/cgroup.py @@ -14,7 +14,7 @@ """Data models for `cgroup.conf` configuration file.""" -from .model import BaseModel, format_key, generate_descriptors +from .model import BaseModel, clean, format_key, generate_descriptors, marshall_content, parse_line from .option import CgroupConfigOptionSet @@ -24,6 +24,26 @@ class CgroupConfig(BaseModel): def __init__(self, **kwargs) -> None: super().__init__(CgroupConfigOptionSet, **kwargs) + @classmethod + def from_str(cls, content: str) -> "CgroupConfig": + """Construct SlurmdbdConfig data model from slurmdbd.conf format.""" + data = {} + lines = content.splitlines() + for index, line in enumerate(lines): + config = clean(line) + if config is None: + continue + + data.update(parse_line(CgroupConfigOptionSet, config)) + + return CgroupConfig.from_dict(data) + + def __str__(self) -> str: + """Return CgroupConfig data model in cgroup.conf format.""" + result = [] + result.extend(marshall_content(CgroupConfigOptionSet, self.dict())) + return "\n".join(result) + for opt in CgroupConfigOptionSet.keys(): setattr(CgroupConfig, format_key(opt), property(*generate_descriptors(opt))) diff --git a/slurmutils/models/model.py b/slurmutils/models/model.py index 53684b4..00dbaf0 100644 --- a/slurmutils/models/model.py +++ b/slurmutils/models/model.py @@ -14,13 +14,21 @@ """Base classes and methods for composing Slurm data models.""" -__all__ = ["BaseModel", "LineInterface", "format_key", "generate_descriptors"] +__all__ = [ + "BaseModel", + "clean", + "format_key", + "generate_descriptors", + "marshall_content", + "parse_line", +] import copy import json import re +import shlex from abc import ABC, abstractmethod -from typing import Any, Callable, Dict, Tuple +from typing import Any, Callable, Dict, List, Optional, Tuple from ..exceptions import ModelError @@ -68,17 +76,61 @@ def deleter(self): return getter, setter, deleter -class LineInterface: - """Interface for data models that can be constructed from a configuration line.""" +def clean(line: str) -> Optional[str]: + """Clean line before further processing. - @classmethod - @abstractmethod - def from_str(cls, line: str): - """Construct data model from configuration line.""" + Returns: + Line with inline comments removed. `None` if line is a comment. + """ + return cleaned if (cleaned := line.split("#", maxsplit=1)[0]) != "" else None + + +def parse_line(options, line: str) -> Dict[str, Any]: + """Parse configuration line. + + Args: + options: Available options for line. + line: Configuration line to parse. + """ + data = {} + opts = shlex.split(line) # Use `shlex.split(...)` to preserve quotation strings. + for opt in opts: + k, v = opt.split("=", maxsplit=1) + if not hasattr(options, k): + raise ModelError( + ( + f"unable to parse configuration option {k}. " + + f"valid configuration options are {list(options.keys())}" + ) + ) + + parse = getattr(options, k).parser + data[k] = parse(v) if parse else v + + return data - @abstractmethod - def __str__(self) -> str: - """Return model as configuration line.""" + +def marshall_content(options, line: Dict[str, Any]) -> List[str]: + """Marshall data model content back into configuration line. + + Args: + options: Available options for line. + line: Data model to marshall into line. + """ + result = [] + for k, v in line.items(): + if not hasattr(options, k): + raise ModelError( + ( + f"unable to marshall configuration option {k}. " + + f"valid configuration options are {[option.name for option in options]}" + ) + ) + + marshall = getattr(options, k).marshaller + result.append(f"{k}={marshall(v) if marshall else v}") + + return result class BaseModel(ABC): @@ -107,6 +159,15 @@ def from_json(cls, obj: str): data = json.loads(obj) return cls.from_dict(data) + @classmethod + @abstractmethod + def from_str(cls, content: str): + """Construct data model from configuration string.""" + + @abstractmethod + def __str__(self) -> str: + """Return model as configuration string.""" + def dict(self) -> Dict[str, Any]: """Return model as dictionary.""" return copy.deepcopy(self.data) diff --git a/slurmutils/models/slurm.py b/slurmutils/models/slurm.py index 52f6753..48b8d2f 100644 --- a/slurmutils/models/slurm.py +++ b/slurmutils/models/slurm.py @@ -26,8 +26,7 @@ import copy from typing import Any, Dict, List, Optional -from ..editors.editor import marshall_content, parse_line -from .model import BaseModel, LineInterface, format_key, generate_descriptors +from .model import BaseModel, clean, format_key, generate_descriptors, marshall_content, parse_line from .option import ( DownNodeOptionSet, FrontendNodeOptionSet, @@ -38,7 +37,7 @@ ) -class Node(BaseModel, LineInterface): +class Node(BaseModel): """`Node` data model.""" def __init__(self, *, NodeName: str, **kwargs) -> None: # noqa N803 @@ -78,7 +77,7 @@ def __str__(self) -> str: return " ".join(line) -class DownNodes(BaseModel, LineInterface): +class DownNodes(BaseModel): """`DownNodes` data model.""" def __init__(self, **kwargs): @@ -95,7 +94,7 @@ def __str__(self) -> str: return " ".join(marshall_content(DownNodeOptionSet, self.data)) -class FrontendNode(BaseModel, LineInterface): +class FrontendNode(BaseModel): """`FrontendNode` data model.""" def __init__(self, *, FrontendName: str, **kwargs) -> None: # noqa N803 @@ -135,7 +134,7 @@ def __str__(self) -> str: return " ".join(line) -class NodeSet(BaseModel, LineInterface): +class NodeSet(BaseModel): """`NodeSet` data model.""" def __init__(self, *, NodeSet: str, **kwargs) -> None: # noqa N803 @@ -175,7 +174,7 @@ def __str__(self) -> str: return " ".join(line) -class Partition(BaseModel, LineInterface): +class Partition(BaseModel): """`Partition` data model.""" def __init__(self, *, PartitionName: str, **kwargs): # noqa N803 @@ -235,6 +234,87 @@ def __init__( self.data["NodeSets"] = NodeSets or {} self.data["Partitions"] = Partitions or {} + @classmethod + def from_str(cls, content: str) -> "SlurmConfig": + """Construct SlurmConfig data model from slurm.conf format.""" + data = {} + lines = content.splitlines() + for index, line in enumerate(lines): + config = clean(line) + if config is None: + continue + + if config.startswith("Include"): + _, v = config.split(maxsplit=1) + data["Include"] = data.get("Include", []) + [v] + elif config.startswith("SlurmctldHost"): + _, v = config.split("=", maxsplit=1) + data["SlurmctldHost"] = data.get("SlurmctldHost", []) + [v] + elif config.startswith("NodeName"): + nodes = data.get("Nodes", {}) + nodes.update(Node.from_str(config).dict()) + data["Nodes"] = nodes + elif config.startswith("DownNodes"): + data["DownNodes"] = data.get("DownNodes", []) + [DownNodes.from_str(config).dict()] + elif config.startswith("FrontendNode"): + frontend_nodes = data.get("FrontendNodes", {}) + frontend_nodes.update(FrontendNode.from_str(config).dict()) + data["FrontendNodes"] = frontend_nodes + elif config.startswith("NodeSet"): + node_sets = data.get("NodeSets", {}) + node_sets.update(NodeSet.from_str(config).dict()) + data["NodeSets"] = node_sets + elif config.startswith("PartitionName"): + partitions = data.get("Partitions", {}) + partitions.update(Partition.from_str(config).dict()) + data["Partitions"] = partitions + else: + data.update(parse_line(SlurmConfigOptionSet, config)) + + return cls.from_dict(data) + + def __str__(self) -> str: + """Return SlurmConfig data model in slurm.conf format.""" + result = [] + data = self.dict() + include = data.pop("Include", []) + slurmctld_host = data.pop("SlurmctldHost", []) + nodes = data.pop("Nodes", {}) + down_nodes = data.pop("DownNodes", []) + frontend_nodes = data.pop("FrontendNodes", {}) + node_sets = data.pop("NodeSets", {}) + partitions = data.pop("Partitions", {}) + + if include: + result.extend([f"Include {i}" for i in include]) + + if slurmctld_host: + result.extend([f"SlurmctldHost={host}" for host in slurmctld_host]) + + result.extend(marshall_content(SlurmConfigOptionSet, data)) + + if nodes: + for k, v in nodes.items(): + result.append(str(Node(NodeName=k, **v))) + + if down_nodes: + for entry in down_nodes: + result.append(str(DownNodes(**entry))) + + if frontend_nodes: + for k, v in frontend_nodes.items(): + result.append(str(FrontendNode(FrontendName=k, **v))) + + if node_sets: + for k, v in node_sets.items(): + result.append(str(NodeSet(NodeSet=k, **v))) + + if partitions: + for k, v in partitions.items(): + result.append(str(Partition(PartitionName=k, **v))) + + return "\n".join(result) + @property def nodes(self): """Get map of all nodes in the Slurm configuration.""" diff --git a/slurmutils/models/slurmdbd.py b/slurmutils/models/slurmdbd.py index c8cb360..4875fe1 100644 --- a/slurmutils/models/slurmdbd.py +++ b/slurmutils/models/slurmdbd.py @@ -14,7 +14,7 @@ """Data models for `slurmdbd.conf` configuration file.""" -from .model import BaseModel, format_key, generate_descriptors +from .model import BaseModel, clean, format_key, generate_descriptors, marshall_content, parse_line from .option import SlurmdbdConfigOptionSet @@ -24,6 +24,26 @@ class SlurmdbdConfig(BaseModel): def __init__(self, **kwargs): super().__init__(SlurmdbdConfigOptionSet, **kwargs) + @classmethod + def from_str(cls, content: str) -> "SlurmdbdConfig": + """Construct SlurmdbdConfig data model from slurmdbd.conf format.""" + data = {} + lines = content.splitlines() + for index, line in enumerate(lines): + config = clean(line) + if config is None: + continue + + data.update(parse_line(SlurmdbdConfigOptionSet, config)) + + return cls.from_dict(data) + + def __str__(self) -> str: + """Return SlurmdbdConfig data model in slurmdbd.conf format.""" + result = [] + result.extend(marshall_content(SlurmdbdConfigOptionSet, self.dict())) + return "\n".join(result) + for opt in SlurmdbdConfigOptionSet.keys(): setattr(SlurmdbdConfig, format_key(opt), property(*generate_descriptors(opt)))