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 "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('--toolchains', type=str, nargs='*', help='List of toolchains to install (use ALL to install all toolchains)', default='ALL')
    parser_install.add_argument('--target-dir', type=str, help='Target directory for installation (use DEFAULT or leave empty to install to default directory)', default='DEFAULT')
    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')

    # 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, required=True, help='Release version to remove')
    parser_remove.add_argument('--toolchains', type=str, nargs='*', help='Toolchains to remove (use ALL to remove all toolchains)')
    parser_install.add_argument('--target-dir', type=str, help='Target directory for removal')

    # 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')
    parser_install.add_argument('--target-dir', type=str, help='Base directory for software environment installation (both source and target have the same base directory)(use DEFAULT or leave empty to install to default directory)', default='DEFAULT')
    
    # 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()