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

idmtools/core/docker_task.py

@@ -0,0 +1,207 @@
+"""
+DockerTask provides a utility to run docker images.
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""
+import os
+import re
+import sys
+import unicodedata
+from dataclasses import dataclass, field
+from datetime import datetime, timezone, timedelta
+from typing import Optional, Type
+from idmtools import __version__ as idmtools_version, IdmConfigParser
+from idmtools.assets import AssetCollection, json
+from idmtools.core.logging import getLogger
+from idmtools.entities.itask import ITask
+from idmtools.entities.platform_requirements import PlatformRequirements
+from idmtools.registry.task_specification import TaskSpecification
+
+logger = getLogger(__name__)
+user_logger = getLogger('user')
+
+
+@dataclass
+class DockerTask(ITask):
+    """
+    Provides a task to run or optionally build a docker container.
+    """
+    image_name: str = field(default=None, metadata={"md": True})
+    # Optional config to build the docker image
+    build: bool = field(default=False, metadata={"md": True})
+    build_path: Optional[str] = field(default=None, metadata={"md": True})
+    # This should in the build_path directory
+    Dockerfile: Optional[str] = field(default=None, metadata={"md": True})
+    pull_before_build: bool = field(default=True, metadata={"md": True})
+    use_nvidia_run: bool = field(default=False, metadata={"md": True})
+
+    __image_built: bool = field(default=False)
+
+    def __post_init__(self):
+        """
+        Set our platform requirements and optionally trigger image build.
+
+        Returns:
+            None
+        """
+        super().__post_init__()
+        self.add_platform_requirement(PlatformRequirements.DOCKER)
+        if self.build:
+            self.build_image()
+
+    def gather_common_assets(self) -> AssetCollection:
+        """
+        Gather common(experiment-level) assets from task.
+
+        Returns:
+            AssetCollection containing all the common assets
+        """
+        if self.image_name is None:
+            raise ValueError("Image Name is required")
+        return self.common_assets
+
+    def gather_transient_assets(self) -> AssetCollection:
+        """
+        Gather transient(simulation-level) assets from task.
+
+        Returns:
+            AssetCollection
+        """
+        return self.transient_assets
+
+    def build_image(self, spinner=None, **extra_build_args):
+        """
+        Build our docker image.
+
+        Args:
+            spinner: Should we display a CLI spinner
+            **extra_build_args: Extra build arguments to pass to docker
+
+        Returns:
+            None
+        """
+        if not self.__image_built:
+            import docker
+            from docker.errors import BuildError
+            if spinner:
+                spinner.text = f"Building {self.image_name}"
+            # if the build_path is none use current working directory
+            if self.build_path is None:
+                self.build_path = os.getcwd()
+
+            client = docker.client.from_env()
+            build_config = dict(
+                path=self.build_path,
+                dockerfile=self.Dockerfile,
+                tag=self.image_name,
+                labels=dict(
+                    uildstamp=f'built-by idmtools {idmtools_version}',
+                    builddate=str(datetime.now(timezone(timedelta(hours=-8)))))
+            )
+            if extra_build_args:
+                build_config.update(extra_build_args)
+            logger.debug(f"Build configuration used: {str(build_config)}")
+            self.__image_built = True
+            if not IdmConfigParser.is_progress_bar_disabled():
+                from tqdm import tqdm
+                prog = tqdm(
+                    desc='Building docker image',
+                    total=10,
+                    bar_format='Building Docker Image: |{bar}| {percentage:3.0f}% [{n_fmt}/{total_fmt}] [{elapsed}] {desc}'
+                )
+            try:
+                build_step = None
+                # regular expression to grab progress
+                progress_grep = re.compile(r'Step ([0-9]+)/([0-9]+) : (.*)')
+                # regular expression to filter out ansi codes
+                ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
+                for line in client.api.build(**build_config):
+                    line = json.loads(line)
+                    if 'stream' in line:
+                        line = line['stream']
+                        line = ansi_escape.sub('', line).strip()
+                        # strip unicode data
+                        line = "".join(ch for ch in line if unicodedata.category(ch)[0] != "C")
+                        logger.debug('Raw Docker Output: %s', line)
+                        if line:
+                            grps = progress_grep.match(line)
+                            if grps:
+                                try:
+                                    step = int(grps.group(1))
+                                    total_steps = int(grps.group(2))
+                                    if prog:
+                                        prog.n = step
+                                        prog.total = total_steps
+                                    line = grps.group(3)
+                                except:  # noqa E722
+                                    pass
+                                if prog:
+                                    prog.set_description(line)
+                                build_step = line
+                            # update build step with output
+                            elif build_step:
+                                if len(line) > 40:
+                                    line = line[:40]
+                                if prog:
+                                    prog.set_description(f'{build_step}: {line}')
+                    elif 'status' in line:
+                        line = line['status'].strip()
+                        if prog:
+                            prog.set_description(line)
+
+                logger.info('Build Successful)')
+            except BuildError as e:
+                logger.info(f"Build failed for {self.image_name} with message {e.msg}")
+                logger.info(f'Build log: {e.build_log}')
+                sys.exit(-1)
+            finally:
+                if prog:
+                    prog.close()
+
+    def reload_from_simulation(self, simulation: 'Simulation'):  # noqa E821
+        """
+        Method to reload task details from simulation object. Currently we do not do this for docker task.
+
+        Args:
+            simulation: Simulation to load data from
+
+        Returns:
+            None
+        """
+        pass
+
+
+class DockerTaskSpecification(TaskSpecification):
+    """
+    DockerTaskSpecification provides the task plugin to idmtools for DockerTask.
+    """
+
+    def get(self, configuration: dict) -> DockerTask:
+        """
+        Get instance of DockerTask with configuration provided.
+
+        Args:
+            configuration: configuration for DockerTask
+
+        Returns:
+            DockerTask with configuration
+        """
+        return DockerTask(**configuration)
+
+    def get_description(self) -> str:
+        """
+        Get description of plugin.
+
+        Returns:
+            Plugin description
+        """
+        return "Defines a docker command"
+
+    def get_type(self) -> Type[DockerTask]:
+        """
+        Get type of task provided by plugin.
+
+        Returns:
+            DockerTask
+        """
+        return DockerTask

idmtools/core/enums.py

@@ -0,0 +1,51 @@
+"""
+Define our common enums to be used through idmtools.
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""
+from pathlib import Path
+
+from enum import Enum
+
+TRUTHY_VALUES = ['1', 'y', 'yes', 'on', 'true', 't', 1, True]
+# Used to store idmtools user specific config/data
+IDMTOOLS_USER_HOME = Path().home().joinpath(".idmtools")
+
+
+class EntityStatus(Enum):
+    """
+    EntityStatus provides status values for Experiment/Simulations/WorkItems.
+    """
+    COMMISSIONING = 'commissioning'
+    CREATED = 'created'
+    RUNNING = 'running'
+    SUCCEEDED = 'succeeded'
+    FAILED = 'failed'
+
+
+class FilterMode(Enum):
+    """
+    Allows user to specify AND/OR for the filtering system.
+    """
+    AND = 0
+    OR = 1
+
+
+class ItemType(Enum):
+    """
+    ItemTypes supported by idmtools.
+    """
+    SUITE = "Suite"
+    EXPERIMENT = "Experiment"
+    SIMULATION = "Simulation"
+    WORKFLOW_ITEM = "WorkItem"  # On Comps this is workitems
+    ASSETCOLLECTION = "Asset Collection"
+
+    def __str__(self):
+        """
+        Returns a string representation of our item type.
+
+        Returns:
+            The string version of our enum value
+        """
+        return str(self.value)

idmtools/core/exceptions.py

@@ -0,0 +1,91 @@
+"""
+Define idmtools common exception as well as idmtools system exception handler.
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""
+import sys
+import typing
+from logging import getLogger
+
+if typing.TYPE_CHECKING:
+    from idmtools.entities.iplatform import TPlatform
+
+user_logger = getLogger('user')
+logger = getLogger(__name__)
+
+
+class ExperimentNotFound(Exception):
+    """
+    Thrown when an experiment cannot be found on a platform.
+    """
+
+    def __init__(self, experiment_id: str, platform: 'TPlatform' = None):
+        """
+        Initialize our ExperimentNotFound.
+
+        Args:
+            experiment_id: Experiment id to say wasn't found
+            platform: Optional platform. Used in error message
+        """
+        if platform:
+            super().__init__(f"Experiment with id '{experiment_id}' could not be retrieved on platform {platform}.")
+        else:
+            super().__init__(f"Experiment with id '{experiment_id}' could not be retrieved.")
+
+
+class UnknownItemException(Exception):
+    """
+    Thrown when an unknown item type is passed to idmtools.
+
+    This usually occurs within the platform operation area.
+    """
+    pass
+
+
+class NoPlatformException(Exception):
+    """
+    Cannot find a platform matching the one requested by user.
+    """
+    pass
+
+
+class TopLevelItem(Exception):
+    """
+    Thrown when a parent of a top-level item is requested by the platform.
+    """
+    pass
+
+
+class UnsupportedPlatformType(Exception):
+    """
+    Occurs when an item is not supported by a platform but is requested.
+    """
+    pass
+
+
+class NoTaskFound(Exception):
+    """
+    Thrown when a simulation has no task defined.
+    """
+    pass
+
+
+def idmtools_error_handler(exctype, value: Exception, tb):
+    """
+    Global exception handler. This will write our errors in a nice format as well as find document links if attached to the exception.
+
+    Args:
+        exctype: Type of exception
+        value: Value of the exception
+        tb: Traceback
+
+    Returns:
+        None
+    """
+    if hasattr(value, 'doc_link'):
+        from idmtools.utils.info import get_help_version_url
+        user_logger.error(f"{value.args[0]}. For more details, see {get_help_version_url(value.doc_link)}")
+
+    logger.exception(value)
+    # Call native exception manager
+    sys.__excepthook__(exctype, value, tb)

idmtools/core/experiment_factory.py

@@ -0,0 +1,71 @@
+"""
+Define ExperimentFactory.
+
+This is used mostly internally. It does allow us to support specialized experiment types when needed.
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""
+from logging import getLogger, DEBUG
+from idmtools.entities.experiment import Experiment
+from idmtools.registry import experiment_specification
+
+logger = getLogger(__name__)
+
+
+class ExperimentFactory:
+    """
+    ExperimentFactory allows creating experiments that could be derived through plugins.
+
+    """
+    DEFAULT_KEY = 'idmtools.entities.experiment.Experiment'
+
+    def __init__(self):
+        """
+        Initialize our factory.
+
+        On initialize, we load our plugin and build a map of ids for experiments.
+        """
+        from idmtools.registry.experiment_specification import ExperimentPlugins
+        self._builders = ExperimentPlugins().get_plugin_map()
+        aliases = dict()
+        # register types as full paths as well
+        for _model, spec in self._builders.items():
+            aliases[f'{spec.get_type().__module__}.{spec.get_type().__name__}'] = spec
+        self._builders.update(aliases)
+
+    def create(self, key, fallback=None, **kwargs) -> Experiment:  # noqa: F821
+        """
+        Create an experiment of type key.
+
+        Args:
+            key: Experiment Type
+            fallback: Fallback type. If none, uses DEFAULT_KEY
+            **kwargs: Options to pass to the experiment object
+
+        Returns:
+            Experiment object that was created
+        """
+        if logger.isEnabledFor(DEBUG):
+            logger.debug(f"Attempting to create experiment of type {key}")
+
+        if not key:
+            key = self.DEFAULT_KEY
+        else:
+            logger.warning(f'Experiment type: {key}')
+
+        if key not in self._builders:
+            if not fallback:
+                raise ValueError(f"The ExperimentFactory could not create an experiment of type {key}")
+            else:
+                logger.warning(f'Could not find experiment type {key}. Using Fallback type of {fallback.__class__}')
+                return fallback()
+
+        model_spec: experiment_specification = self._builders.get(key)
+        result = model_spec.get(kwargs)
+
+        if logger.isEnabledFor(DEBUG):
+            logger.debug(f"Experiment created for type {key}")
+        return result
+
+
+experiment_factory = ExperimentFactory()

idmtools/core/id_file.py

@@ -0,0 +1,70 @@
+"""
+Utility method for writing and reading id files.
+
+ID Files allow us to reload entities like Experiment, Simulations, AssetCollections, etc from a platform through files. This can
+be enabling for workflows to chain steps together, or to self-document remote outputs in the local project directory.
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""
+import json
+from os import PathLike
+from typing import Union, Dict, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from idmtools.core.interfaces.ientity import IEntity
+
+
+def read_id_file(filename: Union[str, PathLike]):
+    """
+    Reads an id from an id file.
+
+    An id file is in the format of
+
+    <id>::<item_type>::<config block>::<extra args>
+    Args:
+        filename:
+
+    Returns:
+        None
+    """
+    if isinstance(filename, PathLike):
+        filename = str(filename)
+    platform_block = None
+    item_id = None
+    item_type = None
+    extra_args = None
+    with open(filename, 'r') as id_in:
+        item_id = id_in.read().strip()
+        if "::" in item_id:
+            parts = item_id.split("::")
+            if len(parts) == 2:
+                item_id, item_type = item_id.split("::")
+            elif len(parts) == 3:
+                item_id, item_type, platform_block = item_id.split("::")
+            elif len(parts) == 4:
+                item_id, item_type, platform_block, extra_args = item_id.split("::")
+    return item_id, item_type, platform_block, extra_args
+
+
+def write_id_file(filename: Union[str, PathLike], item: 'IEntity', save_platform: bool = False, platform_args: Dict = None):
+    """
+    Write an item as and id file.
+
+    Args:
+        filename: Filename to write file to
+        item: Item to write out
+        save_platform: When true, writes platform details to the file
+        platform_args: Platform arguments to write out
+
+    Returns:
+        None
+    """
+    from idmtools.utils.json import IDMJSONEncoder
+    if isinstance(filename, PathLike):
+        filename = str(filename)
+    with open(filename, 'w') as filename:
+        filename.write(f'{item.id}::{item.item_type}')
+        if save_platform and hasattr(item.platform, '_config_block'):
+            filename.write(f"::{item.platform._config_block}")
+            if platform_args:
+                filename.write(json.dumps(platform_args, cls=IDMJSONEncoder))

idmtools/core/interfaces/__init__.py

@@ -0,0 +1,5 @@
+"""
+idmtools List of core abstract interfaces that other core objects derive from.
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""

idmtools/core/interfaces/entity_container.py

@@ -0,0 +1,64 @@
+"""
+EntityContainer definition. EntityContainer provides an envelope for a Parent to container a list of sub-items.
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""
+import typing
+
+
+if typing.TYPE_CHECKING:
+    from idmtools.core.interfaces.ientity import IEntity
+    from idmtools.core.enums import EntityStatus
+    from typing import List
+
+
+class EntityContainer(list):
+    """
+    EntityContainer is a wrapper classes used by Experiments and Suites to wrap their children.
+
+    It provides utilities to set status on entities
+    """
+
+    def __init__(self, children: 'List[IEntity]' = None):
+        """
+        Initialize the EntityContainer.
+
+        Args:
+            children: Children to initialize with
+        """
+        super().__init__()
+        self.extend(children or [])
+
+    def set_status(self, status: 'EntityStatus'):
+        """
+        Set status on all the children.
+
+        Args:
+            status: Status to set
+
+        Returns:
+            None
+        """
+        for entity in self:
+            entity.status = status
+
+    def set_status_for_item(self, item_id, status: 'EntityStatus'):
+        """
+        Set status for specific sub-item.
+
+        Args:
+            item_id: Item id to set status for
+            status: Status to set
+
+        Returns:
+            None
+
+        Raises:
+            ValueError when the item_id is not in the children list
+        """
+        for entity in self:
+            if str(entity.uid) == str(item_id):
+                entity.status = status
+                return
+
+        raise ValueError(f"Item with id {item_id} not found in the container")

idmtools/core/interfaces/iassets_enabled.py

@@ -0,0 +1,58 @@
+"""
+IAssetsEnabled interface definition.
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""
+from abc import ABCMeta, abstractmethod
+from dataclasses import dataclass, field
+from typing import NoReturn, Union
+from idmtools.assets import TAsset, TAssetList
+from idmtools.assets.asset_collection import AssetCollection
+
+
+@dataclass
+class IAssetsEnabled(metaclass=ABCMeta):
+    """
+    Base class for objects containing an asset collection.
+    """
+    assets: AssetCollection = field(default_factory=lambda: AssetCollection(), compare=False,
+                                    metadata={"pickle_ignore": True})
+
+    def __post_init__(self):
+        """
+        Post init. Nothing needed for IAssetsEnabled.
+
+        Returns:
+            None
+        """
+        pass
+
+    @abstractmethod
+    def gather_assets(self) -> NoReturn:
+        """
+        Function called at runtime to gather all assets in the collection.
+        """
+        pass
+
+    def add_assets(self, assets: Union[TAssetList, AssetCollection] = None, fail_on_duplicate: bool = True) -> NoReturn:
+        """
+        Add more assets to :class:`~idmtools.assets.asset_collection.AssetCollection`.
+        """
+        for asset in assets:
+            self.assets.add_asset(asset, fail_on_duplicate)
+
+    def add_asset(self, asset: Union[str, 'TAsset'] = None, fail_on_duplicate: bool = True) -> NoReturn:
+        """
+        Add an asset to our item.
+
+        Args:
+            asset: Asset to add. Asset can be a string in which case it is assumed to be a file path
+            fail_on_duplicate: Should we rain an exception if there is an existing file with same information
+
+        Returns:
+            None
+
+        Raise:
+            DuplicatedAssetError in cases where fail_on_duplicate are true
+        """
+        self.assets.add_asset(asset, fail_on_duplicate)