Decouple Serialization and Deserialization Code for tasks

[kaxil]

🎯 Problem Statement

The Task SDK separation in Airflow 3.1 requires decoupling serialization and deserialization code to eliminate server-side dependencies on client SDK implementations:

1. Task SDK Dependencies: airflow-core deserialization currently depends on Task SDK's BaseOperator for default values and field lists
2. Architectural Coupling: Server components import and use Task SDK classes during deserialization, violating client/server separation
3. Independent Deployment Blocker: Tight coupling prevents independent deployment and upgrade of server vs client components

🚀 Solution Overview

This PR decouples (to a great extent) serialization and deserialization code by removing Task SDK dependencies from airflow-core:

• Remove dynamic SDK calls: Replace get_serialized_fields() calls with hardcoded class methods
• Eliminate import dependencies: Remove OPERATOR_DEFAULTS and other Task SDK imports from server-side code
• Schema-driven defaults: Use schema.json and client_defaults instead of Task SDK classes for default resolution
• Independent deployment: Enable server/client components to be deployed and upgraded separately

📊 Benchmark

As part of this change, I optimised how the defaults are stored and when a field is stored and removed anything that matches defaults, nulls and the bigger impact change to remove storing entire callback functions as strings and instead store a boolean to indicate if a callback was set or not.

The bigger the DAG (more tasks + especially with callbacks), the more savings.

Using actual pre-optimization code:

Scenario	Tasks	Before	After	Saved	Reduction
Basic DAGs	10	7.5 KB	5.7 KB	1.8 KB	24.0%
	50	33.7 KB	24.5 KB	9.2 KB	27.3%
	100	66.4 KB	48.0 KB	18.4 KB	27.7%
Production DAGs (3 callbacks/task)	10	15.5 KB	6.6 KB	8.9 KB	57.4%
	50	73.8 KB	28.9 KB	44.9 KB	60.8%
	100	146.7 KB	56.9 KB	89.8 KB	61.2%

🔥 Callback Optimization Analysis (100 tasks with 3 callbacks each):

Storage Method	Before	After	Saved	Reduction
Callback representation	Lists of function code	Boolean flags	-	-
Per-task callback overhead	822 bytes	91 bytes	731 bytes	89%
Total callback overhead	80.3 KB	8.9 KB	71.4 KB	89%
Per-task total size	1,502 bytes	583 bytes	919 bytes	61%

🎯 Key Optimization Impact:

• ✅ Callback transformation: 89% reduction in callback storage overhead (function code → boolean flags)
• ✅ Production scaling: 57% → 61% reduction as DAG size increases (10 → 100 tasks)
• ✅ Per-task efficiency: 919 bytes saved per task (1,502 → 583 bytes) for callback DAGs
• ✅ Consistent baseline: 24-28% reduction even for basic DAGs without callbacks

🏗️ Architecture Changes

Task Default Resolution

Implements hierarchical defaults during deserialization:

1. Schema defaults (from schema.json) - lowest priority
2. client_defaults.tasks - SDK-specific overrides
3. partial_kwargs - MappedOperator values
4. Explicit task values - highest priority

Serialization Exclusion

Fields matching client_defaults are automatically excluded from task serialization, reducing redundancy while maintaining full information.

fyi: Following the Task Execution API pattern, I aim to add versioned schema contract at Airflow website directly or version docs soon'ish:

Thinking about a URL like: https://airflow.apache.org/schemas/dag-serialization/v2.json

🚦 Migration Path

For Users

• No action required - changes are completely transparent
• Existing DAGs continue working unchanged
• New DAGs automatically benefit from optimizations

---

Appendix (for my own tracking)

TODOs (some might be done in a future PR):

• Add defaults to schema.json
• Exclude defaults in schema from Serialized JSON
• Change on_*_callback on tasks to use has_on_*_callback
• Remove unmap method from scheduler-side #54816
• Implement client_defaults generation in serialization (Task SDK side)
• Verify if the change is backwards compatible. If not, Update serialization version to v3 and add backwards compatibility for v2. Update: It is a backwards-compatible change
• Add tests to ensure defaults in Schema are same as the Server side classes. Or better add prek/pre-commit to autogenerate default from server-side to Schema
• Evaluate the alternative of storing the list of attributes needed for Serialization & De-serialization in schema.JSON

Future Work:

• Move the S10n code over for DAG & Task Group classes to Server-side
• Include schema.json in the calver OpenAPI spec for Execution API and/or in airflow versioned docs
• Move Serialization code to Task SDK
• Remove ui_color & ui_fgcolor

Other points

• If we move serialization to Task SDK and keep de-serializtion to the Server side, how do we handle the following:
  • XCom deserialization -- it currently uses the airflow.serialization module
  • ExtendedJSON - TypeDecorator used in serialization of the following:
    • DagRun.context_carrier
    • TaskInstance.next_kwargs

Benchmark script:

#!/usr/bin/env python3
"""
Script to measure real serialization scaling for different DAG sizes.
Uses actual JSON examination for accurate before/after comparison.
"""

import json
import sys
from datetime import datetime, timedelta
from pathlib import Path

from airflow.sdk import DAG
from airflow.providers.standard.operators.bash import BashOperator
from airflow.serialization.serialized_objects import SerializedDAG

def success_callback(context):
    """Example success callback function."""
    print(f"Task {context['task_instance'].task_id} succeeded!")
    return "success"

def failure_callback(context):
    """Example failure callback function."""
    print(f"Task {context['task_instance'].task_id} failed!")
    # Send notification to Slack
    import requests
    requests.post("https://hooks.slack.com/webhook", json={
        "text": f"❌ Task failed: {context['task_instance'].task_id}"
    })
    return "failure_handled"

def retry_callback(context):
    """Example retry callback function."""
    print(f"Task {context['task_instance'].task_id} will retry!")
    return "retry_scheduled"

def create_test_dag(num_tasks: int, with_callbacks: bool = True) -> DAG:
    """Create a test DAG with specified number of tasks."""
    dag = DAG(
        dag_id=f"scaling_test_dag_{num_tasks}",
        start_date=datetime(2024, 1, 1),
        schedule="@daily",
        default_args={
            "owner": "test_user",
            "retries": 2,
            "retry_delay": timedelta(minutes=5),
            "email": "test@example.com",
            "email_on_failure": True,
            "email_on_retry": True,
        }
    )
    
    with dag:
        for i in range(num_tasks):
            task_kwargs = {
                "task_id": f"task_{i:03d}",
                "bash_command": f"echo 'Processing item {i}'",
                "email_on_failure": True,
                "email_on_retry": True,
            }
            
            if with_callbacks:
                task_kwargs.update({
                    "on_success_callback": success_callback,
                    "on_failure_callback": failure_callback,
                    "on_retry_callback": retry_callback,
                })
            
            BashOperator(**task_kwargs)
    
    return dag

def measure_dag_serialization(num_tasks: int, with_callbacks: bool = True) -> dict:
    """Measure serialization for a DAG with specified parameters."""
    dag = create_test_dag(num_tasks, with_callbacks)
    
    try:
        serialized = SerializedDAG.to_dict(dag)
    except Exception as e:
        return {"error": str(e)}
    
    # Convert to compact JSON
    json_compact = json.dumps(serialized, separators=(',', ':'))
    total_size = len(json_compact.encode('utf-8'))
    
    # Analyze callback fields if callbacks are enabled
    callback_info = {}
    if with_callbacks and "dag" in serialized and "tasks" in serialized["dag"]:
        tasks = serialized["dag"]["tasks"]
        if tasks:
            first_task = tasks[0]["__var"]
            
            # Find callback fields
            callback_fields = [k for k in first_task if "callback" in k.lower()]
            
            # Calculate callback overhead
            task_json = json.dumps(first_task, separators=(',', ':'))
            total_task_size = len(task_json.encode('utf-8'))
            
            task_without_callbacks = {k: v for k, v in first_task.items() if "callback" not in k.lower()}
            no_callback_json = json.dumps(task_without_callbacks, separators=(',', ':'))
            no_callback_size = len(no_callback_json.encode('utf-8'))
            
            callback_overhead_per_task = total_task_size - no_callback_size
            total_callback_overhead = callback_overhead_per_task * len(tasks)
            
            callback_info = {
                "callback_fields": callback_fields,
                "callback_overhead_per_task": callback_overhead_per_task,
                "total_callback_overhead": total_callback_overhead,
                "task_size": total_task_size,
                "task_size_without_callbacks": no_callback_size
            }
    
    return {
        "num_tasks": num_tasks,
        "with_callbacks": with_callbacks,
        "total_size": total_size,
        "size_kb": total_size / 1024,
        "per_task_bytes": total_size / num_tasks,
        "callback_info": callback_info
    }

def run_scaling_measurements():
    """Run measurements for different DAG sizes."""
    print("🔍 Measuring Real Serialization Scaling")
    print("=" * 60)
    
    task_counts = [10, 25, 50, 100]
    
    # Test with callbacks
    print("\n📊 **Production DAGs (With Callbacks)**")
    print("| Tasks | Size | Per Task | Callback Overhead |")
    print("|-------|------|----------|-------------------|")
    
    callback_results = []
    for num_tasks in task_counts:
        result = measure_dag_serialization(num_tasks, with_callbacks=True)
        
        if "error" in result:
            print(f"| {num_tasks} | ERROR: {result['error']} |")
            continue
            
        callback_results.append(result)
        
        callback_overhead_kb = result["callback_info"]["total_callback_overhead"] / 1024 if result["callback_info"] else 0
        
        print(f"| {num_tasks} | {result['size_kb']:.1f} KB | {result['per_task_bytes']:.0f} bytes | {callback_overhead_kb:.1f} KB |")
    
    # Test without callbacks
    print("\n📊 **Basic DAGs (No Callbacks)**")
    print("| Tasks | Size | Per Task |")
    print("|-------|------|----------|")
    
    basic_results = []
    for num_tasks in task_counts:
        result = measure_dag_serialization(num_tasks, with_callbacks=False)
        
        if "error" in result:
            print(f"| {num_tasks} | ERROR: {result['error']} |")
            continue
            
        basic_results.append(result)
        print(f"| {num_tasks} | {result['size_kb']:.1f} KB | {result['per_task_bytes']:.0f} bytes |")
    
    # Show callback field details for largest DAG
    if callback_results:
        largest = callback_results[-1]
        if largest["callback_info"]:
            print(f"\n🔍 **Callback Analysis (100 tasks):**")
            ci = largest["callback_info"]
            print(f"   Callback fields: {ci['callback_fields']}")
            print(f"   Per-task callback overhead: {ci['callback_overhead_per_task']} bytes")
            print(f"   Total callback overhead: {ci['total_callback_overhead']:,} bytes ({ci['total_callback_overhead']/1024:.1f} KB)")
            print(f"   Task size with callbacks: {ci['task_size']} bytes")
            print(f"   Task size without callbacks: {ci['task_size_without_callbacks']} bytes")
    
    # Generate CSV for easy import
    print(f"\n📋 **CSV Format (for PR description):**")
    print("# callbacks")
    for result in callback_results:
        print(f"{result['num_tasks']},{result['size_kb']:.1f},{result['per_task_bytes']:.0f},True")
    
    print("# basic")  
    for result in basic_results:
        print(f"{result['num_tasks']},{result['size_kb']:.1f},{result['per_task_bytes']:.0f},False")
    
    return callback_results, basic_results

def main():
    """Main function."""
    run_scaling_measurements()

if __name__ == "__main__":
    main()