# Airflow Migration Epics

**Feature:** Migrate from Celery to Apache Airflow for DAG-based Orchestration
**Tech Spec:** [airflow-migration-tech-spec.md](../airflow-migration-tech-spec.md)
**Created:** 2025-12-08
**Updated:** 2025-12-18

---

## Epic Overview

This feature migrates the document extraction pipeline from Celery-based task processing to Apache Airflow DAG-based orchestration, following a phased approach to minimize risk.

### Current State
- Celery-based task processing with `extract_document_v2` task
- Monolithic `ExtractionPipelineV2` class orchestrating all extraction logic
- Tight coupling between orchestration and business logic
- Limited observability and workflow visualization

### Target State
- Apache Airflow 3.x DAG-based orchestration
- Decoupled operators for each processing step
- Event-driven triggers via SeaweedFS webhooks
- Enhanced observability with DAG visualization and task monitoring

---

## Epic 1: Airflow Infrastructure & Parallel Deployment

**Goal:** Deploy Airflow alongside existing Celery without disrupting current operations.
**Status:** In Progress (6/7 stories complete)

**Scope:**
- Deploy Airflow infrastructure (webserver, scheduler, DB)
- Create project structure for DAGs and plugins
- Implement base operators and TaskFlow tasks
- Shadow mode testing (Airflow processes without saving)

### Implementation Notes

The implementation evolved from the original plan:
- **Heavy processing tasks** use `@task.external_python` decorator to run in an external Python venv (`/opt/airflow/.venv`) with all required dependencies
- **Core operators** (ValidateEvent, FetchDocument) remain as custom `BaseOperator` subclasses
- **Parse, chunk, embed, and store tasks** are implemented as TaskFlow functions in `dags/tasks/`
- **Large XCom data** handled via pickle files to avoid size limits - ALL parse tasks, chunk tasks, and embedding tasks save results to `/opt/airflow/logs/*.pkl` files
- **Storage simplified**: Only Milvus vector storage is used; PostgreSQL storage task (`store_chunks_task`) was removed as document data is not currently stored in PostgreSQL

### Stories

| ID | Story | Priority | Status |
|----|-------|----------|--------|
| 1.1 | Airflow Docker Infrastructure Setup | High | **done** |
| 1.2 | DAG Project Structure & Configuration | High | **done** |
| 1.3 | Core Operators Implementation (Validate, Fetch, Parse) | High | **done** |
| 1.4 | Chunking & Embedding Tasks | High | **done** |
| 1.5 | Storage Tasks (Milvus only) | High | **done** |
| 1.6 | Main DAG Definition & Branching Logic | High | **done** |
| 1.7 | Shadow Mode Implementation & Metrics | Medium | pending |

**Acceptance Criteria:**
- [x] Airflow webserver accessible at port 8081
- [x] All operators pass unit tests
- [x] DAG loads without import errors
- [ ] Shadow mode processes documents without affecting production data
- [ ] Metrics comparison between Celery and Airflow outputs

---

## Epic 2: Gradual Traffic Migration

**Goal:** Progressively route traffic from Celery to Airflow with rollback capability.

**Scope:**
- Feature flag implementation for traffic splitting
- API integration for DAG triggering
- Monitoring and alerting setup
- Gradual rollout (10% → 30% → 50% → 80% → 100%)

### Stories

| ID | Story | Priority | Status |
|----|-------|----------|--------|
| 2.1 | Feature Flag Infrastructure | High | pending |
| 2.2 | Airflow REST API Client Integration | High | pending |
| 2.3 | API Route Modification for Traffic Splitting | High | pending |
| 2.4 | Monitoring Dashboard & Alerts | Medium | pending |
| 2.5 | Canary Deployment (10% traffic) | High | pending |
| 2.6 | Expanded Rollout (30% → 50% → 80%) | High | pending |
| 2.7 | Full Migration (100% traffic) | High | pending |

**Acceptance Criteria:**
- [ ] Traffic can be split via feature flag percentage
- [ ] Both systems process documents correctly in parallel
- [ ] Metrics show equivalent or better performance
- [ ] Rollback can be executed in < 5 minutes

---

## Epic 3: Celery Deprecation & Cleanup

**Goal:** Remove Celery extraction task and related code after full Airflow migration.

**Scope:**
- Remove `extract_document_v2` Celery task
- Archive `extraction_v2_tasks.py`
- Update documentation
- Clean up Celery beat schedules

### Stories

| ID | Story | Priority | Status |
|----|-------|----------|--------|
| 3.1 | Verify 100% Airflow Stability (1 week soak) | High | pending |
| 3.2 | Remove Celery Extraction Task | Medium | pending |
| 3.3 | Archive Legacy Code | Low | pending |
| 3.4 | Update API & Documentation | Medium | pending |
| 3.5 | Final Cleanup & Validation | Medium | pending |

**Acceptance Criteria:**
- [ ] No Celery extraction tasks remain
- [ ] All tests pass with Airflow-only processing
- [ ] Documentation updated to reflect new architecture
- [ ] No regressions in processing success rate

---

## Dependencies

### External Dependencies
- Apache Airflow 3.1.3+ Docker image
- SeaweedFS 4.02+ for S3-compatible object storage
- PostgreSQL database for Airflow metadata
- SeaweedFS webhook configuration (optional for Phase 2+)

### Internal Dependencies
- Existing `src/extraction_v2/` business logic (reused, not modified)
- PostgreSQL for Airflow metadata only
- Milvus for vector indexing (document vectors)
- Azure OpenAI for embeddings

---

## Risk Mitigation

| Risk | Mitigation |
|------|------------|
| Airflow DB migration issues | Use separate Airflow DB, test in staging |
| XCom data size limits | Store large data in pickle files (`/opt/airflow/logs/*.pkl`), pass file paths via XCom |
| Processing timeout | Configure task-level timeouts via environment variables |
| Traffic split issues | Feature flag with instant rollback |

---

## Related Documentation

| Document | Description |
|----------|-------------|
| [`docs/specs/airflow.md`](../../specs/airflow.md) | **Current architecture tech spec** - Comprehensive documentation of the actual implementation including DAG structure, TaskFlow functions, operators, dependencies, and usage examples |
| [`airflow-migration-tech-spec.md`](../airflow-migration-tech-spec.md) | Original migration planning document |

> **Note:** The story files in `sprint_artifacts/` represent the original planning phase. The implementation evolved to use `@task.external_python` decorators instead of custom operators for heavy processing tasks. See the tech spec for current architecture.

---

## Success Metrics

| Metric | Current (Celery) | Target (Airflow) |
|--------|------------------|------------------|
| Processing Success Rate | 95% | ≥ 95% |
| Average Processing Time | 45s | ≤ 50s |
| P99 Processing Time | 180s | ≤ 200s |
| Developer Debugging Time | 30min avg | ≤ 15min |