Skip to content

Utilities API

Configuration utilities for AID2E Framework.

ArtifactSpec

Bases: BaseModel

Output artifact specification for a stage.

Defines expected output files that stages/jobs produce.

Attributes:

Name Type Description
path str

File path pattern (e.g., "objectives_*.json").

format str

File format ("json", "yaml", "csv", or "root").

Example

artifact = ArtifactSpec(path="objectives_*.json", format="json")

Source code in src/aid2e/utilities/configurations/workflow_config.py
75
76
77
78
79
80
81
82
83
84
85
86
87
88
class ArtifactSpec(BaseModel):
    """Output artifact specification for a stage.

    Defines expected output files that stages/jobs produce.

    Attributes:
        path: File path pattern (e.g., "objectives_*.json").
        format: File format ("json", "yaml", "csv", or "root").

    Example:
        >>> artifact = ArtifactSpec(path="objectives_*.json", format="json")
    """
    path: str = Field(..., description="File path pattern (e.g., 'output_*.json')")
    format: str = Field(default="json", pattern="^(json|yaml|csv|root)$", description="File format")

BaseParameter

Bases: BaseModel

Abstract base class for all parameter types. Subclasses should define specific parameter characteristics.

Source code in src/aid2e/utilities/configurations/base_models.py
 5
 6
 7
 8
 9
10
11
12
13
14
15
class BaseParameter(BaseModel):
    """
    Abstract base class for all parameter types.
    Subclasses should define specific parameter characteristics.
    """
    name: str
    type: str  # Discriminator
    value: Union[float, str, int]  # Generic value field

    class Config:
        extra = "allow"  # Allow subclasses to add fields

BranchDefinition

Bases: BaseModel

Branch definition (optional, for organizing stages in a DAG).

A branch is an optional subgraph inside a workflow, useful when you want multiple independent pipelines under one workflow (e.g., "physics sim" branch + "surrogate" branch). Stages within a branch are executed in topological order.

Attributes:

Name Type Description
name str

Branch name (e.g., "main", "physics_sim", "surrogate").

stages List[StageDefinition]

List of stages in DAG order (assumes simple sequential order; extend with explicit DAG if needed).

scheduler Optional[SchedulerConfiguration]

Optional branch-level scheduler default (used if stage not set).

Example

branch = BranchDefinition( ... name="main", ... stages=[ ... StageDefinition(name="evaluate", ...), ... StageDefinition(name="aggregate", ...) ... ] ... )

Notes
  • Multiple branches in one workflow execute independently
  • For complex DAGs, extend this model with explicit edge definitions
Source code in src/aid2e/utilities/configurations/workflow_config.py
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
class BranchDefinition(BaseModel):
    """Branch definition (optional, for organizing stages in a DAG).

    A branch is an optional subgraph inside a workflow, useful when you want
    multiple independent pipelines under one workflow (e.g., "physics sim" branch
    + "surrogate" branch). Stages within a branch are executed in topological order.

    Attributes:
        name: Branch name (e.g., "main", "physics_sim", "surrogate").
        stages: List of stages in DAG order (assumes simple sequential order; extend with explicit DAG if needed).
        scheduler: Optional branch-level scheduler default (used if stage not set).

    Example:
        >>> branch = BranchDefinition(
        ...     name="main",
        ...     stages=[
        ...         StageDefinition(name="evaluate", ...),
        ...         StageDefinition(name="aggregate", ...)
        ...     ]
        ... )

    Notes:
        - Multiple branches in one workflow execute independently
        - For complex DAGs, extend this model with explicit edge definitions
    """
    name: str = Field(..., description="Branch name")
    stages: List[StageDefinition] = Field(default_factory=list, description="Stages in execution order")
    scheduler: Optional[SchedulerConfiguration] = Field(
        default=None,
        description="Branch-level scheduler default (overrides workflow, used if stage unset)",
    )

ChoiceParameter

Bases: BaseParameter

Categorical parameter with discrete choices.

Source code in src/aid2e/utilities/configurations/base_models.py
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
class ChoiceParameter(BaseParameter):
    """Categorical parameter with discrete choices."""
    value: Union[str, int]
    choices: Union[List[str], List[int]]

    @property
    def type(self) -> Literal["choice"]:
        return "choice"

    @model_validator(mode='after')
    def check_value_choices_consistency(self) -> "ChoiceParameter":
        value = self.value
        choices = self.choices

        for choice in choices:
            is_same = type(value) == type(choice)
            if not is_same:
                raise ValueError("Type of each choice must match type of value ({type(value)}), but {choice} is a {type(choice)}")

        return self

CombinedObjectiveMetric

Bases: BaseModel

Metric emitted by a combined objective plan.

Attributes:

Name Type Description
name str

Objective name (e.g., "f1").

direction ObjectiveDirection

Optimization direction for this metric.

metric_key str

Key in the plan output to extract this metric.

Source code in src/aid2e/utilities/configurations/workflow_config.py
24
25
26
27
28
29
30
31
32
33
34
35
class CombinedObjectiveMetric(BaseModel):
    """Metric emitted by a combined objective plan.

    Attributes:
        name: Objective name (e.g., "f1").
        direction: Optimization direction for this metric.
        metric_key: Key in the plan output to extract this metric.
    """

    name: str = Field(..., description="Objective metric name")
    direction: ObjectiveDirection = Field(..., description="Direction for this metric")
    metric_key: str = Field(..., description="Key in plan output for this metric")

CombinedObjectivePlan

Bases: BaseModel

Combined objective execution producing multiple metrics in one plan.

Attributes:

Name Type Description
name str

Identifier for the combined objective bundle.

objective_plan ObjectivePlanSpec

Plan executed once to emit multiple metrics.

metrics list[CombinedObjectiveMetric]

Metrics extracted from the plan output with their directions.

scheduler Optional[SchedulerConfiguration]

Optional scheduler default for this combined plan.

Source code in src/aid2e/utilities/configurations/workflow_config.py
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
class CombinedObjectivePlan(BaseModel):
    """Combined objective execution producing multiple metrics in one plan.

    Attributes:
        name: Identifier for the combined objective bundle.
        objective_plan: Plan executed once to emit multiple metrics.
        metrics: Metrics extracted from the plan output with their directions.
        scheduler: Optional scheduler default for this combined plan.
    """

    name: str = Field(..., description="Combined objective bundle name")
    objective_plan: ObjectivePlanSpec = Field(..., description="Plan producing multiple metrics")
    metrics: list[CombinedObjectiveMetric] = Field(
        ..., min_items=1, description="Metrics emitted by this plan"
    )
    scheduler: Optional[SchedulerConfiguration] = Field(
        default=None,
        description="Scheduler default for this combined plan",
    )

DesignConfig

Bases: BaseModel

Complete design configuration with parameters and constraints.

Encapsulates a design space including all parameter groups, their bounds/choices, and constraints on valid parameter combinations. Validates constraint syntax at load time and provides validated constraints for optimizer integration.

This is the base class for specialized configurations (e.g., EpicDesignConfig) and supports generic toy problems (DTLZ2, etc.).

Attributes:

Name Type Description
design_parameters DesignParameters

Collection of parameter groups defining the design space. Can be specialized for specific contexts such as EpicDesignConfig.

parameter_constraints Optional[List[ParameterConstraint]]

List of constraints on valid parameter combinations.

key str

YAML key associated with model.

Example

config = DesignConfig( ... design_parameters=DesignParameters(...), ... parameter_constraints=[ParameterConstraint(...)], ... optimization_groups={...}, ... ) names = config.get_parameter_names() bounds = config.get_parameter_bounds('group.param')

Constraints are already syntax-validated

search_space = SearchSpace.from_design_config(config)

Source code in src/aid2e/utilities/configurations/design_config.py
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
class DesignConfig(BaseModel):
    """Complete design configuration with parameters and constraints.

    Encapsulates a design space including all parameter groups, their bounds/choices,
    and constraints on valid parameter combinations. Validates constraint syntax at
    load time and provides validated constraints for optimizer integration.

    This is the base class for specialized configurations (e.g., EpicDesignConfig)
    and supports generic toy problems (DTLZ2, etc.).

    Attributes:
        design_parameters: Collection of parameter groups defining the design space.
                           Can be specialized for specific contexts such as
                           EpicDesignConfig.
        parameter_constraints: List of constraints on valid parameter combinations.
        key: YAML key associated with model.

    Example:
        >>> config = DesignConfig(
        ...     design_parameters=DesignParameters(...),
        ...     parameter_constraints=[ParameterConstraint(...)],
        ...     optimization_groups={...},
        ... )
        >>> names = config.get_parameter_names()
        >>> bounds = config.get_parameter_bounds('group.param')
        >>> # Constraints are already syntax-validated
        >>> search_space = SearchSpace.from_design_config(config)
    """
    design_parameters: DesignParameters
    parameter_constraints: Optional[List[ParameterConstraint]] = Field(default_factory=list)
    optimization_groups: Optional[Dict[str, List[str]]] = Field(default_factory=dict)
    key: ClassVar[str] = 'design_space'

    @model_validator(mode='after')
    def validate_constraints_syntax(self) -> "DesignConfig":
        """Validate all constraint syntax and parameter references.

        Ensures all constraints:
        1. Have valid Python syntax
        2. Reference only parameters that exist in the design space
        3. Are well-formed boolean/comparison expressions

        This validation runs automatically when a DesignConfig is instantiated,
        catching configuration errors early.

        Returns:
            Self (for Pydantic validator chaining).

        Raises:
            ValueError: If any constraint has invalid syntax or references
                       unknown parameters.
        """
        if not self.parameter_constraints:
            return self

        valid_param_names = self.get_parameter_names()

        for constraint in self.parameter_constraints:
            is_valid, error_msg = constraint.validate_syntax(valid_param_names)
            if not is_valid:
                raise ValueError(
                    f"Invalid constraint '{constraint.name}': {error_msg}"
                )

        return self

    def get_flat_parameters(self) -> Dict[str, BaseParameter]:
        """Retrieve all parameters as a flat dictionary.

        Flattens the hierarchical group structure into a single dictionary mapping
        qualified parameter names to parameter objects.

        Returns:
            Dictionary mapping qualified names (e.g., "group.param")
            to BaseParameter objects.

        Example:
            >>> flat = config.get_flat_parameters()
            >>> param = flat['tracker.thickness']
        """
        flat = {}
        for group in self.design_parameters.root.values():
            for param in group.parameters.values():
                flat[param.name] = param
        return flat

    def get_parameter_names(self) -> List[str]:
        """Get all parameter qualified names in the design space.

        Returns a list of all unique qualified parameter names in the format
        'group_name.parameter_name'.

        Returns:
            Sorted list of qualified parameter names.

        Example:
            >>> names = config.get_parameter_names()
            >>> print(names)
            ['group1.param1', 'group1.param2', 'group2.param1']
        """
        return list(self.get_flat_parameters().keys())

    def get_parameter_bounds(self, param_name: str) -> Optional[Tuple[float, float]]:
        """Get bounds for a range parameter.

        Retrieves the lower and upper bounds for a RangeParameter by its
        qualified name. Returns None if the parameter is not found or
        does not have bounds (e.g., ChoiceParameter).

        Args:
            param_name: Qualified parameter name (e.g., "tracker.thickness").

        Returns:
            Tuple of (lower_bound, upper_bound) or None if not applicable.

        Raises:
            KeyError: If parameter name is not found (use get_flat_parameters
                     to verify existence first).

        Example:
            >>> bounds = config.get_parameter_bounds('tracker.thickness')
            >>> if bounds:
            ...     print(f"Range: {bounds[0]} to {bounds[1]}")
        """
        flat = self.get_flat_parameters()
        param = flat.get(param_name)
        if param and hasattr(param, 'bounds'):
            return param.bounds
        return None

    def get_parameter_choices(self, param_name: str) -> Optional[List[str]]:
        """Get choices for a choice parameter.

        Retrieves the list of valid choices for a ChoiceParameter by its
        qualified name. Returns None if the parameter is not found or
        does not have choices (e.g., RangeParameter).

        Args:
            param_name: Qualified parameter name (e.g., "detector.type").

        Returns:
            List of choice strings or None if not applicable.

        Example:
            >>> choices = config.get_parameter_choices('detector.type')
            >>> if choices:
            ...     print(f"Available: {choices}")
        """
        flat = self.get_flat_parameters()
        param = flat.get(param_name)
        if param and hasattr(param, 'choices'):
            return param.choices
        return None

    def get_optimization_group(self, group_name: str) -> Optional[List[str]]:
        """Get parameter names for a specific optimization group.

        Args:
            group_name: name of group to get parameters for

        Returns:
            List of parameters in group
        """
        return self.optimization_groups.get(group_name) if self.optimization_groups else None

    def get_all_optimization_groups(self) -> Dict[str, List[str]]:
        """Get all optimization groups.

        Returns:
            List of all optimization groups
        """
        return self.optimization_groups or {}

    def check_constraints(self, param_values: Dict[str, float]) -> Tuple[bool, List[str]]:
        """Check all constraints against provided parameter values at runtime.

        Evaluates each constraint rule with the given parameter values.
        This is primarily for non-Ax optimizers or manual validation.
        For Ax, constraints are passed directly to the optimizer.

        Args:
            param_values: Dictionary mapping qualified parameter names to
                          numeric values.

        Returns:
            Tuple of (all_valid, failed_constraint_names) where:
            - all_valid: True if all constraints passed, False otherwise.
            - failed_constraint_names: List of constraint names that failed.

        Example:
            >>> param_values = {
            ...     'tracker.thickness': 0.35,
            ...     'magnet.strength': 1.5
            ... }
            >>> is_valid, failures = config.check_constraints(param_values)
            >>> if not is_valid:
            ...     print(f"Failed constraints: {failures}")

        Notes:
            - Constraints are already syntax-validated at load time
            - For Ax optimizer, use config.parameter_constraints directly
        """
        if not self.parameter_constraints:
            return True, []

        failed = []
        for constraint in self.parameter_constraints:
            try:
                if not constraint.evaluate(param_values):
                    failed.append(constraint.name)
            except Exception as e:
                failed.append(f"{constraint.name} (error: {e})")

        return len(failed) == 0, failed

check_constraints(param_values)

Check all constraints against provided parameter values at runtime.

Evaluates each constraint rule with the given parameter values. This is primarily for non-Ax optimizers or manual validation. For Ax, constraints are passed directly to the optimizer.

Parameters:

Name Type Description Default
param_values Dict[str, float]

Dictionary mapping qualified parameter names to numeric values.

required

Returns:

Type Description
bool

Tuple of (all_valid, failed_constraint_names) where:

List[str]
  • all_valid: True if all constraints passed, False otherwise.
Tuple[bool, List[str]]
  • failed_constraint_names: List of constraint names that failed.
Example

param_values = { ... 'tracker.thickness': 0.35, ... 'magnet.strength': 1.5 ... } is_valid, failures = config.check_constraints(param_values) if not is_valid: ... print(f"Failed constraints: {failures}")

Notes
  • Constraints are already syntax-validated at load time
  • For Ax optimizer, use config.parameter_constraints directly
Source code in src/aid2e/utilities/configurations/design_config.py
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
def check_constraints(self, param_values: Dict[str, float]) -> Tuple[bool, List[str]]:
    """Check all constraints against provided parameter values at runtime.

    Evaluates each constraint rule with the given parameter values.
    This is primarily for non-Ax optimizers or manual validation.
    For Ax, constraints are passed directly to the optimizer.

    Args:
        param_values: Dictionary mapping qualified parameter names to
                      numeric values.

    Returns:
        Tuple of (all_valid, failed_constraint_names) where:
        - all_valid: True if all constraints passed, False otherwise.
        - failed_constraint_names: List of constraint names that failed.

    Example:
        >>> param_values = {
        ...     'tracker.thickness': 0.35,
        ...     'magnet.strength': 1.5
        ... }
        >>> is_valid, failures = config.check_constraints(param_values)
        >>> if not is_valid:
        ...     print(f"Failed constraints: {failures}")

    Notes:
        - Constraints are already syntax-validated at load time
        - For Ax optimizer, use config.parameter_constraints directly
    """
    if not self.parameter_constraints:
        return True, []

    failed = []
    for constraint in self.parameter_constraints:
        try:
            if not constraint.evaluate(param_values):
                failed.append(constraint.name)
        except Exception as e:
            failed.append(f"{constraint.name} (error: {e})")

    return len(failed) == 0, failed

get_all_optimization_groups()

Get all optimization groups.

Returns:

Type Description
Dict[str, List[str]]

List of all optimization groups

Source code in src/aid2e/utilities/configurations/design_config.py
415
416
417
418
419
420
421
def get_all_optimization_groups(self) -> Dict[str, List[str]]:
    """Get all optimization groups.

    Returns:
        List of all optimization groups
    """
    return self.optimization_groups or {}

get_flat_parameters()

Retrieve all parameters as a flat dictionary.

Flattens the hierarchical group structure into a single dictionary mapping qualified parameter names to parameter objects.

Returns:

Type Description
Dict[str, BaseParameter]

Dictionary mapping qualified names (e.g., "group.param")

Dict[str, BaseParameter]

to BaseParameter objects.

Example

flat = config.get_flat_parameters() param = flat['tracker.thickness']

Source code in src/aid2e/utilities/configurations/design_config.py
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
def get_flat_parameters(self) -> Dict[str, BaseParameter]:
    """Retrieve all parameters as a flat dictionary.

    Flattens the hierarchical group structure into a single dictionary mapping
    qualified parameter names to parameter objects.

    Returns:
        Dictionary mapping qualified names (e.g., "group.param")
        to BaseParameter objects.

    Example:
        >>> flat = config.get_flat_parameters()
        >>> param = flat['tracker.thickness']
    """
    flat = {}
    for group in self.design_parameters.root.values():
        for param in group.parameters.values():
            flat[param.name] = param
    return flat

get_optimization_group(group_name)

Get parameter names for a specific optimization group.

Parameters:

Name Type Description Default
group_name str

name of group to get parameters for

required

Returns:

Type Description
Optional[List[str]]

List of parameters in group

Source code in src/aid2e/utilities/configurations/design_config.py
404
405
406
407
408
409
410
411
412
413
def get_optimization_group(self, group_name: str) -> Optional[List[str]]:
    """Get parameter names for a specific optimization group.

    Args:
        group_name: name of group to get parameters for

    Returns:
        List of parameters in group
    """
    return self.optimization_groups.get(group_name) if self.optimization_groups else None

get_parameter_bounds(param_name)

Get bounds for a range parameter.

Retrieves the lower and upper bounds for a RangeParameter by its qualified name. Returns None if the parameter is not found or does not have bounds (e.g., ChoiceParameter).

Parameters:

Name Type Description Default
param_name str

Qualified parameter name (e.g., "tracker.thickness").

required

Returns:

Type Description
Optional[Tuple[float, float]]

Tuple of (lower_bound, upper_bound) or None if not applicable.

Raises:

Type Description
KeyError

If parameter name is not found (use get_flat_parameters to verify existence first).

Example

bounds = config.get_parameter_bounds('tracker.thickness') if bounds: ... print(f"Range: {bounds[0]} to {bounds[1]}")

Source code in src/aid2e/utilities/configurations/design_config.py
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
def get_parameter_bounds(self, param_name: str) -> Optional[Tuple[float, float]]:
    """Get bounds for a range parameter.

    Retrieves the lower and upper bounds for a RangeParameter by its
    qualified name. Returns None if the parameter is not found or
    does not have bounds (e.g., ChoiceParameter).

    Args:
        param_name: Qualified parameter name (e.g., "tracker.thickness").

    Returns:
        Tuple of (lower_bound, upper_bound) or None if not applicable.

    Raises:
        KeyError: If parameter name is not found (use get_flat_parameters
                 to verify existence first).

    Example:
        >>> bounds = config.get_parameter_bounds('tracker.thickness')
        >>> if bounds:
        ...     print(f"Range: {bounds[0]} to {bounds[1]}")
    """
    flat = self.get_flat_parameters()
    param = flat.get(param_name)
    if param and hasattr(param, 'bounds'):
        return param.bounds
    return None

get_parameter_choices(param_name)

Get choices for a choice parameter.

Retrieves the list of valid choices for a ChoiceParameter by its qualified name. Returns None if the parameter is not found or does not have choices (e.g., RangeParameter).

Parameters:

Name Type Description Default
param_name str

Qualified parameter name (e.g., "detector.type").

required

Returns:

Type Description
Optional[List[str]]

List of choice strings or None if not applicable.

Example

choices = config.get_parameter_choices('detector.type') if choices: ... print(f"Available: {choices}")

Source code in src/aid2e/utilities/configurations/design_config.py
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
def get_parameter_choices(self, param_name: str) -> Optional[List[str]]:
    """Get choices for a choice parameter.

    Retrieves the list of valid choices for a ChoiceParameter by its
    qualified name. Returns None if the parameter is not found or
    does not have choices (e.g., RangeParameter).

    Args:
        param_name: Qualified parameter name (e.g., "detector.type").

    Returns:
        List of choice strings or None if not applicable.

    Example:
        >>> choices = config.get_parameter_choices('detector.type')
        >>> if choices:
        ...     print(f"Available: {choices}")
    """
    flat = self.get_flat_parameters()
    param = flat.get(param_name)
    if param and hasattr(param, 'choices'):
        return param.choices
    return None

get_parameter_names()

Get all parameter qualified names in the design space.

Returns a list of all unique qualified parameter names in the format 'group_name.parameter_name'.

Returns:

Type Description
List[str]

Sorted list of qualified parameter names.

Example

names = config.get_parameter_names() print(names) ['group1.param1', 'group1.param2', 'group2.param1']

Source code in src/aid2e/utilities/configurations/design_config.py
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
def get_parameter_names(self) -> List[str]:
    """Get all parameter qualified names in the design space.

    Returns a list of all unique qualified parameter names in the format
    'group_name.parameter_name'.

    Returns:
        Sorted list of qualified parameter names.

    Example:
        >>> names = config.get_parameter_names()
        >>> print(names)
        ['group1.param1', 'group1.param2', 'group2.param1']
    """
    return list(self.get_flat_parameters().keys())

validate_constraints_syntax()

Validate all constraint syntax and parameter references.

Ensures all constraints: 1. Have valid Python syntax 2. Reference only parameters that exist in the design space 3. Are well-formed boolean/comparison expressions

This validation runs automatically when a DesignConfig is instantiated, catching configuration errors early.

Returns:

Type Description
DesignConfig

Self (for Pydantic validator chaining).

Raises:

Type Description
ValueError

If any constraint has invalid syntax or references unknown parameters.

Source code in src/aid2e/utilities/configurations/design_config.py
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
@model_validator(mode='after')
def validate_constraints_syntax(self) -> "DesignConfig":
    """Validate all constraint syntax and parameter references.

    Ensures all constraints:
    1. Have valid Python syntax
    2. Reference only parameters that exist in the design space
    3. Are well-formed boolean/comparison expressions

    This validation runs automatically when a DesignConfig is instantiated,
    catching configuration errors early.

    Returns:
        Self (for Pydantic validator chaining).

    Raises:
        ValueError: If any constraint has invalid syntax or references
                   unknown parameters.
    """
    if not self.parameter_constraints:
        return self

    valid_param_names = self.get_parameter_names()

    for constraint in self.parameter_constraints:
        is_valid, error_msg = constraint.validate_syntax(valid_param_names)
        if not is_valid:
            raise ValueError(
                f"Invalid constraint '{constraint.name}': {error_msg}"
            )

    return self

DesignConfigLoader

Load design configurations from YAML files with canonical resolution.

Supports either loading design configuration form an external file or from an inline YAML block. In both cases, data should include a design_space block containing design_parameters and optional parameter_constraints.

Class attributes

space_key: YAML key of the design space instance to parse, such as 'epic_design_space'. param_key: YAML key of the design parameters to extract, such as 'epic_design_parameters'. constrain_key: YAML key of the list of parameter constraints to extract, such as 'parameter_constraints'.

Example

Load design space from an external file

config = DesignConfigLoader.load('./configs/design.params')

Or load from an inlined design space in a YAML block

yaml = { "inline design" : { "design_space" : { "design_parameters" : { "group" : { "parameters" : {...} }, }, "parameter_constraints" : [...] } } } config = DesignConfigLoader.load(yaml)

Source code in src/aid2e/utilities/configurations/design_config.py
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
class DesignConfigLoader:
    """Load design configurations from YAML files with canonical resolution.

    Supports either loading design configuration form an external file or from an
    inline YAML block. In both cases, data should include a `design_space` block
    containing `design_parameters` and optional `parameter_constraints`.

    Class attributes:
        space_key: YAML key of the design space instance to parse, such as
                   'epic_design_space'.
        param_key: YAML key of the design parameters to extract, such as
                   'epic_design_parameters'.
        constrain_key: YAML key of the list of parameter constraints to extract,
                       such as 'parameter_constraints'.

    Example:
        >>> # Load design space from an external file
        >>> config = DesignConfigLoader.load('./configs/design.params')

        >>> # Or load from an inlined design space in a YAML block
        >>> yaml = {
        >>>     "inline design" : {
        >>>         "design_space" : {
        >>>             "design_parameters" : {
        >>>                 "group" : {
        >>>                     "parameters" : {...}
        >>>                 },
        >>>             },
        >>>             "parameter_constraints" : [...]
        >>>         }
        >>>     }
        >>> }
        >>> config = DesignConfigLoader.load(yaml)
    """
    space_key = DesignConfig.key
    param_key = DesignParameters.key
    constrain_key = ParameterConstraint.key

    @classmethod
    def _extract_design_space_payload(cls, raw: Dict[str, Any]) -> Dict[str, Any]:
        """Extract canonical design space payload from loaded data.

        Args:
            raw: Dictionary loaded from YAML file or inline config.

        Returns:
            Extracted design space as a dictionary.

        Raises:
            ValueError: If raw data is not a dict, is not canonical
                        or is missing required blocks.
        """
        if not isinstance(raw, dict):
            raise ValueError("Design space content must be a mapping.")

        payload = raw.get(cls.space_key, raw)
        if not isinstance(payload, dict):
            raise ValueError(f"{cls.space_key} must be a mapping.")
        if cls.param_key not in payload:
            raise KeyError(f"Required block {cls.param_key} not found in {cls.space_key}")

        # throw errors if any legacy configurations are being used
        if "design_constraints" in payload or "design_constraints" in raw:
            raise ValueError(
                "Legacy key 'design_constraints' is no longer supported. "
                "Use 'parameter_constraints'."
            )
        if cls.space_key not in raw:
            if cls.param_key in raw:
                raise ValueError(
                    f"Top-level {cls.param_key} is no longer supported. "
                    f"Wrap design content under {cls.space_key}."
                )

        return payload

    @classmethod
    def _resolve_design_space(cls, file_path: str) -> Dict[str, Any]:
        """Resolve design space from a file path

        Args:
            file_path: Path to the YAML design config file

        Returns:
            Dictionary with ``design_parameters`` and optional
            ``parameter_constraints``.

        Raises:
            ValueError: If both 'path' and inline definitions are present.
            FileNotFoundError: If referenced file does not exist.

        Notes:
            - Absolute paths are used as-is.
            - File not found errors include full resolved path in message.
        """
        if not Path(file_path).exists():
            raise FileNotFoundError(f"Design parameters file not found: {file_path}")

        payload = None
        with open(file_path, 'r') as f:
            loaded_data = yaml.safe_load(f)
            payload = cls._extract_design_space_payload(loaded_data)
        return payload

    @classmethod
    def _process_inputs(cls, file_path: str = None, design_data: Dict[str, Any] = None) -> Dict[str, Any]:
        """Process inputs to load

        Either loads a configuration file and extracts design space config,
        Or processes pre-loaded data to extract design space config. Returns
        the extracted design space config as a dictionary.

        Args:
            file_path: Path to the YAML design config file
            design_data: Loaded data stored in a dictionary

        Returns:
            Extracted data as dictionary mapping keys onto parameter groups and,
            if present, a list of parameter constraints

        Raises:
            RunTimeWarning: If both inline data and a file path were provided.
            FileNotFoundError: If the config file does not exist.
            ValueError: If config structure is invalid or references
                       a non-existent design.params file.
            yaml.YAMLError: If the YAML syntax is invalid.
            RunTimeError: If neither inline data nor a file path were provided

        Notes:
            - The configuration file must be valid YAML.
            - Must contain a top-level ``design_space`` key.
            - Directory of config_file is used as base for relative paths.
        """
        # should EITHER provide data as a dict OR a file path
        # as a string
        is_data_provided = design_data is not None
        is_file_provided = file_path is not None
        if is_data_provided and is_file_provided:
            raise RuntimeWarning(f"Both data and a file path ({file_path}) were provided. Defaulting to data.")

        payload = None
        if is_data_provided:
            payload = cls._extract_design_space_payload(design_data)
        elif is_file_provided:
            payload = cls._resolve_design_space(file_path=file_path)
        else:
            raise RuntimeError("Provide either data as a dictionary or a path to a file")

        data = {cls.param_key: payload[cls.param_key]}
        if cls.constrain_key in payload:
            data[cls.constrain_key] = payload[cls.constrain_key]
        if 'optimization_groups' in payload:
            data['optimization_groups'] = payload['optimization_groups']
        return data

    @staticmethod
    def load(file_path: str = None, design_data: Dict[str, Any] = None) -> "DesignConfig":
        """Load design configuration.

        Args:
            file_path: Path to the YAML design config file.
                       Should be resolved ahead of time.
            design_data: Loaded data stored in a dictionary

        Returns:
            DesignConfig instance ready for use in optimization workflows.

        Example:
            >>> config = DesignConfigLoader.load(file_path='examples/design.yml')
            >>> print(config.get_parameter_names())
            >>> is_valid, failures = config.validate_constraints({...})
        """
        data = DesignConfigLoader._process_inputs(file_path, design_data)
        return DesignConfig(**data)

load(file_path=None, design_data=None) staticmethod

Load design configuration.

Parameters:

Name Type Description Default
file_path str

Path to the YAML design config file. Should be resolved ahead of time.

None
design_data Dict[str, Any]

Loaded data stored in a dictionary

None

Returns:

Type Description
DesignConfig

DesignConfig instance ready for use in optimization workflows.

Example

config = DesignConfigLoader.load(file_path='examples/design.yml') print(config.get_parameter_names()) is_valid, failures = config.validate_constraints({...})

Source code in src/aid2e/utilities/configurations/design_config.py
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
@staticmethod
def load(file_path: str = None, design_data: Dict[str, Any] = None) -> "DesignConfig":
    """Load design configuration.

    Args:
        file_path: Path to the YAML design config file.
                   Should be resolved ahead of time.
        design_data: Loaded data stored in a dictionary

    Returns:
        DesignConfig instance ready for use in optimization workflows.

    Example:
        >>> config = DesignConfigLoader.load(file_path='examples/design.yml')
        >>> print(config.get_parameter_names())
        >>> is_valid, failures = config.validate_constraints({...})
    """
    data = DesignConfigLoader._process_inputs(file_path, design_data)
    return DesignConfig(**data)

DesignParameters

Bases: RootModel[Dict[str, ParameterGroup]]

Collection of parameter groups for generic design spaces.

Manages a hierarchical organization of design parameters grouped by context (subsystems, regions, etc.). Automatically injects fully qualified parameter names in the format "group_name.parameter_name" for unique identification.

The root model contains a dictionary mapping group names to ParameterGroup instances.

Attributes:

Name Type Description
root

Dictionary mapping group names to parameter groups.

key str

YAML key associated with models.

Example

params = DesignParameters(root={ ... 'tracker': ParameterGroup(parameters={...}), ... 'magnet': ParameterGroup(parameters={...}) ... })

Notes
  • Qualified names are injected at validation time.
  • Parameter uniqueness is enforced through qualified naming.
Source code in src/aid2e/utilities/configurations/design_config.py
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
class DesignParameters(RootModel[Dict[str, ParameterGroup]]):
    """Collection of parameter groups for generic design spaces.

    Manages a hierarchical organization of design parameters grouped by context
    (subsystems, regions, etc.). Automatically injects fully qualified parameter names
    in the format "group_name.parameter_name" for unique identification.

    The root model contains a dictionary mapping group names to ParameterGroup instances.

    Attributes:
        root: Dictionary mapping group names to parameter groups.
        key: YAML key associated with models.

    Example:
        >>> params = DesignParameters(root={
        ...     'tracker': ParameterGroup(parameters={...}),
        ...     'magnet': ParameterGroup(parameters={...})
        ... })

    Notes:
        - Qualified names are injected at validation time.
        - Parameter uniqueness is enforced through qualified naming.
    """
    key: ClassVar[str] = 'design_parameters'

    @model_validator(mode="before")
    @classmethod
    def inject_qualified_names(cls, values: Dict[str, dict]) -> Dict[str, dict]:
        """Inject fully qualified names into each parameter.

        Modifies parameter objects in-place to add 'name' attribute in the format
        'group_name.parameter_name' if not already present. This ensures every
        parameter has a globally unique identifier within the design space.

        Args:
            values: Dictionary mapping group names to group data dicts.

        Returns:
            Modified values dict with injected qualified names.

        Notes:
            This validator runs before model instantiation and is critical for
            the qualified naming system used throughout this module.
        """
        for group_name, group_data in values.items():
            param_dict = group_data.get("parameters", {})
            for param_name, param_data in param_dict.items():
                if isinstance(param_data, dict) and "name" not in param_data:
                    param_data["name"] = f"{group_name}.{param_name}"
        return values

inject_qualified_names(values) classmethod

Inject fully qualified names into each parameter.

Modifies parameter objects in-place to add 'name' attribute in the format 'group_name.parameter_name' if not already present. This ensures every parameter has a globally unique identifier within the design space.

Parameters:

Name Type Description Default
values Dict[str, dict]

Dictionary mapping group names to group data dicts.

required

Returns:

Type Description
Dict[str, dict]

Modified values dict with injected qualified names.

Notes

This validator runs before model instantiation and is critical for the qualified naming system used throughout this module.

Source code in src/aid2e/utilities/configurations/design_config.py
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
@model_validator(mode="before")
@classmethod
def inject_qualified_names(cls, values: Dict[str, dict]) -> Dict[str, dict]:
    """Inject fully qualified names into each parameter.

    Modifies parameter objects in-place to add 'name' attribute in the format
    'group_name.parameter_name' if not already present. This ensures every
    parameter has a globally unique identifier within the design space.

    Args:
        values: Dictionary mapping group names to group data dicts.

    Returns:
        Modified values dict with injected qualified names.

    Notes:
        This validator runs before model instantiation and is critical for
        the qualified naming system used throughout this module.
    """
    for group_name, group_data in values.items():
        param_dict = group_data.get("parameters", {})
        for param_name, param_data in param_dict.items():
            if isinstance(param_data, dict) and "name" not in param_data:
                param_data["name"] = f"{group_name}.{param_name}"
    return values

EnvironmentConfig

Bases: ABC, BaseModel

Configures environment variables

Generic base model for configuring environment variables. Must be specialized for specific for specifc contexts such as EpicConfiguration.

Example

class MyEnvConfig(EnvironmentConfiguration): ... geometry_install: str ... def activate(self) -> None: ... os.environ['GEOMETRY_INSTALL'] = self.geometry_install ... print(f"[INFO] Set $GEOMETRY_INSTALL to {self.geometry_install}")

Source code in src/aid2e/utilities/configurations/env_config.py
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
class EnvironmentConfig(ABC, BaseModel):
    """Configures environment variables

    Generic base model for configuring environment variables. Must
    be specialized for specific for specifc contexts such as
    EpicConfiguration.

    Example:
        >>> class MyEnvConfig(EnvironmentConfiguration):
        ...     geometry_install: str
        ...     def activate(self) -> None:
        ...         os.environ['GEOMETRY_INSTALL'] = self.geometry_install
        ...         print(f"[INFO] Set $GEOMETRY_INSTALL to {self.geometry_install}")
    """
    @property
    @abstractmethod
    def key(self) -> str:
        """YAML key associated with model (e.g. epic_environment_config)
        """
        pass

    @abstractmethod
    def activate(self) -> None:
        """
        Activate environment variables. Must be implemented
        by subclasses.
        """
        pass

key abstractmethod property

YAML key associated with model (e.g. epic_environment_config)

activate() abstractmethod

Activate environment variables. Must be implemented by subclasses.

Source code in src/aid2e/utilities/configurations/env_config.py
36
37
38
39
40
41
42
@abstractmethod
def activate(self) -> None:
    """
    Activate environment variables. Must be implemented
    by subclasses.
    """
    pass

EnvironmentConfigLoader

Bases: ABC

Loader for environment variables

Generic base class for loading environment config models. Must be specialized for specific contexts like EnvironmentConfig.

Example

class MyEnvConfigLoader(EnvironmentConfigLoader[MyEnvConfig]): ... @staticmethod ... def load(file_path: str) -> MyEnvConfigLoader: ... with open(file_path, 'r') as file: ... data = yaml.safe_load(file) ... return MyEnvConfigLoader(**data)

Source code in src/aid2e/utilities/configurations/env_config.py
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
class EnvironmentConfigLoader(ABC):
    """Loader for environment variables

    Generic base class for loading environment config
    models.  Must be specialized for specific contexts
    like EnvironmentConfig.

    Example:
        >>> class MyEnvConfigLoader(EnvironmentConfigLoader[MyEnvConfig]):
        ...     @staticmethod
        ...     def load(file_path: str) -> MyEnvConfigLoader:
        ...         with open(file_path, 'r') as file:
        ...             data = yaml.safe_load(file)
        ...         return MyEnvConfigLoader(**data)
    """

    @staticmethod
    @abstractmethod
    def load(env_data: Dict[str, Any] = None, file_path: str = None) -> "EnvironmentConfig":
        """
        Load an environment configuration from a YAML file.
        Must instantiate and return a subclass of
        EnvironmentConfig.
        """
        pass

load(env_data=None, file_path=None) abstractmethod staticmethod

Load an environment configuration from a YAML file. Must instantiate and return a subclass of EnvironmentConfig.

Source code in src/aid2e/utilities/configurations/env_config.py
61
62
63
64
65
66
67
68
69
@staticmethod
@abstractmethod
def load(env_data: Dict[str, Any] = None, file_path: str = None) -> "EnvironmentConfig":
    """
    Load an environment configuration from a YAML file.
    Must instantiate and return a subclass of
    EnvironmentConfig.
    """
    pass

FullConfig

Bases: BaseModel

Complete configuration combining problem, optimizer, scheduler, and workflows.

Source code in src/aid2e/utilities/configurations/full_config.py
19
20
21
22
23
24
class FullConfig(BaseModel):
    """Complete configuration combining problem, optimizer, scheduler, and workflows."""
    problem: ProblemConfiguration
    optimizer: OptimizerConfiguration
    scheduler: Optional[SchedulerConfiguration] = None
    workflows: Optional[WorkflowsConfiguration] = None

InlineObjective

Bases: BaseModel

Objective computed via inline Python function.

The entrypoint should reference a callable that accepts objective-step keyword arguments and returns a scalar value or metric mapping.

Attributes: entrypoint: Module and function reference (format: "module.path:function_name").

Example

inline = InlineObjective(entrypoint="my_objectives:compute_f1")

def compute_f1(*, design_point, inputs, outputs,

extra_args, xcom, work_dir, output_dir):

return {"f1": 0.5}

Source code in src/aid2e/utilities/configurations/objectives.py
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
class InlineObjective(BaseModel):
    """
    Objective computed via inline Python function.

   The entrypoint should reference a callable that accepts objective-step
   keyword arguments and returns a scalar value or metric mapping.

    Attributes:
        entrypoint: Module and function reference (format: "module.path:function_name").

   Example:
    >>> inline = InlineObjective(entrypoint="my_objectives:compute_f1")
    >>> # def compute_f1(*, design_point, inputs, outputs,
    >>> #                extra_args, xcom, work_dir, output_dir):
    >>> #     return {"f1": 0.5}
    """

    entrypoint: str = Field(
        ...,
        description="Module path and function reference (e.g., 'my_objectives:compute_f1')"
    )

    @field_validator('entrypoint')
    @classmethod
    def validate_entrypoint_format(cls, v: str) -> str:
        """Validate entrypoint has 'module.path:function_name' format."""
        if ':' not in v or v.count(':') != 1:
            raise ValueError("entrypoint must be 'module.path:function_name' format")
        module_part, func_part = v.split(':')
        if not module_part or not func_part:
            raise ValueError("entrypoint module and function names cannot be empty")
        if not all(c.isalnum() or c in '_.:-' for c in module_part):
            raise ValueError(f"Invalid module name: {module_part}")
        if not (func_part[0].isalpha() or func_part[0] == '_'):
            raise ValueError(f"Invalid function name: {func_part}")
        return v

validate_entrypoint_format(v) classmethod

Validate entrypoint has 'module.path:function_name' format.

Source code in src/aid2e/utilities/configurations/objectives.py
84
85
86
87
88
89
90
91
92
93
94
95
96
97
@field_validator('entrypoint')
@classmethod
def validate_entrypoint_format(cls, v: str) -> str:
    """Validate entrypoint has 'module.path:function_name' format."""
    if ':' not in v or v.count(':') != 1:
        raise ValueError("entrypoint must be 'module.path:function_name' format")
    module_part, func_part = v.split(':')
    if not module_part or not func_part:
        raise ValueError("entrypoint module and function names cannot be empty")
    if not all(c.isalnum() or c in '_.:-' for c in module_part):
        raise ValueError(f"Invalid module name: {module_part}")
    if not (func_part[0].isalpha() or func_part[0] == '_'):
        raise ValueError(f"Invalid function name: {func_part}")
    return v

JobDefinition

Bases: BaseModel

Single job/task definition within a stage.

A job is the smallest schedulable unit (e.g., one simulation, training run, etc). Jobs can be expanded from a template via job_factory.

Attributes:

Name Type Description
name str

Job name (e.g., "simulate").

command str

Executable command (e.g., "python scripts/dtlz2_problem.py").

payload Dict[str, Any]

Command arguments/payload (free-form dict, supports template substitution).

rule Optional[str]

Optional template for constructing final command from payload. Uses format: "{command} {payload[key1]} {payload[key2]}" etc. If not specified, defaults to "{command}" (just execute command).

resources Dict[str, Any]

Resource requirements (free-form dict, e.g., {"memory": "4GB"}).

outputs List[ArtifactSpec]

Output artifacts this job produces.

Example

job = JobDefinition( ... name="dtlz2_evaluate", ... command="python scripts/dtlz2_problem.py", ... rule="{{command}} {{payload[design_params_file]}} {{payload[output_dir]}} {payload[job_id]}", ... payload={ ... "design_params_file": "{{input_design_params}}", ... "output_dir": "{{output_dir}}", ... "job_id": "{{job_id}}" ... }, ... outputs=[ArtifactSpec(path="objectives_*.json", format="json")] ... )

Notes
  • Payload supports template substitution: {{job_id}}, {{output_dir}}, {{stage_outputs[stage_name]}}
  • Rule template follows experimental_stack.py StackLayer pattern
  • Resources dict is executor-dependent (e.g., JobLibRunner ignores, SlurmRunner uses)
Source code in src/aid2e/utilities/configurations/workflow_config.py
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
class JobDefinition(BaseModel):
    """Single job/task definition within a stage.

    A job is the smallest schedulable unit (e.g., one simulation, training run, etc).
    Jobs can be expanded from a template via job_factory.

    Attributes:
        name: Job name (e.g., "simulate").
        command: Executable command (e.g., "python scripts/dtlz2_problem.py").
        payload: Command arguments/payload (free-form dict, supports template substitution).
        rule: Optional template for constructing final command from payload.
              Uses format: "{command} {payload[key1]} {payload[key2]}" etc.
              If not specified, defaults to "{command}" (just execute command).
        resources: Resource requirements (free-form dict, e.g., {"memory": "4GB"}).
        outputs: Output artifacts this job produces.

    Example:
        >>> job = JobDefinition(
        ...     name="dtlz2_evaluate",
        ...     command="python scripts/dtlz2_problem.py",
        ...     rule="{{command}} {{payload[design_params_file]}} {{payload[output_dir]}} {payload[job_id]}",
        ...     payload={
        ...         "design_params_file": "{{input_design_params}}",
        ...         "output_dir": "{{output_dir}}",
        ...         "job_id": "{{job_id}}"
        ...     },
        ...     outputs=[ArtifactSpec(path="objectives_*.json", format="json")]
        ... )

    Notes:
        - Payload supports template substitution: {{job_id}}, {{output_dir}}, {{stage_outputs[stage_name]}}
        - Rule template follows experimental_stack.py StackLayer pattern
        - Resources dict is executor-dependent (e.g., JobLibRunner ignores, SlurmRunner uses)
    """
    name: str = Field(..., description="Job name")
    command: str = Field(..., description="Executable command")
    payload: Dict[str, Any] = Field(default_factory=dict, description="Command arguments and metadata")
    rule: Optional[str] = Field(
        default=None,
        description="Template rule for command construction (e.g., '{{command}} {{payload[input]}} {{payload[output]}}')"
    )
    resources: Dict[str, Any] = Field(default_factory=dict, description="Resource requirements")
    outputs: List[ArtifactSpec] = Field(default_factory=list, description="Output artifacts")

JobFactory

Bases: BaseModel

Factory for generating multiple jobs from a template.

Enables fan-out: creating N parallel jobs from one job definition. Useful for evaluating multiple design points in parallel.

Attributes:

Name Type Description
type str

Factory type ("range", "enumerate", "Cartesian", etc).

params Dict[str, Any]

Factory-specific parameters (e.g., {"n": 4} for range).

Example

Create 4 parallel design point evaluations

factory = JobFactory(type="range", params={"n": 4})

Notes
  • "range" type: creates N copies with job_id = 0..N-1
  • "enumerate" type: creates one job per item in a list
  • "Cartesian" type: creates N_A * N_B jobs from two parameter sets
Source code in src/aid2e/utilities/configurations/workflow_config.py
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
class JobFactory(BaseModel):
    """Factory for generating multiple jobs from a template.

    Enables fan-out: creating N parallel jobs from one job definition.
    Useful for evaluating multiple design points in parallel.

    Attributes:
        type: Factory type ("range", "enumerate", "Cartesian", etc).
        params: Factory-specific parameters (e.g., {"n": 4} for range).

    Example:
        >>> # Create 4 parallel design point evaluations
        >>> factory = JobFactory(type="range", params={"n": 4})

    Notes:
        - "range" type: creates N copies with job_id = 0..N-1
        - "enumerate" type: creates one job per item in a list
        - "Cartesian" type: creates N_A * N_B jobs from two parameter sets
    """
    type: str = Field(default="range", description="Factory type (range, enumerate, Cartesian, etc)")
    params: Dict[str, Any] = Field(default_factory=dict, description="Factory parameters")

ObjectiveDefinition

Bases: BaseModel

Complete objective specification: name, direction, and objective plan.

This is the unified model used across problem, optimization, and workflow layers. It combines what to optimize (name + direction) with how to execute it through one or more script/inline steps.

Attributes:

Name Type Description
name str

Unique objective identifier (e.g., "f1", "efficiency").

direction ObjectiveDirection

Optimization direction (minimize or maximize).

objective_plan Optional[ObjectivePlanSpec]

How to execute the objective through steps.

scheduler Optional[SchedulerConfiguration]

Reserved objective-level scheduler default. The current runtime executes objective plans inside the DAG executor after workflow stages complete; scheduled objective work should be represented as workflow stages.

metrics_keys List[str]

Optional keys to extract from plan output when it returns a dict. Useful when one plan produces multiple metrics. Example: plan outputs {"f1": 0.5, "f2": 0.3, "runtime": 10.2}, metrics_keys=["f1"] extracts only f1.

Source code in src/aid2e/utilities/configurations/objectives.py
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
class ObjectiveDefinition(BaseModel):
    """Complete objective specification: name, direction, and objective plan.

    This is the unified model used across problem, optimization, and workflow layers.
    It combines what to optimize (name + direction) with how to execute it
    through one or more script/inline steps.

    Attributes:
        name: Unique objective identifier (e.g., "f1", "efficiency").
        direction: Optimization direction (minimize or maximize).
        objective_plan: How to execute the objective through steps.
        scheduler: Reserved objective-level scheduler default. The current runtime
            executes objective plans inside the DAG executor after workflow
            stages complete; scheduled objective work should be represented as
            workflow stages.
        metrics_keys: Optional keys to extract from plan output when it returns a dict.
            Useful when one plan produces multiple metrics.
            Example: plan outputs {"f1": 0.5, "f2": 0.3, "runtime": 10.2},
                    metrics_keys=["f1"] extracts only f1.
    """

    name: str = Field(..., description="Objective name (e.g., 'f1', 'efficiency')")
    direction: ObjectiveDirection = Field(
        ...,
        description="Optimization direction: minimize or maximize"
    )
    objective_plan: Optional[ObjectivePlanSpec] = Field(
        default=None,
        description="How to execute the objective through steps",
    )
    scheduler: Optional[SchedulerConfiguration] = Field(
        default=None,
        description="Reserved objective-level scheduler default",
    )
    metrics_keys: List[str] = Field(
        default_factory=list,
        description="Keys to extract from plan output (if dict)",
    )

    def to_directive(self) -> str:
        """Convert to optimization directive string format.

        Returns:
            String like "minimize:f1" or "maximize:efficiency".
            Useful for OptimizationConfiguration.objectives.

        Example:
            >>> obj = ObjectiveDefinition(name="f1", direction=ObjectiveDirection.MINIMIZE, objective_plan=None)
            >>> obj.to_directive()
            'minimize:f1'
        """
        return f"{self.direction.value}:{self.name}"

    @classmethod
    def from_directive(
        cls,
        directive: str,
        objective_plan: Optional[ObjectivePlanSpec] = None,
    ) -> "ObjectiveDefinition":
        """Create ObjectiveDefinition from directive string.

        Parses strings like "minimize:f1" or "maximize:efficiency".

        Args:
            directive: String in format "minimize:name" or "maximize:name".
            objective_plan: Optional objective step plan.

        Returns:
            ObjectiveDefinition with parsed direction and name.

        Raises:
            ValueError: If directive format is invalid.

        Example:
            >>> directive = "minimize:f1"
            >>> obj = ObjectiveDefinition.from_directive(directive)
        """
        if ':' not in directive or directive.count(':') != 1:
            raise ValueError(f"Invalid directive format: {directive}. Expected 'minimize:name' or 'maximize:name'")

        direction_str, name = directive.split(':')
        try:
            direction = ObjectiveDirection(direction_str.lower())
        except ValueError:
            raise ValueError(f"Invalid direction '{direction_str}'. Must be 'minimize' or 'maximize'")

        if not name.strip():
            raise ValueError("Objective name cannot be empty")

        return cls(
            name=name.strip(),
            direction=direction,
            objective_plan=objective_plan,
        )

from_directive(directive, objective_plan=None) classmethod

Create ObjectiveDefinition from directive string.

Parses strings like "minimize:f1" or "maximize:efficiency".

Parameters:

Name Type Description Default
directive str

String in format "minimize:name" or "maximize:name".

required
objective_plan Optional[ObjectivePlanSpec]

Optional objective step plan.

None

Returns:

Type Description
ObjectiveDefinition

ObjectiveDefinition with parsed direction and name.

Raises:

Type Description
ValueError

If directive format is invalid.

Example

directive = "minimize:f1" obj = ObjectiveDefinition.from_directive(directive)

Source code in src/aid2e/utilities/configurations/objectives.py
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
@classmethod
def from_directive(
    cls,
    directive: str,
    objective_plan: Optional[ObjectivePlanSpec] = None,
) -> "ObjectiveDefinition":
    """Create ObjectiveDefinition from directive string.

    Parses strings like "minimize:f1" or "maximize:efficiency".

    Args:
        directive: String in format "minimize:name" or "maximize:name".
        objective_plan: Optional objective step plan.

    Returns:
        ObjectiveDefinition with parsed direction and name.

    Raises:
        ValueError: If directive format is invalid.

    Example:
        >>> directive = "minimize:f1"
        >>> obj = ObjectiveDefinition.from_directive(directive)
    """
    if ':' not in directive or directive.count(':') != 1:
        raise ValueError(f"Invalid directive format: {directive}. Expected 'minimize:name' or 'maximize:name'")

    direction_str, name = directive.split(':')
    try:
        direction = ObjectiveDirection(direction_str.lower())
    except ValueError:
        raise ValueError(f"Invalid direction '{direction_str}'. Must be 'minimize' or 'maximize'")

    if not name.strip():
        raise ValueError("Objective name cannot be empty")

    return cls(
        name=name.strip(),
        direction=direction,
        objective_plan=objective_plan,
    )

to_directive()

Convert to optimization directive string format.

Returns:

Type Description
str

String like "minimize:f1" or "maximize:efficiency".

str

Useful for OptimizationConfiguration.objectives.

Example

obj = ObjectiveDefinition(name="f1", direction=ObjectiveDirection.MINIMIZE, objective_plan=None) obj.to_directive() 'minimize:f1'

Source code in src/aid2e/utilities/configurations/objectives.py
286
287
288
289
290
291
292
293
294
295
296
297
298
def to_directive(self) -> str:
    """Convert to optimization directive string format.

    Returns:
        String like "minimize:f1" or "maximize:efficiency".
        Useful for OptimizationConfiguration.objectives.

    Example:
        >>> obj = ObjectiveDefinition(name="f1", direction=ObjectiveDirection.MINIMIZE, objective_plan=None)
        >>> obj.to_directive()
        'minimize:f1'
    """
    return f"{self.direction.value}:{self.name}"

ObjectiveDirection

Bases: str, Enum

Direction of optimization for an objective.

Attributes:

Name Type Description
MINIMIZE

Minimize the objective value.

MAXIMIZE

Maximize the objective value.

Source code in src/aid2e/utilities/configurations/objectives.py
25
26
27
28
29
30
31
32
33
class ObjectiveDirection(str, Enum):
    """Direction of optimization for an objective.

    Attributes:
        MINIMIZE: Minimize the objective value.
        MAXIMIZE: Maximize the objective value.
    """
    MINIMIZE = "minimize"
    MAXIMIZE = "maximize"

ObjectivePlanSpec

Bases: BaseModel

Plan for executing an objective (always modeled as steps).

The canonical form is a step plan with one or more stages.

Source code in src/aid2e/utilities/configurations/objectives.py
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
class ObjectivePlanSpec(BaseModel):
    """Plan for executing an objective (always modeled as steps).

    The canonical form is a step plan with one or more stages.
    """

    model_config = ConfigDict(populate_by_name=True)

    steps: StepPlanSpec = Field(
        ...,
        description="DAG-style multi-stage plan for the objective",
    )

    @model_validator(mode="before")
    def reject_legacy_shapes(cls, values: Any) -> Any:
        """Reject retired objective plan schema variants."""
        if not isinstance(values, dict):
            return values
        if "multi-steps" in values or "multi_steps" in values:
            raise ValueError(
                "Legacy objective plan step keys are no longer supported. Use 'steps'."
            )
        if "script" in values or "inline" in values:
            raise ValueError(
                "Single-step objective plans are no longer supported. Wrap the "
                "step under 'steps.stages'."
            )
        return values

    def is_steps(self) -> bool:
        """Return True if this plan is a step DAG."""
        return self.steps is not None

is_steps()

Return True if this plan is a step DAG.

Source code in src/aid2e/utilities/configurations/objectives.py
242
243
244
def is_steps(self) -> bool:
    """Return True if this plan is a step DAG."""
    return self.steps is not None

reject_legacy_shapes(values)

Reject retired objective plan schema variants.

Source code in src/aid2e/utilities/configurations/objectives.py
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
@model_validator(mode="before")
def reject_legacy_shapes(cls, values: Any) -> Any:
    """Reject retired objective plan schema variants."""
    if not isinstance(values, dict):
        return values
    if "multi-steps" in values or "multi_steps" in values:
        raise ValueError(
            "Legacy objective plan step keys are no longer supported. Use 'steps'."
        )
    if "script" in values or "inline" in values:
        raise ValueError(
            "Single-step objective plans are no longer supported. Wrap the "
            "step under 'steps.stages'."
        )
    return values

ObjectivesRegistry

Runtime registry for objective definitions.

Allows objectives to be registered and retrieved by name for use during workflow execution. This enables decoupling objective definitions from their runtime computation.

Example

registry = ObjectivesRegistry() obj_f1 = ObjectiveDefinition(name="f1", direction=ObjectiveDirection.MINIMIZE) registry.register(obj_f1) retrieved = registry.get("f1")

Source code in src/aid2e/utilities/configurations/objectives.py
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
class ObjectivesRegistry:
    """Runtime registry for objective definitions.

    Allows objectives to be registered and retrieved by name for use during
    workflow execution. This enables decoupling objective definitions from
    their runtime computation.

    Example:
        >>> registry = ObjectivesRegistry()
        >>> obj_f1 = ObjectiveDefinition(name="f1", direction=ObjectiveDirection.MINIMIZE)
        >>> registry.register(obj_f1)
        >>> retrieved = registry.get("f1")
    """

    def __init__(self):
        """Initialize empty registry."""
        self._objectives: Dict[str, ObjectiveDefinition] = {}

    def register(self, objective: ObjectiveDefinition) -> None:
        """Register an objective by name.

        Args:
            objective: ObjectiveDefinition to register.

        Raises:
            ValueError: If objective with same name already registered.
        """
        if objective.name in self._objectives:
            raise ValueError(f"Objective '{objective.name}' already registered")
        self._objectives[objective.name] = objective

    def get(self, name: str) -> Optional[ObjectiveDefinition]:
        """Retrieve objective definition by name.

        Args:
            name: Objective name.

        Returns:
            ObjectiveDefinition if found, None otherwise.
        """
        return self._objectives.get(name)

    def list_all(self) -> List[ObjectiveDefinition]:
        """Get all registered objectives.

        Returns:
            List of all registered ObjectiveDefinition instances.
        """
        return list(self._objectives.values())

    def clear(self) -> None:
        """Clear all registered objectives."""
        self._objectives.clear()

__init__()

Initialize empty registry.

Source code in src/aid2e/utilities/configurations/objectives.py
357
358
359
def __init__(self):
    """Initialize empty registry."""
    self._objectives: Dict[str, ObjectiveDefinition] = {}

clear()

Clear all registered objectives.

Source code in src/aid2e/utilities/configurations/objectives.py
393
394
395
def clear(self) -> None:
    """Clear all registered objectives."""
    self._objectives.clear()

get(name)

Retrieve objective definition by name.

Parameters:

Name Type Description Default
name str

Objective name.

required

Returns:

Type Description
Optional[ObjectiveDefinition]

ObjectiveDefinition if found, None otherwise.

Source code in src/aid2e/utilities/configurations/objectives.py
374
375
376
377
378
379
380
381
382
383
def get(self, name: str) -> Optional[ObjectiveDefinition]:
    """Retrieve objective definition by name.

    Args:
        name: Objective name.

    Returns:
        ObjectiveDefinition if found, None otherwise.
    """
    return self._objectives.get(name)

list_all()

Get all registered objectives.

Returns:

Type Description
List[ObjectiveDefinition]

List of all registered ObjectiveDefinition instances.

Source code in src/aid2e/utilities/configurations/objectives.py
385
386
387
388
389
390
391
def list_all(self) -> List[ObjectiveDefinition]:
    """Get all registered objectives.

    Returns:
        List of all registered ObjectiveDefinition instances.
    """
    return list(self._objectives.values())

register(objective)

Register an objective by name.

Parameters:

Name Type Description Default
objective ObjectiveDefinition

ObjectiveDefinition to register.

required

Raises:

Type Description
ValueError

If objective with same name already registered.

Source code in src/aid2e/utilities/configurations/objectives.py
361
362
363
364
365
366
367
368
369
370
371
372
def register(self, objective: ObjectiveDefinition) -> None:
    """Register an objective by name.

    Args:
        objective: ObjectiveDefinition to register.

    Raises:
        ValueError: If objective with same name already registered.
    """
    if objective.name in self._objectives:
        raise ValueError(f"Objective '{objective.name}' already registered")
    self._objectives[objective.name] = objective

OptimizerConfiguration

Bases: BaseModel

Canonical optimizer section model.

Source code in src/aid2e/utilities/configurations/optimizer_config.py
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
class OptimizerConfiguration(BaseModel):
    """Canonical optimizer section model."""

    model_config = ConfigDict(extra="forbid")

    name: str = Field(..., description="Optimizer/backend name, for example 'ax'.")
    type: str = Field(..., description="Optimizer family, for example 'bayesian'.")
    parameters: Dict[str, Any] = Field(
        default_factory=dict,
        description="Optimizer runtime parameters and backend-specific settings.",
    )

    def parse_algorithm_params(self) -> Optional[BaseModel]:
        """Parse registered backend-specific parameters if a model exists."""
        model = get(self.name)
        if model:
            return model(**(self.parameters or {}))
        return None

parse_algorithm_params()

Parse registered backend-specific parameters if a model exists.

Source code in src/aid2e/utilities/configurations/optimizer_config.py
30
31
32
33
34
35
def parse_algorithm_params(self) -> Optional[BaseModel]:
    """Parse registered backend-specific parameters if a model exists."""
    model = get(self.name)
    if model:
        return model(**(self.parameters or {}))
    return None

ParallelismPolicy

Bases: BaseModel

Policy for parallel job execution within a stage.

Attributes:

Name Type Description
max_concurrent int

Maximum jobs to run concurrently in a stage.

retry_max int

Maximum retries on job failure.

timeout_sec int

Timeout per job in seconds.

Example

policy = ParallelismPolicy(max_concurrent=4, retry_max=2, timeout_sec=300)

Source code in src/aid2e/utilities/configurations/workflow_config.py
59
60
61
62
63
64
65
66
67
68
69
70
71
72
class ParallelismPolicy(BaseModel):
    """Policy for parallel job execution within a stage.

    Attributes:
        max_concurrent: Maximum jobs to run concurrently in a stage.
        retry_max: Maximum retries on job failure.
        timeout_sec: Timeout per job in seconds.

    Example:
        >>> policy = ParallelismPolicy(max_concurrent=4, retry_max=2, timeout_sec=300)
    """
    max_concurrent: int = Field(default=4, ge=1, description="Max concurrent jobs in stage")
    retry_max: int = Field(default=2, ge=0, description="Max retries per failed job")
    timeout_sec: int = Field(default=300, ge=1, description="Timeout per job (seconds)")

ParameterConstraint

Bases: BaseModel

Represents a mathematical constraint on design parameters.

Constraints are validated for syntactic correctness at configuration load time. The validated constraints can then be passed to optimizers (e.g., Ax) which handle runtime constraint enforcement internally.

Attributes:

Name Type Description
name str

Unique identifier for the constraint.

description Optional[str]

Human-readable explanation of the constraint intent.

rule str

Mathematical expression using qualified parameter names, e.g., "group.param1 + group.param2 < 10.0".

key str

YAML key associated with a list of models.

Example

constraint = ParameterConstraint( ... name="budget_limit", ... description="Total cost must not exceed budget", ... rule="tracker.cost + magnet.cost < 1000" ... )

Validate syntax (done automatically at load time)

constraint.validate_syntax(['tracker.cost', 'magnet.cost']) (True, None)

Source code in src/aid2e/utilities/configurations/design_config.py
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
class ParameterConstraint(BaseModel):
    """Represents a mathematical constraint on design parameters.

    Constraints are validated for syntactic correctness at configuration load time.
    The validated constraints can then be passed to optimizers (e.g., Ax) which handle
    runtime constraint enforcement internally.

    Attributes:
        name: Unique identifier for the constraint.
        description: Human-readable explanation of the constraint intent.
        rule: Mathematical expression using qualified parameter names,
              e.g., "group.param1 + group.param2 < 10.0".
        key: YAML key associated with a list of models.

    Example:
        >>> constraint = ParameterConstraint(
        ...     name="budget_limit",
        ...     description="Total cost must not exceed budget",
        ...     rule="tracker.cost + magnet.cost < 1000"
        ... )
        >>> # Validate syntax (done automatically at load time)
        >>> constraint.validate_syntax(['tracker.cost', 'magnet.cost'])
        (True, None)
    """
    name: str
    description: Optional[str] = None
    rule: str  # Mathematical expression like "x1 + x2 < 10"
    key: ClassVar[str] = 'parameter_constraints'

    def extract_parameter_names(self) -> Set[str]:
        """Extract parameter names referenced in the constraint rule.

        Parses the constraint expression and extracts all identifiers that
        appear to be qualified parameter names (containing a dot).

        Returns:
            Set of parameter names found in the rule.

        Example:
            >>> constraint = ParameterConstraint(
            ...     name="test",
            ...     rule="tracker.x + magnet.y < detector.limit"
            ... )
            >>> names = constraint.extract_parameter_names()
            >>> print(sorted(names))
            ['detector.limit', 'magnet.y', 'tracker.x']
        """
        # Find all qualified parameter names (e.g., "group.param")
        pattern = r'\b([a-zA-Z_][a-zA-Z0-9_]*\.[a-zA-Z_][a-zA-Z0-9_]*)\b'
        return set(re.findall(pattern, self.rule))

    def validate_syntax(self, valid_param_names: List[str]) -> Tuple[bool, Optional[str]]:
        """Validate constraint syntax and parameter references.

        Checks that:
        1. The constraint rule is valid Python syntax
        2. All referenced parameters exist in the valid_param_names list
        3. The expression is a valid comparison/boolean expression

        This is structural validation done at configuration load time,
        NOT runtime constraint evaluation (which is handled by the optimizer).

        Args:
            valid_param_names: List of valid qualified parameter names
                              from the design configuration.

        Returns:
            Tuple of (is_valid, error_message) where:
            - is_valid: True if constraint is syntactically correct
            - error_message: None if valid, otherwise describes the error

        Example:
            >>> constraint = ParameterConstraint(
            ...     name="test", rule="group.x + group.y < 10"
            ... )
            >>> is_valid, err = constraint.validate_syntax(
            ...     ['group.x', 'group.y']
            ... )
            >>> assert is_valid and err is None
        """
        # 1. Check if rule is parseable Python syntax
        try:
            ast.parse(self.rule, mode='eval')
        except SyntaxError as e:
            return False, f"Invalid syntax in constraint rule: {e}"

        # 2. Extract and validate parameter names
        referenced_params = self.extract_parameter_names()
        valid_set = set(valid_param_names)
        unknown_params = referenced_params - valid_set

        if unknown_params:
            return False, f"Unknown parameters in constraint: {', '.join(sorted(unknown_params))}"

        return True, None

    def evaluate(self, param_values: Dict[str, float]) -> bool:
        """Evaluate constraint against parameter values at runtime.

        Substitutes parameter names in the constraint rule with their values
        and evaluates the resulting mathematical expression. This is used for
        runtime constraint checking when the optimizer doesn't support
        constraint enforcement (e.g., random search, some evolutionary algorithms).

        For optimizers with native constraint support (e.g., Ax), use the
        constraint object directly instead of calling this method.

        Args:
            param_values: Dictionary mapping qualified parameter names
                         (e.g., "group.param") to numeric values.

        Returns:
            True if constraint is satisfied, False otherwise.

        Raises:
            ValueError: If the constraint rule cannot be evaluated
                       (e.g., missing parameters, division by zero).

        Example:
            >>> constraint = ParameterConstraint(
            ...     name="test", rule="DTLZ2.x1 < 1.0"
            ... )
            >>> constraint.evaluate({"DTLZ2.x1": 0.5})
            True
            >>> constraint.evaluate({"DTLZ2.x1": 1.5})
            False

        Notes:
            - Prefer using optimizer's native constraint enforcement when available
            - This method is primarily for optimizer-agnostic validation
            - Uses eval() on sanitized expressions (validated at load time)
        """
        # Replace parameter names with their values
        expr = self.rule
        for param_name, value in param_values.items():
            # Use word boundaries to avoid partial matches
            expr = re.sub(rf'\b{re.escape(param_name)}\b', str(value), expr)

        try:
            # Evaluate the expression
            result = eval(expr)
            return bool(result)
        except Exception as e:
            raise ValueError(f"Failed to evaluate constraint '{self.name}': {e}")

evaluate(param_values)

Evaluate constraint against parameter values at runtime.

Substitutes parameter names in the constraint rule with their values and evaluates the resulting mathematical expression. This is used for runtime constraint checking when the optimizer doesn't support constraint enforcement (e.g., random search, some evolutionary algorithms).

For optimizers with native constraint support (e.g., Ax), use the constraint object directly instead of calling this method.

Parameters:

Name Type Description Default
param_values Dict[str, float]

Dictionary mapping qualified parameter names (e.g., "group.param") to numeric values.

required

Returns:

Type Description
bool

True if constraint is satisfied, False otherwise.

Raises:

Type Description
ValueError

If the constraint rule cannot be evaluated (e.g., missing parameters, division by zero).

Example

constraint = ParameterConstraint( ... name="test", rule="DTLZ2.x1 < 1.0" ... ) constraint.evaluate({"DTLZ2.x1": 0.5}) True constraint.evaluate({"DTLZ2.x1": 1.5}) False

Notes
  • Prefer using optimizer's native constraint enforcement when available
  • This method is primarily for optimizer-agnostic validation
  • Uses eval() on sanitized expressions (validated at load time)
Source code in src/aid2e/utilities/configurations/design_config.py
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
def evaluate(self, param_values: Dict[str, float]) -> bool:
    """Evaluate constraint against parameter values at runtime.

    Substitutes parameter names in the constraint rule with their values
    and evaluates the resulting mathematical expression. This is used for
    runtime constraint checking when the optimizer doesn't support
    constraint enforcement (e.g., random search, some evolutionary algorithms).

    For optimizers with native constraint support (e.g., Ax), use the
    constraint object directly instead of calling this method.

    Args:
        param_values: Dictionary mapping qualified parameter names
                     (e.g., "group.param") to numeric values.

    Returns:
        True if constraint is satisfied, False otherwise.

    Raises:
        ValueError: If the constraint rule cannot be evaluated
                   (e.g., missing parameters, division by zero).

    Example:
        >>> constraint = ParameterConstraint(
        ...     name="test", rule="DTLZ2.x1 < 1.0"
        ... )
        >>> constraint.evaluate({"DTLZ2.x1": 0.5})
        True
        >>> constraint.evaluate({"DTLZ2.x1": 1.5})
        False

    Notes:
        - Prefer using optimizer's native constraint enforcement when available
        - This method is primarily for optimizer-agnostic validation
        - Uses eval() on sanitized expressions (validated at load time)
    """
    # Replace parameter names with their values
    expr = self.rule
    for param_name, value in param_values.items():
        # Use word boundaries to avoid partial matches
        expr = re.sub(rf'\b{re.escape(param_name)}\b', str(value), expr)

    try:
        # Evaluate the expression
        result = eval(expr)
        return bool(result)
    except Exception as e:
        raise ValueError(f"Failed to evaluate constraint '{self.name}': {e}")

extract_parameter_names()

Extract parameter names referenced in the constraint rule.

Parses the constraint expression and extracts all identifiers that appear to be qualified parameter names (containing a dot).

Returns:

Type Description
Set[str]

Set of parameter names found in the rule.

Example

constraint = ParameterConstraint( ... name="test", ... rule="tracker.x + magnet.y < detector.limit" ... ) names = constraint.extract_parameter_names() print(sorted(names)) ['detector.limit', 'magnet.y', 'tracker.x']

Source code in src/aid2e/utilities/configurations/design_config.py
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def extract_parameter_names(self) -> Set[str]:
    """Extract parameter names referenced in the constraint rule.

    Parses the constraint expression and extracts all identifiers that
    appear to be qualified parameter names (containing a dot).

    Returns:
        Set of parameter names found in the rule.

    Example:
        >>> constraint = ParameterConstraint(
        ...     name="test",
        ...     rule="tracker.x + magnet.y < detector.limit"
        ... )
        >>> names = constraint.extract_parameter_names()
        >>> print(sorted(names))
        ['detector.limit', 'magnet.y', 'tracker.x']
    """
    # Find all qualified parameter names (e.g., "group.param")
    pattern = r'\b([a-zA-Z_][a-zA-Z0-9_]*\.[a-zA-Z_][a-zA-Z0-9_]*)\b'
    return set(re.findall(pattern, self.rule))

validate_syntax(valid_param_names)

Validate constraint syntax and parameter references.

Checks that: 1. The constraint rule is valid Python syntax 2. All referenced parameters exist in the valid_param_names list 3. The expression is a valid comparison/boolean expression

This is structural validation done at configuration load time, NOT runtime constraint evaluation (which is handled by the optimizer).

Parameters:

Name Type Description Default
valid_param_names List[str]

List of valid qualified parameter names from the design configuration.

required

Returns:

Type Description
bool

Tuple of (is_valid, error_message) where:

Optional[str]
  • is_valid: True if constraint is syntactically correct
Tuple[bool, Optional[str]]
  • error_message: None if valid, otherwise describes the error
Example

constraint = ParameterConstraint( ... name="test", rule="group.x + group.y < 10" ... ) is_valid, err = constraint.validate_syntax( ... ['group.x', 'group.y'] ... ) assert is_valid and err is None

Source code in src/aid2e/utilities/configurations/design_config.py
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
def validate_syntax(self, valid_param_names: List[str]) -> Tuple[bool, Optional[str]]:
    """Validate constraint syntax and parameter references.

    Checks that:
    1. The constraint rule is valid Python syntax
    2. All referenced parameters exist in the valid_param_names list
    3. The expression is a valid comparison/boolean expression

    This is structural validation done at configuration load time,
    NOT runtime constraint evaluation (which is handled by the optimizer).

    Args:
        valid_param_names: List of valid qualified parameter names
                          from the design configuration.

    Returns:
        Tuple of (is_valid, error_message) where:
        - is_valid: True if constraint is syntactically correct
        - error_message: None if valid, otherwise describes the error

    Example:
        >>> constraint = ParameterConstraint(
        ...     name="test", rule="group.x + group.y < 10"
        ... )
        >>> is_valid, err = constraint.validate_syntax(
        ...     ['group.x', 'group.y']
        ... )
        >>> assert is_valid and err is None
    """
    # 1. Check if rule is parseable Python syntax
    try:
        ast.parse(self.rule, mode='eval')
    except SyntaxError as e:
        return False, f"Invalid syntax in constraint rule: {e}"

    # 2. Extract and validate parameter names
    referenced_params = self.extract_parameter_names()
    valid_set = set(valid_param_names)
    unknown_params = referenced_params - valid_set

    if unknown_params:
        return False, f"Unknown parameters in constraint: {', '.join(sorted(unknown_params))}"

    return True, None

ParameterGroup

Bases: BaseModel

Container for a group of related parameters.

Groups parameters that share common properties or contexts, such as detector subsystems (vertex_barrel, silicon_tracker, etc.). Parameters within a group are accessed via qualified names (group_name.param_name).

Attributes:

Name Type Description
parameters Dict[str, Parameter]

Dictionary mapping parameter names to Parameter objects.

Example

group = ParameterGroup(parameters={ ... 'thickness': RangeParameter(value=0.35, bounds=[0.2, 0.6]), ... 'pitch': RangeParameter(value=25, bounds=[10, 50]) ... })

Source code in src/aid2e/utilities/configurations/design_config.py
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
class ParameterGroup(BaseModel):
    """Container for a group of related parameters.

    Groups parameters that share common properties or contexts, such as detector
    subsystems (vertex_barrel, silicon_tracker, etc.). Parameters within a group
    are accessed via qualified names (group_name.param_name).

    Attributes:
        parameters: Dictionary mapping parameter names to Parameter objects.

    Example:
        >>> group = ParameterGroup(parameters={
        ...     'thickness': RangeParameter(value=0.35, bounds=[0.2, 0.6]),
        ...     'pitch': RangeParameter(value=25, bounds=[10, 50])
        ... })
    """
    parameters: Dict[str, Parameter]

ProblemConfigLoader

Loader for problem YAML/CONFIG files.

Parses files following the schema used by examples/basic/problem.config:

        problem:
            name: "..."
            type: "..."
            output_location: "..."
            work_location: "..."
            design_parameters_file: "./path/to/design.params"
            objectives:
                - name: "f1"
                    direction: "minimize"
                    objective_plan:
                        steps:
                            stages:
                                - name: "evaluate_f1"
                                    script:
                                        path: "scripts/dtlz2_problem.py"
                                        output_file: "objectives_{job_id}.json"
                                    produces_objective: true
                    metrics_keys: ["f1"]
                - name: "f2"
                    direction: "minimize"
Notes

Use ProblemConfigLoader.load() to load from a file path or ProblemConfigLoader.from_dict() to construct from an in-memory dictionary.

Source code in src/aid2e/utilities/configurations/problem_config.py
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
class ProblemConfigLoader:
    """Loader for problem YAML/CONFIG files.

    Parses files following the schema used by `examples/basic/problem.config`:

                problem:
                    name: "..."
                    type: "..."
                    output_location: "..."
                    work_location: "..."
                    design_parameters_file: "./path/to/design.params"
                    objectives:
                        - name: "f1"
                            direction: "minimize"
                            objective_plan:
                                steps:
                                    stages:
                                        - name: "evaluate_f1"
                                            script:
                                                path: "scripts/dtlz2_problem.py"
                                                output_file: "objectives_{job_id}.json"
                                            produces_objective: true
                            metrics_keys: ["f1"]
                        - name: "f2"
                            direction: "minimize"

    Notes:
        Use `ProblemConfigLoader.load()` to load from a file path or
        `ProblemConfigLoader.from_dict()` to construct from an in-memory
        dictionary.
    """
    @staticmethod
    def _build_from_problem_dict(problem: Dict[str, Any], base_dir: Optional[Path]) -> ProblemConfiguration:
        """Build ProblemConfiguration from an inner 'problem' mapping.

        Supports design source via either a file path ('design_parameters_file')
        or inline design payload ('inline_design'). Exactly one must be provided.
        """
        if "type" in problem:
            raise ValueError(
                "Legacy key 'type' is no longer supported in 'problem'. "
                "Use 'problem_type'."
            )
        if "design_space" in problem:
            raise ValueError(
                "Legacy key 'design_space' is no longer supported in 'problem'. "
                "Use 'design_parameters_file' or 'inline_design'."
            )

        # Required scalar fields
        required_scalar = [
            "name",
            "problem_type",
            "output_location",
            "work_location",
            "objectives",
        ]
        missing = [k for k in required_scalar if k not in problem]
        if missing:
            raise ValueError("Invalid problem definition, missing keys: " + ", ".join(missing))

        # Objectives
        objectives_raw = problem.get("objectives", [])
        if not isinstance(objectives_raw, list) or not objectives_raw:
            raise ValueError("'objectives' must be a non-empty list")

        # Design source mutual exclusivity
        has_path = "design_parameters_file" in problem
        has_inline = "inline_design" in problem
        if has_path == has_inline:
            # Either both True or both False → invalid
            raise ValueError("Specify exactly one of 'design_parameters_file' or 'inline_design'")

        design_data = None
        if has_path:
            design_path = resolve_path(problem["design_parameters_file"], base_dir)
            with open(design_path, 'r') as params:
                design_data = yaml.safe_load(params)
        else:
            design_data = problem["inline_design"]

        use_stack_design = False
        design_stack = None
        for stack, components in StackRegistry.list_registered_stacks().items():
            design_loader = components['design_loader']
            if design_loader.space_key in design_data:
                design_config = design_loader.load(design_data=design_data)
                use_stack_design = True
                design_stack = stack
                break

        if not use_stack_design:
             design_config = DesignConfigLoader.load(design_data=design_data)

        # Parse environment config if any present
        env_config = None
        env_stack = None
        for stack, components in StackRegistry.list_registered_stacks().items():
            env_model = components['env_config']
            env_loader = components['env_loader']
            if env_model.key in problem:
                env_config = env_loader.load(env_data=problem)
                env_stack = stack
                break

        # Build ProblemConfiguration
        output_location = Path(problem["output_location"]).expanduser()
        if base_dir and not output_location.is_absolute():
            output_location = (base_dir / output_location).resolve()

        work_location = Path(problem["work_location"]).expanduser()
        if base_dir and not work_location.is_absolute():
            work_location = (base_dir / work_location).resolve()

        config_model = ProblemConfiguration
        if env_config is not None and design_stack == env_stack:
            stack_components = StackRegistry.list_registered_stacks().get(env_stack, {})
            config_model = stack_components.get("problem_config") or ProblemConfiguration

        return config_model(
            name=problem["name"],
            problem_type=problem["problem_type"],
            output_location=str(output_location),
            work_location=str(work_location),
            design_config=design_config,
            objectives=objectives_raw,
            observations=problem.get("observations"),
            environment_config=env_config,
            evaluation_config=problem.get("evaluation_config", {}),
        )

    @staticmethod
    def load(file_path: str) -> ProblemConfiguration:
        path = Path(file_path)
        if not path.exists():
            raise FileNotFoundError(f"Problem file not found: {file_path}")

        with open(path, "r") as f:
            data = yaml.safe_load(f) or {}

        if "problem" not in data or not isinstance(data["problem"], dict):
            raise ValueError("Invalid problem file: missing 'problem' section")

        return ProblemConfigLoader._build_from_problem_dict(data["problem"], base_dir=path.parent)

    @staticmethod
    def from_dict(problem_payload: Dict[str, Any], base_dir: Optional[str] = None) -> ProblemConfiguration:
        """Construct ProblemConfiguration from a dict payload.

        Accepts the inner 'problem' mapping as a Python dict and supports both
        file-based and inline design definitions. Set base_dir for reliable
        relative path resolution when using 'design_parameters_file'.
        """
        return ProblemConfigLoader._build_from_problem_dict(problem_payload, base_dir=Path(base_dir) if base_dir else None)

from_dict(problem_payload, base_dir=None) staticmethod

Construct ProblemConfiguration from a dict payload.

Accepts the inner 'problem' mapping as a Python dict and supports both file-based and inline design definitions. Set base_dir for reliable relative path resolution when using 'design_parameters_file'.

Source code in src/aid2e/utilities/configurations/problem_config.py
330
331
332
333
334
335
336
337
338
@staticmethod
def from_dict(problem_payload: Dict[str, Any], base_dir: Optional[str] = None) -> ProblemConfiguration:
    """Construct ProblemConfiguration from a dict payload.

    Accepts the inner 'problem' mapping as a Python dict and supports both
    file-based and inline design definitions. Set base_dir for reliable
    relative path resolution when using 'design_parameters_file'.
    """
    return ProblemConfigLoader._build_from_problem_dict(problem_payload, base_dir=Path(base_dir) if base_dir else None)

ProblemConfiguration

Bases: BaseModel

Generic problem configuration.

Focuses on core problem attributes, a design configuration, objectives, and optional observations. Environment and scheduler/trial management belong to separate workflow components (e.g., WorkflowManager).

Notes
  • design_config accepts any subclass of DesignConfig.
  • objectives must be non-empty with unique names.
Source code in src/aid2e/utilities/configurations/problem_config.py
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
class ProblemConfiguration(BaseModel):
    """Generic problem configuration.

    Focuses on core problem attributes, a design configuration, objectives, and
    optional observations. Environment and scheduler/trial management belong to
    separate workflow components (e.g., `WorkflowManager`).

    Notes:
        - `design_config` accepts any subclass of `DesignConfig`.
        - `objectives` must be non-empty with unique names.
    """
    name: str
    output_location: str
    work_location: str
    problem_type: str  # e.g., "EPIC_TRACKING", "DTLZ2", "CLOSURE_MOO"

    # Accept stack-specific subclasses through the stack registry.
    design_config: DesignConfig
    objectives: List[ObjectiveDefinition]
    observations: Optional[List[Dict[str, Any]]] = Field(default=None)
    environment_config: Optional[EnvironmentConfig] = Field(default=None)
    evaluation_config: Dict[str, Any] = Field(default_factory=dict)

    @field_validator("objectives", mode="before")
    @classmethod
    def normalize_objectives(cls, raw_objectives: Any) -> List[ObjectiveDefinition]:
        """Normalize various objective payload shapes into ObjectiveDefinition."""
        if not isinstance(raw_objectives, list) or not raw_objectives:
            raise ValueError("'objectives' must be a non-empty list")

        normalized: List[ObjectiveDefinition] = []

        for entry in raw_objectives:
            if isinstance(entry, ObjectiveDefinition):
                normalized.append(entry)
                continue

            if isinstance(entry, dict):
                normalized.append(cls._objective_from_dict(entry))
                continue

            raise ValueError(f"Unsupported objective entry type: {type(entry)}")

        return normalized

    @model_validator(mode="after")
    def validate_paths(self) -> "ProblemConfiguration":
        """Validate directory paths and objective correctness.

        - Ensures existing output and work paths are directories.
        - Ensures `objectives` is non-empty with unique names.
        """
        errors = []

        for label, value in [("output_location", self.output_location),
                             ("work_location", self.work_location)]:
            path = Path(value)
            if path.exists() and not path.is_dir():
                errors.append(f"{label} is not a directory: {value}")

        # Objectives must be provided and unique
        if not self.objectives:
            errors.append("objectives must be provided and non-empty")
        else:
            names = [obj.name for obj in self.objectives]
            if len(set(names)) != len(names):
                errors.append("objective names must be unique")

        if errors:
            raise ValueError("ProblemConfiguration validation failed:\n" + "\n".join(errors))

        return self

    @staticmethod
    def _parse_objective_plan(objective_plan: Any) -> Optional[ObjectivePlanSpec]:
        """Convert an objective_plan payload to ObjectivePlanSpec."""
        if objective_plan is None:
            return None
        if isinstance(objective_plan, ObjectivePlanSpec):
            return objective_plan
        if isinstance(objective_plan, dict):
            return ObjectivePlanSpec(**dict(objective_plan))
        raise ValueError(
            "Invalid objective_plan block for objective; expected mapping or "
            "ObjectivePlanSpec"
        )

    @classmethod
    def _objective_from_dict(cls, payload: Dict[str, Any]) -> ObjectiveDefinition:
        """Build ObjectiveDefinition from a mapping payload."""
        if not isinstance(payload, dict):
            raise ValueError("Objective entry must be a mapping")

        if "name" not in payload:
            raise ValueError("Objective entry missing required field 'name'")

        name = payload["name"]

        if "minimize" in payload:
            raise ValueError(
                "Legacy key 'minimize' is no longer supported. Use 'direction'."
            )
        if "computation" in payload:
            raise ValueError(
                "Legacy key 'computation' is no longer supported. Use "
                "'objective_plan'."
            )
        if "direction" not in payload:
            raise ValueError("Objective entry missing required field 'direction'")

        direction_raw = payload["direction"]
        if isinstance(direction_raw, ObjectiveDirection):
            direction = direction_raw
        else:
            direction = ObjectiveDirection(str(direction_raw).lower())

        objective_plan = cls._parse_objective_plan(payload.get("objective_plan"))
        metrics_keys = payload.get("metrics_keys", []) or []

        return ObjectiveDefinition(
            name=name,
            direction=direction,
            objective_plan=objective_plan,
            metrics_keys=metrics_keys,
        )

normalize_objectives(raw_objectives) classmethod

Normalize various objective payload shapes into ObjectiveDefinition.

Source code in src/aid2e/utilities/configurations/problem_config.py
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
@field_validator("objectives", mode="before")
@classmethod
def normalize_objectives(cls, raw_objectives: Any) -> List[ObjectiveDefinition]:
    """Normalize various objective payload shapes into ObjectiveDefinition."""
    if not isinstance(raw_objectives, list) or not raw_objectives:
        raise ValueError("'objectives' must be a non-empty list")

    normalized: List[ObjectiveDefinition] = []

    for entry in raw_objectives:
        if isinstance(entry, ObjectiveDefinition):
            normalized.append(entry)
            continue

        if isinstance(entry, dict):
            normalized.append(cls._objective_from_dict(entry))
            continue

        raise ValueError(f"Unsupported objective entry type: {type(entry)}")

    return normalized

validate_paths()

Validate directory paths and objective correctness.

  • Ensures existing output and work paths are directories.
  • Ensures objectives is non-empty with unique names.
Source code in src/aid2e/utilities/configurations/problem_config.py
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
@model_validator(mode="after")
def validate_paths(self) -> "ProblemConfiguration":
    """Validate directory paths and objective correctness.

    - Ensures existing output and work paths are directories.
    - Ensures `objectives` is non-empty with unique names.
    """
    errors = []

    for label, value in [("output_location", self.output_location),
                         ("work_location", self.work_location)]:
        path = Path(value)
        if path.exists() and not path.is_dir():
            errors.append(f"{label} is not a directory: {value}")

    # Objectives must be provided and unique
    if not self.objectives:
        errors.append("objectives must be provided and non-empty")
    else:
        names = [obj.name for obj in self.objectives]
        if len(set(names)) != len(names):
            errors.append("objective names must be unique")

    if errors:
        raise ValueError("ProblemConfiguration validation failed:\n" + "\n".join(errors))

    return self

RangeParameter

Bases: BaseParameter

Continuous parameter with min/max bounds.

Source code in src/aid2e/utilities/configurations/base_models.py
18
19
20
21
22
23
24
25
class RangeParameter(BaseParameter):
    """Continuous parameter with min/max bounds."""
    value: float
    bounds: Tuple[float, float]

    @property
    def type(self) -> Literal["range"]:
        return "range"

SchedulerConfigLoader

Loader for scheduler YAML/CONFIG files.

Parses files following the scheduler schema:

scheduler:
  runner_type: "JobLibRunner"
  parameters:
    n_jobs: 4
Notes

Use SchedulerConfigLoader.load() to load from a file path or SchedulerConfigLoader.from_dict() to construct from an in-memory dictionary.

Source code in src/aid2e/utilities/configurations/scheduler_config.py
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
class SchedulerConfigLoader:
    """Loader for scheduler YAML/CONFIG files.

    Parses files following the scheduler schema:

        scheduler:
          runner_type: "JobLibRunner"
          parameters:
            n_jobs: 4

    Notes:
        Use `SchedulerConfigLoader.load()` to load from a file path or
        `SchedulerConfigLoader.from_dict()` to construct from an in-memory dictionary.
    """

    @staticmethod
    def _build_from_scheduler_dict(
        scheduler_payload: Dict[str, Any],
        base_dir: Optional[Path] = None,
    ) -> SchedulerConfiguration:
        """Build SchedulerConfiguration from an inner ``scheduler`` mapping."""
        if not isinstance(scheduler_payload, dict):
            raise ValueError("Invalid scheduler definition: expected a mapping")
        scheduler_payload = dict(scheduler_payload)
        required_keys = ["runner_type", "parameters"]
        missing = [key for key in required_keys if key not in scheduler_payload]
        if missing:
            raise ValueError("Invalid scheduler definition, missing keys: " + ", ".join(missing))
        config = SchedulerConfiguration(**scheduler_payload)
        parameters = dict(config.parameters)
        if config.runner_type == "SlurmRunner":
            template_file = parameters.get("template_file")
            if template_file is None:
                if not parameters:
                    raise ValueError("Invalid SlurmRunner scheduler parameters, provide inline definitions or template_file")
            else:
                template_path = Path(template_file).expanduser()
                if base_dir and not template_path.is_absolute():
                    template_path = (base_dir / template_path).resolve()
                if not template_path.exists():
                    raise FileNotFoundError(f"Slurm template file not found: {template_file}")

                with open(template_path, "r") as f:
                    template_data = json.load(f)
                if not isinstance(template_data, dict):
                    raise ValueError("Invalid Slurm template file: expected a JSON object")

                inline_parameters = dict(parameters)
                inline_parameters.pop("template_file")
                if not template_data and not inline_parameters:
                    raise ValueError("Invalid Slurm template file: expected scheduler parameters")
                scheduler_payload["parameters"] = {
                    **template_data,
                    **inline_parameters,
                }
            config = SchedulerConfiguration(**scheduler_payload)
            parameters = dict(config.parameters)
        if config.runner_type == "PanDAiDDSRunner":
            if not parameters:
                raise ValueError("Invalid PanDAiDDSRunner scheduler parameters, provide PanDA definitions")

        Model = get(config.runner_type)
        if Model is None:
            raise ValueError(f"No scheduler config model registered for {config.runner_type}")
        unknown_keys = sorted(set(parameters) - set(Model.model_fields))
        if unknown_keys:
            raise ValueError(
                f"Invalid {config.runner_type} scheduler parameters, unknown keys: "
                + ", ".join(unknown_keys)
            )
        Model(**parameters)

        return config

    @staticmethod
    def load(file_path: str) -> SchedulerConfiguration:
        """Load a scheduler configuration from a YAML file."""
        path = Path(file_path)
        if not path.exists():
            raise FileNotFoundError(f"Scheduler file not found: {file_path}")

        with open(path, "r") as f:
            data = yaml.safe_load(f) or {}

        if "scheduler" not in data or not isinstance(data["scheduler"], dict):
            raise ValueError("Invalid scheduler file: missing 'scheduler' section")

        return SchedulerConfigLoader._build_from_scheduler_dict(
            data["scheduler"],
            base_dir=path.parent,
        )

    @staticmethod
    def from_dict(
        scheduler_payload: Dict[str, Any],
        base_dir: Optional[str] = None,
    ) -> SchedulerConfiguration:
        """Construct SchedulerConfiguration from an inner scheduler mapping."""
        return SchedulerConfigLoader._build_from_scheduler_dict(
            scheduler_payload,
            base_dir=Path(base_dir) if base_dir else None,
        )

from_dict(scheduler_payload, base_dir=None) staticmethod

Construct SchedulerConfiguration from an inner scheduler mapping.

Source code in src/aid2e/utilities/configurations/scheduler_config.py
192
193
194
195
196
197
198
199
200
201
@staticmethod
def from_dict(
    scheduler_payload: Dict[str, Any],
    base_dir: Optional[str] = None,
) -> SchedulerConfiguration:
    """Construct SchedulerConfiguration from an inner scheduler mapping."""
    return SchedulerConfigLoader._build_from_scheduler_dict(
        scheduler_payload,
        base_dir=Path(base_dir) if base_dir else None,
    )

load(file_path) staticmethod

Load a scheduler configuration from a YAML file.

Source code in src/aid2e/utilities/configurations/scheduler_config.py
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
@staticmethod
def load(file_path: str) -> SchedulerConfiguration:
    """Load a scheduler configuration from a YAML file."""
    path = Path(file_path)
    if not path.exists():
        raise FileNotFoundError(f"Scheduler file not found: {file_path}")

    with open(path, "r") as f:
        data = yaml.safe_load(f) or {}

    if "scheduler" not in data or not isinstance(data["scheduler"], dict):
        raise ValueError("Invalid scheduler file: missing 'scheduler' section")

    return SchedulerConfigLoader._build_from_scheduler_dict(
        data["scheduler"],
        base_dir=path.parent,
    )

SchedulerConfiguration

Bases: BaseModel

Complete scheduler/runner configuration.

Specifies which scheduler backend to use and its parameters. Runner-specific configuration is validated via the scheduler registry.

Attributes:

Name Type Description
runner_type Literal['JobLibRunner', 'SlurmRunner', 'PanDAiDDSRunner']

Type of runner/scheduler to use (JobLibRunner, SlurmRunner, PanDAiDDSRunner).

parameters Dict[str, Any]

Runner-specific parameters as free-form dict.

max_retries int

Global maximum retries for failed jobs.

output_location str

Base directory for scheduler output files.

monitor_interval int

Monitoring interval in seconds for job status checks.

Example

config = SchedulerConfiguration( ... runner_type="JobLibRunner", ... parameters={"n_jobs": -1, "backend": "threading"}, ... output_location="./output" ... )

Source code in src/aid2e/utilities/configurations/scheduler_config.py
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
class SchedulerConfiguration(BaseModel):
    """Complete scheduler/runner configuration.

    Specifies which scheduler backend to use and its parameters.
    Runner-specific configuration is validated via the scheduler registry.

    Attributes:
        runner_type: Type of runner/scheduler to use (JobLibRunner, SlurmRunner, PanDAiDDSRunner).
        parameters: Runner-specific parameters as free-form dict.
        max_retries: Global maximum retries for failed jobs.
        output_location: Base directory for scheduler output files.
        monitor_interval: Monitoring interval in seconds for job status checks.

    Example:
        >>> config = SchedulerConfiguration(
        ...     runner_type="JobLibRunner",
        ...     parameters={"n_jobs": -1, "backend": "threading"},
        ...     output_location="./output"
        ... )
    """

    model_config = ConfigDict(extra="forbid")

    runner_type: Literal["JobLibRunner", "SlurmRunner", "PanDAiDDSRunner"] = Field(
        default="JobLibRunner",
        description="Type of runner/scheduler to use"
    )

    parameters: Dict[str, Any] = Field(
        default_factory=dict,
        description="Runner-specific parameters (validated by scheduler registry)"
    )

    max_retries: int = Field(
        default=3,
        ge=0,
        description="Global maximum retries for failed jobs"
    )
    output_location: str = Field(
        default="./scheduler_output",
        description="Base directory for scheduler output files"
    )
    monitor_interval: int = Field(
        default=30,
        ge=1,
        description="Monitoring interval in seconds for job status checks"
    )

    def parse_runner_params(self) -> Optional[BaseModel]:
        """Parse and validate runner-specific parameters via registry.

        Looks up the registered config model for this runner_type and
        validates the parameters dict against it.

        Returns:
            Validated runner-specific config model instance, or None if 
            runner type not found in registry.

        Raises:
            ValidationError: If parameters don't match the runner's schema.

        Example:
            >>> config = SchedulerConfiguration(
            ...     runner_type="JobLibRunner",
            ...     parameters={"n_jobs": 4}
            ... )
            >>> joblib_config = config.parse_runner_params()
            >>> joblib_config.n_jobs
            4
        """
        Model = get(self.runner_type)
        if Model:
            return Model(**self.parameters)
        return None

parse_runner_params()

Parse and validate runner-specific parameters via registry.

Looks up the registered config model for this runner_type and validates the parameters dict against it.

Returns:

Type Description
Optional[BaseModel]

Validated runner-specific config model instance, or None if

Optional[BaseModel]

runner type not found in registry.

Raises:

Type Description
ValidationError

If parameters don't match the runner's schema.

Example

config = SchedulerConfiguration( ... runner_type="JobLibRunner", ... parameters={"n_jobs": 4} ... ) joblib_config = config.parse_runner_params() joblib_config.n_jobs 4

Source code in src/aid2e/utilities/configurations/scheduler_config.py
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def parse_runner_params(self) -> Optional[BaseModel]:
    """Parse and validate runner-specific parameters via registry.

    Looks up the registered config model for this runner_type and
    validates the parameters dict against it.

    Returns:
        Validated runner-specific config model instance, or None if 
        runner type not found in registry.

    Raises:
        ValidationError: If parameters don't match the runner's schema.

    Example:
        >>> config = SchedulerConfiguration(
        ...     runner_type="JobLibRunner",
        ...     parameters={"n_jobs": 4}
        ... )
        >>> joblib_config = config.parse_runner_params()
        >>> joblib_config.n_jobs
        4
    """
    Model = get(self.runner_type)
    if Model:
        return Model(**self.parameters)
    return None

ScriptObjective

Bases: BaseModel

Objective computed via external script.

Attributes:
    path: Path to executable script (resolved relative to config directory).
    output_file: Expected output file pattern (e.g., "objectives_*.json").
        The script should create a file matching this pattern containing
        the objective value in JSON/YAML format.
    timeout_sec: Timeout in seconds (optional, default: 300).

Example:
    >>> script = ScriptObjective(
    ...     path="scripts/dtlz2_problem.py",
    ...     output_file="objectives_{job_id}.json"
    ... )
Notes

Scripts receive the design-point and output paths through --design_params_file / --output_file and the corresponding AID2E_PARAMS_FILE / AID2E_OUTPUT_FILE environment variables.

Source code in src/aid2e/utilities/configurations/objectives.py
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
class ScriptObjective(BaseModel):
    """Objective computed via external script.

    Attributes:
        path: Path to executable script (resolved relative to config directory).
        output_file: Expected output file pattern (e.g., "objectives_*.json").
            The script should create a file matching this pattern containing
            the objective value in JSON/YAML format.
        timeout_sec: Timeout in seconds (optional, default: 300).

    Example:
        >>> script = ScriptObjective(
        ...     path="scripts/dtlz2_problem.py",
        ...     output_file="objectives_{job_id}.json"
        ... )

Notes:
    Scripts receive the design-point and output paths through
    ``--design_params_file`` / ``--output_file`` and the corresponding
    ``AID2E_PARAMS_FILE`` / ``AID2E_OUTPUT_FILE`` environment variables.
    """
    path: str = Field(..., description="Path to objective computation script")
    output_file: str = Field(..., description="Output file pattern (e.g., objectives_*.json)")
    timeout_sec: int = Field(default=300, ge=1, description="Computation timeout in seconds")

StackBranchDefinition

Bases: BranchDefinition

Definition of a workflow branch narrowed to stages from an experimental software stack.

Source code in src/aid2e/utilities/configurations/experimental_stack_config.py
123
124
125
126
127
128
class StackBranchDefinition(BranchDefinition):
    """
    Definition of a workflow branch narrowed to stages
    from an experimental software stack.
    """
    stages: List[StackStageDefinition] = Field(default_factory=list, description="Software stack stage definitions")

StackJobDefinition

Bases: JobDefinition

Extends the base JobDefinition with a list of stack layers to utilize built-in commands of an experimental stack.

Extensions

script: Name of driver script to generate layers: Layer configurations to run in this job

Source code in src/aid2e/utilities/configurations/experimental_stack_config.py
 99
100
101
102
103
104
105
106
107
108
109
110
111
class StackJobDefinition(JobDefinition):
    """
    Extends the base JobDefinition with a list of stack
    layers to utilize built-in commands of an experimental
    stack.

    Extensions:
        script: Name of driver script to generate
        layers: Layer configurations to run in this job
    """
    command: Optional[str] = Field(default="./{script}", description="Executable command")
    script: Optional[str] = Field(default="do_job_{{context.job_id}}.sh", description="Driver script name")
    layers: List[StackLayerConfig] = Field(default_factory=list, description="Software stack layer configurations")

StackLayerConfig

Bases: BaseModel

Configures layer of an experimental stack

A job may consist of 1 or many layers from an experimental software stack. Sanitizes provided data to make sure singular vs. plural inputs/outputs are handled consistently.

Attributes:

Name Type Description
name Optional[str]

Unique name for this layer instance

layer str

Layer key (e.g. "sim", "rec", "ana")

inputs List[str]

List of inputs to layer

outputs List[str]

List of outputs from layer

arguments Optional[List[str]]

Optional list of any additional arguments to apply

command Optional[str]

Optional command to be run. Can be used to override default of layer.

rule Optional[str]

Optional recipe for combining inputs, outputs, arguments, and command. Can be used to override default of layer.

Notes
  • rule supports template substitutions for {inputs}, {outputs}, {arguments}, and {command}. See StackLayer for more details.
Source code in src/aid2e/utilities/configurations/experimental_stack_config.py
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
class StackLayerConfig(BaseModel):
    """Configures layer of an experimental stack

    A job may consist of 1 or many layers from an experimental
    software stack. Sanitizes provided data to make sure
    singular vs. plural inputs/outputs are handled consistently.

    Attributes:
        name: Unique name for this layer instance
        layer: Layer key (e.g. "sim", "rec", "ana")
        inputs: List of inputs to layer
        outputs: List of outputs from layer
        arguments: Optional list of any additional arguments to apply
        command: Optional command to be run. Can be used to
                 override default of layer.
        rule: Optional recipe for combining inputs, outputs, arguments,
              and command. Can be used to override default of layer.

    Notes:
        - rule supports template substitutions for {inputs}, {outputs},
          {arguments}, and {command}. See StackLayer for more details.
    """
    name: Optional[str] = Field(default=None, description="Unique name of instance")
    layer: str = Field(..., description="Layer key")
    inputs: List[str] = Field(..., description="List of inputs", validation_alias=AliasChoices('inputs', 'input'))
    outputs: List[str] = Field(..., description="List of outputs", validation_alias=AliasChoices('outputs', 'output'))
    arguments: Optional[List[str]] = Field(default=None, description="List of arguments")
    command: Optional[str] = Field(default=None, description="Executable command")
    rule: Optional[str] = Field(default=None, description="Recipe for combining arguments")

    @classmethod
    def pluralize_strings(cls, data, singular, plural):
        """Pluralize strings in data

         Sanitize input by data by ensuring that
         'singular' keys are always 'plural' and
         that their values are wrapped in lists.

         Args:
             data: the to be sanitized
             singular: the singular case of the key
                       (e.g. 'input')
             plural: the plural case of the key
                     (e.g. 'inputs')

         Returns:
             Sanitized data
         """
        if singular in data and plural not in data:
            data[singular] = data[plural]
        if isinstance(data.get(plural), str):
            data[plural] = [data[plural]]
        return data

    @model_validator(mode='before')
    @classmethod
    def handle_input_variants(cls, data):
        """
        Handles cases where (1) 'input' vs. 'inputs' was used in key,
        and (2) only 1 string was provided.
        """
        return cls.pluralize_strings(data, "input", "inputs")

    @model_validator(mode='before')
    @classmethod
    def handle_output_variants(cls, data):
        """
        Handles cases where (1) 'output' vs. 'outputs' was used in key,
        and (2) only 1 string was provided.
        """
        return cls.pluralize_strings(data, "output", "outputs")

handle_input_variants(data) classmethod

Handles cases where (1) 'input' vs. 'inputs' was used in key, and (2) only 1 string was provided.

Source code in src/aid2e/utilities/configurations/experimental_stack_config.py
80
81
82
83
84
85
86
87
@model_validator(mode='before')
@classmethod
def handle_input_variants(cls, data):
    """
    Handles cases where (1) 'input' vs. 'inputs' was used in key,
    and (2) only 1 string was provided.
    """
    return cls.pluralize_strings(data, "input", "inputs")

handle_output_variants(data) classmethod

Handles cases where (1) 'output' vs. 'outputs' was used in key, and (2) only 1 string was provided.

Source code in src/aid2e/utilities/configurations/experimental_stack_config.py
89
90
91
92
93
94
95
96
@model_validator(mode='before')
@classmethod
def handle_output_variants(cls, data):
    """
    Handles cases where (1) 'output' vs. 'outputs' was used in key,
    and (2) only 1 string was provided.
    """
    return cls.pluralize_strings(data, "output", "outputs")

pluralize_strings(data, singular, plural) classmethod

Pluralize strings in data

Sanitize input by data by ensuring that 'singular' keys are always 'plural' and that their values are wrapped in lists.

Parameters:

Name Type Description Default
data

the to be sanitized

required
singular

the singular case of the key (e.g. 'input')

required
plural

the plural case of the key (e.g. 'inputs')

required

Returns:

Type Description

Sanitized data

Source code in src/aid2e/utilities/configurations/experimental_stack_config.py
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
@classmethod
def pluralize_strings(cls, data, singular, plural):
    """Pluralize strings in data

     Sanitize input by data by ensuring that
     'singular' keys are always 'plural' and
     that their values are wrapped in lists.

     Args:
         data: the to be sanitized
         singular: the singular case of the key
                   (e.g. 'input')
         plural: the plural case of the key
                 (e.g. 'inputs')

     Returns:
         Sanitized data
     """
    if singular in data and plural not in data:
        data[singular] = data[plural]
    if isinstance(data.get(plural), str):
        data[plural] = [data[plural]]
    return data

StackRegistry

Unified registry for experimental stack configuration models + interfaces

Source code in src/aid2e/utilities/configurations/stack_registry.py
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
class StackRegistry:
    """
    Unified registry for experimental stack configuration models + interfaces
     """
    _env_configs: Dict[str, Type[BaseModel]] = {}
    _env_loaders: Dict[str, Type[Any]] = {}
    _design_configs: Dict[str, Type[BaseModel]] = {}
    _design_loaders: Dict[str, Type[Any]] = {}
    _workflow_configs: Dict[str, Type[BaseModel]] = {}
    _problem_configs: Dict[str, Type[BaseModel]] = {}
    _experimental_stacks: Dict[str, Type[Any]] = {}

    @classmethod
    def register_stack(
        cls,
        name: str,
        env_config: Type[BaseModel],
        env_loader: Type[Any],
        design_config: Type[BaseModel],
        design_loader: Type[Any],
        workflow_config: Type[BaseModel],
        experimental_stack: Type[Any],
        problem_config: Optional[Type[BaseModel]] = None,
    ) -> None:
        """Register a stack type and its configuration/implementation pair."""
        cls._env_configs[name] = env_config
        cls._env_loaders[name] = env_loader
        cls._design_configs[name] = design_config
        cls._design_loaders[name] = design_loader
        cls._workflow_configs[name] = workflow_config
        if problem_config is not None:
            cls._problem_configs[name] = problem_config
        cls._experimental_stacks[name] = experimental_stack

    @classmethod
    def get_env_config(cls, name: str) -> Type[BaseModel]:
        """Get the environment config model for a stack."""
        if name not in cls._env_configs:
            raise KeyError(f"Stack config model not registered: {name}")
        return cls._env_configs[name]

    @classmethod
    def get_env_loader(cls, name: str) -> Type[Any]:
        """Get the environment config loader for a stack."""
        if name not in cls._env_configs:
            raise KeyError(f"Stack config loader not registered: {name}")
        return cls._env_loaders[name]

    @classmethod
    def get_design_config(cls, name: str) -> Type[BaseModel]:
        """Get the design config model for a stack."""
        if name not in cls._design_configs:
            raise KeyError(f"Stack config model not registered: {name}")
        return cls._design_configs[name]

    @classmethod
    def get_design_loader(cls, name: str) -> Type[Any]:
        """Get the design config loader for a stack."""
        if name not in cls._design_configs:
            raise KeyError(f"Stack config loader not registered: {name}")
        return cls._design_loaders[name]

    @classmethod
    def get_workflow_config(cls, name: str) -> Type[BaseModel]:
        """Get the workflow config model for a stack."""
        if name not in cls._workflow_configs:
            raise KeyError(f"Stack config model not registered: {name}")
        return cls._workflow_configs[name]

    @classmethod
    def get_experimental_stack(cls, name: str) -> Type[Any]:
        """Get the stack implementation class for a stack name."""
        if name not in cls._experimental_stacks:
            raise KeyError(f"Experimental stack not registered: {name}")
        return cls._experimental_stacks[name]

    @classmethod
    def list_registered_stacks(cls) -> Dict[str, Dict[str, Type[Any]]]:
        return {
            name: {
                "env_config": cls._env_configs[name],
                "env_loader": cls._env_loaders[name],
                "design_config" : cls._design_configs[name],
                "design_loader" : cls._design_loaders[name],
                "workflow_config" : cls._workflow_configs[name],
                "problem_config": cls._problem_configs.get(name),
                "experimental_stack": cls._experimental_stacks[name],
            }
            for name in cls._env_configs
        }

get_design_config(name) classmethod

Get the design config model for a stack.

Source code in src/aid2e/utilities/configurations/stack_registry.py
60
61
62
63
64
65
@classmethod
def get_design_config(cls, name: str) -> Type[BaseModel]:
    """Get the design config model for a stack."""
    if name not in cls._design_configs:
        raise KeyError(f"Stack config model not registered: {name}")
    return cls._design_configs[name]

get_design_loader(name) classmethod

Get the design config loader for a stack.

Source code in src/aid2e/utilities/configurations/stack_registry.py
67
68
69
70
71
72
@classmethod
def get_design_loader(cls, name: str) -> Type[Any]:
    """Get the design config loader for a stack."""
    if name not in cls._design_configs:
        raise KeyError(f"Stack config loader not registered: {name}")
    return cls._design_loaders[name]

get_env_config(name) classmethod

Get the environment config model for a stack.

Source code in src/aid2e/utilities/configurations/stack_registry.py
46
47
48
49
50
51
@classmethod
def get_env_config(cls, name: str) -> Type[BaseModel]:
    """Get the environment config model for a stack."""
    if name not in cls._env_configs:
        raise KeyError(f"Stack config model not registered: {name}")
    return cls._env_configs[name]

get_env_loader(name) classmethod

Get the environment config loader for a stack.

Source code in src/aid2e/utilities/configurations/stack_registry.py
53
54
55
56
57
58
@classmethod
def get_env_loader(cls, name: str) -> Type[Any]:
    """Get the environment config loader for a stack."""
    if name not in cls._env_configs:
        raise KeyError(f"Stack config loader not registered: {name}")
    return cls._env_loaders[name]

get_experimental_stack(name) classmethod

Get the stack implementation class for a stack name.

Source code in src/aid2e/utilities/configurations/stack_registry.py
81
82
83
84
85
86
@classmethod
def get_experimental_stack(cls, name: str) -> Type[Any]:
    """Get the stack implementation class for a stack name."""
    if name not in cls._experimental_stacks:
        raise KeyError(f"Experimental stack not registered: {name}")
    return cls._experimental_stacks[name]

get_workflow_config(name) classmethod

Get the workflow config model for a stack.

Source code in src/aid2e/utilities/configurations/stack_registry.py
74
75
76
77
78
79
@classmethod
def get_workflow_config(cls, name: str) -> Type[BaseModel]:
    """Get the workflow config model for a stack."""
    if name not in cls._workflow_configs:
        raise KeyError(f"Stack config model not registered: {name}")
    return cls._workflow_configs[name]

register_stack(name, env_config, env_loader, design_config, design_loader, workflow_config, experimental_stack, problem_config=None) classmethod

Register a stack type and its configuration/implementation pair.

Source code in src/aid2e/utilities/configurations/stack_registry.py
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
@classmethod
def register_stack(
    cls,
    name: str,
    env_config: Type[BaseModel],
    env_loader: Type[Any],
    design_config: Type[BaseModel],
    design_loader: Type[Any],
    workflow_config: Type[BaseModel],
    experimental_stack: Type[Any],
    problem_config: Optional[Type[BaseModel]] = None,
) -> None:
    """Register a stack type and its configuration/implementation pair."""
    cls._env_configs[name] = env_config
    cls._env_loaders[name] = env_loader
    cls._design_configs[name] = design_config
    cls._design_loaders[name] = design_loader
    cls._workflow_configs[name] = workflow_config
    if problem_config is not None:
        cls._problem_configs[name] = problem_config
    cls._experimental_stacks[name] = experimental_stack

StackStageDefinition

Bases: StageDefinition

Definition of a workflow stage narrowed to jobs from an experimental software stack.

Source code in src/aid2e/utilities/configurations/experimental_stack_config.py
114
115
116
117
118
119
120
class StackStageDefinition(StageDefinition):

    """
    Definition of a workflow stage narrowed to jobs from
    an experimental software stack.
    """
    jobs: List[StackJobDefinition] = Field(default_factory=list, description="Software stack job definitions")

StackWorkflowDefinition

Bases: WorkflowDefinition

Definition of an experimental software stack workflow.

Source code in src/aid2e/utilities/configurations/experimental_stack_config.py
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
class StackWorkflowDefinition(WorkflowDefinition):
    """
    Definition of an experimental software stack
    workflow.
    """
    branches: List[StackBranchDefinition] = Field(default_factory=list, description="Software stack workflow branches (optional)")

    # FIXME there should be a way to automate this (incl. grabbing the
    # appropriate branch type) in the superclass
    def get_implicit_branch(self) -> StackBranchDefinition:
        """
        Get or create single implicit branch if branches list is empty.
        Overrides WorkflowDefinition.get_implicit_branch to return
        StackBranchDefinition.
        """
        if self.branches:
            raise ValueError("Branches already defined; cannot use implicit branch")
        return StackBranchDefinition(name="implicit")

get_implicit_branch()

Get or create single implicit branch if branches list is empty. Overrides WorkflowDefinition.get_implicit_branch to return StackBranchDefinition.

Source code in src/aid2e/utilities/configurations/experimental_stack_config.py
140
141
142
143
144
145
146
147
148
def get_implicit_branch(self) -> StackBranchDefinition:
    """
    Get or create single implicit branch if branches list is empty.
    Overrides WorkflowDefinition.get_implicit_branch to return
    StackBranchDefinition.
    """
    if self.branches:
        raise ValueError("Branches already defined; cannot use implicit branch")
    return StackBranchDefinition(name="implicit")

StackWorkflowsConfiguration

Bases: WorkflowsConfiguration

Container for software stack workflows.

Source code in src/aid2e/utilities/configurations/experimental_stack_config.py
151
152
153
154
155
156
157
class StackWorkflowsConfiguration(WorkflowsConfiguration):
    """
    Container for software stack workflows.
    """
    # FIXME ideally we should allow for workflows from separate
    # stacks to be run concurrently
    workflows: List[StackWorkflowDefinition] = Field(..., min_items=1, description="List of workflows")

StageDefinition

Bases: BaseModel

Stage/layer definition with jobs and scheduler.

A stage is a logical step group where multiple jobs run in parallel (fan-out), then their outputs feed into downstream stages (fan-in).

Attributes:

Name Type Description
name str

Stage name (e.g., "evaluate", "aggregate").

jobs List[JobDefinition]

Job definitions to execute (usually one template, expanded via job_factory).

job_factory Optional[JobFactory]

Optional factory for expanding jobs (e.g., N parallel evals).

scheduler Optional[SchedulerConfiguration]

Stage-level scheduler (optional, inherits global if not set).

parallelism ParallelismPolicy

Parallelism policy for this stage.

outputs List[ArtifactSpec]

Output artifact specs produced by this stage.

objective_plan Optional[ObjectivePlanSpec]

Optional plan that computes or collects objective values after stage jobs complete.

Example

stage = StageDefinition( ... name="evaluate", ... jobs=[ ... JobDefinition( ... name="dtlz2_evaluate", ... command="python scripts/dtlz2_problem.py", ... payload={...}, ... outputs=[ArtifactSpec(path="objectives_.json", format="json")] ... ) ... ], ... job_factory=JobFactory(type="range", params={"n": 4}), ... parallelism=ParallelismPolicy(max_concurrent=4, retry_max=2), ... outputs=[ArtifactSpec(path="objectives_.json", format="json")] ... )

Notes
  • job_factory expands the first job in jobs list to N parallel jobs
  • scheduler overrides global scheduler (from WorkflowsConfiguration)
  • outputs are collected after all jobs complete
Source code in src/aid2e/utilities/configurations/workflow_config.py
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
class StageDefinition(BaseModel):
    """Stage/layer definition with jobs and scheduler.

    A stage is a logical step group where multiple jobs run in parallel (fan-out),
    then their outputs feed into downstream stages (fan-in).

    Attributes:
        name: Stage name (e.g., "evaluate", "aggregate").
        jobs: Job definitions to execute (usually one template, expanded via job_factory).
        job_factory: Optional factory for expanding jobs (e.g., N parallel evals).
        scheduler: Stage-level scheduler (optional, inherits global if not set).
        parallelism: Parallelism policy for this stage.
        outputs: Output artifact specs produced by this stage.
        objective_plan: Optional plan that computes or collects objective values
            after stage jobs complete.

    Example:
        >>> stage = StageDefinition(
        ...     name="evaluate",
        ...     jobs=[
        ...         JobDefinition(
        ...             name="dtlz2_evaluate",
        ...             command="python scripts/dtlz2_problem.py",
        ...             payload={...},
        ...             outputs=[ArtifactSpec(path="objectives_*.json", format="json")]
        ...         )
        ...     ],
        ...     job_factory=JobFactory(type="range", params={"n": 4}),
        ...     parallelism=ParallelismPolicy(max_concurrent=4, retry_max=2),
        ...     outputs=[ArtifactSpec(path="objectives_*.json", format="json")]
        ... )

    Notes:
        - job_factory expands the first job in jobs list to N parallel jobs
        - scheduler overrides global scheduler (from WorkflowsConfiguration)
        - outputs are collected after all jobs complete
    """
    name: str = Field(..., description="Stage name")
    jobs: List[JobDefinition] = Field(default_factory=list, description="Job definitions")
    job_factory: Optional[JobFactory] = Field(default=None, description="Job expansion factory")
    scheduler: Optional[SchedulerConfiguration] = Field(default=None, description="Stage-level scheduler override")
    parallelism: ParallelismPolicy = Field(default_factory=ParallelismPolicy, description="Parallelism policy")
    outputs: List[ArtifactSpec] = Field(default_factory=list, description="Output artifacts")
    objective_plan: Optional[ObjectivePlanSpec] = Field(
        default=None,
        description="Plan that computes or collects objective values after this stage",
    )

StepPlanSpec

Bases: BaseModel

DAG-style step plan for an objective.

Replaces the earlier "branch" terminology with a clearer "steps" concept. A step plan is a small DAG of stages where exactly one stage must produce the objective value.

Attributes:

Name Type Description
stages List[StepStage]

Ordered list of stage definitions. Dependencies define the DAG.

produces_from_stage Optional[str]

Optional explicit producing stage name. If omitted, exactly one stage must set produces_objective=True.

Source code in src/aid2e/utilities/configurations/objectives.py
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
class StepPlanSpec(BaseModel):
    """DAG-style step plan for an objective.

    Replaces the earlier "branch" terminology with a clearer "steps"
    concept. A step plan is a small DAG of stages where exactly
    one stage must produce the objective value.

    Attributes:
        stages: Ordered list of stage definitions. Dependencies define the DAG.
        produces_from_stage: Optional explicit producing stage name. If omitted,
            exactly one stage must set ``produces_objective=True``.
    """

    model_config = ConfigDict(populate_by_name=True)

    stages: List[StepStage] = Field(..., min_items=1, description="Stages composing the computation DAG")
    produces_from_stage: Optional[str] = Field(
        default=None,
        description="Explicit stage name that emits the objective (overrides flag)",
        alias="produces_from_stage",
    )

    @model_validator(mode="after")
    def validate_stages(self) -> "StepPlanSpec":
        """Ensure unique names, valid dependencies, and single producer."""
        names = [stage.name for stage in self.stages]
        if len(set(names)) != len(names):
            raise ValueError("Stage names within steps must be unique")

        for stage in self.stages:
            for dep in stage.depends_on:
                if dep not in names:
                    raise ValueError(f"Stage '{stage.name}' depends on unknown stage '{dep}'")

        explicit = self.produces_from_stage
        producing_flags = [s.name for s in self.stages if s.produces_objective]

        if explicit:
            if explicit not in names:
                raise ValueError(f"produces_from_stage '{explicit}' not found in stages")
            chosen = explicit
        else:
            if len(producing_flags) != 1:
                raise ValueError("Exactly one stage must set produces_objective=True when produces_from_stage is not provided")
            chosen = producing_flags[0]

        self.produces_from_stage = chosen
        return self

    def producing_stage(self) -> str:
        """Return the name of the stage that emits the objective value."""
        if not self.produces_from_stage:
            raise ValueError("produces_from_stage was not resolved")
        return self.produces_from_stage

producing_stage()

Return the name of the stage that emits the objective value.

Source code in src/aid2e/utilities/configurations/objectives.py
206
207
208
209
210
def producing_stage(self) -> str:
    """Return the name of the stage that emits the objective value."""
    if not self.produces_from_stage:
        raise ValueError("produces_from_stage was not resolved")
    return self.produces_from_stage

validate_stages()

Ensure unique names, valid dependencies, and single producer.

Source code in src/aid2e/utilities/configurations/objectives.py
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
@model_validator(mode="after")
def validate_stages(self) -> "StepPlanSpec":
    """Ensure unique names, valid dependencies, and single producer."""
    names = [stage.name for stage in self.stages]
    if len(set(names)) != len(names):
        raise ValueError("Stage names within steps must be unique")

    for stage in self.stages:
        for dep in stage.depends_on:
            if dep not in names:
                raise ValueError(f"Stage '{stage.name}' depends on unknown stage '{dep}'")

    explicit = self.produces_from_stage
    producing_flags = [s.name for s in self.stages if s.produces_objective]

    if explicit:
        if explicit not in names:
            raise ValueError(f"produces_from_stage '{explicit}' not found in stages")
        chosen = explicit
    else:
        if len(producing_flags) != 1:
            raise ValueError("Exactly one stage must set produces_objective=True when produces_from_stage is not provided")
        chosen = producing_flags[0]

    self.produces_from_stage = chosen
    return self

StepStage

Bases: BaseModel

Single stage within an objective step plan.

Each stage executes either a script or an inline function, can declare inputs/outputs/extra_args, and may depend on upstream stages. If a plan has only one step, it is represented as a single-element step list.

Attributes:

Name Type Description
name str

Unique stage identifier.

description Optional[str]

Optional human-readable description of the stage intent.

script Optional[ScriptObjective]

Script-based execution for this stage (mutually exclusive with inline).

inline Optional[InlineObjective]

Inline Python callable for this stage (mutually exclusive with script).

inputs Dict[str, Any]

Optional input bindings for this stage (free-form mapping).

outputs Dict[str, Any]

Optional output bindings for this stage (free-form mapping).

extra_args Dict[str, Any]

Additional args/metadata for the stage executor.

produces_objective bool

Whether this stage emits the objective value.

depends_on List[str]

Names of upstream stages this stage depends on.

Source code in src/aid2e/utilities/configurations/objectives.py
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
class StepStage(BaseModel):
    """Single stage within an objective step plan.

    Each stage executes either a script or an inline function, can declare
    inputs/outputs/extra_args, and may depend on upstream stages. If a plan has
    only one step, it is represented as a single-element step list.

    Attributes:
        name: Unique stage identifier.
        description: Optional human-readable description of the stage intent.
        script: Script-based execution for this stage (mutually exclusive with inline).
        inline: Inline Python callable for this stage (mutually exclusive with script).
        inputs: Optional input bindings for this stage (free-form mapping).
        outputs: Optional output bindings for this stage (free-form mapping).
        extra_args: Additional args/metadata for the stage executor.
        produces_objective: Whether this stage emits the objective value.
        depends_on: Names of upstream stages this stage depends on.
    """

    model_config = ConfigDict(populate_by_name=True)

    name: str = Field(..., description="Stage name (unique within steps)")
    description: Optional[str] = Field(default=None, description="Stage description")
    script: Optional[ScriptObjective] = Field(default=None, description="Script execution for this stage")
    inline: Optional[InlineObjective] = Field(default=None, description="Inline callable for this stage")
    inputs: Dict[str, Any] = Field(default_factory=dict, description="Input bindings for this stage")
    outputs: Dict[str, Any] = Field(default_factory=dict, description="Output bindings for this stage")
    extra_args: Dict[str, Any] = Field(default_factory=dict, description="Extra args/metadata for the stage executor")
    produces_objective: bool = Field(default=False, description="Whether this stage emits the objective value")
    depends_on: List[str] = Field(default_factory=list, description="Upstream stage dependencies")

    @model_validator(mode="after")
    def validate_action(self) -> "StepStage":
        """Ensure stage has a valid execution definition.

        A stage must choose exactly one execution method: script or inline.
        """
        has_script = self.script is not None
        has_inline = self.inline is not None

        if has_script == has_inline:
            raise ValueError("Stage must define exactly one of: script or inline")

        return self

    @field_validator('depends_on')
    @classmethod
    def validate_dependencies(cls, depends_on: List[str], info: ValidationInfo) -> List[str]:
        """Ensure stages do not depend on themselves."""
        name = None
        if info and info.data:
            name = info.data.get("name")
        if name and name in depends_on:
            raise ValueError(f"Stage '{name}' cannot depend on itself")
        return depends_on

validate_action()

Ensure stage has a valid execution definition.

A stage must choose exactly one execution method: script or inline.

Source code in src/aid2e/utilities/configurations/objectives.py
131
132
133
134
135
136
137
138
139
140
141
142
143
@model_validator(mode="after")
def validate_action(self) -> "StepStage":
    """Ensure stage has a valid execution definition.

    A stage must choose exactly one execution method: script or inline.
    """
    has_script = self.script is not None
    has_inline = self.inline is not None

    if has_script == has_inline:
        raise ValueError("Stage must define exactly one of: script or inline")

    return self

validate_dependencies(depends_on, info) classmethod

Ensure stages do not depend on themselves.

Source code in src/aid2e/utilities/configurations/objectives.py
145
146
147
148
149
150
151
152
153
154
@field_validator('depends_on')
@classmethod
def validate_dependencies(cls, depends_on: List[str], info: ValidationInfo) -> List[str]:
    """Ensure stages do not depend on themselves."""
    name = None
    if info and info.data:
        name = info.data.get("name")
    if name and name in depends_on:
        raise ValueError(f"Stage '{name}' cannot depend on itself")
    return depends_on

WorkflowDefinition

Bases: BaseModel

Workflow definition with branches, objectives, and scheduler defaults.

A workflow is an end-to-end evaluation unit (e.g., one design point evaluation). It consists of optional branches, each with multiple stages, and defines the objectives to compute from the outputs.

Attributes:

Name Type Description
name str

Workflow name (e.g., "dtlz2_eval").

description Optional[str]

Optional description.

branches List[BranchDefinition]

Workflow branches (optional, defaults to single implicit branch if missing).

objectives List[ObjectiveDefinition]

Objectives to compute (reuses ObjectiveDefinition).

combined_objectives List[CombinedObjectivePlan]

Optional combined plans emitting multiple metrics in one run.

scheduler Optional[SchedulerConfiguration]

Workflow-level scheduler default (used if branch/stage unset).

Notes
  • If branches is empty, executor creates single implicit branch
  • Objectives are unified model (ObjectiveDefinition) for consistency
Source code in src/aid2e/utilities/configurations/workflow_config.py
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
class WorkflowDefinition(BaseModel):
    """Workflow definition with branches, objectives, and scheduler defaults.

    A workflow is an end-to-end evaluation unit (e.g., one design point evaluation).
    It consists of optional branches, each with multiple stages, and defines the
    objectives to compute from the outputs.

    Attributes:
        name: Workflow name (e.g., "dtlz2_eval").
        description: Optional description.
        branches: Workflow branches (optional, defaults to single implicit branch if missing).
        objectives: Objectives to compute (reuses ObjectiveDefinition).
        combined_objectives: Optional combined plans emitting multiple metrics in one run.
        scheduler: Workflow-level scheduler default (used if branch/stage unset).

    Notes:
        - If branches is empty, executor creates single implicit branch
        - Objectives are unified model (ObjectiveDefinition) for consistency
    """
    name: str = Field(..., description="Workflow name")
    description: Optional[str] = Field(default=None, description="Workflow description")
    branches: List[BranchDefinition] = Field(default_factory=list, description="Workflow branches (optional)")
    objectives: List[ObjectiveDefinition] = Field(
        default_factory=list,
        description="Objectives to compute (reuses ObjectiveDefinition)"
    )
    combined_objectives: List[CombinedObjectivePlan] = Field(
        default_factory=list,
        description="Combined objective plans emitting multiple metrics in one run",
    )
    scheduler: Optional[SchedulerConfiguration] = Field(
        default=None,
        description="Workflow-level scheduler default (overrides global, used if branch/stage unset)",
    )
    stack_type: Optional[str] = Field(
        default=None,
        description="Experimental stack type for workflow-level geometry prep",
    )

    def get_implicit_branch(self) -> BranchDefinition:
        """Get or create single implicit branch if branches list is empty.

        Returns:
            Single implicit branch if branches is empty, else raises error.

        Raises:
            ValueError: If branches list is not empty.
        """
        if self.branches:
            raise ValueError("Branches already defined; cannot use implicit branch")
        return BranchDefinition(name="implicit")

get_implicit_branch()

Get or create single implicit branch if branches list is empty.

Returns:

Type Description
BranchDefinition

Single implicit branch if branches is empty, else raises error.

Raises:

Type Description
ValueError

If branches list is not empty.

Source code in src/aid2e/utilities/configurations/workflow_config.py
280
281
282
283
284
285
286
287
288
289
290
291
def get_implicit_branch(self) -> BranchDefinition:
    """Get or create single implicit branch if branches list is empty.

    Returns:
        Single implicit branch if branches is empty, else raises error.

    Raises:
        ValueError: If branches list is not empty.
    """
    if self.branches:
        raise ValueError("Branches already defined; cannot use implicit branch")
    return BranchDefinition(name="implicit")

WorkflowsConfiguration

Bases: BaseModel

Top-level workflows configuration.

Container for multiple independent workflows (e.g., one per objective in a holistic optimization). Each workflow can have its own stages, scheduler, and objective specs.

Attributes:

Name Type Description
workflows List[WorkflowDefinition]

List of independent workflows.

global_scheduler Optional[SchedulerConfiguration]

Default scheduler for all stages (can be overridden per-stage).

Example

config = WorkflowsConfiguration( ... workflows=[ ... WorkflowDefinition(name="dtlz2_eval", ...), ... WorkflowDefinition(name="physics_sim", ...) ... ], ... global_scheduler=SchedulerConfiguration( ... runner_type="JobLibRunner", ... joblib=JobLibRunnerConfig(n_jobs=-1) ... ) ... )

Notes
  • global_scheduler is inherited by all stages unless overridden
  • workflows list must be non-empty
  • Useful for Option B: multiple independent workflows per objective
Source code in src/aid2e/utilities/configurations/workflow_config.py
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
class WorkflowsConfiguration(BaseModel):
    """Top-level workflows configuration.

    Container for multiple independent workflows (e.g., one per objective in a
    holistic optimization). Each workflow can have its own stages, scheduler, and
    objective specs.

    Attributes:
        workflows: List of independent workflows.
        global_scheduler: Default scheduler for all stages (can be overridden per-stage).

    Example:
        >>> config = WorkflowsConfiguration(
        ...     workflows=[
        ...         WorkflowDefinition(name="dtlz2_eval", ...),
        ...         WorkflowDefinition(name="physics_sim", ...)
        ...     ],
        ...     global_scheduler=SchedulerConfiguration(
        ...         runner_type="JobLibRunner",
        ...         joblib=JobLibRunnerConfig(n_jobs=-1)
        ...     )
        ... )

    Notes:
        - global_scheduler is inherited by all stages unless overridden
        - workflows list must be non-empty
        - Useful for Option B: multiple independent workflows per objective
    """
    workflows: List[WorkflowDefinition] = Field(..., min_items=1, description="List of workflows")
    global_scheduler: Optional[SchedulerConfiguration] = Field(
        default=None,
        description="Default scheduler for all stages (can be overridden per-stage)"
    )

    @field_validator('workflows')
    @classmethod
    def validate_unique_workflow_names(cls, workflows: List[WorkflowDefinition]) -> List[WorkflowDefinition]:
        """Ensure all workflow names are unique.

        Args:
            workflows: List of workflow definitions.

        Returns:
            Same list if valid.

        Raises:
            ValueError: If duplicate workflow names found.
        """
        names = [w.name for w in workflows]
        if len(set(names)) != len(names):
            raise ValueError("Workflow names must be unique")
        return workflows

validate_unique_workflow_names(workflows) classmethod

Ensure all workflow names are unique.

Parameters:

Name Type Description Default
workflows List[WorkflowDefinition]

List of workflow definitions.

required

Returns:

Type Description
List[WorkflowDefinition]

Same list if valid.

Raises:

Type Description
ValueError

If duplicate workflow names found.

Source code in src/aid2e/utilities/configurations/workflow_config.py
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
@field_validator('workflows')
@classmethod
def validate_unique_workflow_names(cls, workflows: List[WorkflowDefinition]) -> List[WorkflowDefinition]:
    """Ensure all workflow names are unique.

    Args:
        workflows: List of workflow definitions.

    Returns:
        Same list if valid.

    Raises:
        ValueError: If duplicate workflow names found.
    """
    names = [w.name for w in workflows]
    if len(set(names)) != len(names):
        raise ValueError("Workflow names must be unique")
    return workflows

create_scheduler_context(objective_scheduler=None, workflow_scheduler=None, branch_scheduler=None, stage_scheduler=None, global_scheduler=None)

Create a context dictionary with scheduler information for logging/debugging.

Parameters:

Name Type Description Default
objective_scheduler Optional[SchedulerConfiguration]

Reserved; objective-level scheduler execution is not supported

None
workflow_scheduler Optional[SchedulerConfiguration]

Scheduler at workflow level

None
branch_scheduler Optional[SchedulerConfiguration]

Scheduler at branch level

None
stage_scheduler Optional[SchedulerConfiguration]

Scheduler at stage level

None
global_scheduler Optional[SchedulerConfiguration]

Global scheduler

None

Returns:

Type Description
Dict[str, Any]

Dictionary with scheduler cascade information for each level

Example

context = create_scheduler_context( ... workflow_scheduler=SchedulerConfiguration(...), ... branch_scheduler=SchedulerConfiguration(...), ... ) print(context["effective_scheduler"]) # Will show the effective one

Source code in src/aid2e/utilities/configurations/scheduler_cascade.py
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
def create_scheduler_context(
    objective_scheduler: Optional[SchedulerConfiguration] = None,
    workflow_scheduler: Optional[SchedulerConfiguration] = None,
    branch_scheduler: Optional[SchedulerConfiguration] = None,
    stage_scheduler: Optional[SchedulerConfiguration] = None,
    global_scheduler: Optional[SchedulerConfiguration] = None,
) -> Dict[str, Any]:
    """
    Create a context dictionary with scheduler information for logging/debugging.

    Args:
        objective_scheduler: Reserved; objective-level scheduler execution is not supported
        workflow_scheduler: Scheduler at workflow level
        branch_scheduler: Scheduler at branch level
        stage_scheduler: Scheduler at stage level
        global_scheduler: Global scheduler

    Returns:
        Dictionary with scheduler cascade information for each level

    Example:
        >>> context = create_scheduler_context(
        ...     workflow_scheduler=SchedulerConfiguration(...),
        ...     branch_scheduler=SchedulerConfiguration(...),
        ... )
        >>> print(context["effective_scheduler"])  # Will show the effective one
    """
    effective = resolve_scheduler_cascade(
        stage_scheduler, branch_scheduler, workflow_scheduler, objective_scheduler, global_scheduler
    )

    return {
        "cascade_levels": {
            "stage": stage_scheduler.runner_type if stage_scheduler else None,
            "branch": branch_scheduler.runner_type if branch_scheduler else None,
            "workflow": workflow_scheduler.runner_type if workflow_scheduler else None,
            "global": global_scheduler.runner_type if global_scheduler else None,
        },
        "effective_scheduler": effective.runner_type if effective else None,
        "source": _get_cascade_source(
            stage_scheduler,
            branch_scheduler,
            workflow_scheduler,
            global_scheduler,
        ),
    }

get(name)

Retrieve a registered runner config model by name.

Parameters:

Name Type Description Default
name str

Runner type identifier.

required

Returns:

Type Description
Optional[Type[BaseModel]]

The registered Pydantic model class, or None if not registered.

Source code in src/aid2e/utilities/configurations/scheduler_registry.py
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
def get(name: str) -> Optional[Type[BaseModel]]:
    """Retrieve a registered runner config model by name.

    Args:
        name: Runner type identifier.

    Returns:
        The registered Pydantic model class, or None if not registered.
    """
    if name in _runner_config_registry:
        return _runner_config_registry[name]
    if name in _runner_config_loaders:
        model = _runner_config_loaders[name]()
        _runner_config_registry[name] = model
        return model
    return None

list_registered()

Get all registered runner config models.

Returns:

Type Description
Dict[str, Type[BaseModel]]

Dict mapping runner type names to Pydantic config classes.

Source code in src/aid2e/utilities/configurations/scheduler_registry.py
54
55
56
57
58
59
60
61
62
63
64
65
66
def list_registered() -> Dict[str, Type[BaseModel]]:
    """Get all registered runner config models.

    Returns:
        Dict mapping runner type names to Pydantic config classes.
    """
    for name in list(_runner_config_loaders.keys()):
        if name not in _runner_config_registry:
            try:
                _runner_config_registry[name] = _runner_config_loaders[name]()
            except Exception:
                pass
    return _runner_config_registry.copy()

load_config(config_file)

Load complete configuration from a YAML file.

Parameters:

Name Type Description Default
config_file str

Path to YAML configuration file

required

Returns:

Type Description
FullConfig

FullConfig object with all configurations loaded

Raises:

Type Description
FileNotFoundError

If config file doesn't exist

ValueError

If configuration is invalid

Source code in src/aid2e/utilities/configurations/full_config.py
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
def load_config(config_file: str) -> FullConfig:
    """
    Load complete configuration from a YAML file.

    Args:
        config_file: Path to YAML configuration file

    Returns:
        FullConfig object with all configurations loaded

    Raises:
        FileNotFoundError: If config file doesn't exist
        ValueError: If configuration is invalid
    """
    config_path = Path(config_file)

    if not config_path.exists():
        raise FileNotFoundError(f"Config file not found: {config_file}")

    with open(config_path, 'r') as f:
        data = yaml.safe_load(f)

    normalized = _normalize_full_config_data(data or {}, config_path)

    return FullConfig(**normalized)

load_optimizer_config(config_file)

Load only the optimizer section from a full config file.

Source code in src/aid2e/utilities/configurations/loaders.py
65
66
67
68
def load_optimizer_config(config_file: str) -> OptimizerConfiguration:
    """Load only the optimizer section from a full config file."""
    normalized = _normalize_sections(config_file)
    return normalized["optimizer"]

load_problem_config(config_file)

Load only the problem section from a full config file.

Source code in src/aid2e/utilities/configurations/loaders.py
59
60
61
62
def load_problem_config(config_file: str) -> ProblemConfiguration:
    """Load only the problem section from a full config file."""
    normalized = _normalize_sections(config_file)
    return normalized["problem"]

load_raw_config(config_file)

Load raw YAML/JSON configuration content from disk.

Parameters:

Name Type Description Default
config_file str

Path to YAML or JSON config.

required

Returns:

Type Description
Dict[str, Any]

Raw top-level configuration dictionary.

Raises:

Type Description
FileNotFoundError

If the config file does not exist.

ValueError

If the extension is unsupported.

Source code in src/aid2e/utilities/configurations/loaders.py
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
def load_raw_config(config_file: str) -> Dict[str, Any]:
    """Load raw YAML/JSON configuration content from disk.

    Args:
        config_file: Path to YAML or JSON config.

    Returns:
        Raw top-level configuration dictionary.

    Raises:
        FileNotFoundError: If the config file does not exist.
        ValueError: If the extension is unsupported.
    """
    path = Path(config_file)
    if not path.exists():
        raise FileNotFoundError(f"Config file not found: {config_file}")

    text = path.read_text(encoding="utf-8")
    suffix = path.suffix.lower()
    if suffix in {".yaml", ".yml"}:
        return yaml.safe_load(text) or {}
    if suffix == ".json":
        return json.loads(text)

    raise ValueError(
        f"Unsupported config extension '{suffix}'. "
        "Use .yaml, .yml, or .json."
    )

load_scheduler_config(config_file)

Load only the scheduler section from a full config file.

Accepts only canonical scheduler payloads with parameters.

Source code in src/aid2e/utilities/configurations/loaders.py
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def load_scheduler_config(config_file: str) -> Optional[SchedulerConfiguration]:
    """Load only the scheduler section from a full config file.

    Accepts only canonical scheduler payloads with ``parameters``.
    """
    path = Path(config_file)
    raw = load_raw_config(config_file)
    scheduler_raw = raw.get("scheduler")
    if not scheduler_raw:
        normalized = _normalize_full_config_data(raw, path)
        return normalized.get("scheduler")

    if not isinstance(scheduler_raw, dict):
        raise ValueError("'scheduler' section must be a mapping")

    return SchedulerConfigLoader.from_dict(scheduler_raw, base_dir=str(path.parent))

load_workflow_config(config_file)

Load workflow configuration from full config if present.

Accepted layout
  • workflows: {workflows: [...]}
Source code in src/aid2e/utilities/configurations/loaders.py
89
90
91
92
93
94
95
96
def load_workflow_config(config_file: str) -> Optional[WorkflowsConfiguration]:
    """Load workflow configuration from full config if present.

    Accepted layout:
        - ``workflows: {workflows: [...]}``
    """
    normalized = _normalize_sections(config_file)
    return normalized.get("workflows")

register(name, model)

Register a Pydantic model for a scheduler runner type.

Parameters:

Name Type Description Default
name str

Runner type identifier (e.g., "JobLibRunner", "SlurmRunner", "PanDAiDDSRunner").

required
model Type[BaseModel]

Pydantic model class that validates runner-specific params.

required
Source code in src/aid2e/utilities/configurations/scheduler_registry.py
25
26
27
28
29
30
31
32
33
def register(name: str, model: Type[BaseModel]) -> None:
    """Register a Pydantic model for a scheduler runner type.

    Args:
        name: Runner type identifier (e.g., "JobLibRunner", "SlurmRunner", "PanDAiDDSRunner").
        model: Pydantic model class that validates runner-specific params.
    """
    name_key = name
    _runner_config_registry[name_key] = model

resolve_scheduler_cascade(stage_scheduler=None, branch_scheduler=None, workflow_scheduler=None, objective_scheduler=None, global_scheduler=None)

Resolve the effective scheduler configuration using cascade precedence.

Cascade order (highest to lowest priority): 1. Stage-level scheduler (stage override) 2. Branch-level scheduler (branch default) 3. Workflow-level scheduler (workflow default) 4. Global scheduler (global default)

Parameters:

Name Type Description Default
stage_scheduler Optional[SchedulerConfiguration]

Scheduler at stage level (highest priority)

None
branch_scheduler Optional[SchedulerConfiguration]

Scheduler at branch level

None
workflow_scheduler Optional[SchedulerConfiguration]

Scheduler at workflow level

None
objective_scheduler Optional[SchedulerConfiguration]

Reserved. Objective-level scheduler execution is not supported; scheduled objective work should be modeled as workflow stages.

None
global_scheduler Optional[SchedulerConfiguration]

Global scheduler (lowest priority)

None

Returns:

Type Description
Optional[SchedulerConfiguration]

The first non-None scheduler in the cascade, or None if all are None

Example

stage_sched = SchedulerConfiguration(runner_type="SlurmRunner", parameters={}) branch_sched = SchedulerConfiguration(runner_type="JobLibRunner", parameters={}) effective = resolve_scheduler_cascade(stage_sched, branch_sched) assert effective == stage_sched # Stage overrides branch

effective = resolve_scheduler_cascade(None, branch_sched) assert effective == branch_sched # Branch used if stage is None

Source code in src/aid2e/utilities/configurations/scheduler_cascade.py
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
def resolve_scheduler_cascade(
    stage_scheduler: Optional[SchedulerConfiguration] = None,
    branch_scheduler: Optional[SchedulerConfiguration] = None,
    workflow_scheduler: Optional[SchedulerConfiguration] = None,
    objective_scheduler: Optional[SchedulerConfiguration] = None,
    global_scheduler: Optional[SchedulerConfiguration] = None,
) -> Optional[SchedulerConfiguration]:
    """
    Resolve the effective scheduler configuration using cascade precedence.

    Cascade order (highest to lowest priority):
    1. Stage-level scheduler (stage override)
    2. Branch-level scheduler (branch default)
    3. Workflow-level scheduler (workflow default)
    4. Global scheduler (global default)

    Args:
        stage_scheduler: Scheduler at stage level (highest priority)
        branch_scheduler: Scheduler at branch level
        workflow_scheduler: Scheduler at workflow level
        objective_scheduler: Reserved. Objective-level scheduler execution is
            not supported; scheduled objective work should be modeled as
            workflow stages.
        global_scheduler: Global scheduler (lowest priority)

    Returns:
        The first non-None scheduler in the cascade, or None if all are None

    Example:
        >>> stage_sched = SchedulerConfiguration(runner_type="SlurmRunner", parameters={})
        >>> branch_sched = SchedulerConfiguration(runner_type="JobLibRunner", parameters={})
        >>> effective = resolve_scheduler_cascade(stage_sched, branch_sched)
        >>> assert effective == stage_sched  # Stage overrides branch

        >>> effective = resolve_scheduler_cascade(None, branch_sched)
        >>> assert effective == branch_sched  # Branch used if stage is None
    """
    if objective_scheduler is not None:
        raise ValueError(
            "Objective-level scheduler cascade is not supported yet. "
            "Represent scheduled objective work as workflow stages."
        )

    # Check in order of precedence
    if stage_scheduler is not None:
        return stage_scheduler
    if branch_scheduler is not None:
        return branch_scheduler
    if workflow_scheduler is not None:
        return workflow_scheduler
    if global_scheduler is not None:
        return global_scheduler
    return None

ePIC-specific utilities for AID2E Framework.

EpicAnaLayer dataclass

Bases: AnaLayer

Analysis layer of ePIC stack

Source code in src/aid2e/utilities/epic_utils/epic_stack.py
133
134
135
class EpicAnaLayer(AnaLayer):
    """Analysis layer of ePIC stack"""
    pass

EpicBranchDefinition

Bases: StackBranchDefinition

Definition of an ePIC branch in a workflow. Stages are restricted to ePIC stages.

Source code in src/aid2e/utilities/epic_utils/epic_stack_config.py
45
46
47
48
49
50
class EpicBranchDefinition(StackBranchDefinition):
    """
    Definition of an ePIC branch in a workflow.
    Stages are restricted to ePIC stages.
    """
    stages: List[EpicStageDefinition] = Field(default_factory = list, description="ePIC stack branch definitions")

EpicDesignConfig

Bases: DesignConfig

ePIC-specific design configuration with XML integration. Extends DesignConfig with XML modification capabilities and optimization groups.

Note: Uses 'epic_design_parameters' instead of 'design_parameters' to distinguish from generic configs in YAML files.

Source code in src/aid2e/utilities/epic_utils/epic_design_config.py
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
class EpicDesignConfig(DesignConfig):
    """
    ePIC-specific design configuration with XML integration.
    Extends DesignConfig with XML modification capabilities and optimization groups.

    Note: Uses 'epic_design_parameters' instead of 'design_parameters' to distinguish
    from generic configs in YAML files.
    """
    # Override to use ePIC-specific parameters
    design_parameters: Optional[Any] = None  # Set to None to avoid conflicts
    epic_design_parameters: EpicDesignParameters
    key: ClassVar[str] = 'epic_design_space'

    def get_flat_parameters(self) -> Dict[str, BaseParameter]:
        """Returns a flat dictionary of all parameters keyed by their qualified name."""
        flat = {}
        for group in self.epic_design_parameters.root.values():
            for param in group.parameters.values():
                flat[param.name] = param
        return flat

    def get_parameter_names(self) -> List[str]:
        """Get all parameter qualified names."""
        return list(self.get_flat_parameters().keys())

    def get_xml_modifications(self, param_values: Optional[Dict[str, float]] = None) -> Dict[str, List[Tuple[str, str, str, Any]]]:
        """
        Get XML modifications for given parameter values.

        Args:
            param_values: Dictionary of qualified parameter names to values.
                         If None, uses default values from config.

        Returns:
            Dictionary mapping file_path -> [(xml_path, attribute, unit, new_value), ...]
        """
        if param_values is None:
            # Use default values from config
            param_values = {name: param.value for name, param in self.get_flat_parameters().items()}

        modifications = {}

        for group_name, group in self.epic_design_parameters.root.items():
            # Expand environment variables in file path
            file_path = os.path.expandvars(group.file_path)

            if file_path not in modifications:
                modifications[file_path] = []

            for param_name, param in group.parameters.items():
                qualified_name = f"{group_name}.{param_name}"
                if qualified_name in param_values:
                    new_value = param_values[qualified_name]
                    modifications[file_path].append((
                        param.xml_path,
                        param.attribute,
                        param.unit or "",
                        new_value
                    ))

        return modifications

    def get_file_paths(self) -> List[str]:
        """Get all unique file paths referenced in the configuration."""
        return list(set(
            os.path.expandvars(group.file_path)
            for group in self.epic_design_parameters.root.values()
        ))

get_file_paths()

Get all unique file paths referenced in the configuration.

Source code in src/aid2e/utilities/epic_utils/epic_design_config.py
140
141
142
143
144
145
def get_file_paths(self) -> List[str]:
    """Get all unique file paths referenced in the configuration."""
    return list(set(
        os.path.expandvars(group.file_path)
        for group in self.epic_design_parameters.root.values()
    ))

get_flat_parameters()

Returns a flat dictionary of all parameters keyed by their qualified name.

Source code in src/aid2e/utilities/epic_utils/epic_design_config.py
91
92
93
94
95
96
97
def get_flat_parameters(self) -> Dict[str, BaseParameter]:
    """Returns a flat dictionary of all parameters keyed by their qualified name."""
    flat = {}
    for group in self.epic_design_parameters.root.values():
        for param in group.parameters.values():
            flat[param.name] = param
    return flat

get_parameter_names()

Get all parameter qualified names.

Source code in src/aid2e/utilities/epic_utils/epic_design_config.py
 99
100
101
def get_parameter_names(self) -> List[str]:
    """Get all parameter qualified names."""
    return list(self.get_flat_parameters().keys())

get_xml_modifications(param_values=None)

Get XML modifications for given parameter values.

Parameters:

Name Type Description Default
param_values Optional[Dict[str, float]]

Dictionary of qualified parameter names to values. If None, uses default values from config.

None

Returns:

Type Description
Dict[str, List[Tuple[str, str, str, Any]]]

Dictionary mapping file_path -> [(xml_path, attribute, unit, new_value), ...]

Source code in src/aid2e/utilities/epic_utils/epic_design_config.py
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
def get_xml_modifications(self, param_values: Optional[Dict[str, float]] = None) -> Dict[str, List[Tuple[str, str, str, Any]]]:
    """
    Get XML modifications for given parameter values.

    Args:
        param_values: Dictionary of qualified parameter names to values.
                     If None, uses default values from config.

    Returns:
        Dictionary mapping file_path -> [(xml_path, attribute, unit, new_value), ...]
    """
    if param_values is None:
        # Use default values from config
        param_values = {name: param.value for name, param in self.get_flat_parameters().items()}

    modifications = {}

    for group_name, group in self.epic_design_parameters.root.items():
        # Expand environment variables in file path
        file_path = os.path.expandvars(group.file_path)

        if file_path not in modifications:
            modifications[file_path] = []

        for param_name, param in group.parameters.items():
            qualified_name = f"{group_name}.{param_name}"
            if qualified_name in param_values:
                new_value = param_values[qualified_name]
                modifications[file_path].append((
                    param.xml_path,
                    param.attribute,
                    param.unit or "",
                    new_value
                ))

    return modifications

EpicDesignConfigLoader

Bases: DesignConfigLoader

Loader for ePIC design configurations. Can load either from external files or inline YAML blocks. instantiates EpicDesignConfig objects.

Source code in src/aid2e/utilities/epic_utils/epic_design_config.py
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
class EpicDesignConfigLoader(DesignConfigLoader):
    """
    Loader for ePIC design configurations. Can load either from
    external files or inline YAML blocks. instantiates EpicDesignConfig
    objects.
    """
    space_key = EpicDesignConfig.key
    param_key = EpicDesignParameters.key

    @staticmethod
    def load(file_path: str = None, design_data: Dict[str, Any] = None) -> "EpicDesignConfig":
        """
        Load an ePIC design configuration.
        """
        data = EpicDesignConfigLoader._process_inputs(file_path, design_data)
        return EpicDesignConfig(**data)

load(file_path=None, design_data=None) staticmethod

Load an ePIC design configuration.

Source code in src/aid2e/utilities/epic_utils/epic_design_config.py
157
158
159
160
161
162
163
@staticmethod
def load(file_path: str = None, design_data: Dict[str, Any] = None) -> "EpicDesignConfig":
    """
    Load an ePIC design configuration.
    """
    data = EpicDesignConfigLoader._process_inputs(file_path, design_data)
    return EpicDesignConfig(**data)

EpicDesignParameters

Bases: RootModel[Dict[str, EpicParameterGroup]]

Collection of ePIC parameter groups.

Source code in src/aid2e/utilities/epic_utils/epic_design_config.py
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
class EpicDesignParameters(RootModel[Dict[str, EpicParameterGroup]]):
    """
    Collection of ePIC parameter groups.
    """
    key: ClassVar[str] = 'epic_design_parameters'

    @model_validator(mode="before")
    @classmethod
    def inject_qualified_names(cls, values: Dict[str, dict]):
        """
        Injects full qualified names like 'group.param' into each parameter.
        This ensures parameters are uniquely identified.
        """
        for group_name, group_data in values.items():
            param_dict = group_data.get("parameters", {})
            for param_name, param_data in param_dict.items():
                if isinstance(param_data, dict) and "name" not in param_data:
                    param_data["name"] = f"{group_name}.{param_name}"
        return values

inject_qualified_names(values) classmethod

Injects full qualified names like 'group.param' into each parameter. This ensures parameters are uniquely identified.

Source code in src/aid2e/utilities/epic_utils/epic_design_config.py
63
64
65
66
67
68
69
70
71
72
73
74
75
@model_validator(mode="before")
@classmethod
def inject_qualified_names(cls, values: Dict[str, dict]):
    """
    Injects full qualified names like 'group.param' into each parameter.
    This ensures parameters are uniquely identified.
    """
    for group_name, group_data in values.items():
        param_dict = group_data.get("parameters", {})
        for param_name, param_data in param_dict.items():
            if isinstance(param_data, dict) and "name" not in param_data:
                param_data["name"] = f"{group_name}.{param_name}"
    return values

EpicEnvConfig

Bases: EnvironmentConfig

ePIC-specific environment configuration.

Manages ePIC detector environment variables including singularity image, installation paths, and EIC reconstruction settings. If both singularity image and EIC shell script paths are provided, defaults to singularity image.

Attributes:

Name Type Description
eic_shell Optional[str]

Path to the EIC shell script (usually named eic-shell), either this OR singularity_image must be provided

singularity_image Optional[str]

Path to the EIC shell singularity image, either this OR eic_shell must be provided

epic_install Optional[str]

Optional path to ePIC installation directory, will be used as template for modifying geometry

epic_config Optional[str]

ePIC geometry configuration to use (e.g. epic, epic_full)

geometry_mode Optional[str]

Geometry activation mode, either build or no_build

eic_recon_install Optional[str]

Optional path to EIC reconstruction installation

eic_recon Optional[str]

Optional override for EIC reconstruction command

Source code in src/aid2e/utilities/epic_utils/epic_env_config.py
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
class EpicEnvConfig(EnvironmentConfig):
    """ePIC-specific environment configuration.

    Manages ePIC detector environment variables including singularity image,
    installation paths, and EIC reconstruction settings. If both singularity
    image and EIC shell script paths are provided, defaults to singularity
    image.

    Attributes:
        eic_shell: Path to the EIC shell script (usually named eic-shell),
                   either this OR singularity_image must be provided
        singularity_image: Path to the EIC shell singularity image, either
                           this OR eic_shell must be provided
        epic_install: Optional path to ePIC installation directory, will be
                      used as template for modifying geometry
        epic_config: ePIC geometry configuration to use (e.g. epic, epic_full)
        geometry_mode: Geometry activation mode, either build or no_build
        eic_recon_install: Optional path to EIC reconstruction installation
        eic_recon: Optional override for EIC reconstruction command
    """
    epic_install:Optional[str]
    epic_config: Optional[str]
    geometry_mode: Optional[str] = "build"
    eic_shell: Optional[str] = None
    singularity_image: Optional[str] = None
    eic_recon_install: Optional[str] = None
    eic_recon: Optional[str] = None

    # set key associated with model
    key: ClassVar[str] = "epic_environment"

    @model_validator(mode='before')
    @classmethod
    def ensure_shell_or_image(cls, data):
        """
        Ensure that either eic_shell or singularity_image
        were provided
        """
        is_shell_there = 'eic_shell' in data
        is_image_there = 'singularity_image' in data
        assert is_shell_there or is_image_there
        return data

    def activate(self) -> None:
        """Activate ePIC environment variables and print a summary."""

        # default to singularity image over eic-shell
        if self.singularity_image:
            os.environ["EIC_SINGULARITY_IMAGE"] = self.singularity_image
        elif self.eic_shell:
            os.environ["EIC_SHELL"] = self.eic_shell

        # set other variables
        if self.epic_install:
            os.environ["EPIC_INSTALL"] = self.epic_install
            if not self.eic_recon_install:
                self.eic_recon_install = str(Path(self.epic_install) / "local")
        if self.epic_config:
            os.environ["EPIC_CONFIG"] = self.epic_config
        if self.eic_recon_install:
            os.environ["EIC_RECON_INSTALL"] = self.eic_recon_install
        if self.eic_recon:
            os.environ["EIC_RECON"] = self.eic_recon

        print("[INFO] ePIC environment variables set:")
        for var in ["EPIC_INSTALL", "EIC_RECON_INSTALL", "EIC_SHELL", "EIC_RECON"]:
            if var in os.environ:
                print(f"  {var} = {os.environ[var]}")

activate()

Activate ePIC environment variables and print a summary.

Source code in src/aid2e/utilities/epic_utils/epic_env_config.py
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
def activate(self) -> None:
    """Activate ePIC environment variables and print a summary."""

    # default to singularity image over eic-shell
    if self.singularity_image:
        os.environ["EIC_SINGULARITY_IMAGE"] = self.singularity_image
    elif self.eic_shell:
        os.environ["EIC_SHELL"] = self.eic_shell

    # set other variables
    if self.epic_install:
        os.environ["EPIC_INSTALL"] = self.epic_install
        if not self.eic_recon_install:
            self.eic_recon_install = str(Path(self.epic_install) / "local")
    if self.epic_config:
        os.environ["EPIC_CONFIG"] = self.epic_config
    if self.eic_recon_install:
        os.environ["EIC_RECON_INSTALL"] = self.eic_recon_install
    if self.eic_recon:
        os.environ["EIC_RECON"] = self.eic_recon

    print("[INFO] ePIC environment variables set:")
    for var in ["EPIC_INSTALL", "EIC_RECON_INSTALL", "EIC_SHELL", "EIC_RECON"]:
        if var in os.environ:
            print(f"  {var} = {os.environ[var]}")

ensure_shell_or_image(data) classmethod

Ensure that either eic_shell or singularity_image were provided

Source code in src/aid2e/utilities/epic_utils/epic_env_config.py
47
48
49
50
51
52
53
54
55
56
57
@model_validator(mode='before')
@classmethod
def ensure_shell_or_image(cls, data):
    """
    Ensure that either eic_shell or singularity_image
    were provided
    """
    is_shell_there = 'eic_shell' in data
    is_image_there = 'singularity_image' in data
    assert is_shell_there or is_image_there
    return data

EpicEnvConfigLoader

Bases: EnvironmentConfigLoader

Loader for ePIC environment configuration. Loads YAML files, instantiates EpicEnvConfig.

Source code in src/aid2e/utilities/epic_utils/epic_env_config.py
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
class EpicEnvConfigLoader(EnvironmentConfigLoader):
    """
    Loader for ePIC environment configuration. Loads YAML
    files, instantiates EpicEnvConfig.
    """
    @staticmethod
    def load(env_data: Dict[str, Any] = None, file_path: str = None) -> "EpicEnvConfig":
        """
        Load ePIC environment configuration.

        Args:
            env_data: Loaded data stored in a dictionary
            file_path: Path to YAML configuration file

        Returns:
            EpicEnvConfig instance
        """
        # should EITHER provide data as a dict OR a file path
        # as a string
        is_data_provided = env_data is not None
        is_file_provided = file_path is not None
        if is_data_provided and is_file_provided:
            raise RuntimeWarning(f"Both data and a file path ({file_path}) were provided. Defaulting to data.")

        data = None
        if is_data_provided:
            data = env_data
        elif is_file_provided:
            path = pathlib.Path(file_path)
            if not path.exists():
                raise FileNotFoundError(f"Configuration file not found: {file_path}")
            with open(path, 'r') as file:
                data = yaml.safe_load(file)
        else:
            raise RuntimeError("Must provide either data as a dictionary or a 'path' to a file")

        if EpicEnvConfig.key not in data:
            raise ValueError(f"Invalid data configuration: missing '{EpicEnvConfig.key}' in data")
        return EpicEnvConfig(**data[EpicEnvConfig.key])

load(env_data=None, file_path=None) staticmethod

Load ePIC environment configuration.

Parameters:

Name Type Description Default
env_data Dict[str, Any]

Loaded data stored in a dictionary

None
file_path str

Path to YAML configuration file

None

Returns:

Type Description
EpicEnvConfig

EpicEnvConfig instance

Source code in src/aid2e/utilities/epic_utils/epic_env_config.py
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
@staticmethod
def load(env_data: Dict[str, Any] = None, file_path: str = None) -> "EpicEnvConfig":
    """
    Load ePIC environment configuration.

    Args:
        env_data: Loaded data stored in a dictionary
        file_path: Path to YAML configuration file

    Returns:
        EpicEnvConfig instance
    """
    # should EITHER provide data as a dict OR a file path
    # as a string
    is_data_provided = env_data is not None
    is_file_provided = file_path is not None
    if is_data_provided and is_file_provided:
        raise RuntimeWarning(f"Both data and a file path ({file_path}) were provided. Defaulting to data.")

    data = None
    if is_data_provided:
        data = env_data
    elif is_file_provided:
        path = pathlib.Path(file_path)
        if not path.exists():
            raise FileNotFoundError(f"Configuration file not found: {file_path}")
        with open(path, 'r') as file:
            data = yaml.safe_load(file)
    else:
        raise RuntimeError("Must provide either data as a dictionary or a 'path' to a file")

    if EpicEnvConfig.key not in data:
        raise ValueError(f"Invalid data configuration: missing '{EpicEnvConfig.key}' in data")
    return EpicEnvConfig(**data[EpicEnvConfig.key])

EpicGeoLayer dataclass

Bases: StackLayer

Geometry layer of ePIC stack

Source code in src/aid2e/utilities/epic_utils/epic_stack.py
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
class EpicGeoLayer(StackLayer):
    """Geometry layer of ePIC stack"""
    name = "geo"
    command = "checkOverlaps"
    rule = '{{command}} {{arguments}} -c {{inputs}} {{outputs}}'

    def _make_input_arg(self, inputs: List[str]) -> str:
        """
        Formats inputs for ePIC-specific geometry
        layer. There should be exactly one input,
        the geometry configuration file to run
        overlap check on.
        """
        if len(inputs) != 1:
            raise ValueError(f"EpicGeoLayer takes one input, got {len(inputs)}")
        return inputs[0]

    def _make_output_arg(self, outputs: List[str]) -> str:
        """
        Formats outputs for ePIC-specific geometry
        layer. There should be exactly one output,
        the log file to store the results of the
        check.

        Also adds shell code to check for overlaps/
        extrusions and exit if any found.
        """
        if len(outputs) != 1:
            raise ValueError(f"EpicGeoLayer takes one output, got {len(outputs)}")
        output = outputs[0]

        checks = [
            f' >& {output}',
            "grep -Eq 'Number of illegal overlaps/extrusions[[:space:]]*"
            rf":[[:space:]]*0[[:space:]]*$' {output} || exit 9",
        ]
        return '\n'.join(checks)

EpicJobDefinition

Bases: StackJobDefinition

Definition a job to run 1 or more ePIC stack layers. If no command provided, will set default based on specified evaluator_type.

Source code in src/aid2e/utilities/epic_utils/epic_stack_config.py
28
29
30
31
32
33
34
class EpicJobDefinition(StackJobDefinition):
    """
    Definition a job to run 1 or more ePIC stack
    layers. If no command provided, will set
    default based on specified evaluator_type.
    """
    layers: List[EpicLayerConfig] = Field(default_factory=list, description="Software stack layer configurations")

EpicLayerConfig

Bases: StackLayerConfig

Configuration of a layer of ePIC stack

Source code in src/aid2e/utilities/epic_utils/epic_stack_config.py
23
24
25
class EpicLayerConfig(StackLayerConfig):
    """Configuration of a layer of ePIC stack"""
    pass

EpicParameterGroup

Bases: BaseModel

Group of ePIC parameters that share the same XML file.

Source code in src/aid2e/utilities/epic_utils/epic_design_config.py
49
50
51
52
53
54
class EpicParameterGroup(BaseModel):
    """
    Group of ePIC parameters that share the same XML file.
    """
    file_path: str  # Path to XML file, can include $DETECTOR_PATH
    parameters: Dict[str, EpicParameter]

EpicRecLayer dataclass

Bases: StackLayer

Reconstruction layer of ePIC stack

Source code in src/aid2e/utilities/epic_utils/epic_stack.py
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
class EpicRecLayer(StackLayer):
    """Reconstruction layer of ePIC stack"""
    name = "rec"
    command = "eicrecon"
    rule = '{{command}} {{arguments}} {{outputs}} {{inputs}}'

    def _make_input_arg(self, inputs: List[str]) -> str:
        """
        Formats inputs for ePIC-specific reconstruction
        layer.
        """
        in_arg = ' '.join(inputs)
        return in_arg

    def _make_output_arg(self, outputs: List[str]) -> str:
        """
        Formats outputs for ePIC-specific reconstruction
        layer.
        """
        formatted_outputs = list()
        for out_file in outputs:
            formatted_outputs.append(f"-Ppodio:output_file={out_file}")
        return ' '.join(formatted_outputs)

EpicSimLayer dataclass

Bases: StackLayer

Simulation layer of ePIC stack

Source code in src/aid2e/utilities/epic_utils/epic_stack.py
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
class EpicSimLayer(StackLayer):
    """Simulation layer of ePIC stack"""
    name = "sim"
    command = "npsim"
    rule = '{{command}} --compactFile $DETECTOR_PATH/$DETECTOR_CONFIG.xml {{arguments}} {{inputs}} {{outputs}}'

    def _make_input_arg(self, inputs: List[str]) -> str:
        """
        Formats inputs for ePIC-specific simulation
        layer. Applies appropriate CLI option based
        on file extension of input.
        """
        has_gun = False
        has_macro = False
        formatted_inputs = list()
        for in_file in inputs:
            if in_file.endswith(".py"):
                formatted_inputs.append(f"--steeringFile {in_file}")
                has_gun = True
            if in_file.endswith(".hepmc3.root") or in_file.endswith(".hepmc"):
                formatted_inputs.append(f"-I {in_file}")
            if in_file.endswith(".mac"):
                formatted_inputs.append(f"--macroFile {in_file}")
                has_macro = True

        if has_gun:
            formatted_inputs.insert(0, "-G")
        if has_macro:
            formatted_inputs.insert(0, "--enableG4GPS")
        return ' '.join(formatted_inputs)

    def _make_output_arg(self, outputs: List[str]) -> str:
        """
        Formats outputs for ePIC-specific simulation
        layer.
        """
        out_arg = ' '.join(outputs)
        return f"--outputFile {out_arg}"

EpicStack dataclass

Bases: ExperimentStack

The ePIC software stack

Source code in src/aid2e/utilities/epic_utils/epic_stack.py
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
@dataclass
class EpicStack(ExperimentStack):
    """The ePIC software stack"""
    geo: EpicGeoLayer = field(default_factory = EpicGeoLayer)
    sim: EpicSimLayer = field(default_factory = EpicSimLayer)
    rec: EpicRecLayer = field(default_factory = EpicRecLayer)
    ana: EpicAnaLayer = field(default_factory = EpicAnaLayer)

    def prepare_workflow_geometry(
        self,
        workflow_dir: str,
        design_point: Dict[str, Any],
        problem_config: Any,
        workflow_id: str,
    ) -> str:
        """
        Prepare geometry once for the whole workflow/design point.
        Returns the prepared geometry directory.
        """
        if problem_config is None or problem_config.design_config is None:
            raise AttributeError("DesignConfig not present in workflow context.")
        if not isinstance(problem_config.design_config, EpicDesignConfig):
            raise TypeError("DesignConfig is not an instance of EpicDesignConfig.")

        env_config = problem_config.environment_config
        geometry_mode = getattr(env_config, "geometry_mode", "build") if env_config else "build"
        epic_install = getattr(env_config, "epic_install", None) if env_config else None

        if not epic_install and 'EPIC_INSTALL' not in os.environ:
            raise EnvironmentError("Variable 'EPIC_INSTALL' not set. Must define epic_install.")

        design = problem_config.design_config
        template_geo_dir = epic_install or os.environ['EPIC_INSTALL']

        if geometry_mode == "no_build":
            template_geo_dir = os.path.join(template_geo_dir, "share", "epic")
            trial_geo_dir = os.path.join(workflow_dir, "geometry", "epic")
            if os.path.exists(trial_geo_dir):
                shutil.rmtree(trial_geo_dir)
            shutil.copytree(template_geo_dir, trial_geo_dir)
            os.environ["DETECTOR_PATH"] = trial_geo_dir
            modify_xml_files(design.get_xml_modifications(design_point))
            return trial_geo_dir

        trial_geo_dir = os.path.join(workflow_dir, os.path.basename(template_geo_dir))

        if not os.path.exists(trial_geo_dir):
            shutil.copytree(template_geo_dir, trial_geo_dir)

        original_modifications = design.get_xml_modifications(design_point)

        remapped_modifications = {}
        for src_file, params in original_modifications.items():
            if src_file.startswith(template_geo_dir):
                dst_file = src_file.replace(template_geo_dir, trial_geo_dir, 1)
            else:
                dst_file = src_file
            remapped_modifications[dst_file] = params

        modify_xml_files(remapped_modifications)

        compile_commands = (
            f"cmake -B {trial_geo_dir}/build -S {trial_geo_dir} -DCMAKE_INSTALL_PREFIX={trial_geo_dir}/install\n"
            f"cmake --build {trial_geo_dir}/build\n"
            f"cmake --install {trial_geo_dir}/build\n"
        )
        compile_script = os.path.join(trial_geo_dir, "compile_geo.sh")
        with open(compile_script, "w") as script:
            script.writelines(compile_commands)
        os.chmod(compile_script, 0o777)

        compiled_log = os.path.join(trial_geo_dir, "compiled.log")
        do_compiling = self.make_driver_command(compile_script)
        if not os.path.exists(compiled_log):
            os.system(f"{do_compiling}")
            with open(compiled_log, "w") as f:
                f.write(f"Workflow {workflow_id} geometry compiled\n")

        return trial_geo_dir

    def prepare_for_execution(self, **kwargs) -> Optional[str]:
        context = None
        for _, value in kwargs.items():
            if isinstance(value, JobContext):
                context = value

        if context is None:
            raise RuntimeError("No JobContext provided to EpicStack.prepare_for_execution")

        if context.workflow_context is None:
            raise RuntimeError("No workflow context provided to EpicStack.prepare_for_execution")

        trial_geo_dir = context.workflow_context.parameters.get("prepared_geometry_dir")
        if not trial_geo_dir:
            raise RuntimeError("No prepared geometry directory found in workflow context")

        env_config = context.problem_config.environment_config if context.problem_config else None
        geometry_mode = getattr(env_config, "geometry_mode", "build")
        context.add_log(f"Using {geometry_mode} geometry from {trial_geo_dir}")
        return None

    def make_driver_script(
        self,
        script: str,
        configs: List[StackLayerConfig],
        preparations: str = None,
        **kwargs
    ) -> None:
        """
        Create a driver script to run ePIC layers.
        """
        context = None
        for arg, value in kwargs.items():
            if isinstance(value, JobContext):
                context = value

        # JobContext, WorkflowContext must be provided to access execution, geoemtry dir
        if context is None:
            raise RuntimeError("No JobContext provided to EpicStack.make_driver_script")
        if context.workflow_context is None:
            raise RuntimeError("No workflow context provided to EpicStack.make_driver_script")

        trial_geo_dir = context.workflow_context.parameters.get("prepared_geometry_dir")
        if not trial_geo_dir:
            raise RuntimeError("No prepared geometry directory found in workflow context")

        env_config = context.problem_config.environment_config if context.problem_config else None
        epic_install = getattr(env_config, "epic_install", None)
        epic_config = getattr(env_config, "epic_config", None) or os.environ.get("EPIC_CONFIG")
        geometry_mode = getattr(env_config, "geometry_mode", "build")
        if not epic_config:
            raise EnvironmentError("Variable 'epic_config' not set. Must define epic_config.")

        if geometry_mode == "no_build":
            if not epic_install:
                raise EnvironmentError("Variable 'epic_install' not set. Must define epic_install.")
            detector_setup = (
                f"source \"{epic_install}/bin/thisepic.sh\" {epic_config}\n"
                f"export EPIC_INSTALL=\"{epic_install}\"\n"
                f"export EPIC_CONFIG=\"{epic_config}\"\n"
                f"export DETECTOR_PATH=\"{trial_geo_dir}\"\n"
                f"export DETECTOR_CONFIG=\"{epic_config}\""
            )
        else:
            detector_setup = (
                f"source \"{trial_geo_dir}/install/bin/thisepic.sh\"\n"
                f"export EPIC_CONFIG=\"{epic_config}\"\n"
                f"export DETECTOR_CONFIG=\"{epic_config}\""
            )

        commands = [
            self._determine_shebang(script),
            "set -euo pipefail",
            detector_setup,
        ]
        if preparations != None:
            commands.append(preparations)
        commands.extend(self._make_commands(configs))

        text = "\n\n".join(commands)
        with open(script, 'w') as driver:
            driver.write(text)
        os.chmod(script, 0o777)

    def make_driver_command(self, script: str, **kwargs) -> str:
        """
        Form command to run ePIC driver script.
        """
        if 'EIC_SINGULARITY_IMAGE' in os.environ:
            return f"singularity exec {os.environ['EIC_SINGULARITY_IMAGE']} {script}"
        elif 'EIC_SHELL' in os.environ:
            return f"{os.environ['EIC_SHELL']} -- {script}"
        else:
            raise EnvironmentError("Neither 'EIC_SINGULARITY_IMAGE' nor 'EIC_SHELL' set. Must define eic_shell or singularity image.")

make_driver_command(script, **kwargs)

Form command to run ePIC driver script.

Source code in src/aid2e/utilities/epic_utils/epic_stack.py
302
303
304
305
306
307
308
309
310
311
def make_driver_command(self, script: str, **kwargs) -> str:
    """
    Form command to run ePIC driver script.
    """
    if 'EIC_SINGULARITY_IMAGE' in os.environ:
        return f"singularity exec {os.environ['EIC_SINGULARITY_IMAGE']} {script}"
    elif 'EIC_SHELL' in os.environ:
        return f"{os.environ['EIC_SHELL']} -- {script}"
    else:
        raise EnvironmentError("Neither 'EIC_SINGULARITY_IMAGE' nor 'EIC_SHELL' set. Must define eic_shell or singularity image.")

make_driver_script(script, configs, preparations=None, **kwargs)

Create a driver script to run ePIC layers.

Source code in src/aid2e/utilities/epic_utils/epic_stack.py
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
def make_driver_script(
    self,
    script: str,
    configs: List[StackLayerConfig],
    preparations: str = None,
    **kwargs
) -> None:
    """
    Create a driver script to run ePIC layers.
    """
    context = None
    for arg, value in kwargs.items():
        if isinstance(value, JobContext):
            context = value

    # JobContext, WorkflowContext must be provided to access execution, geoemtry dir
    if context is None:
        raise RuntimeError("No JobContext provided to EpicStack.make_driver_script")
    if context.workflow_context is None:
        raise RuntimeError("No workflow context provided to EpicStack.make_driver_script")

    trial_geo_dir = context.workflow_context.parameters.get("prepared_geometry_dir")
    if not trial_geo_dir:
        raise RuntimeError("No prepared geometry directory found in workflow context")

    env_config = context.problem_config.environment_config if context.problem_config else None
    epic_install = getattr(env_config, "epic_install", None)
    epic_config = getattr(env_config, "epic_config", None) or os.environ.get("EPIC_CONFIG")
    geometry_mode = getattr(env_config, "geometry_mode", "build")
    if not epic_config:
        raise EnvironmentError("Variable 'epic_config' not set. Must define epic_config.")

    if geometry_mode == "no_build":
        if not epic_install:
            raise EnvironmentError("Variable 'epic_install' not set. Must define epic_install.")
        detector_setup = (
            f"source \"{epic_install}/bin/thisepic.sh\" {epic_config}\n"
            f"export EPIC_INSTALL=\"{epic_install}\"\n"
            f"export EPIC_CONFIG=\"{epic_config}\"\n"
            f"export DETECTOR_PATH=\"{trial_geo_dir}\"\n"
            f"export DETECTOR_CONFIG=\"{epic_config}\""
        )
    else:
        detector_setup = (
            f"source \"{trial_geo_dir}/install/bin/thisepic.sh\"\n"
            f"export EPIC_CONFIG=\"{epic_config}\"\n"
            f"export DETECTOR_CONFIG=\"{epic_config}\""
        )

    commands = [
        self._determine_shebang(script),
        "set -euo pipefail",
        detector_setup,
    ]
    if preparations != None:
        commands.append(preparations)
    commands.extend(self._make_commands(configs))

    text = "\n\n".join(commands)
    with open(script, 'w') as driver:
        driver.write(text)
    os.chmod(script, 0o777)

prepare_workflow_geometry(workflow_dir, design_point, problem_config, workflow_id)

Prepare geometry once for the whole workflow/design point. Returns the prepared geometry directory.

Source code in src/aid2e/utilities/epic_utils/epic_stack.py
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
def prepare_workflow_geometry(
    self,
    workflow_dir: str,
    design_point: Dict[str, Any],
    problem_config: Any,
    workflow_id: str,
) -> str:
    """
    Prepare geometry once for the whole workflow/design point.
    Returns the prepared geometry directory.
    """
    if problem_config is None or problem_config.design_config is None:
        raise AttributeError("DesignConfig not present in workflow context.")
    if not isinstance(problem_config.design_config, EpicDesignConfig):
        raise TypeError("DesignConfig is not an instance of EpicDesignConfig.")

    env_config = problem_config.environment_config
    geometry_mode = getattr(env_config, "geometry_mode", "build") if env_config else "build"
    epic_install = getattr(env_config, "epic_install", None) if env_config else None

    if not epic_install and 'EPIC_INSTALL' not in os.environ:
        raise EnvironmentError("Variable 'EPIC_INSTALL' not set. Must define epic_install.")

    design = problem_config.design_config
    template_geo_dir = epic_install or os.environ['EPIC_INSTALL']

    if geometry_mode == "no_build":
        template_geo_dir = os.path.join(template_geo_dir, "share", "epic")
        trial_geo_dir = os.path.join(workflow_dir, "geometry", "epic")
        if os.path.exists(trial_geo_dir):
            shutil.rmtree(trial_geo_dir)
        shutil.copytree(template_geo_dir, trial_geo_dir)
        os.environ["DETECTOR_PATH"] = trial_geo_dir
        modify_xml_files(design.get_xml_modifications(design_point))
        return trial_geo_dir

    trial_geo_dir = os.path.join(workflow_dir, os.path.basename(template_geo_dir))

    if not os.path.exists(trial_geo_dir):
        shutil.copytree(template_geo_dir, trial_geo_dir)

    original_modifications = design.get_xml_modifications(design_point)

    remapped_modifications = {}
    for src_file, params in original_modifications.items():
        if src_file.startswith(template_geo_dir):
            dst_file = src_file.replace(template_geo_dir, trial_geo_dir, 1)
        else:
            dst_file = src_file
        remapped_modifications[dst_file] = params

    modify_xml_files(remapped_modifications)

    compile_commands = (
        f"cmake -B {trial_geo_dir}/build -S {trial_geo_dir} -DCMAKE_INSTALL_PREFIX={trial_geo_dir}/install\n"
        f"cmake --build {trial_geo_dir}/build\n"
        f"cmake --install {trial_geo_dir}/build\n"
    )
    compile_script = os.path.join(trial_geo_dir, "compile_geo.sh")
    with open(compile_script, "w") as script:
        script.writelines(compile_commands)
    os.chmod(compile_script, 0o777)

    compiled_log = os.path.join(trial_geo_dir, "compiled.log")
    do_compiling = self.make_driver_command(compile_script)
    if not os.path.exists(compiled_log):
        os.system(f"{do_compiling}")
        with open(compiled_log, "w") as f:
            f.write(f"Workflow {workflow_id} geometry compiled\n")

    return trial_geo_dir

EpicStageDefinition

Bases: StackStageDefinition

Definition of an ePIC stage in a workflow. Jobs are restricted to ePIC jobs.

Source code in src/aid2e/utilities/epic_utils/epic_stack_config.py
37
38
39
40
41
42
class EpicStageDefinition(StackStageDefinition):
    """
    Definition of an ePIC stage in a workflow.
    Jobs are restricted to ePIC jobs.
    """
    jobs: List[EpicJobDefinition] = Field(default_factory = list, description="ePIC stack job definitions")

EpicWorkflowDefinition

Bases: StackWorkflowDefinition

Definition of an ePIC workflow.

Source code in src/aid2e/utilities/epic_utils/epic_stack_config.py
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
class EpicWorkflowDefinition(StackWorkflowDefinition):
    """
    Definition of an ePIC workflow.
    """
    stack_type: Optional[str] = Field(default="epic", description="Experimental stack type for workflow-level geometry prep")
    branches: List[EpicBranchDefinition] = Field(default_factory=list, description="Software stack workflow branches (optional)")

    def get_implicit_branch(self) -> StackBranchDefinition:
        """
        Get or create single implicit branch if branches list is empty.
        Overrides StackWorkflowDefinition.get_implicit_branch to return
        EpicBranchDefinition.
        """
        if self.branches:
            raise ValueError("Branches already defined; cannot use implicit branch")
        return EpicBranchDefinition(name="implicit")

get_implicit_branch()

Get or create single implicit branch if branches list is empty. Overrides StackWorkflowDefinition.get_implicit_branch to return EpicBranchDefinition.

Source code in src/aid2e/utilities/epic_utils/epic_stack_config.py
60
61
62
63
64
65
66
67
68
def get_implicit_branch(self) -> StackBranchDefinition:
    """
    Get or create single implicit branch if branches list is empty.
    Overrides StackWorkflowDefinition.get_implicit_branch to return
    EpicBranchDefinition.
    """
    if self.branches:
        raise ValueError("Branches already defined; cannot use implicit branch")
    return EpicBranchDefinition(name="implicit")

EpicWorkflowsConfiguration

Bases: StackWorkflowsConfiguration

Container for ePIC workflows.

Source code in src/aid2e/utilities/epic_utils/epic_stack_config.py
71
72
73
74
75
class EpicWorkflowsConfiguration(StackWorkflowsConfiguration):
    """
    Container for ePIC workflows.
    """
    workflows: List[EpicWorkflowDefinition] = Field(..., min_items=1, description="List of workflows")