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": "logs/install-software-environment.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",
}
shared_var = {
    "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(),
    "spe_branch": None,
    "spe_commit_hash": None,
    "available_toolchains": None,
    "toolchain_base_dir": None,
}


# 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(shared_var, msg=None):
    # Create log directory if it doesn't exist
    if not os.path.exists("logs"):
        os.makedirs("logs")
    # 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:
            f.write(msg + "\n")
        else:
            f.write("-" * 50 + "\n")
            cmd_line = " ".join(sys.argv)
            f.write(f"{datetime.datetime.now().isoformat()}, {cmd_line}\n")
            f.write(
                f"Software environment installer branch: {shared_var['script_branch']} (commit hash: {shared_var['script_commit_hash']})\n"
            )
            f.write(
                f"Spack environments branch: {shared_var['spe_branch']} (commit hash: {shared_var['spe_commit_hash']})\n"
            )


def prepare_environment(mpsd_release, script_dir):
    release_base_dir = os.path.join(script_dir, mpsd_release)
    if not os.path.exists(release_base_dir):
        os.makedirs(release_base_dir)
    else:
        # Warn that the target directory already exists.
        print(">Target directory already exists. Continuing...")
    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"])
            subprocess.run(["git", "checkout", mpsd_release])
            subprocess.run(["git", "pull"])

            # Get the branch and commit hash of the spack-environments repo and store them in shared_var
            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")
    setup_log_cmd(shared_var)
    return shared_var


def install_environment(
    release, toolchains, target_dir, force_reinstall, skip_build_cache
):
    print(f"Installing release {release} with toolchains {toolchains} to {target_dir}")
    shared_var = prepare_environment(target_dir, release, force_reinstall, shared_var)


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
    target_dir = os.path.dirname(os.path.realpath(__file__))

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


if __name__ == "__main__":
    main()