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.

This function builds toolchains for toolchains at the appropriate directory for the current system architecture
and MPSD software stack version.
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",
    "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):
    # 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)
    with open(config_vars['cmd_log_file'], "a") as f:
        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"
        )


class builder:
    def __init__(
        self, release, cmd, toolchain_list=None, target_dir=None, skip_build_cache=False
    ) -> None:
        # Variables taken from cli arguments
        self.mpsd_os = os.getenv("MPSD_OS", "UNKNOWN_OS")
        self.mpsd_microarch = os.getenv("MPSD_MICROARCH", "UNKNOWN_MICROARCH")
        self.mpsd_spack_ver = release
        self.toolchain_list = toolchain_list
        if target_dir == "DEFAULT" or target_dir is None:
            target_dir = (
                f"/opt_mpsd/{self.mpsd_os}/{self.mpsd_spack_ver}/{self.mpsd_microarch}"
            )
        self.toolchain_base_dir = Path(target_dir)
        self.skip_build_cache = skip_build_cache
        self.skip_dir_check = False
        self.current_dir = os.getcwd()
        self.run_mode = cmd
        # Spack environments related variables
        self.spe_branch = None  # To be set after cloning the spack-environments repo
        self.spe_commit_hash = (
            None  # To be set after cloning the spack-environments repo
        )
        # script related variables
        self.script_branch = (
            subprocess.run(
                ["git", "rev-parse", "--abbrev-ref", "HEAD"], stdout=subprocess.PIPE
            )
            .stdout.decode()
            .strip()
        )
        self.script_commit_hash = (
            subprocess.run(
                ["git", "rev-parse", "--short", "HEAD"], stdout=subprocess.PIPE
            )
            .stdout.decode()
            .strip()
        )
        # Two log files, one for logging commands and the other for logging the build process
        self.cmd_log_file = "logs/install-software-environment.log"
        self.build_log_file = f"build_toolchains_{self.mpsd_spack_ver}_{time.strftime('%Y%m%d-%H%M%S')}.log"

        # run the script
        # self.run() # better to call this seperately after creating the object


    def run(self):
        if self.run_mode == "remove":
            pass
            # self.remove_toolchains()
        elif self.run_mode == "start_new":
            pass
            # self.start_new_env()
        elif self.run_mode == "install":
            pass
            # self.build_toolchains()
        else:
            self.prepare_env()

    def prepare_env(self):
        # ensure required arguments are not None
        if self.toolchain_base_dir is None:
            raise ValueError("Error: Please specify the target directory")
        if self.mpsd_spack_ver is None:
            raise ValueError("Error: Please specify the release version to install")
        if self.toolchain_base_dir is None:
            raise ValueError(
                "Error: Please specify the target directory for installation"
            )
        if not os.path.exists(self.toolchain_base_dir):
            os.makedirs(self.toolchain_base_dir)
        else:
            if not self.skip_dir_check:
                raise ValueError(
                    f"Error: Target directory {self.toolchain_base_dir} already exists. \n\
                    Please remove it and try again."
                )
        with os_chdir(self.toolchain_base_dir):
            subprocess.run(
                [
                    "git",
                    "clone",
                    "https://gitlab.gwdg.de/mpsd-cs/spack-environments.git",
                ]
            )
            with os_chdir("spack-environments"):
                subprocess.run(["git", "checkout", self.mpsd_spack_ver])
                self.spe_branch = subprocess.run(
                    ["git", "rev-parse", "--abbrev-ref", "HEAD"],
                    stdout=subprocess.PIPE,
                ).stdout.decode()
                self.spe_commit_hash = subprocess.run(
                    ["git", "rev-parse", "--short", "HEAD"], stdout=subprocess.PIPE
                ).stdout.decode()
                # keep the available toolchains in current spack-environments branch
                self.available_toolchains = os.listdir("toolchains")

        # Setup the log file for commands
        self.setup_log_cmd()

    def build_toolchains(
        self,
        release,
        install,
        target_dir,
        force_reinstall,
        skip_build_cache,
        skip_dir_check,
    ):
        # ensure required arguments are not None

        if self.toolchain_list is None:
            raise ValueError("Error: Please specify the toolchains to install")
        self.prepare_env()

    def start_new_env(set_up, from_release):
        pass

    def remove_toolchain(release, remove):
        pass


def main():
    parser = argparse.ArgumentParser(description=about_tool)
    parser.add_argument(
        "--release", type=str, help="Specify the release version to install"
    )
    parser.add_argument(
        "--target-directory",
        type=str,
        help="Specify the target directory for installation (use DEFAULT to use /opt_mpsd/<MPSD_OS>/<MPSD_RELEASE>/<MPSD_MICROARCH)",
    )
    parser.add_argument(
        "--install",
        nargs="+",
        help="Specify toolchain(s) to install eg foss2021a-mpi (use ALL to install all available toolchains)",
    )
    parser.add_argument("--remove", type=str, help="Specify toolchain to remove")
    parser.add_argument(
        "--set-up",
        type=str,
        help="Start a new software environment version, must specify --from <release>",
    )
    parser.add_argument(
        "--from",
        dest="from_release",
        type=str,
        help="Specify the release version to start from",
    )
    parser.add_argument(
        "--force-reinstall",
        action="store_true",
        help="Delete and reinstall an existing toolchain directory",
    )
    parser.add_argument(
        "--skip-build-cache",
        action="store_true",
        help="Skip Spack build cache during installation",
    )
    parser.add_argument(
        "--skip-dir-check",
        action="store_true",
        help="Skip checking if the target directory already exists",
    )

    args = parser.parse_args()

    if args.release is None:
        parser.print_help()
        sys.exit(1)

    if args.remove:
        cmd = "remove"
    elif args.set_up:
        cmd = "start_new"
    else:
        if args.install:
            cmd = "install"
        else:
            cmd = "prepare_env"
    run = builder(
        release=args.release,
        cmd=cmd,
        toolchain_list=args.install,
        target_dir=args.target_directory,
        # force_reinstall=args.force_reinstall,
        skip_build_cache=args.skip_build_cache,
    )
    run.run()


if __name__ == "__main__":
    main()