Design Work Orchestration for Machine-Learning Data Pipelines

Scenario

Modern machine learning teams execute hundreds of complex, multi-stage pipelines daily—from raw telemetry ingestion and feature transformation to model evaluation and artifact promotion. You are tasked with designing a centralized ML workflow orchestrator (similar to an ML-native Flyte or Airflow) capable of scheduling and executing versioned task Directed Acyclic Graphs (DAGs) while enforcing data integrity, fair resource usage, and complete lineage tracking.

Functional Requirements

• DAG Definition & Versioning: Support code-defined DAGs where tasks, dependencies, code versions, and runtime environments (e.g., container images) are immutably versioned.
• Task Execution & Fault Tolerance: Provide durable state transitions, configurable retry policies with exponential backoff, worker heartbeat monitoring, and safe task recovery upon worker death.
• Heterogeneous Compute Pools: Route tasks dynamically to appropriate compute environments (standard CPU, high-memory nodes, GPU clusters) based on resource specifications.
• Data Quality Gates & Atomic Publishing: Support automated validation gates (e.g., schema drift detection, null rate thresholds). Downstream consumers must only access datasets that have passed validation via atomic dataset registration.
• Lineage & Deterministic Backfills: Track input/output dataset hashes and execution parameters to ensure reproducible executions and efficient historical backfilling without duplicate processing.
• Multi-Tenant Fairness: Prevent any single team or pipeline from starving shared cluster resources.

Scale & Operational Assumptions

• Workload Volume: 500 registered pipelines, 10,000 DAG runs per day, averaging 20 tasks per DAG (peak concurrency of ~3,000 parallel tasks).
• Data Scale: Intermediate and final dataset artifacts range from 1 GB to 20 TB.
• Scheduling Latency: Task dispatch latency should be under 5 seconds from dependency satisfaction.

What Your Design Should Cover

1. High-Level Architecture: Control plane (Scheduler, Metadata Store, State Manager) and data plane (Queue, Execution Workers, Storage, Artifact/Metadata Catalog).
2. Task State Machine & Worker Recovery: How task state changes are made durable and how orphaned tasks are reconciled without duplicate execution.
3. Quality Gates & Atomic Publication: Mechanism for running pre-commit checks and switching pointers to newly produced datasets without race conditions.
4. Tenant Fair-Share & Priority Queueing: Resource pool abstraction and admission control across competing tenants.
5. Lineage & Reproducibility: Data model linking code commits, inputs, task run configurations, and output dataset hashes.