mpsd_software.py

        A string representing the MPSD release version.
    package_sets : list of str
        A list of strings representing the package_sets to install
        (e.g., "foss2021a-mpi", "global_generic", "ALL").
    root_dir : pathlib.Path
        A Path object representing the path to the directory where
        the release and package_sets will be installed.
    enable_build_cache : bool, optional
        A boolean indicating whether to build the build cache
        when installing package_sets. Defaults to False.

    Raises
    ------
    ValueError
        If a requested package_set is not available in the specified release.

    Returns
    -------
    None
    """
    logging.info(
        f"Installing release {mpsd_release} with package_sets {package_sets} "
        f"to {root_dir}"
    )

    # Set required variables
    release_base_dir = root_dir / mpsd_release
    microarch = get_native_microarchitecture()
    package_set_dir = release_base_dir / microarch
    package_set_dir.mkdir(parents=True, exist_ok=True)
    spack_setup_script = release_base_dir / "spack-environments" / "spack_setup.sh"
    install_flags = []
    if not enable_build_cache:
        install_flags.append("-b")

    # run the prepare_environment function
    available_package_sets = prepare_environment(mpsd_release, root_dir)
    # Ensure that the requested package_sets are available in the release
    if package_sets == "ALL":
        package_sets = available_package_sets
    elif package_sets == "NONE":
        # No package_sets requested, so we only create the env and print the
        # list of available package_sets
        logging.warning(
            "No package_sets requested. Available package_sets for release "
            f"{mpsd_release} are: \n {available_package_sets}"
        )
        print_log = logging.getLogger("print")
        print_log.info(f"{available_package_sets=}")
        return

    for package_set in package_sets:
        if package_set not in available_package_sets:
            msg = f"Package_Set '{package_set}' is not available"
            msg += f" in release {mpsd_release}. "
            msg += "Use 'available' command to see list of available package_sets."
            logging.error(msg)
            sys.exit(20)

    # Install the package_sets
    with os_chdir(package_set_dir):
        # run spack_setup_script with the package_sets as arguments
        for package_set in package_sets:
            # Set the install log file name from create_log_file_names
            build_log_path = get_log_file_path(
                mpsd_release, "install", root_dir, package_set
            )

            # Log the build_log_path and the package_set_dir
            logging.info(f"Installing package_set {package_set} to {package_set_dir}")
            logging.info(f"> Logging installation of {package_set} at {build_log_path}")

            run(
                f"bash {spack_setup_script} "
                f"{' '.join(install_flags)} {package_set} 2>&1 "
                f"| tee -a {build_log_path} ",
                shell=True,
                check=True,
            )


def remove_environment(mpsd_release, root_dir, package_sets="NONE", force_remove=False):
    """Remove release from installation.

    Handle 3 situations :
    1. remove does not specify what to remove
        -> warn and exit
    2. remove all package_sets from release
        -> remove release folder except logs
    3. remove specific package_sets from release
        -> remove spack environments via spack commands

    Parameters
    ----------
    mpsd_release : str
        A string representing the MPSD release version.
    root_dir : pathlib.Path
        A Path object representing the path to the directory where
        the release and package_sets will be installed.
    package_sets : list of str
        A list of strings representing the package_sets to remove
        (e.g., "foss2021a-mpi", "global_generic", "ALL").
    force_remove : bool, optional
        A boolean indicating whether to force remove the release.
        If False, the user will be prompted to confirm the removal.
        Defaults to False.

    Raises
    ------
    ValueError

    """
    msg = (
        f"Removing release {mpsd_release}"
        f" with package_sets {package_sets} from {root_dir}"
    )
    logging.warning(msg)

    if package_sets == "NONE":
        logging.warning(
            "Please specify package_sets to remove, or 'ALL' to remove all package_sets"
        )
        sys.exit(50)
    # 2nd case: remove the entire release for microarchitecture
    dir_to_remove = root_dir / mpsd_release / get_native_microarchitecture()
    if "ALL" in package_sets:
        # we need to remove the entire release folder
        logging.warning(
            f"Removing release {mpsd_release}"
            f"from {root_dir} for {get_native_microarchitecture()}"
        )
        if not force_remove:
            logging.warning("do you want to continue? [y/n]")
            if input().lower() != "y":
                sys.exit(60)

            # Set the remove log file name from create_log_file_names
        build_log_path = get_log_file_path(mpsd_release, "remove", root_dir, "ALL")

        logging.info(f"> Logging removal of {mpsd_release} at {build_log_path}")
        folders_to_remove = os.listdir(dir_to_remove)
        for folder in folders_to_remove:
            # shutil.rmtree(dir_to_remove / folder) #dosent delete file
            run(
                f"rm -rf {dir_to_remove / folder} 2>&1 | tee -a {build_log_path}",
                shell=True,
                check=True,
            )
        logging.warning(f"Removed release {mpsd_release} from {root_dir}")
        return
    # 3rd case: remove specific package_sets from release
    for package_set in package_sets:
        # we load the spack environment and remove the package_set
        build_log_path = get_log_file_path(
            mpsd_release, "remove", root_dir, package_set
        )
        logging.info(f"> Logging removal of {package_set} at {build_log_path}")
        if package_set not in ["global_generic", "global"]:
            remove_spack_environment(
                dir_to_remove / "spack", package_set, build_log_path
            )
        else:
            # list all specs from the global_packages.list
            spe_folder = root_dir / mpsd_release / "spack-environments"
            package_list_file = (
                spe_folder / "toolchains" / package_set / "global_packages.list"
            )
            with open(package_list_file, "r") as f:
                package_dump = f.read()

            # remove all content from # to the end of the line
            package_dump = re.sub(r"#.*\n", "\n", package_dump)
            # replace \\n with "" to remove line breaks
            package_list = package_dump.replace("\\\n", "").split("\n")
            # remove all empty lines
            package_list = [line for line in package_list if line != ""]

            # remove all packages in package_list
            for package in package_list:
                logging.info(f"Removing package {package} from installation")
                remove_spack_package(dir_to_remove / "spack", package, build_log_path)


def remove_spack_environment(spack_dir, environment_name, build_log_path=None):
    """Remove spack environment including packages exclusive to it.

    First activate the environment,
    then uninstall all packages exclusive to the environment,
    then deactivate the environment,
    remove the environment,
    and finally remove the environment lua file.

    Parameters
    ----------
    spack_dir : pathlib.Path
        A Path object representing the path to the spack directory.
    environment_name : str
        A string representing the name of the spack environment to remove.
    build_log_path : pathlib.Path, optional
        A Path object representing the path to where the logs will be teed
    """
    logging.warning(f"Removing spack environment {environment_name}")
    spack_env = spack_dir / "share" / "spack" / "setup-env.sh"
    commands_to_execute = [
        f"export SPACK_ROOT={spack_dir}",  # need to set SPACK_ROOT in dash and sh
        f". {spack_env}",
        f"spack env activate {environment_name}",
        f"for spec in $(spack -e {environment_name} find"  # this line continues
        r' --format "{name}@{version}%{compiler.name}@{compiler.version}");do'
        " spack uninstall -y $spec; done",
        "spack env deactivate",
        f"spack env remove -y {environment_name}",
    ]
    build_log_path = build_log_path or "/dev/null"
    run(
        "(" + " && ".join(commands_to_execute) + f") 2>&1 |tee -a {build_log_path}",
        shell=True,
        check=True,
    )
    # remove the environment lua file
    lua_file = (
        spack_dir / ".." / "lmod" / "Core" / "toolchains" / f"{environment_name}.lua"
    )
    run(f"rm {lua_file}", shell=True, check=True)


def remove_spack_package(spack_dir, package, build_log_path=None):
    """Remove spack package.

    Used to remove global packages.

    Parameters
    ----------
    spack_dir : pathlib.Path
        A Path object representing the path to the spack directory.
    package : str
        A string representing the name of the spack package to remove.
    build_log_path : pathlib.Path, optional
        A Path object representing the path to where the logs will be teed

    """
    logging.info(f"Removing spack package {package}")
    spack_env = spack_dir / "share" / "spack" / "setup-env.sh"
    commands_to_execute = [
        f"export SPACK_ROOT={spack_dir}",  # need to set SPACK_ROOT in dash and sh
        f". {spack_env}",
        f"spack uninstall -y {package}",
    ]
    run(
        "(" + " && ".join(commands_to_execute) + f") 2>&1 |tee -a {build_log_path}",
        shell=True,
        check=True,
    )


def start_new_environment(release, from_release, target_dir):
    """Start new MPSD software environment version."""
    msg = f"Starting new release {release} from {from_release} to {target_dir}"
    logging.info(msg)
    raise NotImplementedError(msg)


def list_installed_releases(root_dir: Path, print_output: bool = False) -> List[str]:
    """
    List installed releases.

    Parameters
    ----------
    root_dir : pathlib.Path

    Returns
    -------
    None
    """
    plog = logging.getLogger("print")
    list_of_files = os.listdir(root_dir)
    installed_releases = [
        x for x in list_of_files if (root_dir / x / "spack-environments").exists()
    ]
    if print_output:
        plog.info("Available MPSD software releases:")
        for release in installed_releases:
            plog.info(f"    {release}")
    return installed_releases


def list_installed_toolchains(
    mpsd_release: str, root_dir: Path, print_output: bool = False
) -> Union[dict, None]:
    """
    List installed toolchains.

    Parameters
    ----------
    mpsd_release : str
        A string representing the MPSD release version.
    root_dir : pathlib.Path
        A Path object pointing to the root directory of the installation.
        Expect a subfolder root/mpsd_release in which we search for the
        toolchains.
    print_output : bool, optional
        A boolean indicating whether to print the output to the terminal.

    Returns
    -------
    toolchain_map : dict
        A dictionary containing available microarchitectures as keys and
        a list of available package_sets as values for each microarchitecture.
        If the release is not installed/found, None is returned.

        Note: only toolchains can be reported at the moment (i.e. package_sets
        such as global and global_generic are missing, even if installed).
    """
    plog = logging.getLogger("print")
    release_base_dir = root_dir / mpsd_release
    microarch = get_native_microarchitecture()
    toolchain_dir = release_base_dir / microarch
    spack_dir = toolchain_dir / "spack"

    # if the mpsd_release does not exist:
    if not release_base_dir.exists():
        logging.debug(f"Directory {str(release_base_dir)} does not exist.")
        logging.error(f"MPSD release '{mpsd_release}' is not installed.")
        return None

    # if the mpds_release directory exists but the spack repository is not fully
    # cloned - indicates some kind of incomplete installation:
    if not spack_dir.exists():
        logging.info(f"Could not find directory {spack_dir}.")
        logging.error(
            f"MPSD release '{mpsd_release}' has not been completely installed."
        )

        return None

    # find all folders for all microarch in the release directory
    # except for the blacklisted files
    black_listed_files = [
        config_vars["cmd_log_file"],
        "spack-environments",
        "logs",
        "mpsd-spack-cache",
    ]

    list_of_microarchs_candidates = os.listdir(release_base_dir)
    list_of_microarchs = [
        x for x in list_of_microarchs_candidates if x not in black_listed_files
    ]
    logging.debug(f"{list_of_microarchs=}")

    toolchain_map = {}
    for microarch in list_of_microarchs:
        # get a list of all the toolchains in the microarch
        possible_toolchains = (release_base_dir / microarch).glob(
            "lmod/Core/toolchains/*.lua"
        )
        # append toolchain which is the name of the file without the .lua extension
        toolchain_map[microarch] = [toolchain.stem for toolchain in possible_toolchains]

    logging.debug(f"{toolchain_map=}")

    # pretty print the toolchain map key as the heading
    # and the value as the list of toolchains
    if print_output:
        plog.info(f"Installed toolchains ({mpsd_release}):\n")
        for microarch, toolchains in toolchain_map.items():
            plog.info(f"- {microarch}")
            for toolchain in toolchains:
                plog.info(f"    {toolchain}")
            plog.info(
                f"    [module use {str(release_base_dir / microarch / 'lmod/Core')}]"
            )
            plog.info("")
    return toolchain_map


def pretty_print_spec(spec: str) -> None:
    """
    Print the specs with colours using rich.

    - packages in white (everything until first %)
    - compiler in green (everything between % and first+)
    - variants in cyan (everything that starts with +)
    - build_system in yellow (everything that starts with build_system=)
    - architecture in purple (everything that starts with arch=)
    """
    # Note that this implementation necessitates the definition of
    # flags in the order in which we ask spack to format the output
    # also for flags that need the same colour because they are
    # interchangeable (like `+` and `~`) we need to define them together
    colour_map = {
        "%": "green",
        "+": "cyan",
        "~": "cyan",
        "build_system=": "yellow",
        "libs=": "blue",
        "arch=": "purple",
    }

    prev_colour = ""

    for flag in colour_map.keys():
        # If the flag is in the spec string,
        # replace it with: previous closing colour, new colour, flag
        if flag in spec:
            if (
                colour_map[flag] not in prev_colour
            ):  # avoid duplicates for eg when having both ~ and +
                spec = spec.replace(flag, f"{prev_colour}[{colour_map[flag]}]{flag}", 1)
                prev_colour = f"[/{colour_map[flag]}]"  # for next iter

    # Add the final closing tag to the spec string
    spec += prev_colour
    rprint(spec)


def list_installed_packages(
    mpsd_release: str, root_dir: Path, package_set: str, microarch: str
) -> Union[List[str], None]:
    """
    List installed packages and their specs.

    Uses `spack -e package_set find` to list the installed packages,
    in the following format
    "{name}{@versions}{%compiler.name}{@compiler.versions}{compiler_flags}{variants}{arch=architecture}"

    Parameters
    ----------
    mpsd_release : str
        A string representing the MPSD release version.
    root_dir : pathlib.Path
        A Path object pointing to the root directory of the installation.
        Expect a subfolder root/mpsd_release in which we search for the
        toolchains.
    package_set : str
        A string representing the package_sets to show the packages for.
    microarch : str
        A string representing the microarchitecture to show the packages for.

    Returns
    -------
    list
        A list of strings representing the packages installed for the
        specified package_sets and microarch.
        If the release is not installed/found, None is returned.

    """
    plog = logging.getLogger("print")
    plog.info(f"listing packages installed for {package_set=}, {microarch=}")
    spack_dir = root_dir / mpsd_release / microarch / "spack"
    spack_env = spack_dir / "share" / "spack" / "setup-env.sh"
    commands_to_execute = [
        f"export SPACK_ROOT={spack_dir}",  # need to set SPACK_ROOT in dash and sh
        f". {spack_env}",
        f"spack -e {package_set}"
        " find --format "
        r"{name}{@versions}{%compiler.name}{@compiler.versions}{compiler_flags}{variants}{arch=architecture}",
    ]
    process = run(
        " && ".join(commands_to_execute), shell=True, check=True, capture_output=True
    )
    package_list = process.stdout.decode().strip().split("\n")
    for package in package_list:
        pretty_print_spec(package)
    return package_list


def environment_status(
    mpsd_release: str, root_dir: Path, package_set="NONE"
) -> Union[dict, List[str], None]:
    """Show status of release in installation.

    - 1) If no mpsd_release, list available releases
    - 2) If mpsd_release, list available toolchains
    - 3) If mpsd_release and toolchain, list available packages

    Parameters
    ----------
    mpsd_release : str
        A string representing the MPSD release version.
    root_dir : pathlib.Path
        A Path object pointing to the root directory of the installation.
        Expect a subfolder root/mpsd_release in which we search for the
        toolchains.
    package_set : str, optional
        A string representing the package_sets to show the status for.

    Returns
    -------
    installed_release : List[str]
        A list of installed (valid) releases.

    OR
    toolchain_map : dict
        A dictionary containing available microarchitectures as keys and
        a list of available package_sets as values for each microarchitecture.
        If the release is not installed/found, None is returned.

        Note: only toolchains can be reported at the moment (i.e. package_sets
        such as global and global_generic are missing, even if installed).

    OR
    package_list : List[str]
        A list of strings representing the packages installed for the
        specified package_sets and microarch.

    """
    msg = f"Showing status of release {mpsd_release} in {root_dir}"
    logging.info(msg)
    if not mpsd_release:
        # 1) if no mpsd_release is specified, list available releases
        return list_installed_releases(root_dir=root_dir, print_output=True)
    # 2) if mpsd_release is specified, list installed toolchains
    # Test is the mpsd_release is valid
    if mpsd_release not in list_installed_releases(root_dir=root_dir):
        logging.error(f"MPSD release '{mpsd_release}' is not available.")
        return None
    if package_set == "NONE":
        return list_installed_toolchains(
            mpsd_release=mpsd_release, root_dir=root_dir, print_output=True
        )
    # 3) if mpsd_release and toolchain is specified, list installed packages
    # check that the package-set is a valid toolchain
    if (
        package_set
        not in list_installed_toolchains(mpsd_release=mpsd_release, root_dir=root_dir)[
            get_native_microarchitecture()
        ]
    ):
        logging.error(f"Package-set '{package_set}' is not available.")
        return None
    return list_installed_packages(
        mpsd_release=mpsd_release,
        root_dir=root_dir,
        package_set=package_set,
        microarch=get_native_microarchitecture(),
    )


def initialise_environment(root_dir: Path) -> None:
    """Initialize the software environment.

    This creates a hidden file ``.mpsd-software-root`` to tag the location for
    as the root of the installation. All compiled files, logs etc are written in
    or below this subdirectory.

    Parameters
    ----------
    root_dir : pathlib.Path
        A Path object pointing to the current directory where the script was called.

    """
    # check if the root_dir is not already initialized
    init_file = root_dir / config_vars["init_file"]
    if init_file.exists():
        logging.error(f"Directory {str(root_dir)} is already initialised.")
        sys.exit(30)
    else:
        # create the init file
        init_file.touch()
        # note the execution in the execution summary log
        # create the log file and fill it with the headers
        record_script_execution_summary(root_dir=root_dir)
        # record the msg in the log file
        init_log_msg = f"Initialising MPSD software instance at {root_dir}.\n"
        init_log_msg += f"MPSD Software manager version: {__version__}\n"
        write_to_cmd_log(
            root_dir=root_dir,
            msg=init_log_msg,
        )


def get_root_dir() -> Path:
    """Get the root directory of the installation.

    Look for the hidden file ``.mpsd-software-root``
    (defined in config_vars["init_file"])
    in the current directory, or any parent directory.
    If found, return the path to the root directory
    of the MPSD software instance.
    If not found, exit with an error message.

    Returns
    -------
    root_dir : pathlib.Path
        A Path object pointing to the root directory of the installation.
        This folder contains the hidden file ``.mpsd-software-root``,
        ``mpsd_releases`` ( for eg ``dev-23a``) and ``mpsd-spack-cache``.


    """
    # check if the root_dir is not already initialized
    script_call_dir = Path.cwd()
    init_file = script_call_dir / config_vars["init_file"]
    if init_file.exists():
        return script_call_dir

    # if not, look for the init file in the parent directories
    for parent_folder in script_call_dir.parents:
        init_file = parent_folder / config_vars["init_file"]
        if init_file.exists():
            script_call_dir = parent_folder
            return script_call_dir

    # if not found in any parent directory, exit with an error message
    logging.debug(f"Directory {str(script_call_dir)} is not a MPSD software instance.")

    logging.error(
        "Could not find MPSD software instance "
        "in the current directory or any parent directory.\n\n"
        f"The current directory is {script_call_dir}.\n\n"
        "To initialise a MPSD software instance here, "
        "run 'mpsd-software init'.\n\n"
        f"To find the root directory of an existing MPSD software instance, look "
        f"for the directory containing '{config_vars['cmd_log_file']}' "
        + f"and the hidden file '{config_vars['init_file']}'."
    )
    sys.exit(40)


def main():
    """Execute main entry point."""
    parser = argparse.ArgumentParser(
        description=about_intro,
        epilog=about_epilog,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "-l",
        dest="loglevel",
        choices=["warning", "info", "debug"],
        required=False,
        default="warning",
        help="Set the log level",
    )

    parser.add_argument("--version", action="version", version=__version__)

    subparsers = parser.add_subparsers(
        dest="action", title="actions", description="valid actions", required=True
    )
    subparsers.required = True
    list_of_cmds = [
        ("init", "Initialise the MPSD software instance in the current directory"),
        ("available", "What is available for installation?"),
        ("install", "Install a software environment"),
        # ("reinstall", "Reinstall a package_set"),
        ("remove", "Remove a package set"),
        # ("start-new", "Start a new MPSD software release version"),
        ("status", "Show status: what is installed?"),
        ("prepare", "Prepare installation of MPSD-release (dev only)"),
    ]
    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:
            # most commands except need a release version
            if cmd in ["install", "prepare", "reinstall", "remove"]:
                subp.add_argument(
                    "release",
                    type=str,
                    help="Release version to prepare, install, reinstall or remove",
                )
            elif cmd in ["available", "status"]:
                # for some commands the release version is optional
                subp.add_argument(
                    "release",
                    type=str,
                    nargs="?",
                    help="Release version to prepare, install, reinstall or remove",
                )

            if cmd in ["install", "reinstall", "remove"]:
                # "install" command needs additional documentation
                package_set_help = (
                    f"One or more package sets (like toolchains) to be {cmd}ed. "
                    "Use 'ALL' to refer to all available package sets."
                )

                subp.add_argument(
                    "package_set",  # first option defines attribute
                    # name `args.package_set` in `args = parser_args()`
                    type=str,
                    nargs="+",
                    default="NONE",
                    help=package_set_help,
                )
                # TODO Move the enable-build-cache flag to only 'install' cmd
                subp.add_argument(
                    "--enable-build-cache",
                    action="store_true",
                    help=(
                        "Enable Spack build cache. Useful for reinstallation but "
                        "consumes time and disk space."
                    ),
                )
            if cmd in ["status"]:
                subp.add_argument(
                    "package_set",
                    type=str,
                    nargs="?",
                    default="NONE",
                    help="Package set to show status for.",
                )
    # Carry out the action
    args = parser.parse_args()

    # Set up logging without file handle:
    # this is used in the init action and for logging the
    # get_root_dir() function
    set_up_logging(args.loglevel)

    # Check if the action is init
    # if so, call the init function and exit
    if args.action == "init":
        initialise_environment(Path(os.getcwd()))
        sys.exit(0)

    # if a release version is specified:
    if args.release:
        # sanity check for common mistakes in command line arguments
        if args.release.endswith("/"):  # happens easily with autocompletion
            args.release = args.release.removesuffix("/")
            logging.warning(f"Removed trailing slash from release: {args.release}")

    # root_dir is the place where this MPSD software instance has its root
    root_dir = get_root_dir()

    # set up logging filename: we record activities that change the installation
    if args.action in ["init", "install", "prepare", "reinstall", "remove"]:
        apex_log_file = get_log_file_path(
            args.release,
            args.action,
            root_dir,
        )
    # some commands do not write any log_files:
    elif args.action in ["available", "status"]:
        apex_log_file = None
    else:
        # sanity check
        raise NotImplementedError(f"Should never happen: unknown {args.action=}")

    set_up_logging(
        args.loglevel,
        apex_log_file,
    )
    if args.action not in ["status", "available"]:
        # record the script execution summary only if
        # the action is one that changes files on disk
        record_script_execution_summary(root_dir, apex_log_file)
    # Check the command and run related function
    if args.action == "remove":
        remove_environment(args.release, root_dir, args.package_set)
    elif args.action == "start-new":
        start_new_environment(args.from_release, args.to_release, root_dir)
    elif args.action == "install":
        install_environment(
            args.release, args.package_set, root_dir, args.enable_build_cache
        )
    elif args.action == "status":
        _ = environment_status(args.release, root_dir, args.package_set)
    elif args.action == "prepare":
        prepare_environment(args.release, root_dir)
    elif args.action == "available":
        if args.release:
            get_available_package_sets(args.release)
        else:
            get_available_releases(print_result=True)
            sys.exit(0)
    else:
        message = (
            f"No known action found ({args.action=}). Should probably never happen."
        )
        logging.error(message)
        raise NotImplementedError(message)


if __name__ == "__main__":
    main()