install-mpsd-software-environment.py

#!/usr/bin/env python3
import os
import subprocess
import time
import datetime
import argparse
import sys
from pathlib import Path

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": f"build_toolchains_mpsd_spack_ver_{time.strftime('%Y%m%d-%H%M%S')}.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:
    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, script_dir, msg=None, *args, **kwargs):
    release_base_dir = script_dir / mpsd_release

    with os_chdir(release_base_dir):
        # Write to the log file with the following format
        # --------------------------------------------------
        # 2023-02-29T23:32:01, install-software-environment.py --release dev-23a --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 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
                script_branch = (
                    subprocess.run(["git", "rev-parse", "--abbrev-ref", "HEAD"], stdout=subprocess.PIPE)
                    .stdout.decode()
                    .strip()
                )
                script_commit_hash = (
                    subprocess.run(["git", "rev-parse", "--short", "HEAD"], stdout=subprocess.PIPE).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} (commit hash: {script_commit_hash})\n")
                f.write(f"Spack environments branch: {spe_branch} (commit hash: {spe_commit_hash})\n")


def create_dir_structure(mpsd_release, script_dir):
    # 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"],
                ]
            )
        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"])
            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"])


def get_release_info(mpsd_release, script_dir):
    # 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).stdout.decode().strip()
            )
            spe_branch = (
                subprocess.run(["git", "rev-parse", "--abbrev-ref", "HEAD"], stdout=subprocess.PIPE)
                .stdout.decode()
                .strip()
            )
            available_toolchains = os.listdir("toolchains")
    return spe_branch, spe_commit_hash, available_toolchains


def prepare_environment(mpsd_release, script_dir):
    # create the release folder and clone the spack-environments repo
    # get the branch and commit hash of the spack-environments repo and available toolchains
    # log the command usage.

    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, toolchains, script_dir, force_reinstall, skip_build_cache):
    print(f"Installing release {mpsd_release} with toolchains {toolchains} to {script_dir}")

    # Set required variables
    release_base_dir = script_dir / mpsd_release
    mpsd_os = os.environ.get("MPSD_OS", "UNKNOWN_OS")
    mpsd_microarch = os.environ.get("MPSD_MICROARCH", "UNKNOWN_MICROARCH")
    toolchain_dir = release_base_dir / mpsd_microarch
    spack_setup_script = release_base_dir / "spack-environments" / "spack_setup.sh"
    install_flags = []
    if skip_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
    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
        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
            install_log_file = (
                config_vars["install_log_file"]
                .replace("_toolchain_", f"_{toolchain}_")
                .replace("_mpsd_spack_ver_", f"_{mpsd_release}_")
            )
            subprocess.run(
                ["bash", spack_setup_script, *install_flags, toolchain, f" | tee -a {install_log_file} 2>&1"],
            )


def remove_environment(release, toolchains, target_dir):
    print(f"Removing release {release} with toolchains {toolchains} from {target_dir}")


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


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 install or remove")
            if cmd in ["install", "reinstall"]:
                subp.add_argument(
                    "--toolchains",
                    type=str,
                    nargs="*",
                    default="ALL",
                    help="List of toolchains to install (use ALL to install all toolchains)",
                )
                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.skip_build_cache)
    elif args.action == "prepare":
        prepare_environment(args.release, script_dir)


if __name__ == "__main__":
    main()