From c137738acc5e83f15f788a4a849061ccb9509d45 Mon Sep 17 00:00:00 2001 From: Leopold Talirz Date: Mon, 13 Jun 2022 01:50:15 -0700 Subject: [PATCH 1/2] add code registry tools This adds a registry that allows users to place a set of yaml files describing computers and codes installed on them in a folder of their choice (AIIDA_CODE_REGISTRY). When using the `load_code` or `load_computer` functions from aiida_code_registry, these codes are set up automatically on demand. --- .gitignore | 1 + .pre-commit-config.yaml | 56 ++++-- LICENSE | 21 ++ README.md | 9 + aiida_code_registry/__init__.py | 329 ++++++++++++++++++++++++++++++++ configurations/aiidalab.yaml | 27 +++ pyproject.toml | 69 +++++++ 7 files changed, 496 insertions(+), 16 deletions(-) create mode 100644 .gitignore create mode 100644 LICENSE create mode 100644 aiida_code_registry/__init__.py create mode 100644 configurations/aiidalab.yaml create mode 100644 pyproject.toml diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..c18dd8d --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +__pycache__/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d5d540e..fdcb341 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,18 +1,42 @@ repos: - - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.0.1 # Use the ref you want to point at - hooks: - - id: trailing-whitespace - - id: check-yaml +- repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.0.1 # Use the ref you want to point at + hooks: + - id: trailing-whitespace + - id: check-yaml + +- repo: https://github.com/adrienverge/yamllint.git + rev: v1.17.0 + hooks: + - id: yamllint + args: [-c=.yamllint, -s] + # only check yaml files for codes and computers setting + exclude: > + (?x)^( + .github/.* | + .pre-commit-config.yaml | + )$ + +- repo: https://github.com/PyCQA/isort + rev: 5.10.1 + hooks: + - id: isort + +- repo: https://github.com/psf/black + rev: 22.3.0 + hooks: + - id: black + +- repo: local + hooks: + - id: pylint + language: system + types: [file, python] + name: pylint + description: "This hook runs the pylint static code analyzer" + exclude: &exclude_files > + (?x)^( + docs/.*| + )$ + entry: pylint - - repo: https://github.com/adrienverge/yamllint.git - rev: v1.17.0 - hooks: - - id: yamllint - args: [-c=.yamllint, -s] - # only check yaml files for codes and computers setting - exclude: > - (?x)^( - .github/.* | - .pre-commit-config.yaml | - )$ \ No newline at end of file diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..8c6f14c --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2022 AiiDA Team. + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md index 9b0c577..d45fe54 100644 --- a/README.md +++ b/README.md @@ -4,6 +4,15 @@ This repository collects configurations of simulation codes on public compute resources for quick and easy setup in AiiDA. +## Using the AiiDA code registry (python package) + +``` +import aiida_code_registry as registry + +code1 = registry.load_code('base@localhost') +code2 = registry.load_code('qe-pw@localhost', template_vars(dict(executable='pw'))) +``` + ## Using the AiiDA code registry In the following we'll take the example of [Piz Daint](https://www.cscs.ch/computers/piz-daint/), a HPC system at the Swiss National Supercomputing Centre. diff --git a/aiida_code_registry/__init__.py b/aiida_code_registry/__init__.py new file mode 100644 index 0000000..8aa285c --- /dev/null +++ b/aiida_code_registry/__init__.py @@ -0,0 +1,329 @@ +""" +Utilities for dealing with a registry of computers and codes. + +""" +__version__ = "0.1.0" +import copy +import functools +import json +import logging +import os +from pathlib import Path +import pprint +from typing import Union + +from frozendict import frozendict +import jinja2 +from voluptuous import Any, Optional, Required, Schema +import yaml + +from aiida import common, orm +from aiida.orm.utils.builders.code import CodeBuilder +from aiida.orm.utils.builders.computer import ComputerBuilder + +THIS_DIR = Path(__file__).parent.absolute() +CONFIGURATIONS_DIR = THIS_DIR.parent / "configurations" + +NUMBER = Any(float, int) + +logger = logging.getLogger() + +_JINJA_ENV = jinja2.Environment(loader=jinja2.BaseLoader) + +# Code schemas +CODE_SETUP_SCHEMA = { + "label": str, + "description": str, + "computer": str, + "on_computer": bool, + "remote_abs_path": str, + Required("input_plugin"): str, + # Optional("use_double_quotes", default=False): bool, + "prepend_text": str, + "append_text": str, +} + +# Computer schemas +COMPUTER_SETUP_SCHEMA = { + "label": str, + Required("hostname"): str, + "description": str, + Required("transport"): str, + "scheduler": str, + "work_dir": str, + "append_text": str, + "prepend_text": str, + "shebang": str, + # Optional("use_double_quotes", default=False): bool, + "mpirun_command": str, + "mpiprocs_per_machine": int, + Optional("default_memory_per_machine", default=None): Any(int, None), + "extras": dict, +} + +COMPUTER_CONFIGURE_SSH_SCHEMA = { + "timeout": NUMBER, + "safe_interval": NUMBER, + "compress": bool, + "key_policy": str, + "key_filename": str, +} + +COMPUTER_CONFIGURE_LOCAL_SCHEMA = {} + +COMPUTER_CONFIGURE_SCHEMA = { + "core.local": COMPUTER_CONFIGURE_LOCAL_SCHEMA, + "core.ssh": COMPUTER_CONFIGURE_SSH_SCHEMA, +} + + +COMPUTER_SCHEMA = { + Required("label"): str, + Required("setup"): COMPUTER_SETUP_SCHEMA, + "configure": COMPUTER_CONFIGURE_SCHEMA, + "codes": [CODE_SETUP_SCHEMA], +} + +# Registry schema +CODE_REGISTRY_SCHEMA = { + Required("computers"): [COMPUTER_SCHEMA], + "codes": [CODE_SETUP_SCHEMA], +} + +SCHEMA = Schema(CODE_REGISTRY_SCHEMA, required=False) + +# See https://stackoverflow.com/a/53394430/1069467 +def freezeargs(func): + """Transform mutable dictionnary + Into immutable + Useful to be compatible with cache + """ + + @functools.wraps(func) + def wrapped(*args, **kwargs): + args = tuple(frozendict(arg) if isinstance(arg, dict) else arg for arg in args) + kwargs = { + k: frozendict(v) if isinstance(v, dict) else v for k, v in kwargs.items() + } + return func(*args, **kwargs) + + return wrapped + + +@functools.lru_cache(maxsize=3) +def load_code_registry(directory: str) -> dict: + """Load AiiDA code registry data from directory.""" + directory = Path(directory) + if not directory.is_dir(): + raise ValueError( + f"{directory} is not a Path. Use the 'AIIDA_CODE_REGISTRY' variable to point to a directory." + ) + + yamls = list(directory.glob("*.yaml")) + list(directory.glob("*.yml")) + + if not yamls: + raise ValueError(f"No YAML files found in {directory}.") + + registry = {} + for yaml_path in yamls: + with open(yaml_path, encoding="utf8") as handle: + data = yaml.safe_load(handle) + + data = SCHEMA(data) + registry[yaml_path.name] = data + + return registry + + +# @freezeargs +# @functools.lru_cache(maxsize=3) +def _replace_template_vars(template, template_vars): + """Replace template variables in JSON-serializable python object. + + Uses jinja2 templating mechanism plus intermediate serialization to json. + """ + # render even if no template vars provided (will raise if unreplaced template is detected) + if template_vars is None: + template_vars = {} + dict_str = json.dumps(template) + jinja_template = _JINJA_ENV.from_string(dict_str) + dict_str = jinja_template.render(**template_vars) + + return json.loads(dict_str) + + +class CodeRegistry: + """Code registry with convenience functions.""" + + @classmethod + def from_directory(cls, directory: Union[str, Path]): + """Load AiiDA code registry from directory.""" + return cls(raw=load_code_registry(directory)) + + @classmethod + def from_env(cls): + """Load AiiDA code registry from environment. + + Precedence: + - Value of 'AIIDA_CODE_REGISTRY' environment variable + - ../configurations + """ + directory = os.getenv("AIIDA_CODE_REGISTRY") + + if directory is None: + directory = CONFIGURATIONS_DIR + + return cls.from_directory(directory) + + def __init__(self, raw: dict): + """Create AiiDA code registry data from raw dict.""" + self._raw = raw + + # now merge the raw yaml data + computers = {} + for _file_name, file_content in copy.deepcopy(raw).items(): + data = file_content + + for computer in data["computers"]: + label = computer["label"] + if label in computers: + logger.warning( + "Computer '%s' found in registry multiple times, overwriting.", + label, + ) + + # label is optional in 'setup' to avoid duplication + if "label" not in computer["setup"]: + computer["setup"]["label"] = label + + if "codes" not in computer: + computer["codes"] = [] + + # codes specified at top level are added to each computer + if "codes" in data: + computer["codes"] += data["codes"] + + # make sure codes specify the correct computer (can be left out in yaml) + for code in computer["codes"]: + code["computer"] = label + + # transform codes into dictionary + computer["codes"] = {code["label"]: code for code in computer["codes"]} + + computers[label] = computer + + self.computers = computers + + @property + def computer_list(self) -> list: + """Return list of available computers.""" + return list(self.computers.keys()) + + def get_computer(self, label) -> dict: + """Return configuration for given computer.""" + return self.computers[label] + + @property + def code_list(self) -> list: + """Return list of available codes.""" + code_labels = [] + for computer_label, computer in self.computers.items(): + for code_label in computer["codes"]: + code_labels.append(f"{code_label}@{computer_label}") + + return code_labels + + def get_code(self, label) -> dict: + """Return configuration for given code. + + Note: Label needs to be of form 'code@computer'. + """ + if "@" not in label: + raise ValueError(f"Label '{label}' is not of form 'code@computer'.") + + code_label, computer_label = label.split("@") + return self.computers[computer_label]["codes"][code_label] + + def render(self, template_vars: dict): + """Return copy of registry where template variables have been replaced.""" + if not template_vars: + return self + + return self.__class__(raw=_replace_template_vars(self._raw, template_vars)) + + +AIIDA_CODE_REGISTRY = CodeRegistry.from_env() + + +def load_computer(label: str, template_vars: dict = None) -> orm.Computer: + """Load computer; create if it does not exist. + + :param label: computer label + :param template_vars: dictionary with template variables to be replaced + """ + try: + return orm.load_computer(label=label) + except common.NotExistent as exc: + computer_list = AIIDA_CODE_REGISTRY.render(template_vars).computer_list + if label not in computer_list: + raise common.NotExistent( + f"Computer '{label}' found neither in database nor in registry. " + + f"Computers in registry: {pprint.pformat(computer_list)}" + ) from exc + + cfg = AIIDA_CODE_REGISTRY.render(template_vars).get_computer(label) + if "extras" in cfg["setup"]: + # do custom setup... + del cfg["setup"]["extras"] + computer_builder = ComputerBuilder(**cfg["setup"]) + logger.info("Setting up computer '%s'.", label) + computer = computer_builder.new().store() + + user = orm.User.collection.get_default() + + try: + configure_dict = cfg["configure"][cfg["setup"]["transport"]] + except KeyError as exc: + raise KeyError( + f"Computer '{label}' is missing configure info for transport '{cfg['setup']['transport']}'." + ) from exc + + logger.info( + "Configuring computer '%s' for '%s' transport.", + label, + cfg["setup"]["transport"], + ) + computer.configure(user=user, **configure_dict) + + return computer + + +def load_code(label: str, template_vars: dict = None) -> orm.Code: + """Load code; create if it does not exist. + + :param label: code label + :param template_vars: dictionary with template variables to be replaced + """ + try: + return orm.load_code(label=label) + except common.NotExistent as exc: + code_list = AIIDA_CODE_REGISTRY.render(template_vars).code_list + if label not in code_list: + raise common.NotExistent( + f"Code '{label}' found neither in database nor in registry. Codes in registry:\n{pprint.pformat(code_list)}" + ) from exc + + setup_dict = AIIDA_CODE_REGISTRY.render(template_vars).get_code(label).copy() + if setup_dict.pop("on_computer"): + setup_dict["code_type"] = CodeBuilder.CodeType.ON_COMPUTER + else: + setup_dict["code_type"] = CodeBuilder.CodeType.STORE_AND_UPLOAD + + # create computer, if it does not exist + setup_dict["computer"] = load_computer( + label=setup_dict["computer"], template_vars=template_vars + ) + + code_builder = CodeBuilder(**setup_dict) + + return code_builder.new().store() diff --git a/configurations/aiidalab.yaml b/configurations/aiidalab.yaml new file mode 100644 index 0000000..88692d2 --- /dev/null +++ b/configurations/aiidalab.yaml @@ -0,0 +1,27 @@ +--- +computers: + - label: localhost + setup: + hostname: localhost + description: AiiDAlab docker container + transport: core.local + scheduler: core.direct + work_dir: /home/{username}/aiida_run + append_text: " " + prepend_text: "" + shebang: "#!/bin/bash" + mpirun_command: "mpirun -np {tot_num_mpiprocs}" + codes: + - label: base + description: conda base environment in container + computer: localhost + on_computer: true + remote_abs_path: /opt/conda/bin/python + input_plugin: dynamic_workflows.PyCalcJob + prepend_text: "source /opt/conda/etc/profile.d/conda.sh\nconda activate base" + - label: qe-{{ executable }} + description: Quantum ESPRESSO 7.0 + computer: localhost + on_computer: true + remote_abs_path: /a/b/c/{{ executable }}.x + input_plugin: quantumespresso.{{ executable }} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..b8dc055 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,69 @@ +[build-system] +# build the package with [flit](https://flit.readthedocs.io) +requires = ["flit_core >=3.4,<4"] +build-backend = "flit_core.buildapi" + +[project] +# See https://www.python.org/dev/peps/pep-0621/ +name = "aiida-code-registry" +dynamic = ["version"] # read from aiida_code_registry/__init__.py +description = "AiiDA code registry." +authors = [{name = "AiiDA Team"}] +readme = "README.md" +license = {file = "LICENSE"} +classifiers = [ + "Programming Language :: Python", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: MIT License", + "Natural Language :: English", + "Development Status :: 3 - Alpha", + "Framework :: AiiDA" +] +keywords = ["aiida", "plugin"] +requires-python = ">=3.7" +dependencies = [ + "aiida-core>=1.6.0,<3", + "frozendict", +] + +[project.urls] +Source = "https://github.com/aiidateam/aiida-code-registry" + +[project.optional-dependencies] +pre-commit = [ + "pre-commit~=2.2", + "pylint>=2.5.0,<2.9" +] + +[tool.flit.module] +name = "aiida_code_registry" + +[tool.pylint.format] +max-line-length = 125 + +[tool.pylint.messages_control] +disable = [ + "too-many-ancestors", + "invalid-name", + "duplicate-code", + # black compatibility + "C0330", + "C0326", +] + +[tool.pytest.ini_options] +# Configuration for [pytest](https://docs.pytest.org) +python_files = "test_*.py example_*.py" +filterwarnings = [ + "ignore::DeprecationWarning:aiida:", + "ignore::DeprecationWarning:plumpy:", + "ignore::DeprecationWarning:yaml:", +] + +[tool.isort] +# Configuration of [isort](https://isort.readthedocs.io) +line_length = 120 +force_sort_within_sections = true +sections = ['FUTURE', 'STDLIB', 'THIRDPARTY', 'AIIDA', 'FIRSTPARTY', 'LOCALFOLDER'] +known_aiida = ['aiida'] + From d950b2df877be8d792b7a04da0f2b6f1f5c727c6 Mon Sep 17 00:00:00 2001 From: Leopold Talirz Date: Tue, 24 Jan 2023 07:31:03 -0800 Subject: [PATCH 2/2] Update README.md --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index d45fe54..3800cb7 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ This repository collects configurations of simulation codes on public compute re import aiida_code_registry as registry code1 = registry.load_code('base@localhost') -code2 = registry.load_code('qe-pw@localhost', template_vars(dict(executable='pw'))) +code2 = registry.load_code('qe-pw@localhost', template_vars=dict(executable='pw'))) ``` ## Using the AiiDA code registry