Airflow Is Not the Endgame
Linear Airflow DAGs where every step runs sequentially and you wait 12 hours for a pipeline to fail at step 9? That’s 2020 thinking.
Modern orchestration is dynamic. Steps get skipped when nothing changed. Resources get allocated per-step based on what the step actually needs — CPU for preprocessing, big RAM for joins, GPU for training. If something crashes with OOM, the orchestrator doesn’t just retry — it doubles the memory and relaunches automatically.
Let’s build a hyper-scale orchestration system step by step. Python for the pipeline logic. Rust for the speed-critical parts (cache addressing, resource scheduling).
The Problem with Static DAGs
A static Airflow DAG looks like this:
# static_dag.py — Traditional Airflow DAG (the old way)
# Every step runs every time, even if nothing changed
# If preprocessing takes 2 hours and you only changed a learning rate,
# you still wait 2 hours for preprocessing to rerun
from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime
with DAG("ml_pipeline", start_date=datetime(2026, 1, 1), schedule=None) as dag:
# Step 1: Extract — runs EVERY TIME even if data hasn't changed
extract = PythonOperator(task_id="extract", python_callable=extract_data)
# Step 2: Preprocess — reruns even if only hyperparameters changed
preprocess = PythonOperator(task_id="preprocess", python_callable=preprocess_data)
# Step 3: Train — the only step that actually needs to rerun
train = PythonOperator(task_id="train", python_callable=train_model)
# Step 4: Evaluate — pointless rerun if the model didn't change
evaluate = PythonOperator(task_id="evaluate", python_callable=evaluate_model)
extract >> preprocess >> train >> evaluate
# # Problem:
# # Step 1 runs even if the raw data hasn't changed since last run
# # Step 2 runs even if you only changed the learning_rate param
# # Step 4 runs even if the model output is identical to last run
# # Total runtime: fixed, linear, wasteful
Intelligent Caching via Content-Addressable Storage
The fix: each step computes a hash of its inputs (data + code + parameters). If the hash matches a previous run, skip the step and reuse the cached output. This is how Flyte and Kubeflow Pipelines work under the hood.
We implement the cache key generation in Rust for speed, since hashing terabytes of data needs to be fast.
// cache_key.rs — Content-addressable cache key generation in Rust
// This runs at the start of each pipeline step to determine if we can skip it
// Rust is used here because hashing large files needs to be fast (~GB/s)
use std::fs::File;
use std::hash::{Hash, Hasher};
use std::collections::hash_map::DefaultHasher;
use std::io::{Read, BufReader};
use std::path::Path;
use rayon::prelude::*; // Parallel hashing across CPU cores
/// Compute a cache key from multiple inputs:
/// - File paths (data, code, configs)
/// - Parameter values (hyperparameters, dates, flags)
/// - The step name
///
/// If ALL inputs produce the same hash as a previous run, the step is skipped.
/// This is the core of content-addressable caching.
#[derive(Hash)]
struct CacheInput {
/// Name of the pipeline step (e.g. "preprocess", "train")
step_name: String,
/// SHA-256 of all input files concatenated
data_hash: String,
/// Stringified hyperparameters (e.g. "lr=0.001,batch_size=32")
params_hash: String,
}
/// Hash a file using SHA-256 in chunks so we don't load GB files into RAM
fn hash_file(path: &str) -> String {
// Open the file with buffered reading for performance
let file = File::open(path).expect("Failed to open file");
// BufReader reduces syscalls by reading 64KB chunks at a time
let mut reader = BufReader::with_capacity(64 * 1024, file);
let mut hasher = DefaultHasher::new();
// Read the file in chunks and feed each chunk into the hasher
// std::io::copy reads in 64KB blocks, which is efficient for large files
// For multi-GB files, we'd use mmap instead for even better performance
let mut buffer = [0u8; 8192]; // 8KB chunk size — good balance of speed/memory
loop {
let bytes_read = reader.read(&mut buffer).expect("Read error");
if bytes_read == 0 {
break; // End of file
}
// Hash the bytes directly — no need to allocate a full file buffer
buffer[..bytes_read].hash(&mut hasher);
}
// Convert the 64-bit hash to a hex string for use as a cache key
// In production, use SHA-256 instead of DefaultHasher to avoid collisions
format!("{:x}", hasher.finish())
}
/// Hash a directory by hashing all file paths + their contents
/// This detects: new files, deleted files, and modified files
fn hash_directory(path: &str) -> String {
// Walk the directory recursively in parallel using Rayon
// Rayon's parallel iterator distributes files across CPU cores
let entries: Vec<_> = std::fs::read_dir(path)
.expect("Failed to read directory")
.filter_map(|e| e.ok())
.collect();
// Hash all files in parallel — each core hashes a different file
// This is where Rust beats Python: true parallel computation
let mut hashes: Vec<String> = entries
.par_iter()
.map(|entry| {
let path = entry.path();
if path.is_file() {
// For each file, hash both its NAME and its CONTENT
// This catches renames (name changes) AND modifications (content changes)
format!("{}:{}", path.to_string_lossy(), hash_file(path.to_str().unwrap()))
} else {
String::new()
}
})
.collect();
// Sort for deterministic output (filesystem order varies)
hashes.sort();
// Hash the combined list of file hashes
let mut final_hasher = DefaultHasher::new();
for h in &hashes {
h.hash(&mut final_hasher);
}
format!("{:x}", final_hasher.finish())
}
/// Check if a cached result exists for the given cache key
/// Returns the path to cached artifacts if found, None otherwise
fn check_cache(cache_key: &str, cache_dir: &str) -> Option<String> {
let cache_path = format!("{}/{}", cache_dir, cache_key);
if Path::new(&cache_path).exists() {
// Cache hit — we can skip this step entirely
Some(cache_path)
} else {
// Cache miss — need to run the step
None
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_hash_consistency() {
// Same inputs must always produce the same hash
let hash1 = hash_file("Cargo.toml");
let hash2 = hash_file("Cargo.toml");
assert_eq!(hash1, hash2, "Hashes must be deterministic");
}
}
Now the Python side that uses the Rust cache:
```python
# orchestration/cache.py — Python wrapper for the Rust cache key generator
# Calls the Rust binary via subprocess for maximum speed
# Falls back to Python-only hashing if the Rust binary isn't available
import subprocess
import json
import hashlib
import os
from pathlib import Path
from typing import Optional
class ContentAddressableCache:
"""
Content-addressable cache for pipeline steps.
Each step's cache key = hash(inputs + code + parameters).
If hash matches a previous run, skip the step and use cached output.
This is the same pattern Flyte and Kubeflow Pipelines use internally.
"""
def __init__(self, cache_dir: str = "./.cache/pipeline"):
# Directory where cached step outputs are stored
# Structure: {cache_dir}/{step_name}/{cache_key}/output.parquet
self.cache_dir = Path(cache_dir)
self.cache_dir.mkdir(parents=True, exist_ok=True)
def _compute_cache_key(
self,
step_name: str,
input_paths: list[str],
params: dict,
) -> str:
"""
Compute a cache key from step name + input files + parameters.
Uses the Rust binary if available (10x faster for large inputs).
Falls back to Python hashlib if Rust binary isn't found.
"""
# Try the Rust binary first — it's 10x faster for large file sets
try:
result = subprocess.run(
["./target/release/cache-key"],
input=json.dumps({
"step": step_name,
"paths": input_paths,
"params": params,
}),
capture_output=True,
text=True,
timeout=30,
check=True,
)
return result.stdout.strip()
except (FileNotFoundError, subprocess.CalledProcessError, subprocess.TimeoutExpired):
# Fallback to Python — slower but always available
return self._compute_cache_key_python(step_name, input_paths, params)
def _compute_cache_key_python(
self,
step_name: str,
input_paths: list[str],
params: dict,
) -> str:
"""Pure Python fallback for cache key computation."""
hasher = hashlib.sha256()
# Hash the step name
hasher.update(step_name.encode())
# Hash input files — read in chunks to avoid memory issues
for path in input_paths:
if os.path.isfile(path):
with open(path, "rb") as f:
while True:
chunk = f.read(65536) # 64KB chunks
if not chunk:
break
hasher.update(chunk)
# Hash the parameters (sorted for deterministic output)
params_str = json.dumps(params, sort_keys=True)
hasher.update(params_str.encode())
return hasher.hexdigest()
def get_cached_output(self, step_name: str, input_paths: list[str], params: dict) -> Optional[Path]:
"""
Check if a cached output exists for this step + inputs + params.
Returns the path to cached data, or None if cache miss.
"""
cache_key = self._compute_cache_key(step_name, input_paths, params)
cached_path = self.cache_dir / step_name / cache_key
if cached_path.exists():
print(f"🔁 Cache HIT for '{step_name}' — skipping computation")
return cached_path
else:
print(f"⚡ Cache MISS for '{step_name}' — running step")
return None
def store_output(self, step_name: str, input_paths: list[str], params: dict, output_path: Path):
"""
Store a step's output in the cache for future pipeline runs.
The cache key is derived from inputs, so any change invalidates the cache.
"""
cache_key = self._compute_cache_key(step_name, input_paths, params)
cached_dir = self.cache_dir / step_name / cache_key
cached_dir.mkdir(parents=True, exist_ok=True)
# Copy output to cache (hardlink for large files to save disk space)
# If hardlink fails (cross-device), fall back to copy
dest = cached_dir / output_path.name
try:
os.link(output_path, dest) # Hardlink — instant, no disk space used
except OSError:
import shutil
shutil.copy2(output_path, dest) # Copy — slower but always works
# Store metadata about this cache entry
meta = {
"step": step_name,
"params": params,
"input_paths": input_paths,
"created_at": str(datetime.now()),
}
with open(cached_dir / "meta.json", "w") as f:
json.dump(meta, f, indent=2)
print(f"💾 Cached '{step_name}' at {dest}")
Heterogeneous Compute Scheduling
Different steps need different hardware. Preprocessing is CPU-bound. Feature joins need lots of RAM. Training needs GPUs. The orchestrator should schedule each step onto the right instance type automatically.
# orchestration/scheduler.py — Heterogeneous compute scheduler
# Each step declares what resources it needs.
# The scheduler allocates instances dynamically based on the step's requirements.
from dataclasses import dataclass
from enum import Enum
from typing import Optional
class ComputeTier(Enum):
"""
Available compute tiers — each maps to a different instance type.
The orchestrator chooses a tier per step based on the step's resource profile.
"""
CPU_CHEAP = "cpu-cheap" # Spot instances, preemptible CPUs
CPU_STANDARD = "cpu-standard" # On-demand CPUs, moderate RAM
MEMORY_OPTIMIZED = "memory-optimized" # High-RAM instances for joins
GPU_SINGLE = "gpu-single" # Single GPU (A10G, L4) for small models
GPU_MULTI = "gpu-multi" # Multi-GPU (8x A100 H100) with NCCL/EFA
@dataclass
class ResourceProfile:
"""
Resource requirements for a single pipeline step.
Each step declares what it needs, and the scheduler finds the best matching instance.
"""
# CPU cores needed (1-128)
cpu: int = 2
# RAM in GB (1-1024)
memory_gb: int = 8
# GPUs needed (0-64)
gpus: int = 0
# GPU memory per card in GB (0 if no GPU)
gpu_memory_gb: int = 0
# Whether this step can use preemptible/spot instances
# Spot instances are 70% cheaper but can be terminated at any time
spot_ok: bool = True
# Estimated runtime in minutes (for cost calculation)
estimated_minutes: int = 10
def recommended_tier(self) -> ComputeTier:
"""
Automatically determine the best compute tier for this step.
This is the core of heterogeneous scheduling — match step to instance type.
Rules:
- GPU steps → GPU tier
- CPU with high RAM → memory tier
- Everything else → CPU tier (standard or cheap)
"""
if self.gpus > 0:
# GPU step — check if multi-GPU
if self.gpus > 1:
return ComputeTier.GPU_MULTI
return ComputeTier.GPU_SINGLE
if self.memory_gb >= 64:
# High memory footprint — use memory-optimized instances
# These are expensive ($2-4/hr) but necessary for large joins
return ComputeTier.MEMORY_OPTIMIZED
if self.spot_ok and self.cpu <= 8:
# Cheap, interruptible work — use spot instances
# Spot instances cost ~30% of on-demand
return ComputeTier.CPU_CHEAP
# Default: standard CPU — reliable and moderately priced
return ComputeTier.CPU_STANDARD
// scheduler.rs — Fast resource scheduling in Rust
// Used when allocating instances across many parallel pipeline branches
// Rust is faster here because scheduling is a bin-packing problem
// that gets computationally expensive with hundreds of parallel steps
use std::collections::HashMap;
/// Represents a compute instance that can run pipeline steps
#[derive(Clone, Debug)]
struct ComputeNode {
/// Node identifier (instance ID or Kubernetes node name)
id: String,
/// Total CPU cores available on this node
cpu_total: u32,
/// CPU cores currently allocated to running steps
cpu_used: u32,
/// Total RAM in GB
memory_total: u32,
/// RAM currently allocated
memory_used: u32,
/// Number of GPUs available
gpus_total: u32,
/// GPUs currently allocated
gpus_used: u32,
/// Whether this is a spot/preemptible instance (cheaper but can be killed)
is_spot: bool,
/// Estimated cost per hour in dollars
cost_per_hour: f64,
}
impl ComputeNode {
/// Available CPU cores (total - used)
fn cpu_available(&self) -> u32 {
self.cpu_total - self.cpu_used
}
/// Available RAM in GB
fn memory_available(&self) -> u32 {
self.memory_total - self.memory_used
}
/// Available GPUs (total - used)
fn gpus_available(&self) -> u32 {
self.gpus_total - self.gpus_used
}
/// Check if this node can fit the given resource profile
fn can_fit(&self, cpu: u32, memory_gb: u32, gpus: u32) -> bool {
self.cpu_available() >= cpu
&& self.memory_available() >= memory_gb
&& self.gpus_available() >= gpus
}
}
/// A step waiting to be scheduled on a compute node
#[derive(Clone, Debug)]
struct Step {
/// Unique step identifier
id: String,
/// CPU cores required
cpu: u32,
/// RAM required in GB
memory_gb: u32,
/// GPUs required
gpus: u32,
/// Can this step run on spot instances?
spot_ok: bool,
}
/// Schedule steps onto nodes using a best-fit-decreasing algorithm.
/// This minimizes the number of nodes needed (bin packing).
/// For 100+ steps, Rust is ~50x faster than Python for this algorithm.
fn schedule_steps(steps: Vec<Step>, nodes: Vec<ComputeNode>) -> HashMap<String, String> {
// Result: step_id -> node_id mapping
let mut assignment: HashMap<String, String> = HashMap::new();
// Mutable copy of nodes (we'll update resource usage as we assign steps)
let mut available_nodes = nodes.clone();
// Sort steps by resource requirement (largest first = best fit)
// This minimizes wasted resources on each node
let mut sorted_steps: Vec<Step> = steps
.into_iter()
.filter(|s| s.spot_ok) // Only schedule spot-eligible steps here
.collect();
sorted_steps.sort_by(|a, b| {
// Sort by total resource requirements descending
let a_total = a.cpu + a.memory_gb * 4 + a.gpus * 100;
let b_total = b.cpu + b.memory_gb * 4 + b.gpus * 100;
b_total.cmp(&a_total)
});
// Best-fit-decreasing: place each step on the node that fits best
// "Best fit" = the node where the step leaves the least wasted resources
for step in &sorted_steps {
let best_node = available_nodes
.iter_mut()
.filter(|n| n.can_fit(step.cpu, step.memory_gb, step.gpus))
.min_by(|a, b| {
// Find the node where this step leaves the least wasted CPU
let a_waste = a.cpu_available() - step.cpu;
let b_waste = b.cpu_available() - step.cpu;
a_waste.cmp(&b_waste)
});
if let Some(node) = best_node {
// Allocate resources on this node
node.cpu_used += step.cpu;
node.memory_used += step.memory_gb;
node.gpus_used += step.gpus;
assignment.insert(step.id.clone(), node.id.clone());
}
}
assignment
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_schedule_steps() {
let steps = vec![
Step { id: "preprocess".into(), cpu: 4, memory_gb: 16, gpus: 0, spot_ok: true },
Step { id: "train".into(), cpu: 8, memory_gb: 64, gpus: 1, spot_ok: false },
];
let nodes = vec![
ComputeNode {
id: "spot-pool-1".into(),
cpu_total: 32, cpu_used: 0,
memory_total: 128, memory_used: 0,
gpus_total: 0, gpus_used: 0,
is_spot: true, cost_per_hour: 0.12,
},
ComputeNode {
id: "gpu-cluster-1".into(),
cpu_total: 64, cpu_used: 0,
memory_total: 512, memory_used: 0,
gpus_total: 4, gpus_used: 0,
is_spot: false, cost_per_hour: 8.0,
},
];
let result = schedule_steps(steps, nodes);
assert_eq!(result.get("preprocess").unwrap(), "spot-pool-1");
assert_eq!(result.get("train").unwrap(), "gpu-cluster-1");
}
}
Auto-Healing with Semantic Retries
When a step fails, most orchestrators retry with the same resources. That’s stupid — if it failed due to OOM, retrying with the same memory will fail again. Semantic retries analyze the error and adjust resources accordingly.
# orchestration/retry.py — Semantic retry policies for auto-healing
# Instead of blind retries, this analyzes the error and adjusts resources
import time
import subprocess
import psutil # System resource monitoring
import json
from dataclasses import dataclass, field
from typing import Optional, Callable
@dataclass
class RetryPolicy:
"""
Configuration for semantic retry behavior.
Each step can define its own retry policy.
Semantic retry ≠ simple retry:
- Simple retry: try again with the same resources (will fail again)
- Semantic retry: analyze WHY it failed, adjust resources, then retry
"""
# Maximum number of retry attempts before giving up
max_retries: int = 3
# Base delay between retries (doubles each attempt via exponential backoff)
base_delay_seconds: float = 5.0
# Memory multiplier per retry (1.5 = add 50% more RAM each attempt)
memory_multiplier: float = 1.5
# CPU multiplier per retry (1.25 = add 25% more CPU each attempt)
cpu_multiplier: float = 1.25
# Whether to increase retry delay for spot instance preemption
# Spot instances need longer delays (market recovery time)
spot_backoff: bool = False
def analyze_error(error: Exception, log_output: str) -> dict:
"""
Analyze a step failure to determine the root cause.
Returns a dictionary of adjustments to make for the retry.
This is where "semantic" happens — we parse the error message
and decide what to change, not just that we should retry.
"""
adjustments = {}
error_msg = str(error) + log_output[-2000:] # Last 2000 chars of log
# Check for OOM (Out of Memory) — increase memory
if any(kw in error_msg.upper() for kw in ["OOM", "OUT OF MEMORY", "CUDA OUT OF MEMORY",
"KILLED", "EXIT CODE 137", "ALLOC FAILED"]):
adjustments["memory_multiplier"] = 2.0 # Double the memory
adjustments["reason"] = "OOM detected — doubling memory allocation"
adjustments["retry_delay"] = 10 # Give the system time to free resources
# Check for GPU-related failures
if any(kw in error_msg for kw in ["CUDA", "NCCL", "GPU", "cuda"]):
if "NCCL TIMEOUT" in error_msg or "NCCL" in error_msg:
# NCCL timeout — usually network congestion, retry with longer timeout
adjustments["nccl_timeout_multiplier"] = 2.0
adjustments["reason"] = "NCCL timeout — increasing timeout"
elif "CUDA OUT OF MEMORY" in error_msg:
# GPU OOM — reduce batch size instead of adding more GPUs
adjustments["batch_size_divider"] = 2
adjustments["reason"] = "GPU OOM — halving batch size"
else:
# General CUDA error — retry with fresh GPU context
adjustments["reload_model"] = True
adjustments["reason"] = "CUDA error — reloading model"
# Check for spot instance preemption
if "SPOT" in error_msg.upper() or "PREEMPT" in error_msg.upper():
adjustments["spot_ok"] = False # Don't use spot for retry
adjustments["retry_delay"] = 60 # Wait 1 minute for spot market
adjustments["reason"] = "Spot instance preempted — switching to on-demand"
# Check for network timeout
if "TIMEOUT" in error_msg.upper() or "CONNECTION" in error_msg.upper():
adjustments["retry_delay"] = 30
adjustments["reason"] = "Network issue — waiting 30s before retry"
# Check for disk full
if "DISK" in error_msg.upper() or "NO SPACE" in error_msg.upper() or "QUOTA" in error_msg.upper():
adjustments["cleanup_disk"] = True
adjustments["reason"] = "Disk full — cleaning up temp files"
return adjustments
def execute_with_semantic_retry(
step_fn: Callable,
step_name: str,
initial_resources: dict,
policy: RetryPolicy = None,
) -> dict:
"""
Execute a step function with semantic retry logic.
If a step fails, analyze the error, adjust resources, and retry.
Args:
step_fn: The function to execute (e.g. train_model, preprocess_data)
step_name: Name of the step (for logging)
initial_resources: Initial resource allocation {cpu, memory_gb, gpus}
policy: Retry policy configuration
Returns:
The step's output (typically a dict of metrics + artifact paths)
Raises:
RuntimeError: If all retry attempts fail
"""
if policy is None:
policy = RetryPolicy()
resources = initial_resources.copy()
last_error = None
total_attempts = 0
for attempt in range(1, policy.max_retries + 2): # +1 for initial attempt
total_attempts += 1
# Wait for retry with exponential backoff
if attempt > 1:
delay = policy.base_delay_seconds * (2 ** (attempt - 2))
print(f"⏳ Retry {attempt - 1}/{policy.max_retries} after {delay:.0f}s...")
time.sleep(delay)
# Allocate resources for this attempt
memory_gb = int(resources.get("memory_gb", 8) *
(policy.memory_multiplier ** (attempt - 1)))
cpu_cores = int(resources.get("cpu", 2) *
(policy.cpu_multiplier ** (attempt - 1)))
print(f"▶️ Running '{step_name}' attempt {attempt} "
f"(CPU: {cpu_cores}, RAM: {memory_gb}GB)")
try:
# Execute the actual step function
# In production, this submits a job to the compute cluster
start = time.time()
result = step_fn(cpu=cpu_cores, memory_gb=memory_gb)
elapsed = time.time() - start
print(f"✅ '{step_name}' completed in {elapsed:.1f}s")
return result
except Exception as e:
last_error = e
# Capture the last part of log output for error analysis
log_output = getattr(e, "log_output", "")
# Analyze the error semantically
adjustments = analyze_error(e, log_output)
reason = adjustments.get("reason", "Unknown error")
print(f"❌ '{step_name}' failed (attempt {attempt}): {reason}")
# Apply resource adjustments for the next retry
if adjustments.get("memory_multiplier"):
resources["memory_gb"] = int(
resources.get("memory_gb", 8) * adjustments["memory_multiplier"]
)
if adjustments.get("cpu_multiplier"):
resources["cpu"] = int(
resources.get("cpu", 2) * adjustments["cpu_multiplier"]
)
if adjustments.get("batch_size_divider"):
resources["batch_size"] = resources.get("batch_size", 32) // int(
adjustments["batch_size_divider"]
)
if adjustments.get("reload_model"):
resources["fresh_gpu_context"] = True
# If OOM keeps happening even with maxed resources, fail fast
if resources.get("memory_gb", 0) > 512:
raise RuntimeError(
f"Step '{step_name}' needs >512GB RAM — "
f"check for memory leak or data skew"
)
# All retries exhausted
raise RuntimeError(
f"Step '{step_name}' failed after {total_attempts} attempts. "
f"Last error: {last_error}"
)
def simulate_unstable_step(should_fail: bool = True, fail_with_oom: bool = True):
"""
Simulates a step that fails with OOM on the first attempt,
then succeeds after resource adjustments.
Used to demonstrate auto-healing.
"""
import random
# First call: simulate OOM
if should_fail and not hasattr(simulate_unstable_step, "_has_run"):
simulate_unstable_step._has_run = True
error_msg = "CUDA OUT OF MEMORY. Allocated 8GB but needed 12GB."
raise RuntimeError(error_msg)
# Second call: success with the extra memory
return {"accuracy": 0.94, "loss": 0.23}
# Demo
if __name__ == "__main__":
print("=== Auto-Healing Demo ===")
result = execute_with_semantic_retry(
step_fn=simulate_unstable_step,
step_name="train_model",
initial_resources={"memory_gb": 8, "cpu": 4},
policy=RetryPolicy(max_retries=3),
)
print(f"Result: {result}")
The Full Orchestrator
Now combine caching + scheduling + retries into one pipeline:
# orchestration/engine.py — The full orchestration engine
# Combines: content-addressable caching + heterogeneous scheduling + semantic retries
from datetime import datetime
from pathlib import Path
from typing import Callable
import json
class PipelineStep:
"""
A single step in the pipeline with its resource profile and retry policy.
Each step can be cached, scheduled heterogeneously, and auto-healed.
"""
def __init__(
self,
name: str,
fn: Callable,
input_paths: list[str] = None,
params: dict = None,
cpu: int = 2,
memory_gb: int = 8,
gpus: int = 0,
spot_ok: bool = True,
max_retries: int = 3,
):
self.name = name
self.fn = fn
self.input_paths = input_paths or []
self.params = params or {}
self.cpu = cpu
self.memory_gb = memory_gb
self.gpus = gpus
self.spot_ok = spot_ok
self.max_retries = max_retries
class HyperScalePipeline:
"""
A pipeline with intelligent caching, heterogeneous scheduling, and auto-healing.
Steps declare their resource needs. The engine handles the rest.
Usage:
pipeline = HyperScalePipeline()
pipeline.add_step(PipelineStep("preprocess", preprocess_fn, cpu=4, memory_gb=16))
pipeline.add_step(PipelineStep("train", train_fn, gpus=1, memory_gb=64, spot_ok=False))
pipeline.run()
"""
def __init__(self, name: str = "pipeline", cache_dir: str = "./.cache/pipeline"):
# Pipeline metadata
self.name = name
self.steps: list[PipelineStep] = []
# Caching layer for intelligent skip logic
self.cache = ContentAddressableCache(cache_dir)
# Track which steps ran vs were cached
self.executed_steps = []
def add_step(self, step: PipelineStep):
"""Register a step in the pipeline."""
self.steps.append(step)
def run(self):
"""
Execute the pipeline with full orchestration:
1. Check cache for each step → skip if hash matches
2. Schedule each step to the right compute tier
3. Execute with semantic retries if failure occurs
"""
print(f"\n{'='*50}")
print(f"🚀 Running pipeline: {self.name}")
print(f"Steps: {len(self.steps)}")
print(f"{'='*50}\n")
start_time = datetime.now()
for i, step in enumerate(self.steps, 1):
print(f"\n[{i}/{len(self.steps)}] {step.name}")
# Step 1: Check cache — skip if nothing changed
cached = self.cache.get_cached_output(
step_name=step.name,
input_paths=step.input_paths,
params=step.params,
)
if cached:
# Cache hit — skip execution entirely
self.executed_steps.append({
"step": step.name,
"status": "cached",
"output": str(cached),
})
continue
# Step 2: Schedule resources based on step's profile
profile = ResourceProfile(
cpu=step.cpu,
memory_gb=step.memory_gb,
gpus=step.gpus,
spot_ok=step.spot_ok,
)
tier = profile.recommended_tier()
print(f" Compute tier: {tier.value}")
# Step 3: Execute with semantic retries
policy = RetryPolicy(max_retries=step.max_retries)
result = execute_with_semantic_retry(
step_fn=step.fn,
step_name=step.name,
initial_resources={"cpu": step.cpu, "memory_gb": step.memory_gb, "gpus": step.gpus},
policy=policy,
)
# Step 4: Cache the output for future runs
if result and "output_path" in result:
self.cache.store_output(
step_name=step.name,
input_paths=step.input_paths,
params=step.params,
output_path=Path(result["output_path"]),
)
self.executed_steps.append({
"step": step.name,
"status": "executed",
"result": result,
})
elapsed = (datetime.now() - start_time).total_seconds()
print(f"\n{'='*50}")
print(f"✅ Pipeline '{self.name}' complete in {elapsed:.1f}s")
executed = sum(1 for s in self.executed_steps if s["status"] == "executed")
cached_count = sum(1 for s in self.executed_steps if s["status"] == "cached")
print(f"Steps executed: {executed} | Cached: {cached_count}")
print(f"{'='*50}")
def summary(self) -> dict:
"""Return a summary of the pipeline run for monitoring/auditing."""
return {
"pipeline": self.name,
"total_steps": len(self.steps),
"executed": sum(1 for s in self.executed_steps if s["status"] == "executed"),
"cached": sum(1 for s in self.executed_steps if s["status"] == "cached"),
"steps": self.executed_steps,
}
# Example usage
if __name__ == "__main__":
# Define step functions
def preprocess_data(cpu=2, memory_gb=8, **kwargs):
print(f" Preprocessing with {cpu} CPU cores, {memory_gb}GB RAM...")
# In production: read raw data, clean, transform
return {"output_path": "/tmp/preprocessed.parquet", "rows": 1_000_000}
def feature_engineering(cpu=4, memory_gb=32, **kwargs):
print(f" Feature engineering with {cpu} CPU cores, {memory_gb}GB RAM...")
return {"output_path": "/tmp/features.parquet", "features": 128}
def train_model(cpu=8, memory_gb=64, gpus=1, **kwargs):
print(f" Training with {gpus} GPU(s), {memory_gb}GB RAM...")
return {"accuracy": 0.94, "output_path": "/tmp/model.pt"}
# Build pipeline
pipeline = HyperScalePipeline(name="ml-training-v2", cache_dir="./.cache/pipeline")
pipeline.add_step(PipelineStep(
"preprocess", preprocess_data,
input_paths=["data/raw/transactions.parquet"],
cpu=4, memory_gb=16, spot_ok=True,
))
pipeline.add_step(PipelineStep(
"features", feature_engineering,
input_paths=["data/raw/transactions.parquet"],
cpu=8, memory_gb=64, spot_ok=True,
params={"window": "7d"},
))
pipeline.add_step(PipelineStep(
"train", train_model,
input_paths=["/tmp/preprocessed.parquet", "/tmp/features.parquet"],
cpu=16, memory_gb=128, gpus=2, spot_ok=False,
max_retries=3,
params={"lr": 0.001, "batch_size": 64, "epochs": 10},
))
# Run — second run with unchanged inputs will use cache for all steps
pipeline.run()
print("\n=== Pipeline Summary ===")
print(json.dumps(pipeline.summary(), indent=2))
# Second run — everything should be cached
print("\n\n=== Second Run (should be fully cached) ===")
pipeline2 = HyperScalePipeline(name="ml-training-v2")
pipeline2.add_step(PipelineStep("preprocess", preprocess_data, input_paths=["data/raw/transactions.parquet"]))
pipeline2.add_step(PipelineStep("train", train_model, input_paths=["/tmp/model.pt"], gpus=2, spot_ok=False))
pipeline2.run()
Wann brauchst du das?
| Szenario | Ohne Hyper-Scale | Mit Hyper-Scale |
|---|---|---|
| Du änderst einen Hyperparameter | 12h Pipeline, alles läuft neu | 30min, nur Training läuft neu |
| Preprocessing OOM bei 16GB | Crasht, manueller Restart | Verdoppelt RAM automatisch |
| Training braucht GPU | Wartet auf freie GPU | Wird auf GPU-Cluster geroutet |
| Spot Instance stirbt | Pipeline bricht ab | Wechselt zu On-Demand, läuft weiter |
| 100 Steps parallel | Manuell partitioniert | Auto-scheduled via Bin-Packing |
Die drei Konzepte — Caching, Scheduling, Auto-Healing — machen den Unterschied zwischen einer Pipeline, die ständig Aufmerksamkeit braucht, und einer, die sich selbst verwaltet. Der Code dafür ist nicht riesig. Das meiste ist Logik, die einmal geschrieben und dann vergessen wird.
Comments
Leave a comment