Source code for idmtools_platform_comps.cli.comps

"""idmtools comps cli comands.

Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
"""
# flake8: D301
import glob
import json as json_parser
import os
import sys
from typing import Optional, List
import tabulate
from getpass import getpass
from logging import getLogger
from COMPS.CredentialPrompt import CredentialPrompt
import json as je
from idmtools.assets import AssetCollection
from idmtools_platform_comps.utils.singularity_build import SingularityBuildWorkItem

logger = getLogger(__name__)
user_logger = getLogger('user')


[docs]def add_item(assets: AssetCollection, file: str):
    """
    Add Item.

    Args:
        assets: Assets
        file: File or Directory

    Returns:
        None

    Raises:
        FileNotFoundError if file cannot be found.
    """
    if os.path.isdir(file):
        assets.add_directory(file)
    elif os.path.isfile(file):
        assets.add_asset(file)
    else:
        user_logger.error(f"Cannot find file {file}")
        raise FileNotFoundError(f"Cannot find file {file}")


try:

[docs]    class StaticCredentialPrompt(CredentialPrompt):
        """Provides a class to allow login to comps from a username password that is static or provided."""
[docs]        def __init__(self, comps_url, username, password):
            """Constructor."""
            if (comps_url is None) or (username is None) or (password is None):
                raise RuntimeError('Missing comps_url, or username or password')
            self._times_prompted = 0
            self.comps_url = comps_url
            self.username = username
            self.password = password

[docs]        def prompt(self):
            """Return our stores username and password."""
            self._times_prompted = self._times_prompted + 1
            if self._times_prompted > 3:
                raise PermissionError('Failure authenticating')
            return {'Username': self.username, 'Password': self.password}

    os.environ['IDMTOOLS_NO_CONFIG_WARNING'] = '1'
    from idmtools.core.platform_factory import Platform
    import click
    from idmtools_platform_comps.utils.assetize_output.assetize_output import AssetizeOutput
    from idmtools_platform_comps.utils.file_filter_workitem import DEFAULT_EXCLUDES
    from idmtools_platform_comps.comps_platform import COMPSPlatform

    @click.group(short_help="COMPS platform related commands.")
    @click.argument('config-block')
    @click.pass_context
    def comps(ctx: click.Context, config_block):
        """
        Commands related to managing the COMPS platform.

        CONFIG_BLOCK - Name of configuration section or alias to load COMPS connection information from
        """
        ctx.obj = dict(config_block=config_block)

    @comps.command(help="Login to COMPS")
    @click.option('--username', required=True, help="Username")
    @click.option('--password', help="Password")
    @click.pass_context
    def login(ctx: click.Context, username, password):  # noqa D103
        from COMPS import Client
        from idmtools.core.logging import SUCCESS
        os.environ['IDMTOOLS_LOGGING_USER_OUTPUT'] = '0'
        if password:
            user_logger.warning("Password the password via the command line is considered insecure")
        else:
            password = getpass("Password:")
        # make platform object to load info from alias or config but don't login
        platform = Platform(ctx.obj['config_block'], _skip_login=True)

        try:
            Client.login(platform.endpoint,
                         StaticCredentialPrompt(comps_url=platform.endpoint, username=username, password=password))
            user_logger.log(SUCCESS, "Login succeeded")
        except PermissionError:
            user_logger.error(f"Could not loging to {platform.endpoint}")
            sys.exit(-1)

    @comps.command(help="Allows Downloading outputs from the command line")
    @click.option('--pattern', default=[], multiple=True, help="File patterns")
    @click.option('--exclude-pattern', default=DEFAULT_EXCLUDES, multiple=True, help="File patterns")
    @click.option('--experiment', default=[], multiple=True, help="Experiment ids to filter for files to download")
    @click.option('--simulation', default=[], multiple=True, help="Simulation ids to filter for files to download")
    @click.option('--work-item', default=[], multiple=True, help="WorkItems ids to filter for files to download")
    @click.option('--asset-collection', default=[], multiple=True,
                  help="Asset Collection ids to filter for files to download")
    @click.option('--dry-run/--no-dry-run', default=False,
                  help="Gather a list of files that would be downloaded instead of actually downloading")
    @click.option('--wait/--no-wait', default=True, help="Wait on item to finish")
    @click.option('--include-assets/--no-include-assets', default=False,
                  help="Scan common assets of WorkItems and Experiments when filtering")
    @click.option('--verbose/--no-verbose', default=True, help="Enable verbose output in worker")
    @click.option('--json/--no-json', default=False, help="Outputs File list as JSON when used with dry run")
    @click.option('--simulation-prefix-format-str', default=None,
                  help="Simulation Prefix Format str. Defaults to '{simulation.id}'. For no prefix, pass a empty string")
    @click.option('--work-item-prefix-format-str', default=None, help="WorkItem Prefix Format str. Defaults to ''")
    @click.option('--name', default=None, help="Name of Download Workitem. If not provided, one will be generated")
    @click.option('--output-path', default=os.getcwd(), help="Output path to save zip")
    @click.option('--delete-after-download/--no-delete-after-download', default=True,
                  help="Delete the workitem used to gather files after download")
    @click.option('--extract-after-download/--no-extract-after-download', default=True,
                  help="Extract zip after download")
    @click.option('--zip-name', default="output.zip", help="Name of zipfile")
    @click.pass_context
    def download(  # noqa D103
            ctx: click.Context, pattern, exclude_pattern, experiment, simulation, work_item, asset_collection, dry_run,
            wait,
            include_assets, verbose, json, simulation_prefix_format_str, work_item_prefix_format_str, name, output_path,
            delete_after_download,
            extract_after_download, zip_name
    ):
        from idmtools_platform_comps.utils.download.download import DownloadWorkItem

        if json and not dry_run:
            user_logger.error("You cannot return JSON without enabling dry-run mode")
            sys.exit(-1)

        if dry_run and delete_after_download:
            user_logger.warning(
                "You are using dry-run with delete after download. This will most result in an empty file list since "
                "the item will be deleted before the output can be fetched.")

        # ensure no output is enabled when using --json
        if json:
            os.environ['IDMTOOLS_LOGGING_USER_OUTPUT'] = '0'
            os.environ['IDMTOOLS_DISABLE_PROGRESS_BAR'] = '1'

        p: COMPSPlatform = Platform(ctx.obj['config_block'])

        dl_wi = DownloadWorkItem(
            output_path=output_path,
            delete_after_download=delete_after_download,
            extract_after_download=extract_after_download,
            zip_name=zip_name
        )

        if name:
            dl_wi.name = name
        if pattern:
            dl_wi.file_patterns = list(pattern)
        if exclude_pattern:
            dl_wi.exclude_patterns = exclude_pattern if isinstance(exclude_pattern, list) else list(exclude_pattern)

        dl_wi.related_experiments = list(experiment)
        dl_wi.related_simulations = list(simulation)
        dl_wi.related_work_items = list(work_item)
        dl_wi.related_asset_collections = list(asset_collection)
        dl_wi.include_assets = include_assets
        dl_wi.dry_run = dry_run
        dl_wi.verbose = verbose
        if simulation_prefix_format_str is not None:
            if simulation_prefix_format_str.strip() == "":
                dl_wi.no_simulation_prefix = True
            else:
                dl_wi.simulation_prefix_format_str = simulation_prefix_format_str
        if work_item_prefix_format_str is not None:
            dl_wi.work_item_prefix_format_str = work_item_prefix_format_str

        if dl_wi.total_items_watched() == 0:
            user_logger.error("You must specify at least one item to download")

        dl_wi.run(wait_until_done=False, platform=p)
        if not json and not delete_after_download:
            user_logger.info(f"Item can be viewed at {p.get_workitem_link(dl_wi)}")
        if wait:
            dl_wi.wait(wait_on_done_progress=wait)
        if dl_wi.succeeded:
            if dl_wi.dry_run and not delete_after_download:
                file = p.get_files(dl_wi, ['file_list.json'])
                file = file['file_list.json'].decode('utf-8')
                if json:
                    user_logger.info(file)
                else:
                    file = json_parser.loads(file)
                    user_logger.info(tabulate.tabulate([x.values() for x in file], file[0].keys()))
            else:
                pass
        elif dl_wi.failed:
            user_logger.error("Download failed. Check logs in COMPS")
            if dl_wi.failed:
                dl_wi.fetch_error()
            sys.exit(-1)

    @comps.command(help="Allows assetizing outputs from the command line")
    @click.option('--pattern', default=[], multiple=True, help="File patterns")
    @click.option('--exclude-pattern', default=DEFAULT_EXCLUDES, multiple=True, help="File patterns")
    @click.option('--experiment', default=[], multiple=True, help="Experiment ids to assetize")
    @click.option('--simulation', default=[], multiple=True, help="Simulation ids to assetize")
    @click.option('--work-item', default=[], multiple=True, help="WorkItems ids to assetize")
    @click.option('--asset-collection', default=[], multiple=True, help="Asset Collection ids to assetize")
    @click.option('--dry-run/--no-dry-run', default=False,
                  help="Gather a list of files that would be assetized instead of actually assetizing")
    @click.option('--wait/--no-wait', default=True, help="Wait on item to finish")
    @click.option('--include-assets/--no-include-assets', default=False,
                  help="Scan common assets of WorkItems and Experiments when filtering")
    @click.option('--verbose/--no-verbose', default=True, help="Enable verbose output in worker")
    @click.option('--json/--no-json', default=False, help="Outputs File list as JSON when used with dry run")
    @click.option('--simulation-prefix-format-str', default=None,
                  help="Simulation Prefix Format str. Defaults to '{simulation.id}'. For no prefix, pass a empty string")
    @click.option('--work-item-prefix-format-str', default=None, help="WorkItem Prefix Format str. Defaults to ''")
    @click.option('--tag', default=[], type=(str, str), multiple=True,
                  help="Tags to add to the created asset collection as pairs.")
    @click.option('--name', default=None, help="Name of AssetizeWorkitem. If not provided, one will be generated")
    @click.option('--id-file/--no-id-file', default=False, help="Enable or disable writing out an id file")
    @click.option('--id-filename', default=None,
                  help="Name of ID file to save build as. Required when id file is enabled")
    @click.pass_context
    def assetize_outputs(  # noqa D103
            ctx: click.Context, pattern, exclude_pattern, experiment, simulation, work_item, asset_collection, dry_run,
            wait,
            include_assets, verbose, json, simulation_prefix_format_str, work_item_prefix_format_str, tag, name,
            id_file, id_filename
    ):

        if id_file and id_filename is None:
            user_logger.error("--id-filename is required when filename is not provided")
            sys.exit(-1)
        if json:
            os.environ['IDMTOOLS_LOGGING_USER_OUTPUT'] = '0'
            os.environ['IDMTOOLS_DISABLE_PROGRESS_BAR'] = '1'

        p: COMPSPlatform = Platform(ctx.obj['config_block'])
        ao = AssetizeOutput()
        if name:
            ao.name = name
        if pattern:
            ao.file_patterns = list(pattern)
        if exclude_pattern:
            ao.exclude_patterns = exclude_pattern if isinstance(exclude_pattern, list) else list(exclude_pattern)
        ao.related_experiments = list(experiment)
        ao.related_simulations = list(simulation)
        ao.related_work_items = list(work_item)
        ao.related_asset_collections = list(asset_collection)
        ao.include_assets = include_assets
        ao.dry_run = dry_run
        ao.verbose = verbose
        if simulation_prefix_format_str is not None:
            if simulation_prefix_format_str.strip() == "":
                ao.no_simulation_prefix = True
            else:
                ao.simulation_prefix_format_str = simulation_prefix_format_str
        if work_item_prefix_format_str is not None:
            ao.work_item_prefix_format_str = work_item_prefix_format_str
        if tag:
            for name, value in tag:
                ao.asset_tags[name] = value
        if ao.total_items_watched() == 0:
            user_logger.error("You must specify at least one item to assetize")
        ao.run(wait_until_done=False, platform=p)
        if not json:
            user_logger.info(f"Item can be viewed at {p.get_workitem_link(ao)}")
        if wait:
            ao.wait(wait_on_done_progress=wait)
        if ao.succeeded:
            if ao.dry_run:
                file = p.get_files(ao, ['file_list.json'])
                file = file['file_list.json'].decode('utf-8')
                if json:
                    user_logger.info(file)
                else:
                    file = json_parser.loads(file)
                    user_logger.info(tabulate.tabulate([x.values() for x in file], file[0].keys()))
            else:
                if id_file:
                    ao.asset_collection.to_id_file(id_filename)
                if json:
                    result = [i.short_remote_path() for i in ao.asset_collection.assets]
                    user_logger.info(je.dumps(result))
                else:
                    user_logger.info(f"Created {ao.asset_collection.id}")
                    user_logger.info(f"It can be viewed at {p.get_asset_collection_link(ao.asset_collection)}")
                    user_logger.info("Items in Asset Collection")
                    user_logger.info("-------------------------")
                    for asset in ao.asset_collection:
                        user_logger.info(asset.short_remote_path())
        elif ao.failed:
            user_logger.error("Assetized failed. Check logs in COMPS")
            if ao.failed:
                ao.fetch_error()
            sys.exit(-1)

    @comps.command(help="""
        \b
        Create ac from requirement file
        Args:
            asset_tag: tag to be added to ac
            pkg: package name (along with version)
            wheel: package wheel file
        """)
    @click.argument('requirement', type=click.Path(exists=True), required=False)
    @click.option('--asset_tag', multiple=True, help="Tag to be added to AC. Format: 'key:value'")
    @click.option('--pkg', multiple=True, help="Package for override. Format: 'key==value'")
    @click.option('--wheel', multiple=True, help="Local wheel file")
    @click.pass_context
    def req2ac(ctx: click.Context, requirement: str = None, asset_tag: Optional[List[str]] = None,  # noqa D103
               pkg: Optional[List[str]] = None,
               wheel: Optional[List[str]] = None):
        from idmtools_platform_comps.utils.python_requirements_ac.requirements_to_asset_collection import RequirementsToAssetCollection

        pkg_list = list(pkg)
        wheel_list = [os.path.abspath(w) for w in wheel]
        tags = dict()
        for t in asset_tag:
            parts = t.split(':')
            tags[parts[0]] = parts[1]

        p: COMPSPlatform = Platform(ctx.obj['config_block'])
        pl = RequirementsToAssetCollection(p, requirements_path=requirement, pkg_list=pkg_list,
                                           local_wheels=wheel_list, asset_tags=tags)
        ac_id = pl.run()
        print(ac_id)

    @comps.command(help="""
        \b
        Check ac existing based on requirement file
        Args:
            pkg: package name (along with version)
            wheel: package wheel file
        """)
    @click.argument('requirement', type=click.Path(exists=True), required=False)
    @click.option('--pkg', multiple=True, help="Package used for override. Format: say, 'key==value'")
    @click.option('--wheel', multiple=True, help="Local wheel file")
    @click.pass_context
    def ac_exist(ctx: click.Context, requirement: str = None, pkg: Optional[List[str]] = None, wheel: Optional[List[str]] = None):  # noqa D103
        # TODO rename this and move to a subcommand for all the requirements functions
        from idmtools_platform_comps.utils.python_requirements_ac.requirements_to_asset_collection import RequirementsToAssetCollection

        pkg_list = list(pkg)
        wheel_list = [os.path.abspath(w) for w in wheel]
        p: COMPSPlatform = Platform(ctx.obj['config_block'])
        pl = RequirementsToAssetCollection(p, requirements_path=requirement, pkg_list=pkg_list, local_wheels=wheel_list)
        # Check if ac with md5 exists
        ac = pl.retrieve_ac_by_tag()
        if ac:
            print("AC exist: ", ac.id)
        else:
            print("AC doesn't exist")

    @comps.group(help="Singularity commands")
    def singularity():  # noqa D103
        pass

    @singularity.command(help="Build Singularity Image")
    @click.option('--common-input', default=[], multiple=True, help="Files")
    @click.option('--common-input-glob', default=[], multiple=True, help="File patterns")
    @click.option('--transient-input', default=[], multiple=True, help="Transient Files (Paths)")
    @click.option('--transient-input-glob', default=[], multiple=True, help="Transient Files Glob Patterns")
    @click.argument('definition_file')
    @click.option('--wait/--no-wait', default=True, help="Wait on item to finish")
    @click.option('--tag', default=[], type=(str, str), multiple=True,
                  help="Extra Tags as Value Pairs for the Resulting AC")
    @click.option('--workitem-tag', default=[], type=(str, str), multiple=True,
                  help="Extra Tags as Value Pairs for the WorkItem")
    @click.option('--name', default=None, help="Name of WorkItem. If not provided, one will be generated")
    @click.option('--force/--no-force', default=False, help="Force build, ignoring build context")
    @click.option('--image-name', default=None, help="Name of resulting image")
    @click.option('--id-file/--no-id-file', default=True,
                  help="Enable or disable writing out an ID file that points to the created asset collection")
    @click.option('--id-filename', default=None,
                  help="Name of ID file to save build as. If not specified, and id-file is enabled, a name is calculated")
    @click.option('--id-workitem/--no-id-workitem', default=True,
                  help="Enable or disable writing out an id file for the workitem")
    @click.option('--id-workitem-failed/--no-id-workitem-failed', default=False,
                  help="Write id of the workitem even if it failed. You need to enable --id-workitem for this is be active")
    @click.option('--id-workitem-filename', default=None,
                  help="Name of ID file to save workitem to. You need to enable --id-workitem for this is be active")
    @click.pass_context
    def build(ctx: click.Context, common_input, common_input_glob, transient_input, transient_input_glob,  # noqa D103
              definition_file, wait, tag, workitem_tag, name, force, image_name: str,
              id_file: str, id_filename: str, id_workitem: bool, id_workitem_failed: bool, id_workitem_filename: str):
        p: COMPSPlatform = Platform(ctx.obj['config_block'])
        sb = SingularityBuildWorkItem(definition_file=definition_file, name=name, force=force, image_name=image_name)

        if tag:
            for name, value in tag:
                sb.image_tags[name] = value

        if workitem_tag:
            for name, value in tag:
                sb.tags[name] = value

        # Add inputs from files
        for assets, inputs in [(sb.assets, common_input), (sb.transient_assets, transient_input)]:
            for file in inputs:
                add_item(assets, file)

        # And then from glob patterns
        for assets, patterns in [(sb.assets, common_input_glob), (sb.transient_assets, transient_input_glob)]:
            for pattern in patterns:
                for file in glob.glob(pattern):
                    add_item(assets, file)

        sb.run(wait_until_done=wait, platform=p)
        if sb.succeeded and id_file:
            if id_filename is None:
                id_filename = sb.get_id_filename()
            user_logger.info(f"Saving the Asset collection ID that contains the image to {id_filename}")
            sb.asset_collection.to_id_file(id_filename, save_platform=True)
        if id_workitem:
            # TODO when we should use platform id but that need to be updated through the code base
            if sb.succeeded and sb._uid is None:
                user_logger.warning(
                    "Cannot save workitem id because an existing container was found with the same inputs. You can force run using --force, but it is recommended to use the container used.")
            elif id_workitem_failed or sb.succeeded:
                if id_workitem_filename is None:
                    id_workitem_filename = sb.get_id_filename(prefix="builder.")
                user_logger.info(f"Saving the Builder Workitem ID that contains the image to {id_workitem_filename}")
                sb.to_id_file(id_workitem_filename, save_platform=True)
        sys.exit(0 if sb.succeeded else -1)

    @singularity.command(help="Pull Singularity Image")
    @click.argument('image_url')
    @click.option('--wait/--no-wait', default=True, help="Wait on item to finish")
    @click.option('--tag', default=[], type=(str, str), multiple=True,
                  help="Extra Tags as Value Pairs for the Resulting AC")
    @click.option('--workitem-tag', default=[], type=(str, str), multiple=True,
                  help="Extra Tags as Value Pairs for the WorkItem")
    @click.option('--name', default=None, help="Name of WorkItem. If not provided, one will be generated")
    @click.option('--force/--no-force', default=False, help="Force build, ignoring build context")
    @click.option('--image-name', default=None, help="Name of resulting image")
    @click.option('--id-file/--no-id-file', default=True, help="Enable or disable writing out an id file")
    @click.option('--id-filename', default=None,
                  help="Name of ID file to save build as. If not specified, and id-file is enabled, a name is calculated")
    @click.option('--id-workitem/--no-id-workitem', default=True,
                  help="Enable or disable writing out an id file for the workitem")
    @click.option('--id-workitem-failed/--no-id-workitem-failed', default=False,
                  help="Write id of the workitem even if it failed. You need to enable --id-workitem for this is be active")
    @click.option('--id-workitem-filename', default=None,
                  help="Name of ID file to save workitem to. You need to enable --id-workitem for this is be active")
    @click.pass_context
    def pull(ctx: click.Context, image_url, wait, tag, workitem_tag, name, force, image_name: str, id_file: str,  # noqa D103
             id_filename: str,
             id_workitem: bool, id_workitem_failed: bool, id_workitem_filename: str):
        p: COMPSPlatform = Platform(ctx.obj['config_block'])
        sb = SingularityBuildWorkItem(image_url=image_url, force=force, image_name=image_name)
        sb.name = f"Pulling {image_url}" if name is None else name

        if tag:
            for name, value in tag:
                sb.image_tags[name] = value

        if workitem_tag:
            for name, value in tag:
                sb.tags[name] = value

        sb.run(wait_until_done=wait, platform=p)
        if sb.succeeded and id_file:
            if id_filename is None:
                id_filename = sb.get_id_filename()
            user_logger.info(f"Saving ID to {id_filename}")
            sb.asset_collection.to_id_file(id_filename, save_platform=True)

        if id_workitem and sb.succeeded and sb._uid is None:
            user_logger.warning(
                "Cannot save workitem id because an existing container was found with the same inputs. You can force run using --force, but it is recommended to use the container used.")
        elif id_workitem_failed or sb.succeeded:
            if id_workitem_filename is None:
                id_workitem_filename = sb.get_id_filename(prefix="builder.")
            user_logger.info(f"Saving the Builder Workitem ID that contains the image to {id_workitem_filename}")
            sb.to_id_file(id_workitem_filename, save_platform=True)
        sys.exit(0 if sb.succeeded else -1)

except ImportError as e:
    logger.warning(f"COMPS CLI not enabled because a dependency is missing. Most likely it is either click or idmtools cli {e.args}")