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