File indexing completed on 2026-08-12 08:24:55
0001 """Complete DAG Executor usage examples.
0002
0003 Demonstrates how to use the DAG Executor to orchestrate workflows with:
0004 1. Simple sequential workflow
0005 2. Parallel job execution with job_factory
0006 3. Multi-stage workflow with dependencies
0007 4. Container-based evaluators
0008 5. Complete optimizer integration example
0009
0010 Project: AID2E v0.0.0 - AI assisted Detector Design for EIC
0011 """
0012
0013 from aid2e.utilities.workflows import (
0014 DAGExecutor,
0015 WorkflowDefinition,
0016 BranchDefinition,
0017 StageDefinition,
0018 JobDefinition,
0019 JobFactory,
0020 ParallelismPolicy,
0021 ArtifactSpec,
0022 )
0023 from aid2e.utilities.configurations.objectives import (
0024 ObjectiveDefinition,
0025 ObjectiveDirection,
0026 )
0027
0028
0029 def example_1_simple_workflow():
0030 """Example 1: Simple workflow with one stage and one job."""
0031 print("\n" + "="*80)
0032 print("Example 1: Simple Sequential Workflow")
0033 print("="*80)
0034
0035
0036 eval_job = JobDefinition(
0037 name="dtlz2_eval",
0038 command="python -c \"import json; obj={'f1': 0.234, 'f2': 0.876}; print(json.dumps(obj))\"",
0039 payload={"evaluator_type": "bash"},
0040 )
0041
0042
0043 eval_stage = StageDefinition(
0044 name="evaluate",
0045 jobs=[eval_job],
0046 )
0047
0048
0049 main_branch = BranchDefinition(
0050 name="main",
0051 stages=[eval_stage],
0052 )
0053
0054
0055 workflow = WorkflowDefinition(
0056 name="dtlz2_simple",
0057 description="Simple DTLZ2 evaluation workflow",
0058 branches=[main_branch],
0059 objectives=[
0060 ObjectiveDefinition(
0061 name="f1",
0062 direction=ObjectiveDirection.MINIMIZE,
0063 ),
0064 ObjectiveDefinition(
0065 name="f2",
0066 direction=ObjectiveDirection.MINIMIZE,
0067 ),
0068 ],
0069 )
0070
0071
0072 executor = DAGExecutor(
0073 workflow=workflow,
0074 base_output_dir="/tmp/aid2e_examples",
0075 log_level="INFO",
0076 )
0077
0078
0079 design_point = {"x1": 0.5, "x2": 0.7, "x3": 0.3}
0080 print(f"\nExecuting workflow for design point: {design_point}")
0081
0082 objectives = executor.execute(design_point)
0083
0084 print(f"\n✅ Workflow completed!")
0085 print(f"Objectives: {objectives}")
0086 print(f"Output directory: {executor.output_dir}")
0087
0088
0089 def example_2_parallel_jobs():
0090 """Example 2: Workflow with parallel job execution."""
0091 print("\n" + "="*80)
0092 print("Example 2: Parallel Job Execution with JobFactory")
0093 print("="*80)
0094
0095
0096 eval_job = JobDefinition(
0097 name="parallel_eval",
0098 command="echo 'Evaluating design point with job_index={{job_index}}'",
0099 payload={"evaluator_type": "bash"},
0100 )
0101
0102
0103 eval_stage = StageDefinition(
0104 name="parallel_evaluate",
0105 jobs=[eval_job],
0106 job_factory=JobFactory(
0107 type="range",
0108 params={"n": 4},
0109 ),
0110 parallelism=ParallelismPolicy(
0111 max_concurrent=4,
0112 retry_max=2,
0113 timeout_sec=300,
0114 ),
0115 )
0116
0117 branch = BranchDefinition(name="main", stages=[eval_stage])
0118
0119 workflow = WorkflowDefinition(
0120 name="parallel_workflow",
0121 description="Workflow with 4 parallel evaluations",
0122 branches=[branch],
0123 objectives=[],
0124 )
0125
0126 executor = DAGExecutor(workflow, base_output_dir="/tmp/aid2e_examples")
0127 design_point = {"x1": 0.5, "x2": 0.7}
0128
0129 print(f"\nExecuting workflow with 4 parallel jobs...")
0130 objectives = executor.execute(design_point)
0131
0132 print(f"\n✅ Workflow completed with 4 parallel jobs!")
0133 print(f"Total jobs executed: {len(executor.global_xcom)}")
0134
0135
0136 def example_3_multi_stage_workflow():
0137 """Example 3: Multi-stage workflow with sequential dependencies."""
0138 print("\n" + "="*80)
0139 print("Example 3: Multi-Stage Workflow with Dependencies")
0140 print("="*80)
0141
0142
0143 generate_job = JobDefinition(
0144 name="generate_data",
0145 command="echo 'Generating simulation input data'",
0146 payload={"evaluator_type": "bash"},
0147 )
0148
0149 generate_stage = StageDefinition(
0150 name="generate",
0151 jobs=[generate_job],
0152 )
0153
0154
0155 simulate_job = JobDefinition(
0156 name="run_simulation",
0157 command="echo 'Running physics simulation'",
0158 payload={"evaluator_type": "bash"},
0159 )
0160
0161 simulate_stage = StageDefinition(
0162 name="simulate",
0163 jobs=[simulate_job],
0164 )
0165
0166
0167 compute_job = JobDefinition(
0168 name="compute_objectives",
0169 command="python -c \"import json; print(json.dumps({'f1': 0.5, 'f2': 0.8}))\"",
0170 payload={"evaluator_type": "bash"},
0171 )
0172
0173 compute_stage = StageDefinition(
0174 name="compute",
0175 jobs=[compute_job],
0176 )
0177
0178
0179 branch = BranchDefinition(
0180 name="main",
0181 stages=[generate_stage, simulate_stage, compute_stage],
0182 )
0183
0184 workflow = WorkflowDefinition(
0185 name="multi_stage_workflow",
0186 description="Sequential pipeline: generate → simulate → compute",
0187 branches=[branch],
0188 objectives=[],
0189 )
0190
0191 executor = DAGExecutor(workflow, base_output_dir="/tmp/aid2e_examples")
0192 design_point = {"energy": 100, "angle": 45}
0193
0194 print(f"\nExecuting 3-stage sequential workflow...")
0195 objectives = executor.execute(design_point)
0196
0197 print(f"\n✅ Multi-stage workflow completed!")
0198 print(f"Stages executed: generate → simulate → compute")
0199
0200
0201 def example_4_container_evaluators():
0202 """Example 4: Using ContainerEvaluator for Docker-based jobs."""
0203 print("\n" + "="*80)
0204 print("Example 4: Container-Based Workflow (ContainerEvaluator)")
0205 print("="*80)
0206
0207
0208 container_job = JobDefinition(
0209 name="docker_simulation",
0210 command="python /app/simulate.py",
0211 payload={
0212 "evaluator_type": "container",
0213 "image": "python:3.9-slim",
0214 "container_command": [
0215 "/bin/bash", "-c",
0216 "python -c \"import json; print(json.dumps({'f1': 0.3, 'f2': 0.9}))\""
0217 ],
0218 "environment": {
0219 "SIMULATION_MODE": "fast",
0220 "NUM_EVENTS": "1000",
0221 },
0222 "volumes": {
0223 "/tmp/input": "/app/input",
0224 "/tmp/output": "/app/output",
0225 },
0226 },
0227 resources={
0228 "cpu": "2",
0229 "memory": "4GB",
0230 },
0231 )
0232
0233 stage = StageDefinition(
0234 name="containerized_eval",
0235 jobs=[container_job],
0236 )
0237
0238 branch = BranchDefinition(name="main", stages=[stage])
0239
0240 workflow = WorkflowDefinition(
0241 name="container_workflow",
0242 description="Workflow using Docker containers",
0243 branches=[branch],
0244 objectives=[],
0245 )
0246
0247 executor = DAGExecutor(workflow, base_output_dir="/tmp/aid2e_examples")
0248 design_point = {"detector_thickness": 5.0}
0249
0250 print(f"\nExecuting containerized workflow...")
0251 print(f"Container image: python:3.9-slim")
0252 print(f"Environment: SIMULATION_MODE=fast, NUM_EVENTS=1000")
0253
0254 try:
0255 objectives = executor.execute(design_point)
0256 print(f"\n✅ Container workflow completed!")
0257 except Exception as e:
0258 print(f"\n⚠️ Container workflow failed (Docker may not be available): {e}")
0259 print("This is expected if Docker is not installed.")
0260
0261
0262 def example_5_optimizer_integration():
0263 """Example 5: Complete optimizer integration example."""
0264 print("\n" + "="*80)
0265 print("Example 5: Complete Optimizer Integration")
0266 print("="*80)
0267
0268
0269 eval_job = JobDefinition(
0270 name="evaluate_design",
0271 command=(
0272 "python -c \""
0273 "import json, sys; "
0274 "import numpy as np; "
0275 "# DTLZ2 objectives; "
0276 "x = [0.5, 0.7, 0.3]; "
0277 "f1 = x[0]; "
0278 "f2 = (1 + sum([(xi - 0.5)**2 for xi in x[1:]])) * (1 - np.sqrt(x[0]/(1 + sum([(xi - 0.5)**2 for xi in x[1:]])))); "
0279 "print(json.dumps({'f1': float(f1), 'f2': float(f2)})); "
0280 "\""
0281 ),
0282 payload={"evaluator_type": "bash"},
0283 )
0284
0285 eval_stage = StageDefinition(
0286 name="evaluate",
0287 jobs=[eval_job],
0288 )
0289
0290 branch = BranchDefinition(name="main", stages=[eval_stage])
0291
0292 workflow = WorkflowDefinition(
0293 name="optimizer_workflow",
0294 description="Workflow for optimizer integration",
0295 branches=[branch],
0296 objectives=[
0297 ObjectiveDefinition(name="f1", direction=ObjectiveDirection.MINIMIZE),
0298 ObjectiveDefinition(name="f2", direction=ObjectiveDirection.MINIMIZE),
0299 ],
0300 )
0301
0302 executor = DAGExecutor(workflow, base_output_dir="/tmp/aid2e_examples")
0303
0304 print("\nSimulating optimizer loop with 3 design points...")
0305
0306
0307 design_points = [
0308 {"x1": 0.2, "x2": 0.5, "x3": 0.8},
0309 {"x1": 0.5, "x2": 0.7, "x3": 0.3},
0310 {"x1": 0.8, "x2": 0.3, "x3": 0.6},
0311 ]
0312
0313 results = []
0314 for i, design_point in enumerate(design_points):
0315 print(f"\n Iteration {i+1}: Evaluating {design_point}")
0316 objectives = executor.execute(design_point)
0317 results.append({
0318 "design_point": design_point,
0319 "objectives": objectives,
0320 })
0321 print(f" → Objectives: {objectives}")
0322
0323 print(f"\n✅ Optimizer integration complete!")
0324 print(f"\nResults summary:")
0325 for i, result in enumerate(results):
0326 print(f" {i+1}. {result['design_point']} → {result['objectives']}")
0327
0328
0329 def example_6_workflow_from_config():
0330 """Example 6: Load workflow from YAML config file."""
0331 print("\n" + "="*80)
0332 print("Example 6: Loading Workflow from Config File")
0333 print("="*80)
0334
0335 import tempfile
0336 import yaml
0337 from pathlib import Path
0338
0339
0340 config = {
0341 "name": "config_workflow",
0342 "description": "Workflow loaded from YAML config",
0343 "branches": [
0344 {
0345 "name": "main",
0346 "stages": [
0347 {
0348 "name": "evaluate",
0349 "jobs": [
0350 {
0351 "name": "eval_job",
0352 "command": "echo 'Running from config'",
0353 "payload": {"evaluator_type": "bash"},
0354 }
0355 ],
0356 }
0357 ],
0358 }
0359 ],
0360 "objectives": [],
0361 }
0362
0363
0364 with tempfile.NamedTemporaryFile(mode='w', suffix='.yml', delete=False) as f:
0365 yaml.dump(config, f)
0366 config_path = f.name
0367
0368 print(f"\nConfig file created: {config_path}")
0369 print(f"Config contents:")
0370 print(yaml.dump(config, default_flow_style=False, indent=2))
0371
0372
0373 from aid2e.utilities.workflows import create_executor_from_config
0374
0375 executor = create_executor_from_config(
0376 workflow_config_path=config_path,
0377 output_dir="/tmp/aid2e_examples",
0378 )
0379
0380 design_point = {"param1": 1.0}
0381 print(f"\nExecuting workflow from config...")
0382
0383 objectives = executor.execute(design_point)
0384
0385 print(f"\n✅ Config-based workflow completed!")
0386 print(f"Workflow name: {executor.workflow.name}")
0387
0388
0389 Path(config_path).unlink()
0390
0391
0392 if __name__ == "__main__":
0393 print("\n" + "="*80)
0394 print("DAG Executor Complete Examples")
0395 print("="*80)
0396 print("\nThese examples demonstrate the full capabilities of the DAG Executor")
0397 print("for orchestrating multi-stage, multi-objective workflows.\n")
0398
0399
0400 example_1_simple_workflow()
0401 example_2_parallel_jobs()
0402 example_3_multi_stage_workflow()
0403 example_4_container_evaluators()
0404 example_5_optimizer_integration()
0405 example_6_workflow_from_config()
0406
0407 print("\n" + "="*80)
0408 print("All examples completed!")
0409 print("="*80)