mpsd-software-environment.py

#!/usr/bin/env python3

"""mpsd-software-environment: tool for installation of toolchains."""

import argparse
import datetime
import logging
import os
import subprocess
import sys
import time
from pathlib import Path
from typing import List, Tuple, Union
import re
from string import Template

# If 'rich' is available ("pip install rich" or "apt-get install python3-rich"),
# then use coloured output, otherwise proceed as before
try:
    import rich.logging
except ModuleNotFoundError:
    rich_available = False
else:
    rich_available = True


about_tool = """
Build toolchains using Spack.\n

This function builds toolchains for MPSD-HPC at the appropriate directory, \n
for given system architecture and MPSD software stack version.\n
The toolchains
are built using the bash script spack_setup.sh, and the results are logged. """

call_date_iso = datetime.datetime.now().replace(microsecond=0).isoformat()
config_vars = {
    #  kept inside the mpsd_release folder
    "cmd_log_file": "install.log",
    #  Placeholder installer log file name, placed at mpsd_microarch/logs
    "installer_log_template": Template(
        "${mpsd_release}_${mpsd_microarch}_${date}_${action}.log"
    ),
    #  Placeholder build log file name, placed at mpsd_microarch/logs
    "build_log_template": Template(
        "${mpsd_release}_${mpsd_microarch}_${date}_${toolchain}_${action}.log"
    ),
    "metadata_tag_open": "!<meta>",
    "metadata_tag_close": "</meta>!",
    "spack_environments_repo": "https://gitlab.gwdg.de/mpsd-cs/spack-environments.git",
}


def create_log_file_names(
    mpsd_release: str,
    mpsd_microarch: str,
    action: str,
    date: str = call_date_iso,
    toolchain: str = None,
) -> Tuple[str, str]:
    """Create log file names.

    This function creates the log file names for the installer and
    the build log files.
    The installer log file is created
    from the output of this script except
    the toolchain build log that is handled by the spack_setup.sh script.
    The build log `tee`'s the output of the spack_setup.sh script to the log file.
    The log file names are created using the
    template strings defined in config_vars.

    Parameters
    ----------
    mpsd_release : str
        MPSD software stack version
    mpsd_microarch : str
        system architecture
    date : str
        date of the call ins iso format
    action : str
        action performed (install,remove,reinstall,prepare,status)
        only install and remove are valid for build log file.
    toolchain : str
        toolchain name (only for build log file)
    returns : tuple
        tuple containing the strings of installer and build log file names
    """
    installer_log_file = config_vars["installer_log_template"].substitute(
        mpsd_release=mpsd_release,
        mpsd_microarch=mpsd_microarch,
        date=date,
        action=action,
    )
    if toolchain and action in ["install", "remove"]:
        build_log_file = config_vars["build_log_template"].substitute(
            mpsd_release=mpsd_release,
            mpsd_microarch=mpsd_microarch,
            date=date,
            action=action,
            toolchain=toolchain,
        )
    else:
        build_log_file = None
        logging.warning("Incorrect action or toolchain name.")
    return installer_log_file, build_log_file


def log_metadata(key: str, value: str) -> None:
    """Log metadata to the log file.

    This function logs metadata to the log file. The metadata is
    enclosed in a tag, so that it can be easily found in the log file.
    logging module is used to write the metadata to the log file.

    Parameters
    ----------
    key : str
        key of the metadata
    value : str
        value of the metadata
    returns : None
    """
    logging.info(
        f"{config_vars['metadata_tag_open']}{key}:{value}{config_vars['metadata_tag_close']}"
    )


def read_metadata_from_logfile(logfile: Union[str, Path]) -> dict:
    """Read metadata from the log file.

    This function reads metadata from the log file. The metadata is
    enclosed in a tag, so that it can be easily found in the log file.

    Parameters
    ----------
    logfile : str or Path
        log file name
    returns : dict
        dictionary containing the metadata
    """
    with open(logfile, "r") as f:
        log_text = f.read()
    # check for all data that matches the regex
    # metadata_tag_open {key}:{value} metadata_tag_close
    # and return a dictionary with all the matches
    return {
        match.group(1): match.group(2)
        for match in re.finditer(
            f"{config_vars['metadata_tag_open']}(\w+):(\w+){config_vars['metadata_tag_close']}",
            log_text,
        )
    }


def set_up_logging(loglevel="warning", filename=None):
    """Set up logging.

    This function sets up the logging configuration for the script.
    It configures the log level, log format, and log handlers
    for both file and console output.


    Parameters
    ----------
    loglevel : str or int
       Loglevels are:
         - warning (default): only print statements if something is unexpected
         - info (show more detailed progress)
         - debug (show very detailed output)
    filename : str
         - filename to save logging messages into

    If loglevel is 'debug', save line numbers in log messages.
    """
    log_level_numeric = getattr(logging, loglevel.upper(), logging.WARNING)
    assert log_level_numeric
    if not isinstance(log_level_numeric, int):
        raise ValueError("Invalid log level: %s" % loglevel)

    handlers = []
    if filename:
        handlers.append(logging.FileHandler(filename))

    if rich_available:
        # set up logging as recommended for rich, see
        # https://rich.readthedocs.io/en/stable/logging.html
        handlers.append(rich.logging.RichHandler())
        logging_format = "%(message)s"
    else:  # rich not available, define our own output
        # include line numbers in output if level is DEBUG
        linenumbers = " %(lineno)4d" if log_level_numeric == logging.DEBUG else ""
        handlers.append(logging.StreamHandler())
        logging_format = "%(asctime)s %(levelname)7s" + linenumbers + "  |  %(message)s"

    logging.basicConfig(
        level=log_level_numeric,
        format=logging_format,
        datefmt="[%X]",
        handlers=handlers,
        force=True,
    )
    logging.debug(
        f"Logging has been setup, loglevel={loglevel.upper()}"
        + f"{filename=} {rich_available=}"
    )


# Helper class to change directory via context manager
class os_chdir:
    """The os_chdir class is a context manager.

    It changes the current directory to a specified directory
    and returns to the original directory after execution.
    """

    def __init__(self, new_dir):
        """Initialize, save original directory."""
        self.new_dir = new_dir
        self.saved_dir = os.getcwd()

    def __enter__(self):
        """Go to target directory (main action for context)."""
        os.chdir(self.new_dir)

    def __exit__(self, exc_type, exc_val, exc_tb):
        """On exist we return to original directory."""
        os.chdir(self.saved_dir)


def run(*args, counter=[0], **kwargs):
    """
    Run a subprocess and log the call.

    Convenience function to call `subprocess.run` and provide some metadata
    about the call.

    Parameters
    ----------
    args : tuple
        passed on to subprocess.run(*args). For example
        ("ls -l") or (["ls", "-l"])
    counter : TYPE, optional
        list with one integer, starting from [0].
        This is (a Python hack) to count the number of
        calls of this function, so the different calls of subprocess.run
        are easier to follow in the log files.
    kwargs : dict
        keyword-value arguments to be passed to subprocess.run. For example,
        `shell=True`.

    Returns
    -------
    process : subprocess.CompletedProcess
        CompletedProcess object as returned by `subprocess.run` .

    Examples
    --------
    >>> run(['date', '+%Y-%m-%d'])
    ##-03 Starting subprocess.run(['date', '+%Y-%m-%d']) with options
    ##-03   getcwd=/Users/fangohr/git/mpsd-software-environments
    ##-03   COMMAND=date +%Y-%m-%d
    2023-05-30
    ##-03   Completed in 0.0054s.
    ##-03
    CompletedProcess(args=['date', '+%Y-%m-%d'], returncode=0)

    >>> run(['date +%Y-%m-%d'], shell=True)
    ##-04 Starting subprocess.run(['date +%Y-%m-%d']) with options shell=True
    ##-04   getcwd=/Users/fangohr/git/mpsd-software-environments
    ##-04   COMMAND=date +%Y-%m-%d
    2023-05-30
    ##-04   Completed in 0.0069s.
    ##-04
    CompletedProcess(args=['date +%Y-%m-%d'], returncode=0)
    """
    # token is printed in front of every meta-data line - useful for
    # searching the logs. Starts with "##-00", then "##-01", ...
    token = f"##-{counter[0]:02d}"

    counter[0] += 1  # increase counter

    # make command nicely readable: ["ls", "-l"] -> "ls -l"
    assert isinstance(args, tuple)
    assert len(args) == 1
    arg = args[0]
    # either args is a tuple containing a string | Example: ('ls -1',)
    if isinstance(arg, str):
        command = arg
        # or we have a tuple containing a list of strings.
        # Example: (['ls', '-1'],)
    elif isinstance(arg, list):
        command = " ".join(arg)
    else:
        # we do not expect this to happen
        raise NotImplementedError(f"{arg=}, {args=}")

    # make options (such as `shell=True`) nicely readable
    options = ", ".join([f"{key}={value}" for key, value in kwargs.items()])

    # provide information about upcoming subprocess.run call
    logging.info(f"{token} Starting subprocess.run('{command}') with options {options}")
    logging.debug(f"{token}   getcwd={os.getcwd()}")
    logging.debug(f"{token}   exact call: subprocess.run({arg})")

    time_start = time.time()
    process = subprocess.run(*args, **kwargs)
    execution_time = time.time() - time_start

    logging.debug(f"{token}   Completed in {execution_time:.4f}s.")
    logging.debug(f"{token}")  # near-empty line to make reading logs easier
    return process


def setup_log_cmd(
    mpsd_release: str, script_dir: str, msg: str = None, **kwargs
) -> None:
    """
    Log the command used to build the toolchains.

    It also logs information about the software environment installer branch,
    the Spack environments branch, and the commit hashes of each.
    It also logs steps taken
    in the install process using the optional message argument.

    Parameters
    ----------
    - mpsd_release : str
        The name of the release to install toolchains for.
    - script_dir : str
        The path to the directory where the scripts are located.
    - msg : str, optional
        An optional message to log in the command log file.
    - **kwargs : dict
        A dictionary with values for
        - spe_branch : str
            The name of the Spack environments branch.
        - spe_commit_hash : str
            The commit hash of the Spack environments branch.

    Returns
    -------
    - None
    """
    release_base_dir = script_dir / mpsd_release

    # Write to the log file with the following format
    # --------------------------------------------------
    # 2023-02-29T23:32:01, install-software-environment.py --release 23b --install ALL
    # Software environment installer branch: script_branch (commit hash: \
    #    script_commit_hash)
    # Spack environments branch: dev-23a (commit hash: spe_commit_hash)
    # MSGs
    with os_chdir(release_base_dir):
        with open(config_vars["cmd_log_file"], "a") as f:
            if msg:
                # Write the message to the log file
                f.write(msg + "\n")
            else:
                # Write the header
                f.write("-" * 50 + "\n")

                # Gather data to log
                # call statement:
                cmd_line = " ".join(sys.argv)
                # script branch and commit hash
                with os_chdir(script_dir):
                    script_branch = (
                        run(
                            ["git", "rev-parse", "--abbrev-ref", "HEAD"],
                            stdout=subprocess.PIPE,
                            check=True,
                        )
                        .stdout.decode()
                        .strip()
                    )
                    script_commit_hash = (
                        run(
                            ["git", "rev-parse", "--short", "HEAD"],
                            stdout=subprocess.PIPE,
                            check=True,
                        )
                        .stdout.decode()
                        .strip()
                    )
                # spack-environments branch and commit hash from kwargs
                spe_branch = kwargs.get("spe_branch", None)
                spe_commit_hash = kwargs.get("spe_commit_hash", None)

                # Write to log file
                f.write(f"{datetime.datetime.now().isoformat()}, {cmd_line}\n")
                f.write(
                    f"Software environment installer branch: {script_branch} "
                    f"(commit hash: {script_commit_hash})\n"
                )
                f.write(
                    f"Spack environments branch: {spe_branch} "
                    f"(commit hash: {spe_commit_hash})\n"
                )


def create_dir_structure(mpsd_release: str, script_dir: Path) -> None:
    """
    Create the directory structure and clone spack environments repo.

    The create_dir_structure function creates the directory structure for
    the specified release and clones the Spack environments repository if it
    doesn't exist.

    Parameters
    ----------
    - mpsd_release: A string representing the MPSD release version.
    - script_dir: A Path object representing the path to the scripts directory.

    Returns
    -------
    - None
    """
    # Create the directory structure for the release
    release_base_dir = script_dir / mpsd_release
    release_base_dir.mkdir(parents=True, exist_ok=True)

    with os_chdir(release_base_dir):
        # Clone the spack-environments repo if it doesn't exist
        if not os.path.exists("spack-environments"):
            run(
                [
                    "git",
                    "clone",
                    config_vars["spack_environments_repo"],
                ],
                check=True,
            )
        with os_chdir("spack-environments"):
            # Git fetch and checkout the release branch and git pull
            # to be sure that the resulting repo is up to date
            run(["git", "fetch", "--all"], check=True)
            checkout_result = run(["git", "checkout", mpsd_release], check=True)

            if checkout_result.returncode != 0:
                raise Exception(
                    "Release branch does not exist in spack-environment repo \n."
                    "Check for typos."
                )
            run(["git", "pull"], check=True)


def get_release_info(mpsd_release: str, script_dir: Path) -> Tuple[str, str, List[str]]:
    """
    Get information about the specified release.

    Get information about the specified release, such as the branch and commit hash
    of the Spack environments repository and the available toolchains.

    Parameters
    ----------
    mpsd_release : str
        The name of the release to get information for.
    script_dir : pathlib.Path
        The base directory where releases are stored.

    Returns
    -------
    spe_branch : str
        The name of the branch for the Spack environments repository.
    spe_commit_hash : str
        The commit hash for the Spack environments repository.
    available_toolchains : list
        A list of strings representing the available toolchains for the release.

    Raises
    ------
    FileNotFoundError
        If the release directory does not exist. Run `create_dir_structure()` first.
    """
    # Get the info for release
    release_base_dir = script_dir / mpsd_release
    if not os.path.exists(release_base_dir):
        raise FileNotFoundError(
            "Release directory does not exist. Run create_dir_structure() first."
        )
    with os_chdir(release_base_dir):
        with os_chdir("spack-environments"):
            # Get the branch and commit hash of the spack-environments repo
            spe_commit_hash = (
                run(["git", "rev-parse", "HEAD"], stdout=subprocess.PIPE, check=True)
                .stdout.decode()
                .strip()
            )
            spe_branch = (
                run(
                    ["git", "rev-parse", "--abbrev-ref", "HEAD"],
                    stdout=subprocess.PIPE,
                    check=True,
                )
                .stdout.decode()
                .strip()
            )
            available_toolchains = os.listdir("toolchains")
    return spe_branch, spe_commit_hash, available_toolchains


def prepare_environment(mpsd_release: str, script_dir: Path) -> List[str]:
    """
    Create the directory structure for the given MPSD release.

    It does the following steps:
    Clones the spack-environments repository.
    Determines the branch and commit hash of the spack-environments repository
    and the available toolchains.
    Logs the command usage.

    Parameters
    ----------
    mpsd_release : str
        The name of the MPSD release to prepare the environment for.
    script_dir : pathlib.Path
        The base directory to create the release folder and
        clone the spack-environments repository into.

    Returns
    -------
    available_toolchains : list
        A list of available toolchains for the given MPSD release.
    """
    create_dir_structure(mpsd_release, script_dir)
    spe_branch, spe_commit_hash, available_toolchains = get_release_info(
        mpsd_release, script_dir
    )
    setup_log_cmd(
        mpsd_release, script_dir, spe_branch=spe_branch, spe_commit_hash=spe_commit_hash
    )
    return available_toolchains


def get_native_microarchitecture():
    """Return native microarchitecture.

    On MPSD machines, there should be an environment variable "MPSD_MICROARCH".
    We try to read that. If it fails, we use the 'archspec cpu' command.
    If that fails, we ask the user to install it.

    Returns
    -------
    MPSD_MICROARCH : str

    Example
    -------
    >>> get_native_microarchitecture()
    'haswell'
    """
    # attempt to get MICRO_ARCH from environment variable (should work on
    # MPSD_HPC and MPSD linux laptops). If not defined, return
    # "UNKNOWN_MICROARCH"
    microarch = os.environ.get("MPSD_MICROARCH", "UNKNOWN_MICROARCH")

    # if we have not found the microarchitecture environment variable,
    # try calling archspec
    if microarch == "UNKNOWN_MICROARCH":
        logging.debug(
            "Couldn't find MPSD_MICROARCH environment variable. Will try archspec."
        )
        try:
            process = run(["archspec", "cpu"], stdout=subprocess.PIPE, text=True)
        except FileNotFoundError as e:
            logging.debug(f"Call of 'archspec cpu' failed: {e=}")
            # Presumably 'archspec' is not installed.
            msg = "Please install archspec, for example via 'pipx install archspec'.\n"
            msg += "The command we need to execute is 'archspec cpu'.\n"
            msg += "Documentation of package: https://archspec.readthedocs.io/"

            logging.error(msg)
            sys.exit(1)
        else:  # we have found archspec and executed it
            if process.returncode == 0:  # sanity check
                microarch = process.stdout.strip()
                logging.debug(
                    f"Found microarchitecture from 'archspec cpu' to be '{microarch}'"
                )
                assert len(microarch) > 0  # sanity check
            else:
                raise ValueError(
                    f"Some error occurred when calling 'archspec cpu': {process=}"
                )

    # at this point, we have determined the microarchitecture
    log_metadata("microarchitecture", microarch)
    return microarch


def install_environment(
    mpsd_release: str,
    toolchains: List[str],
    script_dir: Path,
    force_reinstall: bool = False,
    enable_build_cache: bool = False,
) -> None:
    """
    Install the specified MPSD release and toolchains.

    The function installs the toolchain to the specified directory, using Spack.

    Parameters
    ----------
    mpsd_release : str
        A string representing the MPSD release version.
    toolchains : list of str
        A list of strings representing the toolchains to install
        (e.g., "foss2021a-mpi", "global_generic", "ALL").
    script_dir : pathlib.Path
        A Path object representing the path to the directory where
        the release and toolchains will be installed.
    force_reinstall : bool, optional
        A boolean indicating whether to force a reinstallation
        even if the release and toolchains already exist. Defaults to False.
    enable_build_cache : bool, optional
        A boolean indicating whether to build the build cache
        when installing toolchains. Defaults to False.

    Raises
    ------
    ValueError
        If a requested toolchain is not available in the specified release.

    Returns
    -------
    None
    """
    logging.info(
        f"Installing release {mpsd_release} with toolchains {toolchains} "
        f"to {script_dir}"
    )

    # Set required variables
    release_base_dir = script_dir / mpsd_release

    toolchain_dir = release_base_dir / mpsd_microarch
    toolchain_dir.mkdir(parents=True, exist_ok=True)
    spack_setup_script = release_base_dir / "spack-environments" / "spack_setup.sh"
    install_flags = []
    if not enable_build_cache:
        install_flags.append("-b")

    # run the prepare_environment function
    available_toolchains = prepare_environment(mpsd_release, script_dir)
    # Ensure that the requested toolchains are available in the release
    if toolchains == "ALL":
        toolchains = available_toolchains
    elif toolchains == "NONE":
        # No toolchains requested, so we only create the env and print the
        # list of available toolchains
        logging.warning(
            "No toolchains requested. Available toolchains for release "
            f"{mpsd_release} are: \n {available_toolchains}"
        )
        return

    for toolchain in toolchains:
        if toolchain not in available_toolchains:
            # TODO: add to message how toolchains can be found
            msg = f"Toolchain '{toolchain}' is not available in release {mpsd_release}."
            logging.error(msg)
            sys.exit(1)

    # Install the toolchains
    with os_chdir(toolchain_dir):
        # run spack_setup_script with the toolchains as arguments
        # if the log folder doesn't exist, create it
        if not os.path.exists("logs"):
            os.mkdir("logs")
        for toolchain in toolchains:
            # Set the install log file name to config_vars["install_log_file"]
            # and replace _toolchains_ with the toolchain name and
            # _mpsd_spack_ver_ with mpsd_release

            logging.info(f"Installing toolchain {toolchain} to {toolchain_dir}")

            # log the command
            setup_log_cmd(
                mpsd_release,
                script_dir,
                msg=f"installing {toolchain} and logging at {install_log_file}",
            )
            setup_log_cmd(
                mpsd_release,
                script_dir,
                msg=(
                    f"CMD: bash {spack_setup_script} {' '.join(install_flags)}"
                    "{toolchain}"
                ),
            )
            run(
                f"bash {spack_setup_script} {' '.join(install_flags)} {toolchain} 2>&1 "
                f"| tee -a {install_log_file} ",
                shell=True,
                check=True,
            )


def remove_environment(release, toolchains, target_dir):
    """Remove release from installation."""
    msg = f"Removing release {release} with toolchains {toolchains} from {target_dir}"
    logging.info(msg)
    raise NotImplementedError(msg)


def start_new_environment(release, from_release, target_dir):
    """Start new MPSD software environment version."""
    msg = f"Starting new release {release} from {from_release} to {target_dir}"
    logging.info(msg)
    raise NotImplementedError(msg)


def main():
    """Execute main entry point."""
    parser = argparse.ArgumentParser(description=about_tool)
    parser.add_argument(
        "--log",
        "-l",
        dest="loglevel",
        choices=["warning", "info", "debug"],
        required=False,
        default="warning",
        help="Set the log level",
    )

    subparsers = parser.add_subparsers(
        dest="action", title="actions", description="valid actions", required=True
    )
    subparsers.required = True
    list_of_cmds = [
        ("prepare", "Prepare the environment for installation on the disk"),
        ("install", "Install a software environment"),
        ("reinstall", "Reinstall a software environment"),
        ("remove", "Remove a software environment or toolchains from an environment"),
        ("start-new", "Start a new software environment version"),
    ]
    for cmd, help_text in list_of_cmds:
        subp = subparsers.add_parser(cmd, help=help_text)

        if cmd == "start-new":
            subp.add_argument(
                "--from-release",
                dest="from_release",
                type=str,
                required=True,
                help="Release version to start from",
            )
            subp.add_argument(
                "--to-release",
                dest="to_release",
                type=str,
                required=True,
                help="Release version to create",
            )

        else:
            subp.add_argument(
                "release",
                type=str,
                help="Release version to prepare, install, reinstall or remove",
            )
            if cmd in ["install", "reinstall", "remove"]:
                # "install" command needs additional documentation
                tool_chain_help = (
                    f"Pass a list of toolchains to command {cmd}. "
                    "Use '--toolchains ALL' to "
                    f"{cmd} all toolchains. If '--toolchain' is not "
                    "specified, list available toolchains for the release "
                    "(after environment has been prepared if not done yet)."
                )

                subp.add_argument(
                    "--toolchains",  # first option defines attribute
                    # name `args.toolchains` in `args = parser_args()`
                    "--toolchain",  # allow singular as alternative
                    # (-> creates attribute `args.toolchains` if used)
                    type=str,
                    dest="toolchains",
                    nargs="+",
                    default="NONE",
                    help=tool_chain_help,
                )
                subp.add_argument(
                    "--enable-build-cache",
                    action="store_true",
                    help=(
                        "Enable Spack build cache. Useful for reinstallation but "
                        "consumes time and disk space."
                    ),
                )
    # Carry out the action
    args = parser.parse_args()

    # Get machine configs
    os.environ.get("MPSD_OS", "UNKNOWN_OS")
    mpsd_microarch = os.environ.get("MPSD_MICROARCH", "UNKNOWN_MICROARCH")
    # release `dev` in script_dir/dev-23a
    script_dir = Path(os.path.dirname(os.path.realpath(__file__)))
    mpsd_release = args.release
    # parse logging first
    # decide the log_file_name
    installer_log_name, build_log_name = create_log_file_names(
        mpsd_release=mpsd_release, mpsd_microarch=mpsd_microarch, action=args.action
    )
    installer_log_file = (
        script_dir / mpsd_release / mpsd_microarch / "logs" / installer_log_name
    )

    set_up_logging(args.loglevel, installer_log_file)

    # target dir is the place where this script exists. the

    # Check the command and run related function
    if args.action == "remove":
        remove_environment(args.release, args.toolchains, script_dir)
    elif args.action == "start-new":
        start_new_environment(args.from_release, args.to_release, script_dir)
    elif args.action == "install":
        install_environment(
            args.release, args.toolchains, script_dir, False, args.enable_build_cache
        )
    elif args.action == "prepare":
        prepare_environment(args.release, script_dir)


if __name__ == "__main__":
    main()