idmtools-platform-comps 0.0.0.dev0__py3-none-any.whl → 0.0.2__py3-none-any.whl

idmtools_platform_comps/ssmt_work_items/comps_work_order_task.py

@@ -0,0 +1,29 @@
+"""idmtools CompsWorkOrderTask.
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""
+from dataclasses import dataclass, field
+from idmtools.assets import AssetCollection
+from idmtools.entities.itask import ITask
+from idmtools.entities.simulation import Simulation
+from idmtools_platform_comps.ssmt_work_items.work_order import IWorkOrder
+
+
+@dataclass
+class CompsWorkOrderTask(ITask):
+    """
+    Defines a task that is purely work order driven, like Singularity build.
+    """
+    work_order: IWorkOrder = field(default=None)
+
+    def gather_common_assets(self) -> AssetCollection:
+        """Gather common assets."""
+        pass
+
+    def gather_transient_assets(self) -> AssetCollection:
+        """Gather transient assets."""
+        pass
+
+    def reload_from_simulation(self, simulation: Simulation):
+        """Reload simulation."""
+        pass

idmtools_platform_comps/ssmt_work_items/comps_workitems.py

@@ -0,0 +1,113 @@
+"""idmtools SSMTWorkItem. This is the base of most comps workitems.
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""
+from logging import getLogger, DEBUG
+from dataclasses import dataclass, field, InitVar
+from uuid import UUID
+from idmtools.assets.file_list import FileList
+from idmtools.entities.command_task import CommandTask
+from idmtools_platform_comps.ssmt_work_items.icomps_workflowitem import ICOMPSWorkflowItem
+
+logger = getLogger(__name__)
+user_logger = getLogger(__name__)
+
+
+@dataclass
+class SSMTWorkItem(ICOMPSWorkflowItem):
+    """
+    Defines the SSMT WorkItem.
+
+    Notes:
+        - We have lots of workitem bases. We need to consolidate these a bit.
+    """
+    docker_image: str = field(default=None)
+    command: InitVar[str] = None
+
+    def __post_init__(self, item_name: str, asset_collection_id: UUID, asset_files: FileList, user_files: FileList, command: str):
+        """Constructor."""
+        if command and not self.task:
+            self.task = CommandTask(command)
+        else:
+            user_logger.warning("You provided both a task and a command. Only using the task")
+        super().__post_init__(item_name, asset_collection_id, asset_files, user_files)
+
+        self.work_item_type = self.work_item_type or 'DockerWorker'
+
+    def get_base_work_order(self):
+        """
+        Builder basic work order.
+
+        Returns: work order as a dictionary
+        """
+        base_wo = {
+            "WorkItem_Type": self.work_item_type,
+            "Execution": {
+                "ImageName": self.get_comps_ssmt_image_name(),
+                "Command": str(self.task.command)
+            }
+        }
+
+        return base_wo
+
+    def get_comps_ssmt_image_name(self):
+        """
+        Build comps ssmt docker image name.
+
+        Returns: final validated name
+        """
+        if self.docker_image:
+            return self.docker_image
+
+        if self.platform.docker_image:
+            return self.platform.docker_image
+
+        SSMT_PRODUCTION_IMAGE = 'docker-production.packages.idmod.org/idmtools/comps_ssmt_worker'
+        SSMT_STAGING_IMAGE = 'docker-staging.packages.idmod.org/idmtools/comps_ssmt_worker'
+
+        from idmtools_platform_comps.utils.package_version import get_latest_ssmt_image_version_from_artifactory
+        from idmtools_platform_comps import __version__
+
+        # Determine the default ssmt docker image
+        if "comps.idmod.org" in self.platform.endpoint.lower():
+            ssmt_image = SSMT_PRODUCTION_IMAGE
+        else:
+            ssmt_image = SSMT_STAGING_IMAGE
+
+        release = get_latest_ssmt_image_version_from_artifactory(base_version=__version__)
+
+        docker_image = f'{ssmt_image}:{release}'
+        if logger.isEnabledFor(DEBUG):
+            logger.debug(f'docker_image in use: {docker_image}')
+
+        return docker_image
+
+
+@dataclass
+class InputDataWorkItem(ICOMPSWorkflowItem):
+    """
+    Idm InputDataWorkItem.
+
+    Notes:
+        - TODO add examples
+    """
+
+    def __post_init__(self, item_name: str, asset_collection_id: UUID, asset_files: FileList, user_files: FileList):
+        """Constructor."""
+        super().__post_init__(item_name, asset_collection_id, asset_files, user_files)
+        self.work_item_type = self.work_item_type or 'InputDataWorker'
+
+
+@dataclass
+class VisToolsWorkItem(ICOMPSWorkflowItem):
+    """
+    Idm VisToolsWorkItem.
+
+    Notes:
+        - TODO add examples
+    """
+
+    def __post_init__(self, item_name: str, asset_collection_id: UUID, asset_files: FileList, user_files: FileList):
+        """Constructor."""
+        super().__post_init__(item_name, asset_collection_id, asset_files, user_files)
+        self.work_item_type = self.work_item_type or 'VisTools'

idmtools_platform_comps/ssmt_work_items/icomps_workflowitem.py

@@ -0,0 +1,71 @@
+"""idmtools ICOMPSWorkflowItem.
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""
+from uuid import UUID
+import json
+from abc import ABC
+from dataclasses import field, dataclass
+from idmtools.assets.file_list import FileList
+from idmtools.entities.iworkflow_item import IWorkflowItem
+
+
+@dataclass
+class ICOMPSWorkflowItem(IWorkflowItem, ABC):
+    """
+    Interface of idmtools work item.
+    """
+
+    name: str = field(default="idmtools workflow item")
+    work_order: dict = field(default_factory=lambda: {})
+    work_item_type: str = field(default=None)
+    plugin_key: str = field(default="1.0.0.0_RELEASE")
+
+    def __post_init__(self, item_name: str, asset_collection_id: UUID, asset_files: FileList, user_files: FileList):
+        """Constructor."""
+        super().__post_init__(item_name, asset_collection_id, asset_files, user_files)
+        self.work_order = self.work_order or {}
+
+    def __repr__(self):
+        """Workitem as string."""
+        return f"<WorkItem {self.uid}>"
+
+    def get_base_work_order(self):
+        """Get the base work order."""
+        base_wo = {"WorkItem_Type": self.work_item_type}
+        return base_wo
+
+    def load_work_order(self, wo_file):
+        """Load work order from a file."""
+        self.work_order = json.load(open(wo_file, 'rb'))
+
+    def set_work_order(self, wo):
+        """
+        Update wo for the name with value.
+
+        Args:
+            wo: user wo
+        Returns: None
+        """
+        self.work_order = wo
+
+    def update_work_order(self, name, value):
+        """
+        Update wo for the name with value.
+
+        Args:
+            name: wo arg name
+            value: wo arg value
+
+        Returns: None
+        """
+        self.work_order[name] = value
+
+    def clear_wo_args(self):
+        """
+        Clear all existing wo args.
+
+        Returns: None
+
+        """
+        self.work_order = {}

idmtools_platform_comps/ssmt_work_items/work_order.py

@@ -0,0 +1,54 @@
+"""idmtools WorkOrder classes.
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""
+from abc import ABC
+from dataclasses import dataclass, field
+from typing import List, Dict
+
+
+@dataclass()
+class IWorkOrder(ABC):
+    """Base workorder type."""
+    WorkItem_Type: str = field()
+
+
+@dataclass()
+class ExecutionDefinition:
+    """Define the execution definition for workorders."""
+    Command: str = field()
+    ImageName: str = field(default='DockerWorker')
+
+
+@dataclass()
+class DockerWorkOrder(IWorkOrder):
+    """Define the docker worker."""
+    WorkItem_Type: str = field(default='DockerWorker')
+    Execution: ExecutionDefinition = field(default_factory=ExecutionDefinition)
+
+
+@dataclass()
+class BuildFlags:
+    """Define build flags."""
+    section: List[str] = field(default_factory=lambda: ['all'])
+    library: str = field(default='https://library.sylabs.io')
+    Switches: List['str'] = field(default_factory=lambda: [])
+
+
+@dataclass()
+class BuildDefinition:
+    """Define options for build definitions."""
+    Type: str = field(default='singularity')
+    Input: str = field(default=None)
+    Flags: BuildFlags = field(default_factory=BuildFlags)
+
+
+@dataclass()
+class ImageBuilderWorkOrder(IWorkOrder):
+    """Defines our Image Builder service workorder."""
+    WorkItem_Type: str = field(default='ImageBuilderWorker')
+    Build: str = field(default=BuildDefinition)
+    Output: str = field(default='image.sif')
+    Tags: Dict[str, str] = field(default_factory=lambda: dict(type='singularity'))
+    AdditionalMounts: List[str] = field(default_factory=list)
+    StaticEnvironment: Dict[str, str] = field(default_factory=dict)

idmtools_platform_comps/utils/__init__.py

@@ -0,0 +1,4 @@
+"""idmtools comps utils.
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""

idmtools_platform_comps/utils/assetize_output/__init__.py

@@ -0,0 +1,4 @@
+"""idmtools assetize output.
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""

idmtools_platform_comps/utils/assetize_output/assetize_output.py

@@ -0,0 +1,125 @@
+"""idmtools assetize output work item.
+
+Notes:
+    - TODO add example heres
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""
+from logging import getLogger
+from pathlib import PurePath
+from uuid import UUID
+from dataclasses import dataclass, field
+from typing import Union, Dict
+from idmtools.assets import AssetCollection
+from idmtools.assets.file_list import FileList
+from idmtools.entities.iplatform import IPlatform
+from idmtools.entities.relation_type import RelationType
+from idmtools.core.enums import EntityStatus
+from idmtools_platform_comps.utils.file_filter_workitem import FileFilterWorkItem
+
+logger = getLogger(__name__)
+user_logger = getLogger('user')
+
+
+@dataclass(repr=False)
+class AssetizeOutput(FileFilterWorkItem):
+    """
+    AssetizeOutput allows creating assets from previously ran items in COMPS.
+
+    Notes:
+        - TODO link examples here.
+    """
+    # Dictionary of tags to apply to the results asset collection
+    asset_tags: Dict[str, str] = field(default_factory=dict)
+    #: The asset collection created by Assetize
+    asset_collection: AssetCollection = field(default=None)
+
+    def __post_init__(self, item_name: str, asset_collection_id: UUID, asset_files: FileList, user_files: FileList, command: str):
+        """Constructor AssetizeOutput init."""
+        self._ssmt_script = str(PurePath(__file__).parent.joinpath("assetize_ssmt_script.py"))
+        super().__post_init__(item_name, asset_collection_id, asset_files, user_files, command)
+
+    def _extra_command_args(self, command: str):
+        """Add our tags to the command."""
+        for name, value in self.asset_tags.items():
+            command += f' --asset-tag "{name}={value}"'
+        return command
+
+    def _filter_workitem_pre_creation(self, platform):
+        """
+        Callback to allow for pre-creation calls.
+
+        In our case, we check if tags exist. If no tags exists, we use our defalts ones.
+
+        Args:
+            platform: Platform we are creating on.
+
+        Returns:
+            None
+        """
+        super(AssetizeOutput, self)._filter_workitem_pre_creation(platform)
+        if len(self.asset_tags) == 0:
+            self.__generate_tags()
+
+    def __generate_tags(self):
+        """
+        Add the defaults tags to the WorkItem.
+
+        Returns:
+            None
+        """
+        for experiment in self.related_experiments:
+            self.asset_tags['AssetizedOutputfromFromExperiment'] = str(experiment.id)
+        for simulation in self.related_simulations:
+            self.asset_tags['AssetizedOutputfromFromSimulation'] = str(simulation.id)
+        for work_item in self.related_work_items:
+            self.asset_tags['AssetizedOutputfromFromWorkItem'] = str(work_item.id)
+        for ac in self.related_asset_collections:
+            self.asset_tags['AssetizedOutputfromAssetCollection'] = str(ac.id)
+
+    def run(self, wait_until_done: bool = False, platform: 'IPlatform' = None, wait_on_done_progress: bool = True, **run_opts) -> Union[AssetCollection, None]:
+        """
+        Run the AssetizeOutput.
+
+        Args:
+            wait_until_done: Wait until Done will wait for the workitem to complete
+            platform: Platform Object
+            wait_on_done_progress: When set to true, a progress bar will be shown from the item
+            **run_opts: Additional options to pass to Run on platform
+
+        Returns:
+            AssetCollection created if item succeeds
+        """
+        p = super()._check_for_platform_from_context(platform)
+        if 'wait_on_done' in run_opts:
+            raise TypeError("The 'wait_on_done' parameter has been removed in idmtools 1.8.0. Please update your code with 'wait_until_done'.")
+        p.run_items(self, wait_on_done_progress=wait_on_done_progress, **run_opts)
+        if wait_until_done:
+            return self.wait(wait_on_done_progress=wait_on_done_progress, platform=p)
+
+    def wait(self, wait_on_done_progress: bool = True, timeout: int = None, refresh_interval=None, platform: 'IPlatform' = None) -> Union[AssetCollection, None]:
+        """
+        Waits on Assetize Workitem to finish. This first waits on any dependent items to finish(Experiment/Simulation/WorkItems).
+
+        Args:
+            wait_on_done_progress: When set to true, a progress bar will be shown from the item
+            timeout: Timeout for waiting on item. If none, wait will be forever
+            refresh_interval: How often to refresh progress
+            platform: Platform
+
+        Returns:
+            AssetCollection created if item succeeds
+        """
+        # wait on related items before we wait on our item
+        p = super()._check_for_platform_from_context(platform)
+        opts = dict(wait_on_done_progress=wait_on_done_progress, timeout=timeout, refresh_interval=refresh_interval, platform=p)
+        self._wait_on_children(**opts)
+
+        super().wait(**opts)
+        if self.status == EntityStatus.SUCCEEDED and not self.dry_run:
+            # If we succeeded, get our AC
+            comps_workitem = self.get_platform_object(force=True)
+            acs = comps_workitem.get_related_asset_collections(RelationType.Created)
+            if acs:
+                self.asset_collection = AssetCollection.from_id(acs[0].id, platform=p)
+                return self.asset_collection

idmtools_platform_comps/utils/assetize_output/assetize_ssmt_script.py

@@ -0,0 +1,144 @@
+"""idmtools ssmt script.
+
+This script is used on server side only and not meant to be ran on a local machine.
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""
+# flake8: noqa F405 F403
+import sys
+from argparse import Namespace
+from logging import getLogger
+from COMPS.Data.WorkItem import RelationType
+from typing import List, Dict
+import os
+from COMPS.Data import AssetCollectionFile, WorkItem
+
+# we have to support two ways to load the utils. The first is load from Assets
+# the fallback is the installed package
+# this allows us to move forward with changes to utils without need for new images
+try:
+    from file_filter import *
+    from common import *
+except (FileNotFoundError, ImportError):
+    from idmtools_platform_comps.utils.ssmt_utils.file_filter import *
+    from idmtools_platform_comps.utils.ssmt_utils.common import *
+
+logger = getLogger(__name__)
+user_logger = getLogger('user')
+
+
+class NoFileFound(Exception):
+    doc_link: str = "platforms/comps/errors.html#errors"
+
+
+def create_asset_collection(file_list: SetOfAssets, ac_files: List[AssetCollectionFile], tags: Dict[str, str]):  # pragma: no cover
+    """
+
+    Args:
+        file_list:
+        ac_files: AC Files
+        tags: Tags to add
+
+    Returns:
+
+    """
+    asset_collection = AssetCollection()
+    asset_collection.set_tags(tags)
+    # Maps checksum to AssetCollectionFile and the path on disk to file
+    asset_collection_map: Dict[uuid.UUID, Tuple[AssetCollectionFile, str, int]] = dict()
+    for file in file_list:
+        fn = os.path.basename(file[1])
+        acf = AssetCollectionFile(file_name=fn, relative_path=file[1].replace(fn, "").strip("/"), md5_checksum=file[2])
+        asset_collection_map[file[2]] = (acf, file[0], file[3])
+        asset_collection.add_asset(acf)
+
+    # add files from acs
+    for f in ac_files:
+        asset_collection.add_asset(f)
+
+    # do initial save to see what assets are in comps
+    missing_files = asset_collection.save(return_missing_files=True)
+    if missing_files:
+        user_logger.info(f"{len(missing_files)} not currently in COMPS as Assets")
+
+        ac2 = AssetCollection()
+        new_files = 0
+        total_to_upload = 0
+        for checksum, asset_details in asset_collection_map.items():
+            if checksum in missing_files:
+                new_files += 1
+                total_to_upload += asset_details[2]
+                acf = AssetCollectionFile(file_name=asset_details[0].file_name, relative_path=asset_details[0].relative_path)
+                ac2.add_asset(acf, asset_details[1])
+            else:
+                ac2.add_asset(asset_details[0])
+        user_logger.info(f"Saving {new_files} totaling {humanfriendly.format_size(total_to_upload)} assets to comps")
+        ac2.set_tags(tags)
+        ac2.save()
+        asset_collection = ac2
+    return asset_collection
+
+
+def get_argument_parser():
+    p = get_common_parser("AssetizeOutputs")
+    p.add_argument("--asset-tag", action='append', help="List of tags as value pairs. Eg: ABC=123")
+    return p
+
+
+def build_asset_tags(parsed_args: Namespace) -> Dict[str, str]:
+    """
+    Builds our Asset tag dic from tags
+
+    Args:
+        parsed_args: Parse Arg
+
+    Returns:
+        Dict of tags
+    """
+    asset_tags = dict()
+    if parsed_args.asset_tag:
+        for tag in parsed_args.asset_tag:
+            name, value = tag.split("=")
+            asset_tags[name] = value
+    return asset_tags
+
+
+if __name__ == "__main__":  # pragma: no cover
+    # build our argument parser and then parse the command line
+    parser = get_argument_parser()
+    args = parser.parse_args()
+
+    # Set our JOB config global with config provided
+    JOB_CONFIG = vars(args)
+    # register our error handler
+    sys.excepthook = get_error_handler_dump_config_and_error(JOB_CONFIG)
+    print(args)
+
+    # Parse the common arguments common to filter scripts
+    entity_filter_func, fn_format_func = parse_filter_args_common(args)
+    # Parse our Tags
+    asset_tags = build_asset_tags(args)
+
+    # login
+    client = login_to_env()
+
+    # Load our workitem
+    wi = WorkItem.get(os.environ['COMPS_WORKITEM_GUID'])
+
+    # Gather all our files in Experiments, Simulations, and Asset Collections
+    files, files_from_ac = filter_files_and_assets(args, entity_filter_func, wi, fn_format_func)
+    if len(files) == 0 and len(files_from_ac) == 0:
+        raise NoFileFound("No files found with patterns specified. Please verify your filters.")
+
+    ensure_no_duplicates(files_from_ac, files)
+
+    if args.dry_run:
+        print_results(files_from_ac, files)
+    else:
+        ac = create_asset_collection(files, files_from_ac, tags=asset_tags)
+
+        with open('asset_collection.id', 'w') as o:
+            user_logger.info(ac.id)
+            o.write(str(ac.id))
+
+        wi.add_related_asset_collection(ac.id, RelationType.Created)

idmtools_platform_comps/utils/base_singularity_work_order.json

@@ -0,0 +1,6 @@
+{
+  "WorkItem_Type": "ImageBuilderWorker",
+  "Build": {
+    "Type": "singularity"
+  }
+}

idmtools_platform_comps/utils/download/__init__.py

@@ -0,0 +1,4 @@
+"""idmtools download workitem.
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""