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 {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")
    setup_log_cmd(shared_var)
    return shared_var

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

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



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 "prepare" action
    parser_prepare = subparsers.add_parser('prepare', help='Prepare the environment for installation on the disk')
    parser_prepare.add_argument('release', type=str, help='Release version to install')

    # Add parser for "install" action
    parser_install = subparsers.add_parser('install', help='Install a software environment')
    parser_install.add_argument('release', type=str, help='Release version to install')
    parser_install.add_argument('--toolchains', type=str, nargs='*', help='List of toolchains to install (use ALL to install all toolchains)', default='ALL')
    parser_install.add_argument('--enable-build-cache', action='store_true', help='Enable Spack build cache. Useful for reinstallation but consumes time and disk space')

    # Add parser for "reinstall" action
    parser_reinstall = subparsers.add_parser('reinstall', help='Reinstall a software environment')
    parser_reinstall.add_argument('release', type=str, help='Release version to install')
    parser_reinstall.add_argument('--toolchains', type=str, nargs='*', help='List of toolchains to install (use ALL to install all toolchains)', default='ALL')
    parser_reinstall.add_argument('--enable-build-cache', action='store_true', help='Enable Spack build cache. Useful for reinstallation but consumes time and disk space')

    # Add parser for "remove" action
    parser_remove = subparsers.add_parser('remove', help='Remove a software environment or toolchains from an environment')
    parser_remove.add_argument('release', type=str, help='Release version to remove')
    parser_remove.add_argument('--toolchains', type=str, nargs='*', help='Toolchains to remove (use ALL to remove all toolchains)')
    parser_remove.add_argument('--enable-build-cache', action='store_true', help='Enable Spack build cache. Useful for reinstallation but consumes time and disk space')

    # Add parser for "start-new" action
    parser_start_new = subparsers.add_parser('start-new', help='Start a new software environment version')
    parser_start_new.add_argument('--from-release', dest='from_release', type=str, required=True, help='Release version to start from')
    parser_start_new.add_argument('--to-release', dest='to_release', type=str, required=True, help='Release version to create')

    
    # Carry out the action
    args = parser.parse_args()
    if args.action == 'remove':
        remove(args.release, args.toolchains,args.target_dir)
    elif args.action == 'start-new':
        start_new(args.from_release, args.to_release, args.target_dir)
    else:
        if args.target_dir:
            target_dir = args.target_dir
        else:
            mpsd_os = os.getenv('MPSD_OS','UNKNOWN_OS')
            mpsd_arch = os.getenv('MPSD_ARCH','UNKNOWN_ARCH')
            if os.path.exists('/opt_mpsd'):
                root_dir = '/opt_mpsd'
            else:
                root_dir = os.path.expanduser('~') + '/mpsd_toolchain'
            target_dir = os.path.join(root_dir, mpsd_os , args.release, mpsd_arch)

        if args.action == 'install':
            install(args.release, args.toolchains, target_dir, args.force_reinstall, args.skip_build_cache)
        else:
            prepare_env(target_dir,args.release, False, shared_var)

if __name__ == "__main__":
    main()