Source code for idmtools_platform_comps.utils.singularity_build

"""idmtools singularity build workitem.

Notes:
    - TODO add examples here.

Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
"""
import hashlib
import io
import json
import os
import re
import uuid
from dataclasses import dataclass, field, InitVar
from logging import getLogger, DEBUG
from os import PathLike
from pathlib import PurePath
from typing import List, Dict, Union, Optional, TYPE_CHECKING
from urllib.parse import urlparse
from uuid import UUID
from COMPS.Data import QueryCriteria
from jinja2 import Environment
from idmtools import IdmConfigParser
from idmtools.assets import AssetCollection, Asset
from idmtools.assets.file_list import FileList
from idmtools.core import EntityStatus, NoPlatformException
from idmtools.core.logging import SUCCESS
from idmtools.entities.command_task import CommandTask
from idmtools.entities.relation_type import RelationType
from idmtools.utils.hashing import calculate_md5_stream
from idmtools_platform_comps.ssmt_work_items.comps_workitems import InputDataWorkItem
from idmtools_platform_comps.utils.general import save_sif_asset_md5_from_ac_id
from idmtools_platform_comps.utils.package_version import get_docker_manifest, get_digest_from_docker_hub

if TYPE_CHECKING:
    from idmtools.entities.iplatform import IPlatform

SB_BASE_WORKER_PATH = os.path.join(os.path.dirname(__file__), 'base_singularity_work_order.json')

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


[docs]@dataclass(repr=False)
class SingularityBuildWorkItem(InputDataWorkItem):
    """
    Provides a wrapper to build utilizing the COMPS build server.

    Notes:
        - TODO add references to examples
    """
    #: Path to definition file
    definition_file: Union[PathLike, str] = field(default=None)
    #: definition content. Alternative to file
    definition_content: str = field(default=None)
    #: Enables Jinja parsing of the definition file or content
    is_template: bool = field(default=False)
    #: template_args
    template_args: Dict[str, str] = field(default_factory=dict)
    #: Image Url
    image_url: InitVar[str] = None
    #: Destination image name
    image_name: str = field(default=None)
    #: Name of the workitem
    name: str = field(default=None)
    #: Tages to add to container asset collection
    image_tags: Dict[str, str] = field(default_factory=dict)
    #: Allows you to set a different library. (The default library is “https://library.sylabs.io”). See https://sylabs.io/guides/3.5/user-guide/cli/singularity_build.html
    library: str = field(default=None)
    #: only run specific section(s) of definition file (setup, post, files, environment, test, labels, none) (default [all])
    section: List[str] = field(default_factory=lambda: ['all'])
    #: build using user namespace to fake root user (requires a privileged installation)
    fix_permissions: bool = field(default=False)
    # AssetCollection created by build
    asset_collection: AssetCollection = field(default=None)
    #: Additional Mounts
    additional_mounts: List[str] = field(default_factory=list)
    #: Environment vars for remote build
    environment_variables: Dict[str, str] = field(default_factory=dict)
    #: Force build
    force: bool = field(default=False)
    #: Don't include default tags
    disable_default_tags: bool = field(default=None)
    # ID that is added to work item and then results collection that can be used to tied the items together
    run_id: uuid.UUID = field(default_factory=uuid.uuid4)

    #: loaded if url is docker://. Used to determine if we need to re-run a build
    __digest: Dict[str, str] = field(default=None)
    __image_tag: str = field(default=None)

    #: rendered template. We have to store so it is calculated before RUN which means outside our normal pre-create hooks
    __rendered_template: str = field(default=None)

    def __post_init__(self, item_name: str, asset_collection_id: UUID, asset_files: FileList, user_files: FileList, image_url: str):
        """Constructor."""
        self.work_item_type = 'ImageBuilderWorker'
        self._image_url = None
        # Set this for now. Later it should be replace with some type of Specialized worker identifier
        self.task = CommandTask("ImageBuilderWorker")
        super().__post_init__(item_name, asset_collection_id, asset_files, user_files)

        self.image_url = image_url if isinstance(image_url, str) else None
        if isinstance(self.definition_file, PathLike):
            self.definition_file = str(self.definition_file)

[docs]    def get_container_info(self) -> Dict[str, str]:
        """Get container info.

        Notes:
            - TODO remove this
        """
        pass

    @property
    def image_url(self):  # noqa: F811
        """Get the image url."""
        return self._image_url

    @image_url.setter
    def image_url(self, value: str):
        """
        Set the image url.

        Args:
            value: Value to set value to

        Returns:
            None
        """
        url_info = urlparse(value)
        if url_info.scheme == "docker":
            if "packages.idmod.org" in value:
                full_manifest, self.__image_tag = get_docker_manifest(url_info.path)
                self.__digest = full_manifest['config']['digest']
            else:
                self.__image_tag = url_info.netloc + ":latest" if ":" not in value else url_info.netloc
                image, tag = url_info.netloc.split(":")
                self.__digest = get_digest_from_docker_hub(image, tag)
            if self.fix_permissions:
                self.__digest += "--fix-perms"
            if self.name is None:
                self.name = f"Load Singularity image from Docker {self.__image_tag}"
        # TODO how to do this for shub
        self._image_url = value

[docs]    def context_checksum(self) -> str:
        """
        Calculate the context checksum of a singularity build.

        The context is the checksum of all the assets defined for input, the singularity definition file, and the environment variables

        Returns:
            Conext checksum.
        """
        file_hash = hashlib.sha256()
        # ensure our template is set
        self.__add_common_assets()
        for asset in sorted(self.assets + self.transient_assets, key=lambda a: a.short_remote_path()):
            if asset.absolute_path:
                with open(asset.absolute_path, mode='rb') as ain:
                    calculate_md5_stream(ain, file_hash=file_hash)
            else:
                self.__add_file_to_context(json.dumps([asset.filename, asset.relative_path, str(asset.checksum)], sort_keys=True) if asset.persisted else asset.bytes, file_hash)

        if len(self.environment_variables):
            contents = json.dumps(self.environment_variables, sort_keys=True)
            self.__add_file_to_context(contents, file_hash)

        if logger.isEnabledFor(DEBUG):
            logger.debug(f'Context: sha256:{file_hash.hexdigest()}')
        return f'sha256:{file_hash.hexdigest()}'

    def __add_file_to_context(self, contents: Union[str, bytes], file_hash):
        """
        Add a specific file content to context checksum.

        Args:
            contents: Contents
            file_hash: File hash to add to

        Returns:
            None
        """
        item = io.BytesIO()
        item.write(contents.encode('utf-8') if isinstance(contents, str) else contents)
        item.seek(0)
        calculate_md5_stream(item, file_hash=file_hash)

[docs]    def render_template(self) -> Optional[str]:
        """
        Render template. Only applies when is_template is True. When true, it renders the template using Jinja to a cache value.

        Returns:
            Rendered Template
        """
        if self.is_template:
            # We don't allow re-running template rendering
            if self.__rendered_template is None:
                if logger.isEnabledFor(DEBUG):
                    logger.debug("Rendering template")
                contents = None
                # try from file first
                if self.definition_file:
                    with open(self.definition_file, mode='r') as ain:
                        contents = ain.read()
                elif self.definition_content:
                    contents = self.definition_content

                if contents:
                    env = Environment()
                    template = env.from_string(contents)
                    self.__rendered_template = template.render(env=os.environ, sbi=self, **self.template_args)
            return self.__rendered_template
        return None

[docs]    @staticmethod
    def find_existing_container(sbi: 'SingularityBuildWorkItem', platform: 'IPlatform' = None) -> Optional[AssetCollection]:
        """
        Find existing container.

        Args:
            sbi: SingularityBuildWorkItem to find existing container matching config
            platform: Platform To load the object from

        Returns:
            Existing Asset Collection
        """
        if platform is None:
            from idmtools.core.context import CURRENT_PLATFORM
            if CURRENT_PLATFORM is None:
                raise NoPlatformException("No Platform defined on object, in current context, or passed to run")
            platform = CURRENT_PLATFORM
        ac = None
        if not sbi.force:  # don't search if it is going to be forced
            qc = QueryCriteria().where_tag(['type=singularity']).select_children(['assets', 'tags']).orderby('date_created desc')
            if sbi.__digest:
                qc.where_tag([f'digest={sbi.__digest}'])
            elif sbi.definition_file or sbi.definition_content:
                qc.where_tag([f'build_context={sbi.context_checksum()}'])
            if len(qc.tag_filters) > 1:
                if logger.isEnabledFor(DEBUG):
                    logger.debug("Searching for existing containers")
                ac = platform._assets.get(None, query_criteria=qc)
            if ac:
                if logger.isEnabledFor(DEBUG):
                    logger.debug(f"Found: {len(ac)} previous builds")
                ac = platform._assets.to_entity(ac[0])
                if logger.isEnabledFor(DEBUG):
                    logger.debug(f'Found existing container in {ac.id}')
            else:
                ac = None

        return ac

    def __add_tags(self):
        """
        Add default tags to the asset collection to be created.

        The most important part of this logic is the digest/run_id information we add. This is what enables the build/pull-cache through comps.

        Returns:
            None
        """
        self.image_tags['type'] = 'singularity'
        # Disable all tags but image name and type
        if not self.disable_default_tags:
            if self.platform is not None and hasattr(self.platform, 'get_username'):
                self.image_tags['created_by'] = self.platform.get_username()
            # allow users to override run id using only the tag
            if 'run_id' in self.tags:
                self.run_id = self.tags['run_id']
            else:
                # set the run id on the workitem and resulting tags
                self.tags['run_id'] = str(self.run_id)
            self.image_tags['run_id'] = self.tags['run_id']
            # Check for the digest
            if self.__digest and isinstance(self.__digest, str):
                self.image_tags['digest'] = self.__digest
                self.image_tags['image_from'] = self.__image_tag
                if self.image_name is None:
                    self.image_name = self.__image_tag.strip(" /").replace(":", "_").replace("/", "_") + ".sif"
            # If we are building from a file, add the build context
            elif self.definition_file:
                self.image_tags['build_context'] = self.context_checksum()
                if self.image_name is None:
                    bn = PurePath(self.definition_file).name
                    bn = str(bn).replace(".def", ".sif")
                    self.image_name = bn
            elif self.definition_content:
                self.image_tags['build_context'] = self.context_checksum()
            if self.image_url:
                self.image_tags['image_url'] = self.image_url

        # Final fall back for image name
        if self.image_name is None:
            self.image_name = "image.sif"
        if self.image_name and not self.image_name.endswith(".sif"):
            self.image_name = f'{self.image_name}.sif'
        # Add image name to the tags
        self.image_tags['image_name'] = self.image_name

    def _prep_work_order_before_create(self) -> Dict[str, str]:
        """
        Prep work order before creation.

        Returns:
            Workorder for singularity build.
        """
        self.__add_tags()
        self.load_work_order(SB_BASE_WORKER_PATH)
        if self.definition_file or self.definition_content:
            self.work_order['Build']['Input'] = "Assets/Singularity.def"
        else:
            self.work_order['Build']['Input'] = self.image_url
        if len(self.environment_variables):
            self.work_order['Build']['StaticEnvironment'] = self.environment_variables
        if len(self.additional_mounts):
            self.work_order['Build']['AdditionalMounts'] = self.additional_mounts
        self.work_order['Build']['Output'] = self.image_name if self.image_name else "image.sif"
        self.work_order['Build']['Tags'] = self.image_tags
        self.work_order['Build']['Flags'] = dict()
        if self.fix_permissions:
            self.work_order['Build']['Flags']['Switches'] = ["--fix-perms"]
        if self.library:
            self.work_order['Build']['Flags']['--library'] = self.library
        if self.section:
            self.work_order['Build']['Flags']['--section'] = self.section
        return self.work_order

[docs]    def pre_creation(self, platform: 'IPlatform') -> None:
        """
        Pre-Creation item.

        Args:
            platform: Platform object

        Returns:
            None
        """
        if self.name is None:
            self.name = "Singularity Build"
            if self.definition_file:
                self.name += f" of {PurePath(self.definition_file).name}"
        super(SingularityBuildWorkItem, self).pre_creation(platform)
        self.__add_common_assets()
        self._prep_work_order_before_create()

    def __add_common_assets(self):
        """
        Add common assets which in this case is the singularity definition file.

        Returns:
            None
        """
        self.render_template()
        if self.definition_file:
            opts = dict(content=self.__rendered_template) if self.is_template else dict(absolute_path=self.definition_file)
            self.assets.add_or_replace_asset(Asset(filename="Singularity.def", **opts))
        elif self.definition_content:
            opts = dict(content=self.__rendered_template if self.is_template else self.definition_content)
            self.assets.add_or_replace_asset(Asset(filename="Singularity.def", **opts))

    def __fetch_finished_asset_collection(self, platform: 'IPlatform') -> Union[AssetCollection, None]:
        """
        Fetch the Singularity asset collection we created.

        Args:
            platform: Platform to fetch from.

        Returns:
            Asset Collection or None
        """
        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=platform if platform else self.platform)
            if IdmConfigParser.is_output_enabled():
                user_logger.log(SUCCESS, f"Created Singularity image as Asset Collection: {self.asset_collection.id}")
                user_logger.log(SUCCESS, f"View AC at {self.platform.get_asset_collection_link(self.asset_collection)}")
            return self.asset_collection
        return None

[docs]    def run(self, wait_until_done: bool = True, platform: 'IPlatform' = None, wait_on_done_progress: bool = True, **run_opts) -> Optional[AssetCollection]:
        """
        Run the build.

        Args:
            wait_until_done: Wait until build completes
            platform: Platform to run on
            wait_on_done_progress: Show progress while waiting
            **run_opts: Extra run options

        Returns:
            Asset collection that was created if successful
        """
        p = super()._check_for_platform_from_context(platform)
        opts = dict(wait_on_done_progress=wait_on_done_progress, wait_until_done=wait_until_done, platform=p, wait_progress_desc=f"Waiting for build of Singularity container: {self.name}")
        ac = self.find_existing_container(self, platform=p)
        if ac is None or self.force:
            super().run(**opts)
            ac = self.asset_collection

        else:
            if IdmConfigParser.is_output_enabled():
                user_logger.log(SUCCESS, f"Existing build of image found with Asset Collection ID of {ac.id}")
                user_logger.log(SUCCESS, f"View AC at {self.platform.get_asset_collection_link(ac)}")
            # Set id to None
            self.uid = None
            if ac:
                self.image_tags = ac.tags
            self.asset_collection = ac
            # how do we get id for original work item from AC?
            self.status = EntityStatus.SUCCEEDED

        save_sif_asset_md5_from_ac_id(ac.id)
        return self.asset_collection

[docs]    def wait(self, wait_on_done_progress: bool = True, timeout: int = None, refresh_interval=None, platform: 'IPlatform' = None, wait_progress_desc: str = None) -> Optional[AssetCollection]:
        """
        Waits on Singularity Build Work item to finish and fetches the resulting asset collection.

        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
            wait_progress_desc: Wait Progress Description Text

        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, wait_progress_desc=wait_progress_desc)

        super().wait(**opts)
        if self.status == EntityStatus.SUCCEEDED:
            return self.__fetch_finished_asset_collection(p)
        return None

[docs]    def get_id_filename(self, prefix: str = None) -> str:
        """
        Determine the id filename. Mostly used when use does not provide one.

        The logic is combine prefix and either
        * definition file minus extension
        * image url using with parts filtered out of the name.

        Args:
            prefix: Optional prefix.

        Returns:
            id file name

        Raises:
            ValueError - When the filename cannot be calculated
        """
        if prefix is None:
            prefix = ''
        if self.definition_file:
            base_name = PurePath(self.definition_file).name.replace(".def", ".id")
            if prefix:
                base_name = f"{prefix}{base_name}"
            filename = str(PurePath(self.definition_file).parent.joinpath(base_name))
        elif self.image_url:
            filename = re.sub(r"(docker|shub)://", "", self.image_url).replace(":", "_")
            if filename:
                filename = f"{prefix}{filename}"
        else:
            raise ValueError("Could not calculate the filename. Please specify one")
        if not filename.endswith(".id"):
            filename += ".id"

        return filename

[docs]    def to_id_file(self, filename: Union[str, PathLike] = None, save_platform: bool = False):
        """
        Create an ID File.

        If the filename is not provided, it will be calculate for definition files or for docker pulls

        Args:
            filename: Filename
            save_platform: Save Platform info to file as well

        Returns:
            None
        """
        if filename is None:
            filename = self.get_id_filename(prefix='builder.')
        super(SingularityBuildWorkItem, self).to_id_file(filename, save_platform)