File indexing completed on 2026-08-12 08:24:56
0001 """Unified objective definition models.
0002
0003 This module defines the single source of truth for objectives across AID2E:
0004 - How objectives are specified in problems (name + direction)
0005 - How they're executed in workflows (script, inline, or steps/DAG)
0006 - How they're optimized by algorithms (directives like "minimize:f1")
0007
0008 Key concepts:
0009 ObjectiveDirection: MINIMIZE or MAXIMIZE
0010 ObjectivePlanSpec: How to compute (script path, inline function, or multi-step plan)
0011 ObjectiveDefinition: Complete spec (name + direction + objective plan)
0012 ObjectivesRegistry: Runtime mapping of objective names to definitions
0013
0014 Project: AID2E v0.0.0 - AI assisted Detector Design for EIC
0015 Homepage: https://aid2e.github.io/aid2e-framework
0016 Repository: https://github.com/aid2e/AID2E-framework.git
0017 """
0018
0019 from enum import Enum
0020 from typing import Optional, List, Union, Dict, Any
0021 from pydantic import BaseModel, Field, field_validator, model_validator, ConfigDict, ValidationInfo
0022 from aid2e.utilities.configurations.scheduler_config import SchedulerConfiguration
0023
0024
0025 class ObjectiveDirection(str, Enum):
0026 """Direction of optimization for an objective.
0027
0028 Attributes:
0029 MINIMIZE: Minimize the objective value.
0030 MAXIMIZE: Maximize the objective value.
0031 """
0032 MINIMIZE = "minimize"
0033 MAXIMIZE = "maximize"
0034
0035
0036 class ScriptObjective(BaseModel):
0037 """Objective computed via external script.
0038
0039 Attributes:
0040 path: Path to executable script (resolved relative to config directory).
0041 output_file: Expected output file pattern (e.g., "objectives_*.json").
0042 The script should create a file matching this pattern containing
0043 the objective value in JSON/YAML format.
0044 timeout_sec: Timeout in seconds (optional, default: 300).
0045
0046 Example:
0047 >>> script = ScriptObjective(
0048 ... path="scripts/dtlz2_problem.py",
0049 ... output_file="objectives_{job_id}.json"
0050 ... )
0051 """
0052 path: str = Field(..., description="Path to objective computation script")
0053 output_file: str = Field(..., description="Output file pattern (e.g., objectives_*.json)")
0054 timeout_sec: int = Field(default=300, ge=1, description="Computation timeout in seconds")
0055
0056
0057 class InlineObjective(BaseModel):
0058 """Objective computed via inline Python function.
0059
0060 The entrypoint should reference a callable that accepts design parameters
0061 and returns the objective value.
0062
0063 Attributes:
0064 entrypoint: Module and function reference (format: "module.path:function_name").
0065
0066 Example:
0067 >>> inline = InlineObjective(entrypoint="my_objectives:compute_f1")
0068 >>> # Expects function: def compute_f1(design_params: Dict[str, float]) -> float
0069
0070 Notes:
0071 - The function is imported at runtime (lazy loading).
0072 - Must accept design_params: Dict[str, float] as argument.
0073 - Must return a single float value.
0074 """
0075 entrypoint: str = Field(
0076 ...,
0077 description="Module:function reference (e.g., 'my_objectives:compute_f1')"
0078 )
0079
0080 @field_validator('entrypoint')
0081 @classmethod
0082 def validate_entrypoint_format(cls, v: str) -> str:
0083 """Validate entrypoint has 'module:function' format."""
0084 if ':' not in v or v.count(':') != 1:
0085 raise ValueError("entrypoint must be 'module:function' format")
0086 module_part, func_part = v.split(':')
0087 if not module_part or not func_part:
0088 raise ValueError("entrypoint module and function names cannot be empty")
0089 if not all(c.isalnum() or c in '_.:-' for c in module_part):
0090 raise ValueError(f"Invalid module name: {module_part}")
0091 if not (func_part[0].isalpha() or func_part[0] == '_'):
0092 raise ValueError(f"Invalid function name: {func_part}")
0093 return v
0094
0095
0096 class StepStage(BaseModel):
0097 """Single stage within an objective step plan.
0098
0099 Each stage executes either a script or an inline function, can declare
0100 inputs/outputs/extra_args, and may depend on upstream stages. If a plan has
0101 only one step, it is represented as a single-element step list.
0102
0103 Attributes:
0104 name: Unique stage identifier.
0105 description: Optional human-readable description of the stage intent.
0106 script: Script-based execution for this stage (mutually exclusive with inline).
0107 inline: Inline Python callable for this stage (mutually exclusive with script).
0108 inputs: Optional input bindings for this stage (free-form mapping).
0109 outputs: Optional output bindings for this stage (free-form mapping).
0110 extra_args: Additional args/metadata for the stage executor.
0111 produces_objective: Whether this stage emits the objective value.
0112 depends_on: Names of upstream stages this stage depends on.
0113 """
0114
0115 model_config = ConfigDict(populate_by_name=True)
0116
0117 name: str = Field(..., description="Stage name (unique within steps)")
0118 description: Optional[str] = Field(default=None, description="Stage description")
0119 script: Optional[ScriptObjective] = Field(default=None, description="Script execution for this stage")
0120 inline: Optional[InlineObjective] = Field(default=None, description="Inline callable for this stage")
0121 inputs: Dict[str, Any] = Field(default_factory=dict, description="Input bindings for this stage")
0122 outputs: Dict[str, Any] = Field(default_factory=dict, description="Output bindings for this stage")
0123 extra_args: Dict[str, Any] = Field(default_factory=dict, description="Extra args/metadata for the stage executor")
0124 produces_objective: bool = Field(default=False, description="Whether this stage emits the objective value")
0125 depends_on: List[str] = Field(default_factory=list, description="Upstream stage dependencies")
0126
0127 @model_validator(mode="after")
0128 def validate_action(self) -> "StepStage":
0129 """Ensure stage has a valid execution definition.
0130
0131 A stage must choose exactly one execution method: script or inline.
0132 """
0133 has_script = self.script is not None
0134 has_inline = self.inline is not None
0135
0136 if has_script == has_inline:
0137 raise ValueError("Stage must define exactly one of: script or inline")
0138
0139 return self
0140
0141 @field_validator('depends_on')
0142 @classmethod
0143 def validate_dependencies(cls, depends_on: List[str], info: ValidationInfo) -> List[str]:
0144 """Ensure stages do not depend on themselves."""
0145 name = None
0146 if info and info.data:
0147 name = info.data.get("name")
0148 if name and name in depends_on:
0149 raise ValueError(f"Stage '{name}' cannot depend on itself")
0150 return depends_on
0151
0152
0153 class StepPlanSpec(BaseModel):
0154 """DAG-style step plan for an objective.
0155
0156 Replaces the earlier "branch" terminology with a clearer "steps"
0157 concept. A step plan is a small DAG of stages where exactly
0158 one stage must produce the objective value.
0159
0160 Attributes:
0161 stages: Ordered list of stage definitions. Dependencies define the DAG.
0162 produces_from_stage: Optional explicit producing stage name. If omitted,
0163 exactly one stage must set ``produces_objective=True``.
0164 """
0165
0166 model_config = ConfigDict(populate_by_name=True)
0167
0168 stages: List[StepStage] = Field(..., min_items=1, description="Stages composing the computation DAG")
0169 produces_from_stage: Optional[str] = Field(
0170 default=None,
0171 description="Explicit stage name that emits the objective (overrides flag)",
0172 alias="produces_from_stage",
0173 )
0174
0175 @model_validator(mode="after")
0176 def validate_stages(self) -> "StepPlanSpec":
0177 """Ensure unique names, valid dependencies, and single producer."""
0178 names = [stage.name for stage in self.stages]
0179 if len(set(names)) != len(names):
0180 raise ValueError("Stage names within steps must be unique")
0181
0182 for stage in self.stages:
0183 for dep in stage.depends_on:
0184 if dep not in names:
0185 raise ValueError(f"Stage '{stage.name}' depends on unknown stage '{dep}'")
0186
0187 explicit = self.produces_from_stage
0188 producing_flags = [s.name for s in self.stages if s.produces_objective]
0189
0190 if explicit:
0191 if explicit not in names:
0192 raise ValueError(f"produces_from_stage '{explicit}' not found in stages")
0193 chosen = explicit
0194 else:
0195 if len(producing_flags) != 1:
0196 raise ValueError("Exactly one stage must set produces_objective=True when produces_from_stage is not provided")
0197 chosen = producing_flags[0]
0198
0199 self.produces_from_stage = chosen
0200 return self
0201
0202 def producing_stage(self) -> str:
0203 """Return the name of the stage that emits the objective value."""
0204 if not self.produces_from_stage:
0205 raise ValueError("produces_from_stage was not resolved")
0206 return self.produces_from_stage
0207
0208
0209 class ObjectivePlanSpec(BaseModel):
0210 """Plan for executing an objective (always modeled as steps).
0211
0212 The canonical form is a step plan with one or more stages. As a
0213 convenience, users may supply a single script or inline definition; it will
0214 be wrapped into a single-step plan automatically.
0215 """
0216
0217 model_config = ConfigDict(populate_by_name=True)
0218
0219 steps: StepPlanSpec = Field(
0220 ...,
0221 description="DAG-style multi-stage plan for the objective",
0222 )
0223
0224 @model_validator(mode="before")
0225 def reject_legacy_shapes(cls, values: Any) -> Any:
0226 """Reject retired objective plan schema variants."""
0227 if not isinstance(values, dict):
0228 return values
0229 if "multi-steps" in values or "multi_steps" in values:
0230 raise ValueError(
0231 "Legacy objective plan step keys are no longer supported. Use 'steps'."
0232 )
0233 if "script" in values or "inline" in values:
0234 raise ValueError(
0235 "Single-step objective plans are no longer supported. Wrap the "
0236 "step under 'steps.stages'."
0237 )
0238 return values
0239
0240 def is_steps(self) -> bool:
0241 """Return True if this plan is a step DAG (always true for canonical form)."""
0242 return self.steps is not None
0243
0244 class ObjectiveDefinition(BaseModel):
0245 """Complete objective specification: name, direction, and objective plan.
0246
0247 This is the unified model used across problem, optimization, and workflow layers.
0248 It combines what to optimize (name + direction) with how to execute it
0249 (script, inline function, or multi-step DAG).
0250
0251 Attributes:
0252 name: Unique objective identifier (e.g., "f1", "efficiency").
0253 direction: Optimization direction (minimize or maximize).
0254 objective_plan: How to execute (script, inline, or steps).
0255 scheduler: Optional objective-level scheduler default (cascades to stages).
0256 metrics_keys: Optional keys to extract from plan output when it returns a dict.
0257 Useful when one plan produces multiple metrics.
0258 Example: plan outputs {"f1": 0.5, "f2": 0.3, "runtime": 10.2},
0259 metrics_keys=["f1"] extracts only f1.
0260 """
0261 name: str = Field(..., description="Objective name (e.g., 'f1', 'efficiency')")
0262 direction: ObjectiveDirection = Field(
0263 ...,
0264 description="Optimization direction: minimize or maximize"
0265 )
0266 objective_plan: Optional[ObjectivePlanSpec] = Field(
0267 default=None,
0268 description="How to execute the objective (script, inline, or steps)",
0269 )
0270 scheduler: Optional[SchedulerConfiguration] = Field(
0271 default=None,
0272 description="Default scheduler for this objective; cascades to its stages",
0273 )
0274 metrics_keys: List[str] = Field(
0275 default_factory=list,
0276 description="Keys to extract from plan output (if dict)",
0277 )
0278
0279 def to_directive(self) -> str:
0280 """Convert to optimization directive string format.
0281
0282 Returns:
0283 String like "minimize:f1" or "maximize:efficiency".
0284 Useful for OptimizationConfiguration.objectives.
0285
0286 Example:
0287 >>> obj = ObjectiveDefinition(name="f1", direction=ObjectiveDirection.MINIMIZE, objective_plan=None)
0288 >>> obj.to_directive()
0289 'minimize:f1'
0290 """
0291 return f"{self.direction.value}:{self.name}"
0292
0293 @classmethod
0294 def from_directive(
0295 cls,
0296 directive: str,
0297 objective_plan: Optional[ObjectivePlanSpec] = None,
0298 ) -> "ObjectiveDefinition":
0299 """Create ObjectiveDefinition from directive string.
0300
0301 Parses strings like "minimize:f1" or "maximize:efficiency".
0302
0303 Args:
0304 directive: String in format "minimize:name" or "maximize:name".
0305 objective_plan: Optional ObjectivePlanSpec (script, inline, or steps).
0306
0307 Returns:
0308 ObjectiveDefinition with parsed direction and name.
0309
0310 Raises:
0311 ValueError: If directive format is invalid.
0312
0313 Example:
0314 >>> directive = "minimize:f1"
0315 >>> obj = ObjectiveDefinition.from_directive(directive)
0316 """
0317 if ':' not in directive or directive.count(':') != 1:
0318 raise ValueError(f"Invalid directive format: {directive}. Expected 'minimize:name' or 'maximize:name'")
0319
0320 direction_str, name = directive.split(':')
0321 try:
0322 direction = ObjectiveDirection(direction_str.lower())
0323 except ValueError:
0324 raise ValueError(f"Invalid direction '{direction_str}'. Must be 'minimize' or 'maximize'")
0325
0326 if not name.strip():
0327 raise ValueError("Objective name cannot be empty")
0328
0329 return cls(
0330 name=name.strip(),
0331 direction=direction,
0332 objective_plan=objective_plan,
0333 )
0334
0335
0336 class ObjectivesRegistry:
0337 """Runtime registry for objective definitions.
0338
0339 Allows objectives to be registered and retrieved by name for use during
0340 workflow execution. This enables decoupling objective definitions from
0341 their runtime computation.
0342
0343 Example:
0344 >>> registry = ObjectivesRegistry()
0345 >>> obj_f1 = ObjectiveDefinition(name="f1", direction=ObjectiveDirection.MINIMIZE)
0346 >>> registry.register(obj_f1)
0347 >>> retrieved = registry.get("f1")
0348 """
0349
0350 def __init__(self):
0351 """Initialize empty registry."""
0352 self._objectives: Dict[str, ObjectiveDefinition] = {}
0353
0354 def register(self, objective: ObjectiveDefinition) -> None:
0355 """Register an objective by name.
0356
0357 Args:
0358 objective: ObjectiveDefinition to register.
0359
0360 Raises:
0361 ValueError: If objective with same name already registered.
0362 """
0363 if objective.name in self._objectives:
0364 raise ValueError(f"Objective '{objective.name}' already registered")
0365 self._objectives[objective.name] = objective
0366
0367 def get(self, name: str) -> Optional[ObjectiveDefinition]:
0368 """Retrieve objective definition by name.
0369
0370 Args:
0371 name: Objective name.
0372
0373 Returns:
0374 ObjectiveDefinition if found, None otherwise.
0375 """
0376 return self._objectives.get(name)
0377
0378 def list_all(self) -> List[ObjectiveDefinition]:
0379 """Get all registered objectives.
0380
0381 Returns:
0382 List of all registered ObjectiveDefinition instances.
0383 """
0384 return list(self._objectives.values())
0385
0386 def clear(self) -> None:
0387 """Clear all registered objectives."""
0388 self._objectives.clear()
