Source code for idmtools.core.experiment_factory

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__)

[docs]class ExperimentFactory: """ ExperimentFactory allows creating experiments that could be derived through plugins. """ DEFAULT_KEY = 'idmtools.entities.experiment.Experiment'
[docs] 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)
[docs] 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()