# DAG Execution Best Practices

Best practices for building robust async DAG-based workflow engines.

## Core Principles

### 1. Atomicity and Idempotency

**Atomic Tasks:** Each task should be self-contained and either complete entirely or not at all.

```python
async def atomic_stage(data: dict) -> dict:
    """Atomic stage - all or nothing."""
    # Start transaction
    async with database.transaction():
        result = await process_data(data)
        await save_result(result)
        return {"processed": True}
    # If any step fails, transaction rolls back
```

**Idempotent Operations:** Tasks should produce the same result if re-executed with the same input.

```python
async def idempotent_stage(data: dict) -> dict:
    """Idempotent - safe to retry."""
    record_id = data["record_id"]
    
    # Check if already processed
    existing = await db.get(record_id)
    if existing and existing.status == "complete":
        return existing.result
    
    # Process and store result with ID
    result = await process(data)
    await db.upsert(record_id, result)
    return result
```

### 2. Task Design

**Keep Tasks Small & Focused:**
```python
# BAD - Too many responsibilities
async def do_everything(data: dict) -> dict:
    validated = validate(data)
    enriched = enrich(validated)
    processed = process(enriched)
    stored = store(processed)
    notified = notify(stored)
    return notified

# GOOD - Single responsibility
stages = [
    PipelineStage(name="validate", handler="app.validate"),
    PipelineStage(name="enrich", handler="app.enrich", depends_on=["validate"]),
    PipelineStage(name="process", handler="app.process", depends_on=["enrich"]),
    PipelineStage(name="store", handler="app.store", depends_on=["process"]),
    PipelineStage(name="notify", handler="app.notify", depends_on=["store"]),
]
```

### 3. Data Flow

**Pass Data via Context, Not Globals:**
```python
# BAD - Shared state
_global_result = None

async def stage_a(data: dict) -> dict:
    global _global_result
    _global_result = await compute()
    return {}

# GOOD - Explicit context
async def stage_a(data: dict) -> dict:
    result = await compute()
    return {"stage_a_result": result}  # Merged into context

async def stage_b(data: dict) -> dict:
    # Access via context
    stage_a_result = data["stage_a_result"]
    return {"stage_b_result": transform(stage_a_result)}
```

### 4. Parallelism Patterns

**Fan-Out/Fan-In:**
```python
# Define independent stages that can run in parallel
pipeline = Pipeline(
    name="parallel_processing",
    stages=[
        # Fan-out: these run in parallel
        PipelineStage(id="fetch_a", name="Fetch A", handler="app.fetch_a"),
        PipelineStage(id="fetch_b", name="Fetch B", handler="app.fetch_b"),
        PipelineStage(id="fetch_c", name="Fetch C", handler="app.fetch_c"),
        # Fan-in: waits for all fetches
        PipelineStage(
            id="aggregate",
            name="Aggregate",
            handler="app.aggregate",
            depends_on=["fetch_a", "fetch_b", "fetch_c"]
        ),
    ],
    max_parallel=3,  # Run all fetches concurrently
)
```

**Diamond Dependency:**
```python
#     A
#    / \
#   B   C
#    \ /
#     D

stages = [
    PipelineStage(id="a", name="A", handler="app.a"),
    PipelineStage(id="b", name="B", handler="app.b", depends_on=["a"]),
    PipelineStage(id="c", name="C", handler="app.c", depends_on=["a"]),
    PipelineStage(id="d", name="D", handler="app.d", depends_on=["b", "c"]),
]
```

### 5. Error Handling

**Retry with Backoff:**
```python
stage = PipelineStage(
    name="api_call",
    handler="app.call_api",
    retry_count=3,
    retry_delay=2.0,  # Exponential backoff
)
```

**Skip on Dependency Failure:**
```python
stage = PipelineStage(
    name="optional_step",
    handler="app.optional",
    depends_on=["required_step"],
    skip_on_failure=True,  # Continue pipeline if dependency failed
)
```

**Fail Fast vs Continue:**
```python
# Stop on first failure
pipeline = Pipeline(stages=stages, fail_fast=True)

# Continue despite failures
pipeline = Pipeline(stages=stages, fail_fast=False)
```

### 6. Timeouts

**Always Set Timeouts:**
```python
# Stage-level timeout
stage = PipelineStage(
    name="external_api",
    handler="app.call_external",
    timeout=30.0,  # 30 seconds max
)

# Pipeline-level timeout
pipeline = Pipeline(
    stages=stages,
    timeout=300.0,  # 5 minutes total
)
```

### 7. Minimize Import Overhead

**Keep DAG Definitions Lightweight:**
```python
# BAD - Heavy imports at module level
import pandas as pd
import tensorflow as tf

# GOOD - Import inside handler
async def ml_stage(data: dict) -> dict:
    import tensorflow as tf  # Lazy import
    model = tf.keras.models.load_model(...)
    return {"prediction": model.predict(...)}
```

### 8. Dynamic DAG Generation

**Generate Tasks Based on Data:**
```python
async def build_dynamic_pipeline(file_list: list[str]) -> Pipeline:
    """Generate stages dynamically based on input."""
    stages = []
    
    for i, file in enumerate(file_list):
        stages.append(PipelineStage(
            id=f"process_{i}",
            name=f"Process {file}",
            handler="app.process_file",
            metadata={"file": file}
        ))
    
    # Aggregation stage depends on all processing stages
    stages.append(PipelineStage(
        id="aggregate",
        name="Aggregate Results",
        handler="app.aggregate",
        depends_on=[f"process_{i}" for i in range(len(file_list))]
    ))
    
    return Pipeline(name="dynamic_pipeline", stages=stages)
```

## Testing DAGs

**Unit Test Handlers:**
```python
@pytest.mark.asyncio
async def test_process_stage():
    result = await process_handler({"input": "test"})
    assert result["status"] == "success"
```

**Test DAG Structure:**
```python
def test_dag_valid():
    pipeline = create_pipeline()
    # Validation happens at construction
    assert len(pipeline.stages) == 5
    assert pipeline.get_root_stages() == [stages[0]]

def test_dag_rejects_cycle():
    with pytest.raises(ValueError, match="cycle"):
        Pipeline(stages=[
            PipelineStage(id="a", depends_on=["b"], ...),
            PipelineStage(id="b", depends_on=["a"], ...),
        ])
```

**Integration Test Execution:**
```python
@pytest.mark.asyncio
async def test_pipeline_execution():
    executor = PipelineExecutor()
    executor.register_handler("test", mock_handler)
    
    result = await executor.execute(pipeline, initial_data={"key": "value"})
    
    assert result.status == PipelineStatus.COMPLETED
    assert result.stages_completed == 3
```

## References

- [Apache Airflow Best Practices](https://www.astronomer.io/docs/learn/dag-best-practices/)
- [Cloudflare Workflows DAG](https://developers.cloudflare.com/workflows/python/dag/)
- [Async DAG Traversal Patterns](https://tomasfarias.dev/articles/async-patterns-in-python-dag-traversal/)
- [Writing Efficient DAGs](https://www.sparkcodehub.com/airflow/best-practices/efficient-dags)
