# Tech-Spec: Airflow DAG & Extraction Pipeline Improvements

**Created:** 2025-12-30
**Status:** Completed
**Completed:** 2026-01-06
**Priority:** High
**Estimated Effort:** 2-3 days

---

## Overview

### Problem Statement

The current Airflow DAG implementation for document extraction has several performance, reliability, and maintainability issues identified during code review:

1. **Performance:** Documents are converted twice (parse + chunk), doubling processing time
2. **Throughput:** `max_active_runs=1` limits system to processing one document at a time
3. **Reliability:** Orphaned pickle files accumulate on task failures; no retry granularity for embeddings
4. **Maintainability:** Hardcoded thresholds scattered across files; flat DAG structure without task groups

### Solution

Implement targeted improvements across 4 phases:
1. **Phase 1:** Quick wins (throughput, cleanup) - 4 hours
2. **Phase 2:** Performance optimization (single conversion) - 8 hours
3. **Phase 3:** Reliability improvements (retry granularity) - 3 hours
4. **Phase 4:** DAG structure improvements (task groups, pools, SLAs) - 6 hours

### Scope

**In Scope:**
- Fix double conversion issue for text documents
- Increase max_active_runs and add resource pools
- Add cleanup task for orphaned pickle files
- Implement checkpoint/resume for embedding task
- Centralize configuration constants
- Add task groups for readability
- Add SLAs and failure callbacks

**Out of Scope:**
- File validation for corrupted/password-protected files (not needed)
- New file format support (CSV, HTML, etc.)
- Complex format handling (nested tables, pivot tables)
- Progress reporting UI (separate story)
- Record deduplication in Milvus (separate story)

---

## Context for Development

### Codebase Patterns

**Current Architecture:**
```
airflow/
├── dags/
│   ├── config.py                    # Shared configuration
│   ├── document_extraction_dag.py   # Main DAG definition
│   └── tasks/
│       ├── parse_tasks.py           # @task.external_python parse functions
│       └── processing_tasks.py      # Chunking, embedding, storage tasks
└── plugins/
    └── operators/                   # ValidateEventOperator, FetchDocumentOperator
```

**Key Patterns:**
- `@task.external_python(python='/opt/airflow/.venv/bin/python')` for heavy processing
- Pickle files in `/opt/airflow/logs/` for large data transfer (XCom limit: 48KB)
- Lazy imports inside task functions
- `trigger_rule='none_failed_min_one_success'` for branch convergence

### Files to Modify

| File | Changes |
|------|---------|
| `airflow/dags/config.py` | Add centralized constants, pool definitions |
| `airflow/dags/document_extraction_dag.py` | Task groups, pools, SLAs, cleanup task, max_active_runs |
| `airflow/dags/tasks/parse_tasks.py` | Store DoclingDocument in pickle, not just markdown |
| `airflow/dags/tasks/processing_tasks.py` | Use stored DoclingDocument, checkpoint embeddings, file validation |
| `airflow/plugins/operators/validate_event.py` | Add file validation hooks |

### Technical Decisions

| Decision | Choice | Rationale |
|----------|--------|-----------|
| DoclingDocument serialization | Pickle the `result.document` object | Docling's DoclingDocument is pickle-serializable; avoids re-conversion |
| Cleanup mechanism | DAG-level cleanup task with `trigger_rule='all_done'` | Runs regardless of success/failure; simple implementation |
| Embedding checkpoints | Save progress to pickle after each batch | Resume from last successful batch on retry |
| Configuration centralization | Single `config.py` with dataclass | Type-safe, IDE-friendly, single source of truth |
| Task groups | Logical grouping: excel_processing, text_processing, storage | Cleaner UI, easier debugging |
| Pools | `llm_api_pool` (3 slots), `embedding_pool` (2 slots) | Prevent API rate limiting |

---

## Implementation Plan

### Phase 1: Quick Wins (4 hours)

#### Task 1.1: Increase max_active_runs

**File:** `airflow/dags/document_extraction_dag.py`

```python
# Change from:
max_active_runs=1,

# To:
max_active_runs=10,
```

**Acceptance Criteria:**
- [x] DAG allows 10 concurrent document processing runs
- [x] No race conditions in pickle file naming (use document_id prefix)

---

#### Task 1.2: Add Cleanup Task

**File:** `airflow/dags/document_extraction_dag.py`

```python
@task(trigger_rule='all_done')
def cleanup_temp_files(document_id: str) -> Dict[str, Any]:
    """Clean up all temporary pickle files for this document."""
    import os
    import glob

    patterns = [
        f'/opt/airflow/logs/parse_*_{document_id}.pkl',
        f'/opt/airflow/logs/chunks_*_{document_id}.pkl',
        f'/opt/airflow/logs/embedding_data/embeddings_*.pkl',  # Match by content
    ]

    cleaned = 0
    for pattern in patterns:
        for f in glob.glob(pattern):
            try:
                os.remove(f)
                cleaned += 1
            except Exception:
                pass

    return {'cleaned_files': cleaned}
```

**Wire into DAG:**
```python
cleanup = cleanup_temp_files(doc_id)

# Both paths lead to cleanup
store_vectors_excel >> cleanup
store_vectors_text >> cleanup
unsupported >> cleanup
```

**Acceptance Criteria:**
- [x] Cleanup task runs after all paths complete (success or failure)
- [x] All pickle files for document_id are removed
- [x] No errors if files already deleted

---

#### Task 1.3: Centralize Configuration

**File:** `airflow/dags/config.py`

```python
from dataclasses import dataclass
from datetime import timedelta
import os

@dataclass(frozen=True)
class ExtractionConfig:
    """Centralized configuration for extraction pipeline."""

    # Thresholds
    LARGE_TABLE_ROW_THRESHOLD: int = 1000
    MAX_TOKENS_PER_CHUNK: int = 8000
    MIN_TOKENS_FOR_MERGE: int = 100

    # Batch sizes
    EMBEDDING_BATCH_SIZE: int = 16
    MAX_EMBEDDING_API_BATCH: int = 100
    LARGE_TABLE_CHUNK_SIZE: int = 100

    # Timeouts
    PARSE_TIMEOUT: timedelta = timedelta(minutes=int(os.getenv('AIRFLOW_PARSE_TIMEOUT_MINUTES', '120')))
    EMBEDDING_TIMEOUT: timedelta = timedelta(minutes=int(os.getenv('AIRFLOW_EMBEDDING_TIMEOUT_MINUTES', '60')))
    STORE_TIMEOUT: timedelta = timedelta(minutes=int(os.getenv('AIRFLOW_STORE_TIMEOUT_MINUTES', '30')))

    # XCom limits
    MAX_XCOM_SIZE: int = 48 * 1024  # 48KB

    # Paths
    TEMP_DATA_DIR: str = '/opt/airflow/logs'
    EXTERNAL_PYTHON: str = '/opt/airflow/.venv/bin/python'

# Singleton instance
CONFIG = ExtractionConfig()

# Pool definitions
POOLS = {
    'llm_api_pool': 3,      # Max 3 concurrent LLM API calls
    'embedding_pool': 2,    # Max 2 concurrent embedding tasks
}
```

**Acceptance Criteria:**
- [x] All hardcoded values moved to CONFIG
- [x] `large_table_extractor.py` and `detect_border.py` import from config
- [x] Processing tasks use CONFIG values

---

### Phase 2: Performance - Single Conversion (8 hours)

#### Task 2.1: Store DoclingDocument in Parse Tasks

**File:** `airflow/dags/tasks/parse_tasks.py`

**Change for `parse_word_document`, `parse_pdf_document`, `parse_powerpoint_document`:**

```python
@task.external_python(python=EXTERNAL_PYTHON)
def parse_word_document(fetch_result: Dict[str, Any]) -> Dict[str, Any]:
    from docling.document_converter import DocumentConverter
    from src.extraction_v2.document_converter import DocumentToHtmlConverter
    import pickle
    import logging

    log = logging.getLogger(__name__)

    # ... existing file handling code ...

    # Convert ONCE using Docling
    converter = DocumentConverter()
    result = converter.convert(file_to_convert)

    # Extract markdown for backward compatibility
    markdown = result.document.export_to_markdown()

    # Store BOTH markdown AND DoclingDocument
    parse_data = {
        'doc_type': 'word',
        'markdown_content': markdown,
        'docling_document': result.document,  # NEW: Store full document
        'local_path': local_path,
        'converted_path': converted_path or '',
        'document_id': document_id,
        'original_filename': original_filename,
    }

    data_file = f"/opt/airflow/logs/parse_word_{document_id}.pkl"
    with open(data_file, 'wb') as f:
        pickle.dump(parse_data, f)

    return {
        'data_file': data_file,
        'doc_type': 'word',
        'document_id': document_id,
        'original_filename': original_filename,
        'has_docling_document': True,  # NEW: Flag for chunk_text_task
        'content_length': len(markdown),
    }
```

---

#### Task 2.2: Use Stored DoclingDocument in chunk_text_task

**File:** `airflow/dags/tasks/processing_tasks.py`

```python
@task.external_python(python=EXTERNAL_PYTHON)
def chunk_text_task(
    document_id: str,
    parse_result: Dict[str, Any],
    max_tokens: int = 8000
) -> Dict[str, Any]:
    from uuid import UUID
    from src.extraction_v2.semantic_chunker import SemanticChunker
    import logging
    import pickle
    import os

    log = logging.getLogger(__name__)
    doc_uuid = UUID(document_id)

    # Load data from pickle file
    data_file = parse_result.get('data_file')
    if data_file and os.path.exists(data_file):
        with open(data_file, 'rb') as f:
            parse_data = pickle.load(f)
    else:
        raise ValueError("No parse data file found")

    original_filename = parse_data.get('original_filename', 'document')
    doc_type = parse_data.get('doc_type', 'word')

    # USE STORED DOCLING DOCUMENT - NO RE-CONVERSION!
    docling_doc = parse_data.get('docling_document')

    if docling_doc is None:
        # Fallback for backward compatibility
        log.warning("No stored DoclingDocument, falling back to re-conversion")
        from docling.document_converter import DocumentConverter
        local_path = parse_data.get('converted_path') or parse_data.get('local_path')
        converter = DocumentConverter()
        result = converter.convert(local_path)
        docling_doc = result.document

    # Chunk using the document
    chunker = SemanticChunker(max_tokens=max_tokens)
    chunks = chunker.chunk_document(
        doc=docling_doc,
        file_id=doc_uuid,
        original_filename=original_filename,
    )

    # ... rest of function unchanged ...
```

**Acceptance Criteria:**
- [x] Text documents are only converted ONCE in parse task
- [x] chunk_text_task uses stored DoclingDocument
- [x] Fallback to re-conversion if DoclingDocument missing (backward compat)
- [x] Performance test: Word/PDF processing time reduced by ~40-50%

---

### Phase 3: Reliability - Embedding Checkpoints (3 hours)

#### Task 3.1: Embedding Checkpoint/Resume

**File:** `airflow/dags/tasks/processing_tasks.py`

```python
@task.external_python(python=EXTERNAL_PYTHON)
def generate_embeddings_task(
    chunks_result: Dict[str, Any],
    batch_size: int = 16
) -> Dict[str, Any]:
    from src.knowledge.embeddings import EmbeddingService
    import asyncio
    import logging
    import pickle
    import os
    from uuid import uuid4

    log = logging.getLogger(__name__)

    # ... load chunks from pickle ...

    # Generate unique run ID for checkpointing
    run_id = uuid4().hex[:8]
    checkpoint_file = f'/opt/airflow/logs/embedding_checkpoint_{run_id}.pkl'

    async def embed_with_checkpoints():
        service = EmbeddingService()
        all_embeddings = []
        start_batch = 0

        # Check for existing checkpoint (retry scenario)
        if os.path.exists(checkpoint_file):
            log.info(f"Found checkpoint, resuming...")
            with open(checkpoint_file, 'rb') as f:
                checkpoint = pickle.load(f)
                all_embeddings = checkpoint['embeddings']
                start_batch = checkpoint['next_batch']
            log.info(f"Resuming from batch {start_batch}, {len(all_embeddings)} embeddings recovered")

        effective_batch_size = min(batch_size, 100)
        total_batches = (len(texts) - 1) // effective_batch_size + 1

        for i in range(start_batch * effective_batch_size, len(texts), effective_batch_size):
            batch = texts[i:i + effective_batch_size]
            batch_num = i // effective_batch_size + 1

            log.info(f"Processing batch {batch_num}/{total_batches}")
            batch_embeddings = await service.embed_batch(batch)
            all_embeddings.extend(batch_embeddings)

            # Save checkpoint after each batch
            checkpoint = {
                'embeddings': all_embeddings,
                'next_batch': batch_num,
                'total_batches': total_batches,
            }
            with open(checkpoint_file, 'wb') as f:
                pickle.dump(checkpoint, f)

            log.info(f"Checkpoint saved: {len(all_embeddings)}/{len(texts)} embeddings")

        # Clean up checkpoint on success
        if os.path.exists(checkpoint_file):
            os.remove(checkpoint_file)

        return all_embeddings

    all_embeddings = asyncio.run(embed_with_checkpoints())

    # ... rest of function unchanged ...
```

**Acceptance Criteria:**
- [x] Progress saved after each batch to checkpoint file
- [x] On retry, resumes from last successful batch
- [x] Checkpoint file cleaned up on success
- [x] Test: Kill task mid-embedding, verify resume works

---

### Phase 4: DAG Structure Improvements (6 hours)

#### Task 4.1: Add Task Groups

**File:** `airflow/dags/document_extraction_dag.py`

```python
from airflow.utils.task_group import TaskGroup

with DAG(...) as dag:

    # Validation group
    with TaskGroup("validation", tooltip="Validate and fetch document") as validation_group:
        validate = ValidateEventOperator(task_id='validate_event')
        fetch = FetchDocumentOperator(task_id='fetch_document')
        validate >> fetch

    # Branch
    branch = detect_format(validate.output)

    # Excel processing group
    with TaskGroup("excel_processing", tooltip="Excel extraction pipeline") as excel_group:
        parse_excel = parse_excel_document.override(
            task_id='parse_excel',
            execution_timeout=CONFIG.PARSE_TIMEOUT,
        )(fetch_result=fetch.output)

        detect_headers = detect_headers_task.override(
            task_id='detect_headers',
            pool='llm_api_pool',  # Limit concurrent LLM calls
        )(parse_result=parse_excel)

        chunk_excel = chunk_excel_task.override(
            task_id='chunk_excel',
        )(document_id=doc_id, parse_result=parse_excel, headers_result=detect_headers)

        embeddings_excel = generate_embeddings_task.override(
            task_id='generate_embeddings',
            execution_timeout=CONFIG.EMBEDDING_TIMEOUT,
            pool='embedding_pool',
        )(chunks_result=chunk_excel)

        store_excel = store_vectors_task.override(
            task_id='store_vectors',
            execution_timeout=CONFIG.STORE_TIMEOUT,
        )(document_id=doc_id, embedded_result=embeddings_excel)

    # Text processing group
    with TaskGroup("text_processing", tooltip="Word/PDF/PPT extraction pipeline") as text_group:
        # ... similar structure ...

    # Dependencies
    validation_group >> branch
    branch >> excel_group
    branch >> text_group
    branch >> unsupported

    [excel_group, text_group, unsupported] >> cleanup
```

---

#### Task 4.2: Add Pools (Admin Setup)

**Create pools via Airflow CLI or UI:**

```bash
# Create LLM API pool (limit concurrent header detection)
airflow pools set llm_api_pool 3 "Limit concurrent LLM API calls"

# Create embedding pool (limit concurrent embedding generation)
airflow pools set embedding_pool 2 "Limit concurrent embedding tasks"
```

**Or via Python (in DAG file):**

```python
from airflow.models import Pool
from airflow.utils.db import create_session

def create_pools():
    with create_session() as session:
        for pool_name, slots in POOLS.items():
            pool = session.query(Pool).filter_by(pool=pool_name).first()
            if not pool:
                session.add(Pool(pool=pool_name, slots=slots))
        session.commit()
```

---

#### Task 4.3: Add SLAs and Callbacks

**File:** `airflow/dags/document_extraction_dag.py`

```python
from airflow.exceptions import AirflowSensorTimeout

def on_failure_callback(context):
    """Log failure details for debugging."""
    import logging
    log = logging.getLogger(__name__)

    task_instance = context['task_instance']
    log.error(
        f"Task failed: {task_instance.task_id}",
        extra={
            'dag_id': context['dag'].dag_id,
            'run_id': context['run_id'],
            'execution_date': str(context['execution_date']),
            'exception': str(context.get('exception', 'Unknown')),
        }
    )

    # Could add: Slack notification, PagerDuty alert, etc.

def sla_miss_callback(dag, task_list, blocking_task_list, slas, blocking_tis):
    """Handle SLA misses."""
    import logging
    log = logging.getLogger(__name__)
    log.warning(f"SLA missed for tasks: {[t.task_id for t in task_list]}")

# Update default_args
default_args = {
    'owner': 'textiq',
    'depends_on_past': False,
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 2,
    'retry_delay': timedelta(minutes=2),
    'retry_exponential_backoff': True,
    'max_retry_delay': timedelta(minutes=30),
    'on_failure_callback': on_failure_callback,
    'sla': timedelta(hours=2),  # 2-hour SLA for all tasks
}

with DAG(
    dag_id='document_extraction_dag',
    default_args=default_args,
    sla_miss_callback=sla_miss_callback,
    # ... rest of config ...
) as dag:
```

**Acceptance Criteria:**
- [x] Task groups visible in Airflow UI
- [x] Pools limit concurrent LLM/embedding calls
- [x] SLA alerts triggered if processing exceeds 2 hours
- [x] Failure callback logs detailed error context

---

## Testing Strategy

### Unit Tests

| Test | File | Description |
|------|------|-------------|
| test_config_centralized | `tests/airflow/test_config.py` | Verify all constants in CONFIG |
| test_cleanup_task | `tests/airflow/test_dag.py` | Verify cleanup runs on all paths |
| test_embedding_checkpoint | `tests/airflow/test_processing_tasks.py` | Verify checkpoint save/resume |

### Integration Tests

| Test | Description |
|------|-------------|
| test_single_conversion_performance | Time text document processing, verify ~40% improvement |
| test_concurrent_runs | Trigger 5 documents simultaneously, verify parallel processing |
| test_failure_cleanup | Kill task mid-run, verify pickle files cleaned |
| test_embedding_retry | Kill embedding task, verify resume from checkpoint |

### Manual Validation

- [ ] Airflow UI shows task groups correctly
- [ ] Pool limits respected under load
- [ ] SLA miss notifications work
- [ ] Cleanup task visible in DAG graph

---

## Dependencies

### Python Packages

None - all changes use existing dependencies.

### External

- None - all changes are internal to codebase

---

## Rollout Plan

1. **Phase 1** (Quick Wins): Deploy immediately, low risk
2. **Phase 2** (Performance): Deploy after unit tests pass, medium risk
3. **Phase 3** (Reliability): Deploy after integration tests, low risk
4. **Phase 4** (Structure): Deploy after DAG visualization verified, low risk

### Rollback

Each phase is independent. Rollback by reverting specific commits:
- Phase 1: Revert max_active_runs and cleanup task
- Phase 2: Revert DoclingDocument storage (fallback code ensures compatibility)
- Phase 3: Revert checkpoint code (no state persisted between deploys)
- Phase 4: Revert task groups (purely cosmetic changes to DAG)

---

## Notes

### Future Improvements (Out of Scope)

1. **Record Deduplication:** Use deterministic IDs based on content hash
2. **Progress Reporting:** Emit XCom progress metrics for UI consumption
3. **Complex Format Handling:** Nested tables, pivot tables, embedded objects
4. **Parallel Table Processing:** Use dynamic task mapping for large table extraction

### Reference Documents

- [Deep-Dive: Airflow & Extraction V2 Integration](../deep-dive-airflow-extraction-v2-integration.md)
- [Airflow Component Spec](../specs/airflow.md)
- [Airflow Migration Tech Spec](../feats/airflow-migration-tech-spec.md)

---

*Generated by BMAD Tech-Spec Workflow*
*Date: 2025-12-30*