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."""
    msg = (
        f"Removing release {mpsd_release}"
        f" with package_sets {package_sets} from {root_dir}"
    )
    logging.warning(msg)
    # Handle 3 situations :
    # 1. remove dosent 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

    if package_sets == "NONE":
        logging.warning(
            "Please specify package_sets to remove, or 'ALL' to remove all toolchains"
        )
        sys.exit(1)
    # 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(50)  # TODO document this code.

        folders_to_remove = os.listdir(dir_to_remove)
        # skip logs folder
        # if "logs" in folders_to_remove:
        #     folders_to_remove.remove("logs")
        for folder in folders_to_remove:
            shutil.rmtree(dir_to_remove / folder)
        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
        if package_set not in ["global_packages", "global"]:
            remove_spack_environment(dir_to_remove / "spack", package_set)
        else:
            # TODO remove global packages by calling remove_spack_package
            pass


def remove_spack_environment(spack_dir, environment_name):
    """Remove spack environment."""
    logging.warning(f"Removing spack environment {environment_name}")
    spack_env = spack_dir / "share" / "spack" / "setup-env.sh"
    commands_to_execute = [
        f"source {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}",
    ]
    run(" && ".join(commands_to_execute), 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 environment_status(mpsd_release: str, root_dir: Path) -> Union[dict, None]:
    """Show status of release in installation.

    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.

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

    """
    msg = f"Showing status of release {mpsd_release} in {root_dir}"
    logging.info(msg)
    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
    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 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", "status"]:
                subp.add_argument(
                    "release",
                    type=str,
                    help="Release version to prepare, install, reinstall or remove",
                )
            elif cmd in ["available"]:
                # 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,
                )
                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()

    # 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,
    )

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