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",
    "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_env(toolchain_base_dir, mpsd_spack_ver, skip_dir_check, shared_var):

    if not os.path.exists(toolchain_base_dir):
        os.makedirs(toolchain_base_dir)
    else:
        if not skip_dir_check:
            raise ValueError(
                f"Error: Target directory {shared_var['toolchain_base_dir']} already exists. \n\
                Please remove it and try again."
            )
    with os_chdir(toolchain_base_dir):
        subprocess.run(
            [
                "git",
                "clone",
                config_vars['spack_environments_repo'],
            ]
        )
        with os_chdir("spack-environments"):
            subprocess.run(
                [
                    "git",
                    "checkout",
                    mpsd_spack_ver,
                ]
            )

            # Get the branch and commit hash of the spack-environments repo and store them in shared_var
            shared_var['spe_commit_hash'] = (
                subprocess.run(
                    ["git", "rev-parse", "HEAD"], stdout=subprocess.PIPE
                )
                .stdout.decode()
                .strip()
            )
            shared_var['spe_branch'] = (
                subprocess.run(
                    ["git", "rev-parse", "--abbrev-ref", "HEAD"], stdout=subprocess.PIPE
                )
                .stdout.decode()
                .strip()
            )
            shared_var['available_toolchains'] = os.listdir("toolchains")


def main():
    parser = argparse.ArgumentParser(description=about_tool)
    # Add subparsers for different actions
    subparsers = parser.add_subparsers(dest='action', title='actions', description='valid actions', required=True)

    # Add parser for "install" action
    parser_install = subparsers.add_parser('install', help='Install a software environment')
    parser_install.add_argument('--release', type=str, required=True, help='Release version to install')
    parser_install.add_argument('--target-dir', type=str, help='Target directory for installation (use DEFAULT to install to default directory)')
    parser_install.add_argument('--force-reinstall', action='store_true', help='Force reinstall an existing toolchain directory')
    parser_install.add_argument('--skip-build-cache', action='store_true', help='Skip Spack build cache during installation')


if __name__ == "__main__":
    main()