File indexing completed on 2026-08-12 08:24:56
0001 """Experimental software stack configuration models
0002
0003 Defines models necessary for wiring experimental software stacks into workflows.
0004 Allows users to specify layers of a pre-defined stack as part of a workflow.
0005
0006 Key Classes:
0007 - StackLayerConfig, configures a stack layer
0008
0009 Project: AID2E v0.0.0 - AI assisted Detector Design for EIC
0010 Homepage: https://aid2e.github.io/aid2e-framework
0011 Repository: https://github.com/aid2e/AID2E-framework.git
0012 """
0013
0014 from pydantic import AliasChoices, BaseModel, Field, model_validator
0015 from typing import List, Optional
0016
0017 from aid2e.utilities.configurations.workflow_config import (
0018 BranchDefinition,
0019 JobDefinition,
0020 StageDefinition,
0021 WorkflowDefinition,
0022 WorkflowsConfiguration,
0023 )
0024
0025
0026 class StackLayerConfig(BaseModel):
0027 """Configures layer of an experimental stack
0028
0029 A job may consist of 1 or many layers from an experimental
0030 software stack. Sanitizes provided data to make sure
0031 singular vs. plural inputs/outputs are handled consistently.
0032
0033 Attributes:
0034 name: Unique name for this layer instance
0035 layer: Layer key (e.g. "sim", "rec", "ana")
0036 inputs: List of inputs to layer
0037 outputs: List of outputs from layer
0038 arguments: Optional list of any additional arguments to apply
0039 command: Optional command to be run. Can be used to
0040 override default of layer.
0041 rule: Optional recipe for combining inputs, outputs, arguments,
0042 and command. Can be used to override default of layer.
0043
0044 Notes:
0045 - rule supports template substitutions for {inputs}, {outputs},
0046 {arguments}, and {command}. See StackLayer for more details.
0047 """
0048 name: Optional[str] = Field(default=None, description="Unique name of instance")
0049 layer: str = Field(..., description="Layer key")
0050 inputs: List[str] = Field(..., description="List of inputs", validation_alias=AliasChoices('inputs', 'input'))
0051 outputs: List[str] = Field(..., description="List of outputs", validation_alias=AliasChoices('outputs', 'output'))
0052 arguments: Optional[List[str]] = Field(default=None, description="List of arguments")
0053 command: Optional[str] = Field(default=None, description="Executable command")
0054 rule: Optional[str] = Field(default=None, description="Recipe for combining arguments")
0055
0056 @classmethod
0057 def pluralize_strings(cls, data, singular, plural):
0058 """Pluralize strings in data
0059
0060 Sanitize input by data by ensuring that
0061 'singular' keys are always 'plural' and
0062 that their values are wrapped in lists.
0063
0064 Args:
0065 data: the to be sanitized
0066 singular: the singular case of the key
0067 (e.g. 'input')
0068 plural: the plural case of the key
0069 (e.g. 'inputs')
0070
0071 Returns:
0072 Sanitized data
0073 """
0074 if singular in data and plural not in data:
0075 data[singular] = data[plural]
0076 if isinstance(data.get(plural), str):
0077 data[plural] = [data[plural]]
0078 return data
0079
0080 @model_validator(mode='before')
0081 @classmethod
0082 def handle_input_variants(cls, data):
0083 """
0084 Handles cases where (1) 'input' vs. 'inputs' was used in key,
0085 and (2) only 1 string was provided.
0086 """
0087 return cls.pluralize_strings(data, "input", "inputs")
0088
0089 @model_validator(mode='before')
0090 @classmethod
0091 def handle_output_variants(cls, data):
0092 """
0093 Handles cases where (1) 'output' vs. 'outputs' was used in key,
0094 and (2) only 1 string was provided.
0095 """
0096 return cls.pluralize_strings(data, "output", "outputs")
0097
0098
0099 class StackJobDefinition(JobDefinition):
0100 """
0101 Extends the base JobDefinition with a list of stack
0102 layers to utilize built-in commands of an experimental
0103 stack.
0104
0105 Extensions:
0106 script: Name of driver script to generate
0107 layers: Layer configurations to run in this job
0108 """
0109 command: Optional[str] = Field(default="./{script}", description="Executable command")
0110 script: Optional[str] = Field(default="do_job_{{context.job_id}}.sh", description="Driver script name")
0111 layers: List[StackLayerConfig] = Field(default_factory=list, description="Software stack layer configurations")
0112
0113
0114 class StackStageDefinition(StageDefinition):
0115
0116 """
0117 Definition of a workflow stage narrowed to jobs from
0118 an experimental software stack.
0119 """
0120 jobs: List[StackJobDefinition] = Field(default_factory=list, description="Software stack job definitions")
0121
0122
0123 class StackBranchDefinition(BranchDefinition):
0124 """
0125 Definition of a workflow branch narrowed to stages
0126 from an experimental software stack.
0127 """
0128 stages: List[StackStageDefinition] = Field(default_factory=list, description="Software stack stage definitions")
0129
0130
0131 class StackWorkflowDefinition(WorkflowDefinition):
0132 """
0133 Definition of an experimental software stack
0134 workflow.
0135 """
0136 branches: List[StackBranchDefinition] = Field(default_factory=list, description="Software stack workflow branches (optional)")
0137
0138
0139
0140 def get_implicit_branch(self) -> StackBranchDefinition:
0141 """
0142 Get or create single implicit branch if branches list is empty.
0143 Overrides WorkflowDefinition.get_implicit_branch to return
0144 StackBranchDefinition.
0145 """
0146 if self.branches:
0147 raise ValueError("Branches already defined; cannot use implicit branch")
0148 return StackBranchDefinition(name="implicit")
0149
0150
0151 class StackWorkflowsConfiguration(WorkflowsConfiguration):
0152 """
0153 Container for software stack workflows.
0154 """
0155
0156
0157 workflows: List[StackWorkflowDefinition] = Field(..., min_items=1, description="List of workflows")
