mpsd-software-environment.py

#!/usr/bin/env python3
import argparse
import datetime
import os
import subprocess
import sys
from pathlib import Path
from typing import List, Tuple

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. """
config_vars = {
    "cmd_log_file": "install.log",
    "build_log_file": (
        "logs/mpsd_spack_ver_toolchains_"
        f"{datetime.datetime.now().replace(microsecond=0).isoformat()}.log"
    ),
    # TODO: modify toolchains,mpsd_spack_ver when the variable is available
    "spack_environments_repo": "https://gitlab.gwdg.de/mpsd-cs/spack-environments.git",
}


# Helper class to change directory via context manager
class os_chdir:
    """
    The os_chdir class is a context manager that
    changes the current directory to a specified directory
    and returns to the original directory after execution.
    """

    def __init__(self, new_dir):
        self.new_dir = new_dir
        self.saved_dir = os.getcwd()

    def __enter__(self):
        os.chdir(self.new_dir)

    def __exit__(self, exc_type, exc_val, exc_tb):
        os.chdir(self.saved_dir)


def setup_log_cmd(
    mpsd_release: str, script_dir: str, msg: str = None, *args, **kwargs
) -> None:
    """The setup_log_cmd function logs the command used to build the
    toolchains, along with information about the software environment installer
    branch, the Spack environments branch, and the commit hashes of each. It
    also logs steps taken in install process using the optional message
    argument.

    Args:
        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.

    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 = (
                        subprocess.run(
                            ["git", "rev-parse", "--abbrev-ref", "HEAD"],
                            stdout=subprocess.PIPE,
                            check=True,
                        )
                        .stdout.decode()
                        .strip()
                    )
                    script_commit_hash = (
                        subprocess.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:
    """The create_dir_structure function creates the directory structure for
    the specified release and clones the Spack environments repository if it
    doesn't exist.

    Args:
    - 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"):
            subprocess.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
            subprocess.run(["git", "fetch", "--all"], check=True)
            checkout_result = subprocess.run(["git", "checkout", mpsd_release])
            if checkout_result.returncode != 0:
                raise Exception(
                    "Release branch does not exist in spack-environment repo \n."
                    "Check for typos."
                )
            subprocess.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, such as the branch and commit hash
    of the Spack environments repository and the available toolchains.

    Args:
    - 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:
    - Exception: 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 Exception(
            "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 = (
                subprocess.run(
                    ["git", "rev-parse", "HEAD"], stdout=subprocess.PIPE, check=True
                )
                .stdout.decode()
                .strip()
            )
            spe_branch = (
                subprocess.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]:
    """
    - Creates the directory structure for the given MPSD release and clones the
      spack-environments repository.
    - Determines the branch and commit hash of the spack-environments repository and
      the available toolchains.
    - Logs the command usage.

    Args:
    - 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 install_environment(
    mpsd_release: str,
    toolchains: List[str],
    script_dir: Path,
    force_reinstall: bool = False,
    enable_build_cache: bool = False,
) -> None:
    """Installs the specified MPSD release and toolchains to the specified
    directory using Spack.

    Args:
        mpsd_release: A string representing the MPSD release version.

        toolchains: A list of strings representing the toolchains to install
                    (e.g., "foss2021a-mpi", "global_generic", "ALL").

        script_dir: A Path object representing the path to the directory where
                    the release and toolchains will be installed.

        force_reinstall: A boolean indicating whether to force a reinstallation
                    even if the release and toolchains already exist. Defaults to False.

        enable_build_cache: A boolean indicating whether to build the build
                    cache when installing toolchains. Defaults to False.

    Raises:
        Exception: If a requested toolchain is not available in the specified release.

    Returns:
        None

    """
    print(
        f"Installing release {mpsd_release} with toolchains {toolchains} "
        f"to {script_dir}"
    )

    # Set required variables
    release_base_dir = script_dir / mpsd_release
    os.environ.get("MPSD_OS", "UNKNOWN_OS")
    mpsd_microarch = os.environ.get("MPSD_MICROARCH", "UNKNOWN_MICROARCH")
    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
        print(
            "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:
            raise Exception(
                f"Toolchain '{toolchain}' is not available in release {mpsd_release}."
            )

    # 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

            print(f"Installing toolchain {toolchain} to {toolchain_dir}")
            install_log_file = (
                config_vars["build_log_file"]
                .replace("mpsd_spack_ver_", f"{mpsd_release}_")
                .replace("_toolchains_", f"_{toolchain}_")
            )
            # 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}"
                ),
            )
            subprocess.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):
    msg = f"Removing release {release} with toolchains {toolchains} from {target_dir}"
    print(msg)
    raise NotImplementedError(msg)


def start_new_environment(release, from_release, target_dir):
    msg = f"Starting new release {release} from {from_release} to {target_dir}"
    print(msg)
    raise NotImplementedError(msg)


def main():
    parser = argparse.ArgumentParser(description=about_tool)
    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()

    # target dir is the place where this script exists. the
    # release `dev` in script_dir/dev-23a
    script_dir = Path(os.path.dirname(os.path.realpath(__file__)))

    # 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()