Back to home page

EIC code displayed by LXR

 
 

    


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     # Define a simple job that evaluates DTLZ2
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     # Define a stage containing the job
0043     eval_stage = StageDefinition(
0044         name="evaluate",
0045         jobs=[eval_job],
0046     )
0047     
0048     # Define a branch containing the stage
0049     main_branch = BranchDefinition(
0050         name="main",
0051         stages=[eval_stage],
0052     )
0053     
0054     # Define the workflow
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     # Create executor
0072     executor = DAGExecutor(
0073         workflow=workflow,
0074         base_output_dir="/tmp/aid2e_examples",
0075         log_level="INFO",
0076     )
0077     
0078     # Execute for a design point
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     # Job template that will be expanded to N parallel jobs
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     # Stage with job_factory to expand to 4 parallel jobs
0103     eval_stage = StageDefinition(
0104         name="parallel_evaluate",
0105         jobs=[eval_job],
0106         job_factory=JobFactory(
0107             type="range",
0108             params={"n": 4},  # Create 4 parallel jobs
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     # Stage 1: Generate input data
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     # Stage 2: Run simulation (depends on generate)
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     # Stage 3: Compute objectives (depends on simulate)
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     # Stages are executed sequentially: generate → simulate → compute
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     # Job using ContainerEvaluator
0208     container_job = JobDefinition(
0209         name="docker_simulation",
0210         command="python /app/simulate.py",  # Command inside container
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     # Define workflow for optimizer
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     # Simulate optimizer providing design points
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     # Create sample config file
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     # Write config to temp file
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     # Load executor from config
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     # Cleanup
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     # Run all examples
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)