Source code for idmtools_platform_comps.comps_operations.suite_operations

"""idmtools comps suite operations.

Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
from dataclasses import dataclass, field
from typing import Any, List, Dict, Tuple, Union, Type, TYPE_CHECKING, Optional
from uuid import UUID
from logging import getLogger
from COMPS.Data import Suite as COMPSSuite, QueryCriteria, Experiment as COMPSExperiment, WorkItem
from idmtools.core import ItemType
from idmtools.entities import Suite
from idmtools.entities.iplatform_ops.iplatform_suite_operations import IPlatformSuiteOperations

if TYPE_CHECKING:  # pragma: no cover
    from idmtools_platform_comps.comps_platform import COMPSPlatform

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

[docs]@dataclass class CompsPlatformSuiteOperations(IPlatformSuiteOperations): """ Provides Suite operation to the COMPSPlatform. """ platform: 'COMPSPlatform' # noqa F821 platform_type: Type = field(default=COMPSSuite)
[docs] def get(self, suite_id: UUID, columns: Optional[List[str]] = None, load_children: Optional[List[str]] = None, query_criteria: Optional[QueryCriteria] = None, **kwargs) -> COMPSSuite: """ Get COMPS Suite. Args: suite_id: Suite id columns: Optional list of columns. Defaults to id and name load_children: Optional list of children to load. Defaults to "tags", "configuration" query_criteria: Optional query criteria **kwargs: Returns: COMPSSuite """ columns = columns or ["id", "name"] children = load_children if load_children is not None else ["tags", "configuration"] # Comps doesn't like getting uuids for some reason query_criteria = query_criteria or QueryCriteria().select(columns).select_children(children) s = COMPSSuite.get(id=str(suite_id), query_criteria=query_criteria) return s
[docs] def platform_create(self, suite: Suite, **kwargs) -> Tuple[COMPSSuite, UUID]: """ Create suite on COMPS. Args: suite: Suite to create **kwargs: Returns: COMPS Suite object and a UUID """ self.platform._login() # Create suite comps_suite = COMPSSuite(, description=suite.description) comps_suite.set_tags(suite.tags) # Update suite uid suite.uid = return comps_suite, suite.uid
[docs] def get_parent(self, suite: COMPSSuite, **kwargs) -> Any: """ Get parent of suite. We always return None on COMPS. Args: suite:Suite to get parent of **kwargs: Returns: None """ return None
[docs] def get_children(self, suite: COMPSSuite, **kwargs) -> List[Union[COMPSExperiment, WorkItem]]: """ Get children for a suite. Args: suite: Suite to get children for **kwargs: Any arguments to pass on to loading functions Returns: List of COMPS Experiments/Workitems that are part of the suite """ cols = kwargs.get("cols") children = kwargs.get("children") cols = cols or ["id", "name", "suite_id"] children = children if children is not None else ["tags"] children = suite.get_experiments(query_criteria=QueryCriteria().select(cols).select_children(children)) return children
[docs] def refresh_status(self, suite: Suite, **kwargs): """ Refresh the status of a suite. On comps, this is done by refreshing all experiments. Args: suite: Suite to refresh status of **kwargs: Returns: None """ for experiment in suite.experiments: self.platform.refresh_status(experiment)
[docs] def to_entity(self, suite: COMPSSuite, children: bool = True, **kwargs) -> Suite: """ Convert a COMPS Suite to an IDM Suite. Args: suite: Suite to Convert children: When true, load simulations, false otherwise **kwargs: Returns: IDM Suite """ # Creat a suite obj = Suite() # Set its correct attributes obj.uid = = obj.description = suite.description obj.tags = suite.tags obj.comps_suite = suite # Convert all experiments if children: comps_exps = suite.get_experiments() obj.experiments = [] for exp in comps_exps: self.platform._experiments.to_entity(exp, parent=obj, **kwargs) return obj
[docs] def create_sim_directory_map(self, suite_id: str) -> Dict: """ Build simulation working directory mapping. Args: suite_id: suite id Returns: Dict of simulation id as key and working dir as value """ # s = Suite.get(suite_id) comps_suite = self.platform.get_item(suite_id, ItemType.SUITE, raw=True, force=True) comps_exps = comps_suite.get_experiments(QueryCriteria().select('id')) sims_map = {} for exp in comps_exps: r = self.platform._experiments.create_sim_directory_map( sims_map = {**sims_map, **r} return sims_map
[docs] def platform_delete(self, suite_id: str) -> None: """ Delete platform suite. Args: suite_id: platform suite id Returns: None """ comps_suite = self.platform.get_item(suite_id, ItemType.SUITE, raw=True) comps_exps = comps_suite.get_experiments() for comps_exp in comps_exps: try: comps_exp.delete() except RuntimeError:"Could not delete the associated experiment ({})...") return try: comps_suite.delete() except RuntimeError:"Could not delete suite ({suite_id})...") return