-
Notifications
You must be signed in to change notification settings - Fork 1k
Description
Objective
The goal of this proposal is to redesign the Mesa batch runner into a modular, flexible system that separates the batch run process into three stages: Preparation, Running, and Processing. The focus will be on the Preparation stage, where different experimental designs can be used to generate run configurations. These configurations will be encapsulated in a dataclass that includes the model class and all relevant parameters, ensuring reusability in the Running stage.
Design Overview
- Preparation Stage:
- Use a dataclass (
RunConfiguration
) to store the model class, run parameters, and configuration details (e.g.,max_steps
,data_collection_period
). - Implement different configuration generators (e.g., full factorial, sparse grids, manual) to allow for flexible experiment designs.
- Use a dataclass (
- Running Stage:
- The batch runner will execute all configurations using multiprocessing when necessary. It will take a list of
RunConfiguration
objects and execute each run independently. - Results will be collected during execution and processed after all runs are completed.
- The batch runner will execute all configurations using multiprocessing when necessary. It will take a list of
- Processing Stage:
- Results from the batch run will be processed into a usable format (e.g., a list of dictionaries, pandas DataFrames) for further analysis.
Key Components
1. RunConfiguration
Dataclass
This dataclass stores all the information required to run a single configuration of the experiment.
from dataclasses import dataclass
from typing import Any, Dict
@dataclass
class RunConfiguration:
model_cls: type[Model]
run_id: int
iteration: int
parameters: Dict[str, Any]
max_steps: int
data_collection_period: int
2. Configuration Generators
Provide different strategies for generating configurations:
- Full Factorial: Generate all combinations of parameters.
- Sparse Grid: Sample a subset of parameter space.
- Base case: Start with a reference scenario and vary parameters from there.
- Manual Configuration: Allow users to specify configurations explicitly.
Each generator will output a list of RunConfiguration
objects.
3. Batch Runner Class
The BatchRunner
class will manage the execution of all runs using the RunConfiguration
objects. It will handle multiprocessing, progress tracking, and result collection.
class BatchRunner:
def __init__(
self,
configurations: List[RunConfiguration],
number_processes: int | None = 1,
display_progress: bool = True,
):
self.configurations = configurations
self.number_processes = number_processes
self.display_progress = display_progress
def run_all(self) -> List[Dict[str, Any]]:
# Core logic to run all configurations in parallel or serially