"""Classes for running parameter sweeps over scenarios."""
from __future__ import annotations
import traceback
from collections.abc import Callable, Iterable
from concurrent.futures import BrokenExecutor, ProcessPoolExecutor, as_completed
from functools import partial
from typing import TYPE_CHECKING, Any
import pandas as pd
from modelly.scenarios.exceptions import (
FailureInfo,
FailureOrigin,
ModelInstantiationException,
ModelRunException,
OutcomeExtractionException,
RunStageException,
)
from modelly.scenarios.store import InMemoryStore, RunId
if TYPE_CHECKING:
from concurrent.futures import Executor
from modelly.core.model import Model
from modelly.scenarios.scenario import Scenario
from modelly.scenarios.store import Reference, Store, Writer
[docs]
class RunConfiguration:
"""Defines how a single Scenario is executed and what is extracted from it.
Can be used as is for simple use cases or subclassed by overriding one or more of the following
methods
- ``instantiate_model`` — construct a Model from a Scenario (default:
``model_class(*model_args, scenario=scenario, **model_kwargs)``).
- ``run_model`` — advance the model. Default delegates to ``model.run_until`` based on the ``until`` attribute.
Override for alternative run control
- ``extract_output`` — return a dict with outcome names as key and dataframes as values
Stopping is the model's responsibility. ``RunConfiguration`` only chooses
which run primitive to call.
"""
def __init__(
self,
model_class: type[Model],
until: float | int,
model_args: list[Any] | None = None,
model_kwargs: dict[str, Any] | None = None,
outcomes: str | list[str] | None = None,
data_recorder_attr_name="data_recorder",
):
"""Initialize a RunConfiguration object.
Args:
model_class: the model class to instantiate
until: until which time to run the model
model_args: any additional model arguments
model_kwargs: any additional model keyword arguments
outcomes: the outcomes to extract. If None, extract all outcomes.
data_recorder_attr_name : the name of the data recorder attribute to use on the model
"""
# we need to avoid circular imports
from modelly.core.model import Model # noqa: PLC0415
if not (isinstance(model_class, type) and issubclass(model_class, Model)):
raise TypeError("model_class must be a subclass of Model")
if not isinstance(until, (int, float)):
raise TypeError("until must be an int or float")
if until <= 0:
raise ValueError("until must be positive")
self.model_class = model_class
self.model_args = [] if model_args is None else model_args
self.model_kwargs = {} if model_kwargs is None else model_kwargs
self.until = until
# fixme:: this code leaves it to the user to set the attribute to which the recorder is assigned
# this is probably a a convention that needs to be pinned down explicitly on the model class
self.data_recorder_attr_name = data_recorder_attr_name
if isinstance(outcomes, str):
outcomes = [outcomes]
self.outcomes = outcomes
[docs]
def instantiate_model(self, scenario: Scenario) -> Model:
"""Instantiate the model."""
return self.model_class(
*self.model_args, scenario=scenario, **self.model_kwargs
)
[docs]
def run_model(self, model: Model) -> None:
"""Run the model."""
model.run_until(self.until)
def __call__(self, scenario: Scenario) -> dict[str, pd.DataFrame]:
"""Run the scenario and extract output."""
try:
model = self.instantiate_model(scenario)
except Exception as e:
raise ModelInstantiationException(
self.model_class, self.model_args, self.model_kwargs, scenario
) from e
try:
self.run_model(model)
except Exception as e:
raise ModelRunException(
RunId(scenario.scenario_id, scenario.replication_id)
) from e
try:
output = self.extract_output(model)
except Exception as e:
raise OutcomeExtractionException(
RunId(scenario.scenario_id, scenario.replication_id), self.outcomes
) from e
return output
def _safe_call(
config: RunConfiguration,
scenario: Scenario,
writer: Writer,
) -> tuple[Reference, None] | tuple[None, FailureInfo]:
"""Run one scenario and persist its outcome. Runs in the worker.
Args:
config: a RunConfiguration instance
scenario: a Scenario instance
writer: a Writer instance
Returns (reference, None) on success or (None, failure_info) on a failure
raised inside the run or the writer. Catching here means a model
error becomes data (a FailureInfo instance) rather than an exception
crossing the process boundary, so one failed run never aborts a parameter sweep.
"""
run_id = RunId(scenario.scenario_id, scenario.replication_id)
try:
outcome = config(scenario)
except RunStageException as e:
cause = e.__cause__ or e
return None, FailureInfo(
origin=e.origin,
exception_type=type(cause).__name__,
message=str(cause),
traceback="".join(traceback.format_exception(e)),
)
try:
ref = writer.to_reference(run_id, outcome)
except Exception as e:
return None, FailureInfo(
origin=FailureOrigin.WRITING,
exception_type=type(e).__name__,
message=str(e),
traceback="".join(traceback.format_exception(e)),
)
return ref, None
[docs]
def run_scenarios(
scenarios: Iterable[Scenario],
config: RunConfiguration,
*,
executor: Executor | None = None,
store: Store | None = None,
progress: bool = True,
workers: int | None = None,
) -> Store:
"""Run the scenarios and return a Store object.
Args:
scenarios: an iterable of scenarios to run
Scenarios to execute. For replications, construct these via
``MyScenario.from_dataframe(df, replications=n)`` — replication is
handled at scenario construction, not here.
config: a RunConfiguration instance
Per-scenario execution unit. Must be picklable when using a
distributed executor (e.g., ProcessPoolExecutor).
executor: an executor to run the scenarios
Execution backend. If None, scenarios run sequentially in the
calling thread (useful for debugging and small experiments).
Otherwise, pass a user-constructed executor; its lifetime is the
caller's responsibility (use a ``with`` block).
store: the Storage backend to use
progress: whether to display the progress
Display a progress bar via ``tqdm`` if installed.
workers: run the scenarios in this many processes
A shortcut for passing a ``ProcessPoolExecutor``: the pool is created
and closed here. Ignored when ``executor`` is given. The model class,
the scenarios and the config must be importable (picklable).
Returns: a Store instance
"""
if workers is not None and executor is None:
with ProcessPoolExecutor(max_workers=workers) as pool:
return run_scenarios(
scenarios, config, executor=pool, store=store, progress=progress
)
if store is None:
store = InMemoryStore()
scenarios = list(scenarios)
writer = store.writer()
store.write_scenarios(scenarios)
bar = partial(_progress_bar, progress=progress, total=len(scenarios))
if executor is None:
# Sequential: run in the loop so the bar advances per scenario.
for scenario in bar(scenarios):
_record(store, scenario, _safe_call(config, scenario, writer))
return store
futures = {
executor.submit(_safe_call, config, scenario, writer): scenario
for scenario in scenarios
}
try:
_collect_futures(futures, store, bar)
except BrokenExecutor as e:
_abort_pending(store, e)
return store
def _progress_bar(iterable: Iterable, progress: bool, total: int) -> Iterable:
"""Wrap iterable in a tqdm progress bar when progress is on and tqdm is installed."""
if not progress:
return iterable
try:
from tqdm.auto import tqdm # noqa: PLC0415
return tqdm(iterable, total=total, desc="Running scenarios")
except ImportError:
return iterable
def _record(
store: Store,
scenario: Scenario,
result: tuple[Reference, None] | tuple[None, FailureInfo],
) -> None:
"""Mark the run of scenario as succeeded or failed in store."""
ref, failure = result
if failure is None:
store.mark_succeeded(ref)
else:
run_id = RunId(scenario.scenario_id, scenario.replication_id)
store.mark_failed(run_id, failure)
def _failure_info(
origin: FailureOrigin, cause: BaseException, exc: BaseException
) -> FailureInfo:
"""Describe a failure from its cause, with the traceback of exc."""
return FailureInfo(
origin=origin,
exception_type=type(cause).__name__,
message=str(cause),
traceback="".join(traceback.format_exception(exc)),
)
def _collect_futures(futures: dict, store: Store, bar: Callable) -> None:
"""Record each future's result as it completes; a broken executor propagates."""
for future in bar(as_completed(futures)):
scenario = futures[future]
try:
result = future.result()
except BrokenExecutor:
raise
except Exception as e:
# pickling failure or CancelledError on the return trip; record as failed and continue
result = (None, _failure_info(FailureOrigin.WRITING, e, e))
_record(store, scenario, result)
def _abort_pending(store: Store, error: BrokenExecutor) -> None:
"""Mark every run still pending in store as aborted by a broken executor."""
cause = error.__cause__ or error
for entry in list(store.pending()):
store.mark_aborted(entry, _failure_info(FailureOrigin.ABORTED, cause, error))
