How to solve Airflow local development problems, testing challenges, and complex setup requirements with a simpler alternative

Setting up Apache Airflow 3 for local development requires at least 4 separate services, a minimum of 4GB RAM (8GB recommended), and setup times that can stretch from hours to days — before writing a single line of workflow code. If you’re searching for “Airflow local development setup,” “Apache Airflow testing problems,” or “Airflow development environment issues,” you’re not alone.

Local development should be straightforward: write code, test it, iterate quickly. However, setting up Airflow 3 locally reveals significant architectural complexity that affects every developer who works with the framework. Despite being the biggest release in Airflow’s history, Airflow 3 maintains design decisions that prioritize production deployment over development experience.

The challenges are systematic rather than incidental. Airflow was built for ops teams managing static workflows on dedicated servers, not for developers who want to iterate quickly. Prefect, on the other hand, was built with modern expectations in mind: teams routinely report 3x improvements in development speed and over 50% reduction in cost when migrating from Airflow to Prefect.

Airflow Local Setup Requirements: The Multi-Component Infrastructure and Configuration Problem

Airflow 3 requires multiple services for even basic local development, each demanding careful configuration. The official Airflow documentation lists the required components:

• Webserver: For the UI interface
• Scheduler: Responsible for triggering DAG runs and submitting tasks to executors
• DAG Processor: A separate process for parsing DAG files
• Database: PostgreSQL or similar for metadata storage

The quickstart documentation attempts to simplify this with airflow standalone, but this approach has limitations for realistic development scenarios. When developers need to test concurrent runs, external dependencies, or realistic configurations, they must configure the full multi-component setup through extensive airflow.cfg files, environment variables, and command-line parameters.

The official Docker Compose setup demonstrates the true scope of this complexity. The docker-compose.yaml file requires multiple separate containers with careful orchestration of startup dependencies, plus configuration management for database connections, executor selection (LocalExecutor, CeleryExecutor, or KubernetesExecutor), resource limits, timeouts, parallelism settings, and security management. The documentation recommends allocating at least 4GB of memory to Docker, with 8GB recommended for stability. Compare this with the far simpler Prefect docker compose setup that supports production work.

Because of this complexity, small configuration mistakes in Airflow can lead to silent failures, tasks hanging indefinitely, or error messages that require deep framework knowledge to decode. This configuration sprawl compounds issues when developers are attempting to debug issues with their production Airflow DAGs locally.

Much of this infrastructure complexity exists because DAGs in Airflow are not functions: they are complicated objects that the Airflow scheduler must parse and execute regularly, requiring persistent services and extensive configuration to coordinate their interactions.

In Prefect, by contrast, flows can be executed directly as functions within unit test suites, tasks can be executed individually outside of their flows, and deployments can be served locally without spinning up any Prefect services. Even schedule objects can be imported and unit tested.

Simple Alternative: Zero Infrastructure Development

Prefect eliminates these problems entirely: no scheduler, database, webserver, or Docker containers are required. Crucially, the exact same code path executes identically in local development and remote execution environments, eliminating the debugging complexity that comes from architectural differences between development and deployment. Workflows run as standard Python processes with zero additional infrastructure requirements and no configuration files, allowing you to focus on workflow logic rather than service orchestration:

# /// script
# dependencies = ["prefect"]
# ///

from prefect import flow, task


@task
def extract_data():
    return {"key": "value"}


@task
def transform_data(data):
    return data["key"].upper()


@flow
def my_workflow():
    data = extract_data()
    result = transform_data(data)
    return result


if __name__ == "__main__":
    my_workflow()

💡 Try it now: Save this code as workflow.py and run uv run workflow.py - no setup, configuration, or infrastructure required!

Airflow Docker Development Issues: Windows Compatibility and Resource Problems

For Windows users these problems are further exacerbated. Airflow treats Windows as an afterthought and requires Linux-based distributions for production use. Because of this, Windows users face particular challenges as noted in community discussions: “running Airflow in Windows natively is dead in the water”.

Common Airflow Docker development issues include:

• Windows compatibility: Airflow has limited Windows support, forcing Windows developers to use WSL or Docker with their own complications
• Resource consumption: Docker setups often consume more resources than the actual work being performed
• Development workflow friction: Mounting volumes, managing container lifecycle, and debugging across containerized services adds overhead
• Network complexity: Getting proper communication between scheduler, webserver, and database requires careful configuration
• Iteration speed: Code changes require container rebuilds or volume sync delays

The Docker approach introduces architectural differences between local development environments and typical deployment scenarios, making it difficult to debug environment-specific issues.

Prefect runs natively on all platforms, including full Windows support. No Docker, WSL, or virtual machines are required for development.

Airflow Task Isolation Problems: Development Friction and Debugging Challenges

While separating orchestration from business logic is critical for workflow systems, Airflow 3 takes this to an extreme that complicates development. It encourages submitting every task through an executor, adding infrastructure overhead even for simple development scenarios. Moreover, users cannot submit tasks directly themselves — the DAG processor and scheduler must parse the code for submission on their respective loops.

Airflow’s task isolation means:

• Executor dependency: Even basic testing requires configuring and running an executor
• Process boundaries: Tasks run in separate processes or containers, making debugging complex
• Resource overhead: Each task submission involves scheduler communication and process spawning
• Development friction: Simple changes require full infrastructure restarts

Prefect balances isolation with development simplicity. Tasks can run in the same process for development and testing, while still supporting distributed execution when needed. The isolation level is configurable rather than mandatory.

Additionally, Prefect supports direct submission to infrastructure without managing persistent processes. Infrastructure is provisioned on-demand, executes the workflow, and tears down automatically.

Configuration-Free Development Alternative

Prefect requires zero configuration for local development. It uses sensible defaults and configuration-as-code patterns that eliminate setup overhead. The same code that runs locally executes identically in remote environments, ensuring consistent behavior across development and deployment contexts.

As an example, the following code will reliably execute the simple workflow on a cron schedule with zero configuration:

# /// script
# dependencies = ["prefect"]
# ///

from prefect import flow


@flow(log_prints=True)
def hello_world():
    print('hello world!')


if __name__ == "__main__":
    hello_world.serve(name="example-deployment", cron="* * * * *")

💡 Try it now: Save this as scheduled_workflow.py and run uv run scheduled_workflow.py - it will start running on a cron schedule immediately!

Airflow Parameter Limitations: Dynamic Workflow and Testing Constraints

Airflow 3’s approach to workflow parametrization reveals a fundamental design limitation that affects both development and testing. While Airflow provides DAG Params, the implementation has significant constraints.

Airflow’s parametrization challenges include:

• Parse-time limitations: DAG parameters cannot be used in the DAG body during parsing, as demonstrated in community discussions. Parameters are only available at runtime through templating.
• Limited dynamic behavior: Parameters cannot drive dynamic task generation or DAG structure changes, requiring complex workarounds.
• Testing difficulties: The parametrization model makes dependency injection patterns for testing nearly impossible.

These limitations make Airflow parameters unsuitable for the dependency injection patterns common in modern software development. Testing scenarios that require different parameter values must rely on external configuration systems or complex templating approaches.

Flexible Parametrization Alternative

Prefect’s parametrization enables proper dependency injection for both runtime flexibility and comprehensive testing. Parameters can drive workflow logic, task generation, and even dynamically chosen execution paths, making workflows truly testable and maintainable. Additionally, parametrization works identically in local development and remote execution, eliminating another source of environment-specific complexity.

# /// script
# dependencies = ["prefect"]
# ///

from prefect import flow, task
from typing import List


@task
def process_item(item: str, multiplier: int = 1) -> str:
    return f"Processed {item} x{multiplier}"


@flow
def dynamic_workflow(items: List[str], processing_multiplier: int = 2):
    results = []
    # Parameters can drive dynamic task generation
    for item in items:
        result = process_item(item, processing_multiplier)
        results.append(result)
    return results


if __name__ == "__main__":
    # Test with different parameters easily
    test_result = dynamic_workflow(["apple", "banana"], processing_multiplier=1)
    print("Test result:", test_result)
    # Run with production parameters
    prod_result = dynamic_workflow(["data1", "data2", "data3"], processing_multiplier=5)
    print("Production result:", prod_result)

💡 Try it now: Save as dynamic_workflow.py and run uv run dynamic_workflow.py - see how parameters drive workflow behavior instantly!

Solutions for Airflow Development Problems

The challenges outlined above reflect fundamental architectural decisions that affect every developer working with Airflow. While the community continues to develop workarounds and improvements, these solutions often add additional complexity rather than addressing the root causes.

For teams experiencing persistent productivity issues with local development, testing, and debugging workflows, Prefect offers a fundamentally different approach that prioritizes developer experience. Recent migration stories reveal dramatic improvements: after switching to Prefect, Endpoint increased their turnaround time by 3x while also achieving a 73% reduction in orchestration costs. LiveEO reported similar results — they tripled their development speed and cut AWS costs by 63%.

Conclusion: Streamlining Workflow Development

Airflow 3’s approach to local development reflects its design heritage as a production-focused batch processing system. While version 3 introduces significant architectural improvements, it maintains the assumption that developers will work within complex infrastructure environments.

The architectural differences create measurably different development experiences. Time spent managing infrastructure and debugging environment issues becomes time available for actual development work. When infrastructure complexity is removed from the development process, teams can focus on building sophisticated data applications rather than managing deployment environments.

Prefect’s architecture prioritizes developer experience as a primary design consideration. The result is faster iteration cycles, higher code quality, and development teams that can concentrate on business logic rather than operational concerns. Most importantly, the identical execution paths between local development and remote deployment eliminate a major source of environment-specific bugs and debugging complexity.

Ready to experience streamlined workflow development? Try any of the code examples above with uv run filename.py and see the difference for yourself - no configuration, setup, or infrastructure required.

To learn more about Prefect:

• visit our website
• visit us on GitHub to open issues and pull requests
• follow us on LinkedIn for regular updates
• join our active Slack community

Happy Engineering!

Airflow Local Development Sucks was originally published in The Prefect Blog on Medium, where people are continuing the conversation by highlighting and responding to this story.

Data teams have embraced Kubernetes for many reasons. It provides dynamic resource allocation for workloads that swing from lightweight data ingestion to GPU-hungry ML training. Its native support for auto-scaling matches infrastructure to demand. Declarative configuration codifies requirements alongside pipelines and multiple teams can share infrastructure safely with workload isolation.

But there’s a fundamental economics problem lurking beneath this technical elegance: the same dynamic allocation that makes Kubernetes perfect for data workloads also makes it expensive. Your ETL pipeline needs 16 cores for 2 hours daily but pays for them 24/7. Your ML training job scales to 100 nodes for a few hours weekly, burning budget on unused capacity the other 160+ hours.

This is precisely why spot instances exist — and why they’re uniquely suited to data workloads running on Kubernetes.

The Spot Instance Opportunity

Spot instances represent cloud providers’ excess capacity, offered at steep discounts when demand is low. The numbers are compelling: using Amazon EC2 Spot Instances can result in up to a 90% discount compared to on-demand prices, with typical savings of 60–75% according to AWS’s Spot Instance Advisor. Google Cloud’s Spot VMs offer similar discounts of 60–91% off standard on-demand prices for most machine types.

For data teams running substantial compute workloads — ML training, ETL pipelines, analytical queries — these savings compound fast. A team spending $50,000 monthly on compute could dramatically reduce costs down to $5,000-$15,000 by migrating appropriate workloads to spot instances.

The economics work because cloud providers need spare capacity for demand spikes but don’t want it sitting idle. Instead of wasting this capacity, they offer it at reduced rates with the caveat that it can be reclaimed with minimal notice when needed elsewhere. AWS provides a two-minute interruption notice, while GCP typically provides one minute notice.

The State Management Problem

Spot instances create a tension for data workflows that doesn’t exist for typical application workloads. When your web server gets terminated, users might see a brief error, but the request can be retried elsewhere, as there is no piece of state that ties an individual request to a specific container instance. When a batch processing job gets terminated halfway through processing 10TB of customer transaction data, the consequences are very different.

Unlike stateless web applications that can be killed and restarted elsewhere with minimal impact, data processing jobs are inherently stateful and often tied to specific business outcomes. Consider a typical ETL pipeline that processes customer transactions:

• Pipeline kickoff requires precise parametrized inputs
• Loading data from multiple sources, possibly using costly queries
• Transforming and validating the data through several interdependent stages
• Storing results to a data warehouse that downstream systems depend on
• Triggering downstream reporting processes that have SLA requirements

If this pipeline is interrupted midway through processing, the naive approach of “just restart everything” can lead to duplicate processing if your pipeline is not idempotent. This results in inconsistent states and the kind of data quality issues that keep data engineers awake at night.

Traditional Kubernetes approaches handle spot instance interruptions through basic restart mechanisms, but these aren’t designed for the complex state management that data workflows require. A pod that gets terminated during a critical transformation step may restart from the beginning, losing hours of computation and risks data inconsistencies that propagate through your entire data ecosystem.

This is where workflow orchestrators come into the picture. Workflow orchestration is fundamentally about managing workflow state, from scheduling and triggering rules to intermediate checkpoints that turn fragile workflows into durable multi-step processes.

How Prefect enables Spot Instances

Prefect’s architecture makes it uniquely suited for running data-intensive workloads on spot instances because it was designed from the ground up with failure recovery and dynamic infrastructure as core principles.

Expressive Task Caching for Idempotent Pipelines

When your pipeline gets interrupted halfway through processing — whether from spot instance termination or any other failure — you need to resume from the last successful task rather than starting over from scratch. Authoring truly idempotent code is difficult, but appropriate caching can give your workflows functional idempotency.

Sophistication here lies in configurability: Prefect offers many dimensions for expressing how workflow tasks should cache their results, and when:

• Cache Policies allow you to define precisely what inputs (if any) determine whether tasks use cached results or execute fresh computation.
• Result Storage controls how and where task results are serialized and stored. This is critical for spot instance deployments because your cache needs to survive infrastructure failures.
• Transactions provide the ability to group tasks as atomic units. This prevents the scenario where some tasks in a dependent chain use cached results while others don’t, creating inconsistent data states.

This level of caching sophistication is what enables workflow restarts to occur from meaningful checkpoints rather than complete do-overs. Instead of a 3-hour pipeline restarting from zero at hour 2.5, it resumes from checkooint and finishes in 30 minutes.

Infrastructure-Aware Orchestration

Most workflow orchestrators treat infrastructure as a black box — they submit jobs and at best retry them wholesale when they fail. Prefect’s Kubernetes worker, by contrast, understands the infrastructure it’s running on and can react intelligently to the specific failure modes that spot instances introduce.

When a pod receives SIGTERM (imminent termination), Prefect handles the signal and reschedules the workflow run with full context preservation. The worker is aware of each job’s backoff limit, and either manages rescheduling the work through Prefect’s backend — allowing any other Kubernetes worker to pick it up — or by letting Kubernetes reschedule the pod and automatically monitoring the new pod. Both patterns ensure no work is lost and the workflow continues seamlessly using task caches where necessary.

Stateless Worker Design for Maximum Resilience

The Prefect Kubernetes worker is stateless by design. Rather than storing run state in memory (where it can be lost during worker restarts or spot interruptions), Prefect stores all necessary run metadata in job and pod labels within the Kubernetes cluster itself. This metadata propagates to the underlying run, where the Prefect client can access task caches and the Prefect API for more granular state information.

Workers can be interrupted and restarted without losing track of running workflows. A terminated worker can restart on different infrastructure and immediately resume orchestrating flows based on the state recorded in the cluster. This resilience is crucial for spot instance deployments, where even the worker’s infrastructure can be reclaimed with minimal notice.

Dynamic Infrastructure Configuration

Unlike traditional orchestrators that require infrastructure to be defined statically at deployment time, Prefect allows for per-run infrastructure configuration.

This capability is useful for spot instance strategies because it enables your data engineers to:

• Specify different instance types for different stages of a pipeline based on current spot pricing
• Automatically fall back to on-demand instances for time-sensitive workloads when spot capacity is limited
• Dynamically adjust resource requirements based on data volume and available spot instance types

This flexibility is useful for production spot instance deployments, where available instance types and pricing can vary significantly across availability zones and over time. Recent analysis shows that spot instance prices have increased by 21% on AWS and decreased by 26% on GCP over the past year, with significant regional variations. A static infrastructure configuration can’t adapt to this pricing volatility, but Prefect’s dynamic approach can.

The Economics of Reliability

Prefect’s robust failure handling and spot instance cost savings creates a compelling economic argument beyond simple cost reduction. Consider a typical data team scenario using D-type ECS instances:

• Traditional approach: $50,000/month on on-demand instances, no rescheduling of work
• Prefect + Spot instances: $12,500/month, some rescheduling of work

The $37,500 monthly savings enables teams to invest in additional tooling, hire more engineers, or expand their data infrastructure significantly, going well beyond simple cost minimization. The slight reliability trade-off is usually acceptable for most data workloads, especially when Prefect’s durability mechanisms ensure that jobs will complete successfully.

The traditional approach to data infrastructure often involves significant over-provisioning to handle peak workloads, leading to resource utilization rates of 30–40%. Spot instances with intelligent orchestration can achieve utilization rates of 80–90% while maintaining reliability, materially changing the economics of data processing at scale.

The Broader Implications

The combination of Kubernetes spot instances and intelligent orchestration can drive a significant shift in your team’s data infrastructure economics. When compute costs drop by 80–90%, new possibilities emerge — experimental workloads, deeper analytics, and more frequent model training all become possible within existing budget constraints.

The technical sophistication required to make spot instances work reliably for data workloads is considerable, but the right orchestration platform can abstract this complexity while providing the control and visibility that data engineers need.

While over 20% of Prefect Cloud users currently run workflows in Kubernetes, this percentage has been slowly declining as teams migrate toward even more dynamic infrastructure patterns. Fully serverless compute platforms like AWS Fargate and Google Cloud Run are becoming more and more common, along with emerging players like Modal and Coiled. Even as the ecosystem evolves toward serverless architectures, the cost optimization and durability strategies explored here remain relevant across different compute paradigms. For teams running substantial workloads, investigating this approach allows teams to get the most value out of their compute.

Further Reading

For more information check out Prefect’s Kubernetes documentation, visit us on GitHub, or start deploying Prefect flows for free.

Stop Paying 24/7 for Batch Pipeline Compute was originally published in The Prefect Blog on Medium, where people are continuing the conversation by highlighting and responding to this story.

Prefect is built on a lot of great ideas, but even more importantly, it’s built on great tools. This post is part of a series of love letters to the libraries and technologies that make Prefect possible.

We’re kicking things off with AnyIO — the async concurrency library that brought sanity, structure, and elegance to our workflow engine. If you’ve ever tried to wrangle asyncio and lived to tell the tale, you’ll understand why this one deserves a bit more than just a passing mention. This is our thank-you, our technical deep dive, and maybe just a little bit of a love letter.

In large-scale data processing and workflow orchestration, concurrency is crucial for keeping systems responsive and utilizing resources efficiently. Yet managing concurrency with raw asyncio can quickly turn into a mind-bending maze of callbacks, race conditions, and partial cancellations. That’s where AnyIO comes in, bridging Python’s main async frameworks under one carefully designed API — and freeing developers (us included!) to focus on building reliable, high-scale applications.

Below, we’ll explore why Prefect relies so heavily on AnyIO, show how we use it to solve real-world concurrency challenges, and share the lessons we’ve learned that might save you from your own async pitfalls. Let’s dive in.

The Challenge of Asynchronous Programming in Python

Asynchronous programming has become indispensable for building responsive, scalable systems. By allowing tasks to run concurrently — whether you’re dealing with heavy network I/O, CPU-bound computations, or file reads — you can make sure your program isn’t just sitting idle waiting for one operation to finish. In other words, async code helps you squeeze the most out of your CPU cycles, memory, and network connections so your application can handle more work without grinding to a halt.

Python’s async journey started with asyncio in Python 3.4, standardizing cooperative multitasking but also introducing a fair bit of complexity. That complexity paved the way for Trio, designed with a focus on simplicity, better debugging, and correctness. However, it also meant developers had to pick sides between asyncio and Trio, leaving the ecosystem fragmented.

In short, as the ecosystem evolved, fragmentation and complexity became the norm. So how did AnyIO unify the landscape? Let’s look under the hood.

What is AnyIO? A Technical Overview

AnyIO is an asynchronous concurrency library that lets you seamlessly write code compatible with Python’s two main async frameworks: asyncio and trio. But it’s more than just a compatibility layer — it’s designed to blend the best ideas from both frameworks into a clear, unified API that smooths out their quirks and differences.

Here’s why this matters in practice:

import asyncio

async def asyncio_example():
    await asyncio.sleep(1)
    return "completed with asyncio"

# trio version
import trio

async def trio_example():
    await trio.sleep(1)
    return "completed with trio"

But with AnyIO, you only write your async code once and run it anywhere:

import anyio

async def anyio_example():
    await anyio.sleep(1)
    return "completed with AnyIO"

# Run your code with the asyncio backend
anyio.run(anyio_example, backend="asyncio")

# Or swap to trio just as easily
anyio.run(anyio_example, backend="trio")

Beyond just simplifying your choices, AnyIO provides clean and consistent interfaces for common async needs:

• Task Groups: Manage groups of tasks easily, with automatic error handling and cleanup.
• Cancellation Scopes: Cancel tasks gracefully and reliably, ensuring resources are cleaned up properly.
• Networking and Streams: Use the same APIs for TCP/UDP sockets, pipes, and more across both frameworks.
• Synchronization Tools: Get locks, events, semaphores, and conditions — everything you need to coordinate tasks without headaches.
• File Operations: Perform file reads and writes asynchronously without blocking your event loop.
• Subprocesses: Run external commands efficiently within your async workflows.

The beauty of AnyIO is how it guides you toward structured concurrency — helping you write async code that’s predictable, easy to reason about, and resistant to subtle bugs like resource leaks or deadlocks.

AnyIO’s Position in the Python Async Ecosystem

AnyIO sits at a unique intersection in Python’s async ecosystem. Unlike most async libraries, which typically commit exclusively to either asyncio or trio, AnyIO bridges the gap by offering a universal abstraction that works seamlessly with both frameworks. This flexibility becomes even more crucial as the ecosystem grows and developers seek solutions that aren’t locked into one specific tool.

AnyIO shines especially bright for middleware and framework developers, enabling them to build software that remains agnostic to their users’ async framework preferences. This versatility is why many prominent projects have embraced AnyIO:

• Starlette: To power backend-agnostic web services.
• Prefect: For robust, async-powered workflow orchestration.
• OpenAI, Anthropic, and Google Gemini Python SDKs: For reliable async file uploads in interactions with large language models.
• DsPY: To efficiently handle concurrent calls to LLMs while managing resources effectively.

By providing a unified and easy-to-use abstraction, AnyIO helps maintain cohesion and compatibility across the diverse landscape of Python’s asynchronous ecosystem.

But who is behind AnyIO’s success? Naturally, you might wonder who’s guiding and shaping its evolution.

The Minds Behind AnyIO: Maintainers and Community

AnyIO is primarily maintained by Alex Grönholm, an experienced software developer deeply involved in Python’s async community. Alex also created asphalt, an application framework, and actively contributed to the ASGI specification, which significantly shaped Python web frameworks.

The project thrives on contributions from developers throughout the Python community, including core maintainers from both asyncio and trio. This collaborative approach ensures that AnyIO embodies a consensus on asynchronous programming best practices rather than reflecting just one developer’s vision.

Why Prefect Chose AnyIO: Technical Decision Points

At Prefect, concurrency isn’t just a feature — it’s at the heart of our workflow engine. Here’s what we needed:

1. Parallel Task Execution: Running many tasks simultaneously while respecting their dependencies.
2. Graceful Cancellation: Easily and reliably stopping tasks when users halt their workflows.
3. Reliable Timeouts: Preventing runaway tasks from consuming resources indefinitely.
4. Clear, Debuggable Execution Paths: Ensuring transparency and ease of debugging in complex concurrent scenarios.

We chose AnyIO for a few key reasons:

• Structured Concurrency: AnyIO’s structured concurrency model clearly defines task lifetimes and relationships. This aligns perfectly with our workflows, ensuring reliable cleanup even when things go wrong.
• Backend Flexibility: Currently, we rely on asyncio, but the option to seamlessly switch to trio without rewriting code gives us strategic flexibility. It also allows us to leverage libraries from either ecosystem.
• Robust Cancellation and Timeout Handling: AnyIO’s powerful primitives for handling cancellations and timeouts surpass raw asyncio capabilities, essential for a workflow system where task cancellation needs to be dependable.
• Clean and Thoughtful API Design: AnyIO’s APIs are carefully designed to avoid common async pitfalls. Its task groups, for example, handle exceptions cleanly and prevent resource leaks — critical features for running user-defined workflows safely and reliably.

AnyIO in Action: How Prefect Leverages the Library

Below, we’ll show two major ways AnyIO underpins Prefect’s concurrency story. First, let’s look at how Prefect’s worker architecture uses structured concurrency for reliable task execution. Then we’ll see how capacity limiters solve a common real-world problem — SQLite concurrency issues.

Flow Run Orchestration Using TaskGroups

In Prefect, workers form the backbone of our orchestration lifecycle, continuously polling for scheduled workflows and delegating their execution. At the core of this lifecycle is structured concurrency, a powerful pattern implemented via AnyIO’s TaskGroups.

Structured concurrency explicitly models task dependencies, creating a clear hierarchical structure of upstream and downstream relationships. This ensures that cancellations and errors propagate swiftly and coherently through the execution graph, gracefully terminating affected tasks while maintaining system integrity.

Early deployments revealed several concurrency challenges, including deadlocks and cancellation race conditions. TaskGroups addressed these issues elegantly by offering predictable cleanup, clear error propagation, resource management, and simpler concurrency overall.

Here’s how we integrate TaskGroups into our BaseWorker class:

async def setup(self) -> None:
    """Prepares the worker to run."""
    self._logger.debug("Setting up worker...")
    self._runs_task_group = anyio.create_task_group()

    # ... other setup code ...

    await self._exit_stack.enter_async_context(self._runs_task_group)

    # ... additional setup code ...

    self.is_setup = True

By initializing a task group within an exit stack during setup, we guarantee all tasks within the group are properly awaited or canceled upon teardown, eliminating orphaned processes.

When submitting flow runs concurrently, we use TaskGroup’s start_soon method:

# From the _submit_scheduled_flow_runs method
self._runs_task_group.start_soon(self._submit_run, flow_run)

For tasks returning values, we prefer the .start() method:

# From the _submit_run method
readiness_result = await self._runs_task_group.start(self._submit_run_and_capture_errors, flow_run)

Handling cancellations cleanly was among our most challenging problems. With structured concurrency, cancellations naturally propagate through task hierarchies. We use AnyIO’s cancellation scopes to define precise boundaries:

async def _schedule_task(
    self, __in_seconds: int, fn: Callable[..., Any], *args: Any, **kwargs: Any
):
    """Schedule a background task to start after a delay."""
    async def wrapper(task_status: anyio.abc.TaskStatus[Any]):
        if self.is_setup:
            with anyio.CancelScope() as scope:
                self._scheduled_task_scopes.add(scope)
                task_status.started()
                await anyio.sleep(__in_seconds)
                self._scheduled_task_scopes.remove(scope)
        else:
            task_status.started()

        result = fn(*args, **kwargs)
        if asyncio.iscoroutine(result):
            await result

    await self._runs_task_group.start(wrapper)

During teardown, cancellation is straightforward:

async def teardown(self, *exc_info: Any) -> None:
    """Cleans up resources after the worker is stopped."""
    self._logger.debug("Tearing down worker...")
    self.is_setup = False
    for scope in self._scheduled_task_scopes:
        scope.cancel()

    # ... other cleanup code ...

In a production setting, dozens or hundreds of tasks can spawn simultaneously. AnyIO’s structured concurrency ensures tasks are never orphaned if the parent is canceled. The predictable cleanup model spares developers from headaches like race conditions and stuck processes.

Database Protection via Capacity Limiters

Beyond orchestrating flow runs, concurrency also impacts our database usage. Prefect uses a database to track workflow states, initially prioritizing SQLite for simplicity. However, SQLite struggles with concurrent writes, quickly leading to “database is locked” errors under high concurrency.

In real-world scenarios, these resource bottlenecks can cripple applications. To protect SQLite from overload, we introduced a simple yet effective middleware using AnyIO’s CapacityLimiter:

class RequestLimitMiddleware:
    """
    Middleware limiting concurrent API requests to protect SQLite.
    """

    def __init__(self, app: Any, limit: float):
        self.app = app
        self._limiter = anyio.CapacityLimiter(limit)

    async def __call__(self, scope: Any, receive: Any, send: Any) -> None:
        async with self._limiter:
            await self.app(scope, receive, send)

This middleware ensures:

• A controlled maximum number of concurrent operations.
• Fair request handling through FIFO queuing.
• Requests waiting in a queue if the concurrency limit is reached, resuming as soon as capacity becomes available.

The result? SQLite remains responsive, and lock errors are drastically reduced. This simple pattern can work for any shared resource — like an external API — where too many concurrent requests can lead to throttling, rate limits, or outright failures.

Performance and Reliability Benefits for Prefect Users

From a performance perspective, AnyIO adds minimal overhead compared to direct asyncio usage, typically less than 1ms per operation. The library is also memory-efficient, avoiding unnecessary allocations and object churn that can impact systems running thousands of concurrent workflows.

Looking Forward: AnyIO’s Future and Its Impact on Prefect

The future of AnyIO looks promising. The library continues to mature, with recent improvements focused on:

• Better memory efficiency for large-scale concurrent operations.
• Enhanced networking primitives for modern protocols.
• More comprehensive support for structured concurrency patterns.

For Prefect, these improvements directly translate to a more capable and reliable workflow engine. As we scale to support more complex orchestration scenarios, AnyIO’s abstractions provide the foundation we need. Looking ahead, we see several exciting possibilities:

1. Enhanced Cancellation Semantics: Future versions of AnyIO may provide even more nuanced cancellation controls, allowing for partial workflow cancellation with sophisticated cleanup logic.
2. Backend Optimizations: As both asyncio and Trio evolve, AnyIO will provide access to performance improvements from either ecosystem without requiring code changes.
3. Ecosystem Growth: As more libraries adopt AnyIO, Prefect will be able to integrate with a wider range of async tools regardless of their underlying implementation.

AnyIO may not be the most well-known Python library, but it’s a critical piece of infrastructure that enables Prefect to deliver a reliable, performant workflow orchestration engine. By providing unified abstractions over Python’s async landscape, it allows us to focus on building the best possible orchestration experience rather than wrestling with low-level concurrency management.

For data engineers already using Prefect, understanding AnyIO helps explain how we achieve reliable concurrent execution. For those building their own async systems, AnyIO represents a battle-tested approach to managing concurrency that’s worth consideration.

We’re grateful to Alex Grönholm and the AnyIO community for their dedication to improving Python’s async story. If you’re curious to try AnyIO for yourself, check out the official AnyIO documentation and Prefect’s own docs on concurrency and orchestration. You’ll be well on your way to writing clear, resilient async code.

Next Steps

• Sign up for Prefect Cloud
• Join our Slack Community
• Give us a star us on GitHub

How AnyIO Powers Prefect’s Async Architecture was originally published in The Prefect Blog on Medium, where people are continuing the conversation by highlighting and responding to this story.

At Prefect, we’ve seen how a platform approach to data engineering can massively increase developer velocity. This approach abstracts away the complexity of infrastructure while giving engineers the power to build resilient, complex workflows. For developers this means reduced mental overhead, less development time, and fewer headaches to get their code up and running quickly.

For platform teams, it means protecting your infrastructure with a minimally invasive architecture while meeting the needs of the teams you support.

Prefect is perfect for platform people

Prefect is a workflow orchestration tool empowering developers to build, observe, and react to whatever workflows they are building. It just so happens that Prefect is also a highly composable application framework that can be customized by platform teams to better serve developers and stakeholders.

With Prefect, platform teams can:

• Customize a platform that developers can use to build and observe their workflows
• Set appropriate levels of access to infrastructure and other resources
• Create and scale opinionated interfaces to infrastructure, without significant overhead
• Provide infrastructure abstractions appropriate for the skill level of the team they support

Developers, the workflow builders

“Everything fails, all the time” -Amazon CTO, Werner Vogels

Code breaks, it is a fact of life. APIs change, infrastructure fails unexpectedly, and any other of an infinite number of problems can occur. Every developer wants to make their workflows resilient to failure.

Developers…

• want to schedule and observe their code
• don’t want to rewrite
• need help with infrastructure abstractions such as containers, Kubernetes, AWS ECS, etc.

With no standardization or platform team helping them, developers are left figuring out how to deploy on their own. While this is possible for some, it isn’t for many.

Platform teams, the infrastructure owners

Enter platform teams, engineers responsible for the infrastructure of an organization.

Compared to a data engineer who is often building one-off data pipelines or trying to piece things together themselves, platform teams care about ensuring that the business has the tools needed to achieve its goals at scale. To this end, platform teams build and expose foundational frameworks to help other teams build faster (particularly with respect to infrastructure) while maintaining compliance and access control standards.

Platform teams…

• are responsible for the sanctity of a company’s infrastructure
• help developer teams put their work into production using this infrastructure

By providing a standardized surface area for how developers interact with data, infrastructure, and other systems, platform teams can help workflow builders create faster and scale safely.

At Prefect, we see this platform approach to data engineering frequently, with organizations ranging from Fortune 100 companies to smaller tech firms who want to move as quickly as possible.

Creating a platform with Prefect

Platform teams can use the modular architecture of Prefect to create a custom platform that helps data engineers, data scientists, or anyone creating workflows to build faster while meeting infrastructure and access requirements.

Give developers code interfaces and a UI

Prefect turns every object into discrete API automatically, accessible by your users. Create deployment templates for your team to submit work via an API, or give them the ability to build their own highly custom workflows.

The Prefect UI gives teams a UI to track and observe their work, or even to submit templated workflows. Even non-technical teams can submit workflows with parameters, and track their progress.

Customize your infrastructure concepts

Prefect workers and work pools allow platform teams to customize their infrastructure definitions in Prefect:

• Create opinionated infrastructure patterns
• Tightly control who uses what infrastructure with RBAC + deployments
• Dynamically assign infrastructure, switching between cloud providers and services easily

Keep it minimally invasive

Prefect’s hybrid model means that execution and orchestration are separate concepts, meaning you can keep your data and code on your own premises. With Prefect Cloud, isolate and scale the infrastructure behind your workflow, without needing to worry about the orchestration server.

Platform user patterns

Let’s see some examples of how Prefect provides clear interfaces to create guardrails for different developer personas.

Deployment creators

Work pools can be used to lock down deployment infrastructure options. For example, you can remove a field you don’t want deployment creators to adjust or hard code a value for it.

Prefect Block types are Python classes that store configuration and code. Blocks have a structured HTML form for input and can be used in workflow code. Alternatively, they can be defined in Python code. Blocks are great for sharing configuration, as they live on the server.

Blocks are extremely flexible. Custom block types can be created to do whatever your team needs. Instead of needing to fork Prefect, you can just create a block with the capabilities you need and have your team use it. We see many platform teams using blocks to create

With Prefect’s team management, granular infrastructure control options, and RBAC options including custom roles, platform teams can make sure right people have access to the right resources.

Deployment runners

Through role-based access controls (RBAC), Prefect Cloud teams can limit users to run workflows but not create them.

Additionally, workflow authors can specify the data types of parameters required for the entry point function. This way at runtime, the deployment runner must provide a parameter value of an acceptable type to run the workflow.

Under the hood, Prefect uses Pydantic, so parameters can be limited to a list of choices through enums. Authors can even create a custom Pydantic class to require that a value in a valid email address format is provided or ensure a maximum or minimum numeric value.

Deployment viewers

Prefect Cloud administrators can further limit user permissions to viewing the results of runs, but not running them. The principle of least privilege states that each entity should be granted the minimum resources and authorizations needed to perform its function. Sometimes the ability to check in is all you need.

Prefect Cloud also allows for automations that fire PagerDuty notifications, emails, Slack messages, or other notifications in response to event triggers you specify. So if a flow run fails or data doesn’t look how you expect, the right people can be automatically alerted.

Additionally, if a Prefect team is using error summaries powered by MarvinAI, they can quickly see the source of workflow failures at a glance in the UI.

As you can see, Prefect provides data platform teams the ability to create self-service guardrails for personas of any type.

Build a platform for your developers

We see more and more companies adopting a platform approach to data engineering. This approach empowers developers to build faster while providing infrastructure and access safeguards.

A platform team can tailor Prefect however they choose: from adding permissions to infrastructure environments, to heavily opinionated deployment patterns. Developer teams will move faster, and your infrastructure will run smoother than ever.

For examples of teams using Prefect as a platform for a variety of engineering teams, check out the following case studies:

• Rent the Runway
• Modern Health
• Rec Room

Prefect makes complex workflows simpler, not harder. Try Prefect Cloud for free for yourself, download our open source package, join our Slack community, or talk to one of our engineers to learn more.

A Platform Approach to Workflow Orchestration was originally published in The Prefect Blog on Medium, where people are continuing the conversation by highlighting and responding to this story.

Forget about Time

Time is, generally speaking, weird. It’s the least-worst way to describe when something has occurred, is occurring, or will occur, but it’s kind of a messy abstraction. In the realm of data pipelines, it’s an abstraction, or a guess, of when something else has happened that demands a pipeline be executed.

Running a refresh pipeline every hour represents a guess that in the span of an hour, new data thats worth moving or displaying elsewhere was produced. Scheduling a pipeline to run weekly is similarly a guess that the supply of or demand for data changes about once a week.

But modern data practitioners need to take action on data as it is produced — emailing users as soon as they sign up, updating fulfillment statuses as soon as an order is placed, and displaying sensor data as soon as it makes a reading.

Event-driven architectures allow us to stop guessing — rather than guessing that some customer data came in in the past hour, we should just run the pipeline when new customer data comes in, immediately, with no extra steps and no educated guesses.

Time is, of course, useful — pipeline runs almost always cost money, and you might not want to spend money every time an event occurs — but it shouldn’t be the only way to say when something should run. Prefect makes event-driven data pipelines simple.

Event-Driven Flows with Kafka, Event Webhooks and Deployment Triggers

Imagine you have a web app, and luckily, people are signing up! You know that your onboarding email series is most effective if you send it immediately after signup.

In this post, I’ll wire up a system that subscribes to a Kafka topic with my new user signups, pipes those events into Prefect Cloud, and triggers the new user onboarding pipeline on receipt of those events. It assumes basic familiarity with Prefect Cloud. If you’re not quite there yet, check out our docs or hit me up in the Slack Community.

New User Signups → Prefect Events

When a new user signs up for my app, my app produces a message to a Kafka topic, which is an event store commonly used in distributed applications. I have a Kafka topic hosted on Confluent Cloud that contains these messages.

I know that I want to use a Prefect Flow to reach out to those new users as soon as possible after they sign up. To do this, I’m going to need to get these signup messages into Prefect, where I can observe and react to them.

To get messages from my topic into Prefect Cloud, we’re going to take advantage of Prefect Cloud event webhooks. Webhooks are an ultra-lightweight way to expose an endpoint that accepts HTTP traffic, and are used for simple integrations across the modern web. When an HTTP call is made to a Prefect webhook, it emits a Prefect event and captures data from the HTTP call for use downstream.

Using Confluent’s Httpsink connector, I can automatically hit a Prefect webhook every time a message appears in the topic.

So if I get a new user signup message in the form of:

"[
  {
    \\"registertime\\":1494287144756,
    \\"gender\\":\\"OTHER\\",
    \\"regionid\\":\\"Region_4\\",
    \\"userid\\":\\"User_5\\"
  }
]"

Kafka will fire out to Prefect and save the message as a Prefect event in real time. We use a Jinja template in our webhook configuration to do some light transformation — instructing prefect to construct an event name from the user id of the message:

Upon save, we immediately see user events start streaming into Prefect Cloud

New User Events → Deployment Triggers

Now that I have my new user signups streaming into Prefect, it’s time to do something with them. As it happens (😉) I have the following Python code, which sends a welcome email, checks to ensure that the email didn’t bounce, and updates an internal marketing database. I’ve turned this code into a flow and my functions into tasks with Prefect decorators and deployed it to Prefect Cloud.

from prefect import flow, task

@task(log_prints=True)
def send_welcome_email(name: str="will") -> None:
    print(f"Sending onboarding campaign to {name}!")
    ...
    pass

@task(log_prints=True)
def check_for_bounce(email_id: int=12345) -> None:
    print(f"checking for bounce on email {email_id}")
    ...
    pass

@task(log_prints=True)
def update_marketing_tracker(email_id: int=12345):
    print(f"updating marketing tracker for email {email_id}")
    ...
    pass

@flow
def new_user_flow(name: str) -> None:
    send_welcome_email(name)
    check_for_bounce()
    update_marketing_tracker()

if __name__ == "__main__":
    new_user_flow(name="will")

I need this flow to run when a new user event hits my system, using the user id from the event payload as the input parameter to the flow. there are a number of ways to do this both in code and the UI, but the simplest way is to navigate to an event like the ones you want to use to trigger your flow run, and click automate.

This bootstraps a trigger definition that listens for customer events. Then we can proceed to the next step of the wizard to tell it what to go when the trigger fires.

You’ll notice some interesting syntax in the user_id field. Automations capture information from the event that triggered them, and you can use Jinja syntax to pass that information into your flow run at runtime. {{ event.payload.userid }} says to take the userid from the event we captured, and use that as the main parameter value for the flow. On creation, this automation is now linked to the deployment and visible on the deployment details page under “Triggers”.

Putting it all Together

On the creation of this link, flows start kicking off in response to the Kafka new user messages:

And our flows and events tell the story of a new user signing up and getting their welcome email just seconds after their signup:

This is just the tip of the iceberg on event-driven flows with Prefect. Triggers can fire on a lack of observed events, webhooks can use jinja to flexibly transform payloads into Prefect events from CloudEvents or custom calls, and you can specify trigger policies directly from deployments without having to create automations manually. Give event-driven flows a try and let me know what you think in the Slack Community!

Beyond Scheduling: Event-Driven Flows with Prefect was originally published in The Prefect Blog on Medium, where people are continuing the conversation by highlighting and responding to this story.

Since March 14, when I posted about instrumenting LLM tools with Prefect, a whole lot has happened in the AI tooling landscape.

Most significant to my world, Marvin has taken on a life of its own. Leveraging the diverse LLM support from the open-source community, the data engineering amenities offered by Prefect, and the type-safety of Pydantic — Marvin has grown up to be a pythonic LLM toolbox that you can drop right into traditional software contexts to solve problems. If you want to learn more about how Marvin works, read the docs!

expanding the original use case

Making a Slack-based assistant for our Prefect Community was the start of all things Marvin, and it has been the use-case where we prove out our concepts to make sure they work in practice.

Today we’ll look at a concrete example of observability in an AI system — namely how we can track Marvin’s global token usage with the Events and Automations systems offered by Prefect Cloud (generous free tier btw).

background

You can chat with our Prefect assistant bot (affectionately, also Marvin) in the #ask-marvin channel in our Slack community. There’s a parallel Marvin bot we run internally to assist with tasks like support, debugging and summarizing internal process documents .

Each time you or I tag any of our Marvins, Slack will hit that Marvin app’s FastAPI endpoints with the Slack event data and prompt a Prefect-aware Bot with the user question:

Each time, the bot will leverage an LLM (like GPT 4) and may use an arbitrary number of tools in order to answer the user question — you can learn more about the state of Marvin’s OpenAI tool use here.

Between LLM calls and plugin execution, this can be time and token intensive stuff!

Prefect Cloud offers the idea of an Event , which is just a named and timestamped payload that we can emit from anywhere and then react to.

emitting events

Once Marvin settles on a response and sends it to you, we’ll run a little prefect routine from inside the endpoint called emit_event:

from prefect.events import emit_event

...

async def _slackbot_response(event: SlackEvent):

    # other things happening ...

    response = await bot.say(text)

    slack_response = await _post_message(
        channel=event.channel, message=response.content, thread_ts=event.ts
    )

    prompt_tokens = count_tokens(text)
    response_tokens = count_tokens(response.content)

    # this will do nothing if Prefect credentials are not configured
    emit_event(
        event=f"bot.{bot.name}.responded",
        resource={"prefect.resource.id": f"bot.{bot.name}"},
        payload={
            "user": event.user,
            "channel": event.channel,
            "thread_ts": thread,
            "text": text,
            "response": response.content,
            "prompt_tokens": prompt_tokens,
            "response_tokens": response_tokens,
            "total_tokens": prompt_tokens + response_tokens,
        },
    )
    
    # more things happening ...

… basically saying, “hey! Marvin responded in slack” to Prefect Cloud.

As a technical note, our community Marvin bot is running as a GCP Cloud Run service and has two important environment variables set:

PREFECT_API_KEY=pnu_XXX
PREFECT_API_URL=https://api.prefect.cloud/api/accounts/<UUID/workspaces/<UUID>

which allow it to talk to our Prefect Cloud workspace.

So after I emit my event, I can pop into the marvin-bot workspace’s Event Feed and see some details about it:

Now this event lives in the Events DB of my Prefect Cloud workspace, so I can query events over time for token usage by user or whatever else.

Being able to track aggregate token usage is great, but most exciting to me is the idea of using this event to immediately trigger other actions.

the automations system

Automations in Prefect Cloud have the form: Trigger > Action

So, with Prefect Cloud as a listener, I can have the receipt of certain events trigger an additional action. So here’s my trigger:

… which reads as, trigger this automation once I receive 1 event where the resource id looks like bot.* .

Now, we decide what our action will be. Let’s choose Run a deployment , since I already implemented a small flow to increment a token tracker:

Once I finish configuring the automation, a flow will run each time a user tags Marvin in Slack, incrementing my token tracker.

Now, if I only ever wanted to increment a token tracker, I could have just done that directly from within the _slackbot_response method from the Marvin app itself, without emitting an event that triggers a Prefect flow.

However, the larger idea here is that I can easily do arbitrary, observable work by eventing directly off of the user’s interaction with my application.

I can maintain global state with Prefect Cloud and use it in many ways:

• To govern the state transitions of tasks within my application, and observe each unit of work performed (especially when the LLM is the one deciding how and when to use additional tools)
• I could moderate user-bot interactions with my application(s), in a way that I can modify without re-deploying the Marvin service (i.e., edit an ACL Block in Cloud; the app reads from the Block in Cloud)
• if I were using my own self-hosted LLM for my Marvin app, I could have this interaction trigger a flow to update a domain-specific training set
• I could service user requests received by the Marvin app which require work from different applications, using the Prefect Cloud Events and Automations systems as the router of that work

going forward

I’m excited to expand the set of events that are emitted by Marvin, and in turn expand the set of actionable information I can leverage as a builder.

Having spent a lot of time in batch-schedule-land historically, its lovely to enter a fully event-driven world without breaking out a whole new streaming solution / tool.

As for Marvin’s future, we’ll be working on a first-class and opt-in Event instrumentation of Marvin’s LLM utilities to make as it as easy as possible to build transparent and reactive systems to control our use of AI in code.

Let us know if you have any questions in our Prefect Slack community or our Marvin Discord.

—

Happy Engineering!

Keeping 👀 on AI tools with Marvin was originally published in The Prefect Blog on Medium, where people are continuing the conversation by highlighting and responding to this story.

The single Postgres setting that almost took down our database

Beginning in March, Prefect Cloud saw a significant increase in volume and traffic driven by both new user growth as well as existing teams scaling up — by most measures we were doubling approximately every two months. This is great! But unsurprisingly this level of growth comes with a whole host of challenges for our team, especially when it comes to service reliability and resiliency.

Importantly, some of our customers suffered from service interruptions related to the issues described here. First and foremost, we are incredibly sorry for the negative impact and stress these incidents caused. We are committed to doing better.

Sometimes configuring a database at scale can seem like a dark art. While the Postgres documentation is fantastic, there are no hard-and-fast rules for tuning given the diversity of workloads out there and much of what we learned came from sparse references and experimentation. We hope that compiling our findings will help future engineers get ahead of the challenges we faced — if you want to skip the narrative, feel free to jump ahead for a summary of our findings.

The Context: Scaling Up

In anticipation of continued Prefect Cloud growth, we accelerated our plans for a proactive scale-up of our backend. The first phase of this took place in the first two weeks of April, during which we vertically scaled our primary Postgres 14 instance (running via CloudSQL in GCP).

The second phase of this work involved partitioning our most heavily used tables; large tables can exhibit poor query performance and latency due to indexes that are too large to fit into memory and unreliable table statistics. Partitioning a table essentially breaks the table up into many smaller virtual tables that the query planner can take advantage of, helping to avoid these issues. At the time of our decision, our fastest-growing table had 400 million rows and we expected it to have over 1 billion rows by the time the partitioning work would be completed. In order to partition our tables without incurring downtime, we rolled out our changes incrementally over several weeks (we’ll share more detail in an upcoming post). This allowed us to incrementally update our application to handle partitions, monitor for expected performance improvements, and adjust the plan as we went along.

We felt confident in our partition calculations and were optimistic that we’d begin seeing performance gains very soon…

The Investigation: High Memory Utilization

All the while that we were rolling out table partitions, our memory usage remained relatively constant and high. This was especially surprising because we had just vertically scaled the database, but it was difficult to understand the root cause because of the overhead of the in-flight partitioning work. It was of course also possible that our math was wrong and the partitions were inefficient, or worse, that the memory usage was an independent problem altogether. We remained confident and focused on rolling out the last two table partitions before reaching any definitive conclusions.

On April 28th, 2023 our hand was forced — during a particularly bursty period of traffic, the database unexpectedly restarted with an out-of-memory (OOM) event when memory utilization hit 90%.

The first thing we observed was that memory utilization roughly correlated with the number of active connections — this isn’t unique or surprising, as individual connections do have a memory footprint. To reduce memory pressure we scaled down our services to reduce the number of connections at the cost of increased API latency. Connections shouldn’t take up more than a few hundred megabytes of memory at most, and we needed to eventually free up almost 100GB of memory to feel long-term stable — this is all to say that we knew we were only helping the system on the margin, and reducing the number of connections wasn’t a true solution.

While still in the midst of our investigations, on May 1st, 2023 the database restarted a second time due to another OOM event. It was painfully clear that we couldn’t wait for the partition work to resolve this issue (and even that wouldn’t be guaranteed) — something needed to be done immediately. We began iterating through various Postgres configuration options such as shared_buffers and work_mem to try and see if any moved the needle significantly (our services are able to handle anticipated DB restarts without disruption). All the while we were collectively trying to construct a meaningful theory for why any of this was happening - there was clearly a missing piece of the puzzle. Maybe it was a silent config change in our vertical scale up? Maybe the partition work was quietly introducing something nefarious? The team floated and tested many hypotheses.

It was in the midst of this experimentation that a new theory emerged — each webserver that communicated with the database managed its own in-process connection pool. In the absence of a pod restart, the connections within each pool were recycled every 2 hours. It was not out of the realm of possibility that these connections might be accruing an increasingly large cache that needed flushing. This seemed all the more likely given the relatively tight correlations we were seeing with scale ups / scale downs and large changes to memory utilization.

We updated our connection pools to recycle their connections every 10 minutes and quickly released the change to prod. Coincidentally and seemingly unrelated, we had noticed that an entirely separate database used for our events system was not being traced properly. So in addition to the pool recycle change, we released a small change that ensured our SQLAlchemy instrumentor was turned on for all engines (essentially replacing the line SQLAlchemyInstrumentor().instrument(engine=engine.sync_engine) with SQLAlchemyInstrumentor().instrument() ).

We released these changes and were all on the edge of our seats, monitoring our dashboards to see if we could discern any change. It didn’t take long to see something noticeable:

It seemed like we had hit the jackpot! The database seemed healthier than ever — we carefully scaled our services back up with almost no effect on memory. On top of that, API latency was better than it had been in weeks. Despite feeling less acutely stressed, we remained unsatisfied with this empirical solution that lacked an airtight explanation; we began trying to refine our original theory to explain why this seemingly simple change brought us back to the healthiest state we had seen in weeks. However, we couldn’t seem to break past the hand-waving phase to anything more specific and testable…

The Twist: Load-Bearing Comments

As we reflected on the prior few days and were organizing our insights, we noticed a curiosity - a different database that manages account configuration saw an unprecedented increase in received bytes (an increase of almost 400%) at exactly the same time that our primary database recovered from its memory woes. This increase wasn’t inherently concerning as this database doesn’t serve complex workloads, but it was interesting and was not easily explained by our current story.

It turns out that enabling SQLAlchemyInstrumentor globally for all engines has the effect of adding comments to every query, for example:

SELECT *
FROM auth.account
WHERE auth.account.id IN (%s) 

/*db_driver='asyncpg',db_framework='sqlalchemy%%3A0.34b0',
traceparent='00-93d35b06052db659c4eb5f5be2db2548-a48d7cb6e8de48f7-01'*/

Because queries to our account database were generally small, this had the effect of increasing query size (in bytes) noticeably. You might ask: why is this relevant? We soon discovered that when we explicitly turned comments off across all engines, memory in our primary database began increasing back to the problematic levels.

This initially made no sense - production-grade tracers shouldn’t have that much impact on performance! We then remembered - we had seen this once before. After we successfully partitioned our events database many months back, we found that executing the exact same prepared statement six times in a row resulted in an indefinite hang on the sixth execution. Adding a unique comment (such as a trace ID) to each execution prevented the hang from ever occurring. This is because Postgres computes a “generic plan” (a plan independent of parameter values) on the sixth execution of the same prepared statement and caches it for future use - but generic plans are typically very costly for queries over partitioned tables (hence the hang). By default, Postgres uses heuristics to decide whether to use the generic plan vs. a custom plan but those heuristics don’t always work.

And this was the final piece of the puzzle.

The Resolution: Configuring plan_cache_mode

Ultimately, here’s what was happening:

• The asyncpg driver for SQLAlchemy prepares every statement
• After seeing the same prepared statement 5 times, Postgres creates a “generic plan” that is then cached for future use
• Generic plans for prepared statements do not take into account specific parameter values and are therefore costly for queries over partitioned tables
• Using SQL comments for tracing served as a “cache buster” — apparently the generic plan cache is keyed on the raw text of the query!
• Postgres exposes a way to configure this behavior; setting plan_cache_mode=force_custom_plan prevents generic plans from being computed and cached entirely

If you needed yet another story for why lines-of-code is not an effective measure of impact, the final resolution of our memory issues was to set this single value on each connection to our primary database (easily done by passing connect_args when creating the engine in SQLAlchemy).

Conclusion

Along with our new knowledge of cache plans in Postgres, our team also emerged with a few learnings that apply to taming the dark art of scaling a production database during periods of rapid iteration; in no particular order:

• incidents can be chaotic, with many moving pieces — maintaining a level head and constantly questioning your assumptions are incredibly valuable skills
• having an airtight release process for quick iteration in prod is critical (thank you to our platform team!)
• Google Cloud SQL reserves 10% of instance memory for other services running on the underlying instance, so database processes are likely to run out of memory when usage exceeds 90%
• If a solution or explanation can’t be deconstructed and reconstructed from multiple perspectives, it’s possible that it isn’t an actual solution
• Observability and tracing are key. Surprisingly, these tools are not always passive participants in your stack!
• plan caches in Postgres are very useful for most standard queries (we definitely saw latency impact on our simple accounts database when toggling this setting), and are destructive for queries over partitioned tables — plan_cache_mode is your friend in this situation
• caching continues to be the root of all evil

Most importantly we want to thank our users and customers for their confidence and trust — which we do not take for granted — as we work harder than ever to strengthen Prefect’s resiliency far ahead of ongoing growth.

The hard-won knowledge represented above was uncovered and compiled by many engineers at Prefect — in particular Zach Angell, Jake Kaplan, Eddie Park, Chris Guidry, Andrew Brookins, and Jonathan Yu spent many restless days and nights experimenting, monitoring and learning to solve the issue described here and significantly improve our overall service reliability and scalability.

To learn more about Prefect:

• visit our website
• join our active Slack community
• ask questions on our Discourse
• follow us on Twitter for updates
• visit us on GitHub to open issues and pull requests

Happy Engineering!

More Memory, More Problems was originally published in The Prefect Blog on Medium, where people are continuing the conversation by highlighting and responding to this story.

I’m invited to several March Madness bracket pools every year because of my love of college basketball (and how much I talk about it). To be honest, I’m not very good at them because navigating March Madness is chaos. As an engineer, I needed to add some logic to this situation to help me figure it out and at least give myself a shot at winning these bracket pools. As Ted Mosby from How I Met Your Mother once said: “This is not March Madness…”

And what better tool than Prefect to help me orchestrate my first go at adding some logic to my decision process. The custom predictive algorithm is simplistic, but by gosh, it is orchestrated and observable thanks to Prefect.

All of my Prefect flow code and infrastructure configuration can be found at the Github link here.

Note: This is just code for the Men’s tournament, but changing the RapidAPI ID’s to those for the Women’s tournament will yield similar predictions.

The Code

Using Prefect flows, subflows, and task decorators, I took the following steps to help me pick the winners of the brackets.

1. I gathered the statistics for the 68 teams in the tournament. I used RapidAPI’s Basketball API to gather data for each team. I was able to securely store the API Key using a Prefect secret block.
2. I then picked the winners for each team's first set of games. I used the statistics from each team to calculate a score using win percentage, points scored, and points allowed. Those scores were compared to determine a winner for each of the first individual matchups.
3. Then, I ranked the predicted first game winners. After calculating the score, I used Python native sort and lambdas to sort the teams by score. I then used these rankings to fill out the rest of my bracket. These rankings were printed to the screen using Prefect’s log_prints configuration.

The Infrastructure

The following infrastructure was put together using Prefect deployments and blocks.

• GitHub for Version Control
• S3 for Prefect flow code storage, connected with a Prefect S3 storage block.
• EKS Cluster created using eksctl.
• ECR to store the image that the job will use to run the Prefect flow.
• Prefect KubernetesJob to define the job configuration for the flow run.

All of these components were combined through Prefect deployments and blocks. Instructions for configuring all of these components can be found in the GitHub repository above.

While these flows can be run locally or on any infrastructure, I chose to run this flow using EKS because of how flexible Kubernetes is with scaling resources and applying updates. Prefect’s integration with Kubernetes is completely customizable, and running these deployments on scalable infrastructure allows me to share this flow with everyone that has access to my workspace, without them having to worry about spinning up new infrastructure. It also provides the flexibility to scale when running more complex models for next year's tournament!

The Results

The following results were outputted from the flow code. The algorithm used to calculate these results is SUPER simple, so take these results with a grain of salt. I look forward to adding more data science and machine learning concepts to next year’s tournament.

11:23:13.032 | INFO    | Flow run 'lavender-asp' - 1. Houston
11:23:13.033 | INFO    | Flow run 'lavender-asp' - 2. Florida Atlantic
11:23:13.034 | INFO    | Flow run 'lavender-asp' - 3. Charleston
11:23:13.036 | INFO    | Flow run 'lavender-asp' - 4. Oral Roberts
11:23:13.040 | INFO    | Flow run 'lavender-asp' - 5. Gonzaga
11:23:13.041 | INFO    | Flow run 'lavender-asp' - 6. Alabama
11:23:13.042 | INFO    | Flow run 'lavender-asp' - 7. UCLA
11:23:13.043 | INFO    | Flow run 'lavender-asp' - 8. Purdue
11:23:13.043 | INFO    | Flow run 'lavender-asp' - 9. Arizona
11:23:13.044 | INFO    | Flow run 'lavender-asp' - 10. Kent State
11:23:13.044 | INFO    | Flow run 'lavender-asp' - 11. Marquette
11:23:13.045 | INFO    | Flow run 'lavender-asp' - 12. Iona
11:23:13.045 | INFO    | Flow run 'lavender-asp' - 13. Drake
11:23:13.046 | INFO    | Flow run 'lavender-asp' - 14. Furman
11:23:13.047 | INFO    | Flow run 'lavender-asp' - 15. St. Marys (CA)
11:23:13.047 | INFO    | Flow run 'lavender-asp' - 16. Louisiana Lafayette
11:23:13.048 | INFO    | Flow run 'lavender-asp' - 17. Texas
11:23:13.049 | INFO    | Flow run 'lavender-asp' - 18. Kansas
11:23:13.052 | INFO    | Flow run 'lavender-asp' - 19. UC Santa Barbara
11:23:13.052 | INFO    | Flow run 'lavender-asp' - 20. Utah State
11:23:13.052 | INFO    | Flow run 'lavender-asp' - 21. Kennesaw State
11:23:13.053 | INFO    | Flow run 'lavender-asp' - 22. Montana State
11:23:13.053 | INFO    | Flow run 'lavender-asp' - 23. Boise State
11:23:13.054 | INFO    | Flow run 'lavender-asp' - 24. Texas A&M
11:23:13.054 | INFO    | Flow run 'lavender-asp' - 25. Texas A&M-CC
11:23:13.054 | INFO    | Flow run 'lavender-asp' - 26. NC State
11:23:13.055 | INFO    | Flow run 'lavender-asp' - 27. Nevada
11:23:13.055 | INFO    | Flow run 'lavender-asp' - 28. USC
11:23:13.055 | INFO    | Flow run 'lavender-asp' - 29. Kentucky
11:23:13.056 | INFO    | Flow run 'lavender-asp' - 30. Pittsburgh
11:23:13.056 | INFO    | Flow run 'lavender-asp' - 31. TCU
11:23:13.057 | INFO    | Flow run 'lavender-asp' - 32. Maryland
11:23:13.058 | INFO    | Flow run 'lavender-asp' - 33. Illinois
11:23:13.058 | INFO    | Flow run 'lavender-asp' - 34. Auburn
11:23:13.059 | INFO    | Flow run 'lavender-asp' - 35. Iowa State
11:23:13.059 | INFO    | Flow run 'lavender-asp' - 36. Fairleigh Dickinson

I will submit these results to an ESPN bracket, provide an update on how the bracket does, and predict the next Sweet 16 and Elite 8 matchups using Prefect’s flow deployment structure! Stay tuned.

Conclusion

All in all, college basketball is fun and chaotic. Filling out brackets is all a guessing game, but attempting to predict using simple math can help. One thing is certain: orchestrating and observing this math on your infrastructure using Prefect is the slam dunk your bracket needs.

Picking March Madness Winners Using Prefect was originally published in The Prefect Blog on Medium, where people are continuing the conversation by highlighting and responding to this story.

Like everyone else, I’ve been playing with AI over the last couple of months.

Projects like OpenAI and Hugging Face have made it simpler than ever to interact with LLMs like GPT3 — whether it’s popping open a new tab in my browser for ChatGPT or making a simple API call with Python.

What is an LLM?
“Large Language Model” refers to a type of AI model that’s been trained on massive amounts of text to generate human-like language. There are open questions about the use of LLMs.

One very interesting way to leverage this new LLM accessibility is being explored by a project called LangChain, the growth of which has exploded.

LangChain is open-source tooling that adds some useful abstractions for interacting with these LLMs, whether that’s:

• callable LLMs to simplify ad-hoc API calls
• successive calls to LLMs + tools to refine answers to questions
• building customized chatbots by persisting information in vector embedding databases, and querying them to provide LLMs with whatever context it needs to answer a question

Through these LLM interactions + others, LangChain is in the business of “enabling developers to build applications that they previously could not.”

I’ll explore some functionality LangChain makes easy, and then present a first pass at a plugin powered by Prefect, which can help us observe our use of LangChain, named langchain-prefect.

Exploring LangChain

Starting barebones, if I have an OPENAI_API_KEY on my machine, I can go ahead and pip install langchain openai and do a random ol’ LLM call:

from langchain.llms import OpenAI

llm = OpenAI(temperature=0.9)
answer: str = llm(
    "What would be a good company name for "
    "a company that makes colorful socks?"
)

print(answer) # prints something like `Socktastic!'

This is cool. But, to quote LangChain’s README:

the real power comes when you can combine (LLMs) with other sources of computation or knowledge

A tool like LangChain starts to gets interesting when I integrate other tools, APIs, and embedding databases — but it also starts to get complex enough that orchestration is non-trivial.

As with any tool I’d want to build on top of…

Before I dive headlong into building some LLM-based application using LangChain, it’s important to think pragmatically about a few things.

Let’s say I want to build a chatbot powered by LangChain and some LLM with domain-specific knowledge I can curate and regularly update.

I’d want the chatbot to handle a conversation with a user, answering their questions using the LLM + context + tools at its disposal.

Implementation details aside, I’d need to answer a couple of questions before deploying this chatbot in a production use case.

Ballpark, how much will it cost?

The main currency with these LLMs are tokens — for example, OpenAI charges $0.0200 / 1K tokens to use their DaVinci model. 1 token roughly corresponds to one word, depending on who you ask.

You pay for total tokens, which is prompt tokens + completion tokens, i.e. text you send the LLM + text you get back from it.

If I’m using LangChain agents to make successive LLM calls / leverage tools, and they’re stuffing extra context (tokens) in those calls, how much will that cost me?

A quick experiment and some math:

Let’s invoke a LangChain agent like my eventual chatbot might.

It will use the llm-math tool to do the math and fall back to asking the LLM when it doesn’t know the answer (you could add serpapi as a tool here).

from langchain.agents import initialize_agent, load_tools
from langchain.llms import OpenAI

llm = OpenAI(temperature=0)
tools = load_tools(["llm-math"], llm=llm)
agent = initialize_agent(tools, llm)

answer: str = agent.run(
    "How old is the current Dalai Lama? "
    "What is his age divided by 2 (rounded to the nearest integer)?"
)
print(answer) # The current Dalai Lama's age divided by 2 is 43

When I ran this, the agent made 7 calls to the LLM:

let total, p, c = total_tokens, prompt_tokens, completion_tokens

• total: 205 = p: 172 + c: 33
• total: 254 = p: 242 + c: 12
• total: 243 = p: 224 + c: 19
• total: 244 = p: 239 + c: 5
• total: 282 = p: 256 + c: 26
• total: 245 = p: 242 + c: 3
• total: 316 = p: 293 + c: 23

That shakes out to a grand total of 1789 tokens = $0.03578 = 3.578 cents

Not bad, but this is only one invocation of a barebones agent. If I used other tools (which could be chains/agents) or pulled more context from a vectorstore to stuff in my prompts, that number would rise pretty dramatically — depending on the math you’d have to do for your app.

If you’ve followed along and run the example above, you might wonder how I even figured out my total token usage. It didn’t come back as a return from agent.run 🧐 — but we’ll get to that in just a bit.

So how much does it cost?

It depends. It depends on how well I write my application, what I want it to do, and how well LangChain is leveraging my context when talking to LLMs for me. It certainly seems like something I want to keep track of.

What I haven’t considered here yet at all is infrastructure costs, although, yeah, we will need a scalable and custom runtime for this app.

How do I know that things are working the way I want?

On its face, this seems obvious. It’s working if my users get a useful answer from the chatbot, right?

Well, that’s part of the story for sure. But, this question strikes me as one that’d be most difficult to nail down at production scale — where processes running LLM calls are just one part of a complex workflow.

As abstractions for interfacing with LLMs mature, it feels like it’s going get difficult to keep track of what’s going on under the hood — just like it’s non-trivial to keep track of modern ELT / ETL at production scale.

Hypothetical question: A user opens a support ticket for my chatbot and says my chatbot left them in the cold. How do I know what went wrong?

• Is OpenAI having an incident? did they choke my API calls?
• Is there a problem with my embeddings db?
• Am I orchestrating my LangChain agents/chains incorrectly?

Building tooling to frame and answer these sorts of questions sounds like a lot. But then again, Prefect does kinda do a lot.

Making a Prefect plugin for LangChain

I’ll share some progress on a plugin for LangChain to address some of the pain points I’ve mentioned above, dubbed langchain-prefect.

Let’s first go ahead and pip install langchain-prefect…

and return to that first example, with a new friend named RecordLLMCalls:

from langchain.llms import OpenAI
from langchain_prefect.plugins import RecordLLMCalls

llm = OpenAI(temperature=0.9)

with RecordLLMCalls():
  answer: str = llm(
      "What would be a good company name for "
      "a company that makes colorful socks?"
  )
  
  print(answer)

17:45:19.782 | INFO | prefect.engine - Created flow run 'scarlet-chital' for flow 'Execute LLM Call'
17:45:21.319 | INFO | Flow run 'Calling langchain.llms.openai' - Sending 'What would be a good company name for a company that makes colorful socks?' to langchain.llms.openai via <function BaseLLM.generate at 0x11e567a30>
17:45:22.211 | INFO | Flow run 'Calling langchain.llms.openai' - Recieved: NotAnArtifact(name='LLM Result', description='The result of the LLM invocation.', content=LLMResult(generations=[[Generation(text='\n\nHappy Feet Socks.', generation_info={'finish_reason': 'stop', 'logprobs': None})]], llm_output={'token_usage': {'completion_tokens': 7, 'total_tokens': 22, 'prompt_tokens': 15}}))
17:45:22.336 | INFO | Flow run 'scarlet-chital' - Finished in state Completed()

Happy Feet Socks.

Uhh, what just happened?

Let’s break it down by looking at the logs.

With langchain-prefect you get prefect. With Prefect, you can wrap arbitrary Python functions in a @flow decorator so when you run one of them, metadata on that execution is tracked (see footnotes).

From inside of RecordLLMCalls, the LangChain methods that actually make the API call to the LLM are wrapped in a @flow so that Prefect can track the execution when they are invoked. Prefect gives fun names like scarlet-chital to each flow run, so we see something like:

prefect.engine - Created flow run 'scarlet-chital' for flow 'Execute LLM Call'

Then we see what is actually being sent to the LLM, how the LLM is being prompted by LangChain, the LLM output, and how many tokens we used.

Flow run 'Calling langchain.llms.openai' - Sending 'What would be a good company name for a company that makes colorful socks?' to langchain.llms.openai via <function BaseLLM.generate at 0x11e567a30>
Flow run 'Calling langchain.llms.openai' - Recieved: NotAnArtifact(name='LLM Result', description='The result of the LLM invocation.', content=LLMResult(generations=[[Generation(text='\n\nHappy Feet Socks.', generation_info={'finish_reason': 'stop', 'logprobs': None})]], llm_output={'token_usage': {'completion_tokens': 7, 'total_tokens': 22, 'prompt_tokens': 15}}))

Finally, since our Python code (or LangChain’s) ran without errors, we see:

Flow run 'scarlet-chital' - Finished in state Completed()

Here’s where it gets fun.

If we pop into a terminal and run prefect server start, you’ll be prompted to Check out the dashboard at http://127.0.0.1:4200

… where you’ll see a fancy dashboard just waiting to fill up with colors:

If we click into our flow run, we see:

The logs and a bunch of metadata on when, what, and how things happened.

Prefect is uniquely interesting here.

In Prefect, you can call flows from inside other flows — it’s called a subflow.
Look what happens if I run the agent example with RecordLLMCalls and add a flow wrapping my agent.run call:

from langchain.agents import initialize_agent, load_tools
from langchain.llms import OpenAI
from langchain_prefect.plugins import RecordLLMCalls
from prefect import flow

llm = OpenAI(temperature=0)
tools = load_tools(["llm-math"], llm=llm)
agent = initialize_agent(tools, llm)

@flow
def my_flow():
    """Flow wrapping any LLM calls made by the agent."""
    return agent.run(
        "How old is the current Dalai Lama? "
        "What is his age divided by 2 (rounded to the nearest integer)?"
    )
with RecordLLMCalls(tags={"agent"}):
    result = my_flow()

Remember how we ran this earlier, and it made 7 LLM calls? Here they are.

For each one, we’re still tracking LLM input / output and token usage, and now we know where these came from. I used tags={‘agent’} to apply an agent tag to any LLM calls that agent.run created on my behalf.

The nice thing about Prefect flows is that my_flow that wraps these LLM calls could itself be called in the context of a larger workflow (even nested as a subflow of another wrapping parent flow).

Here’s what the usage page looks like in OpenAI:

Where do we go from here?

Okay, hopefully, that was enough to get the blood flowing. Let’s sum up.

In general, LLMs are interesting and useful — building apps that use them responsibly feels like a no-brainer. Tools like LangChain make it easier to build LLM apps. We still need to know details about how our apps work, even when we want to use tools with convenient abstractions that can obfuscate those details.

Prefect is built to help data people run, scale, and observe event-driven work wherever they want. It provides a framework for deployments on a whole slew of runtime environments (from Lambda to Kubernetes) on all the main cloud providers. Productionizing LLMs is right up Prefect’s alley.

Among the more advanced concepts in Prefect are results, which are applicable here. Caching common LLM results, aggregating token usage across many actors, and routing good and bad LLM results for use in fine-training all feel like practical choices that results could help implement.

I’m biased, but frankly I’m excited about Prefect’s potential to address problems in LLM observability. As the AI hype train pushes forward, and we try to balance easy abstractions with control, I believe it will be increasingly important to understand what exactly is happening with applications built on LLMs — the same way it is now for apps built on data warehouses. Real outcomes depend on our tools working as expected.

Stay tuned for further exploration of tools like LangChain! Questions on langchain-prefect or want to contribute? Find and fork the code.

Happy Engineering!

—

Footnotes

How is metadata tracked?
By default (with no configuration outside of pip install prefect), Prefect flows spin up a lightweight ephemeral API (FastAPI) and use it to start writing details about flow runs to a SQLite db file ~/.prefect/prefect.db.

You can inspect flow runs with the dashboard after prefect server start.

It’s also super simple to let Prefect Cloud handle UI, API, and DB management for you, for free — but for dev’ing on things, the ephemeral API / local SQLite DB / local UI dashboard is great.

Keeping Your Eyes on Your AI Tools was originally published in The Prefect Blog on Medium, where people are continuing the conversation by highlighting and responding to this story.

When wrangling data on a small team, it’s hard not to rely on tools that help us get event data from the world to our data consumers — they’re just too useful to ignore!

Though, as a data practice scales, it can get difficult to keep track of your tools’ interactions that are not automated and visible.

This tutorial is a quick start to a simple, scalable, and maintainable ELT workflow using an open-source stack and a free-tier data warehouse. We’ll be using the following tools:

• Airbyte to extract data sources and load them to a data warehouse (like Snowflake)
• dbt to transform the data
• Prefect to automate and observe our ecosystem of tools
• GitHub as a source of event data and a place to store our code

For the purpose of illustrating an ELT pattern, we choose a fun question that serves as the motivation for the workflow:

Who are the common contributors to airbyte, dbt and prefect?

Alrighty, let’s get into it!

Table of Contents

· Set up an Airbyte instance
· Create some Airbyte connections
∘ Create a GitHub source connection
∘ Create a Snowflake destination connection
· Create some dbt models to transform the raw data
· Write a Prefect flow to orchestrate our tools
∘ Let’s use some Prefect collections to make our lives easier
· Define a Prefect flow
· Running Airbyte connections
· Run dbt models
· Run some queries against the transformed data
· The whole enchilada
· Running our flow
· … and that’s it!
· Extending this example
· Next steps

Set up an Airbyte instance

This tutorial is based on Open-Source Airbyte. If you are using Airbyte Cloud, you’ll need to wait for the Airbyte Cloud API to programmatically interact with Airbyte from Prefect.

To set up a local Airbyte instance, follow the Airbyte Quickstart guide or use a helm chart.

Create some Airbyte connections

For this tutorial, we’ll use Airbyte’s GitHub source connector to pull data from the GitHub organizations whose repositories are used in this tutorial:

• airbytehq/airbyte
• dbt-labs/dbt-core
• prefecthq/prefect

Create a GitHub source connection

All Airbyte needs to connect to GitHub is a valid personal access token.

Let’s go ahead and create a new connection in the Airbyte UI by clicking on the Sources tab and click + New Source. Select the GitHub source connector and pass along your token (or comma-separated list of tokens to distribute load across multiple actors):

For this tutorial, we’ll create a connection for each of the three GitHub organizations we’re interested in, grabbing the repository of interest for each organization.

Create a Snowflake destination connection

To tell Airbyte where to drop our raw data in Snowflake, we’ll need to provide the following information:

• account: the name of your Snowflake account
• username: the name of the user that Airbyte will use in Snowflake
• password: the password for this user
• database: the name of the database that Airbyte will use to write data to
• schema: the name of the schema that Airbyte will use to write data to

Let’s go ahead and create a new connection in the Airbyte UI by clicking on the Destinations tab and then clicking on + New Destination. Select the Snowflake destination connector and pass in the information above.

💡 Note: While you’ve got this info handy, you can also use it to connect your dbt CLI to Snowflake. See the dbt docs for more info — psst… dbt debug is your friend!

Create some dbt models to transform the raw data

Now that we’ve given Airbyte the ability to pull data from GitHub and write it to Snowflake, we can use dbt to transform the raw data into a format that’s more useful for our downstream consumers (like BI tools, fun dashboards and downstream dbt models).

While not the focus of this tutorial, the source code for the dbt models we’ll be using can be found in the linked repository.

In this tutorial, we’ll create two models: common_commiters and common_issue_submitters. These will create views in Snowflake containing the committers and issues submitters that are common across all three repositories. We’ll then be able to query these views and check out the overlap between the three communities.

Write a Prefect flow to orchestrate our tools

It’s time to connect our tools with Prefect!

Let’s use some Prefect collections to make our lives easier

pip install prefect-airbyte prefect-dbt prefect-snowflake

… and now let’s import some of their useful functionality for our flow:

from prefect_airbyte import AirbyteConnection, AirbyteServer
from prefect_dbt.cli.commands import trigger_dbt_cli_command
from prefect_snowflake.database import SnowflakeConnector, snowflake_multiquery

Define a Prefect flow

It’s often useful to take a step back and think about each step we want to take in our flow. In this case, we want to:

• trigger existing Airbyte connections to drop raw data into our Snowflake database
• run dbt models to transform the raw data into something more useful
• query the transformed data in Snowflake to find the answer to our original question

Let’s dive into the code for each of these steps!

Running Airbyte connections

Lucky for us, the AirbyteConnection block defined in prefect-airbyte takes care of the hard work for us. If we tell it which connection to trigger and which AirbyteServer to use, it'll do the rest!

Let's create our AirbyteConnection blocks, all on the same AirbyteServer:

airbyte_server = AirbyteServer(server_host="localhost", server_port=8000)

airbyte_connections = [
    AirbyteConnection(
        connection_id=connection_id,
        airbyte_server=airbyte_server
    )
    for connection_id in [
        "fa8d5164-ca22-47da-83e0-829cc86a70b8", # Airbyte Github Stats
        "8e3a3ef1-6c2d-4255-99d4-21e9e542d853", # DBT Github Stats
        "980c14d9-2992-46a7-a1dc-2b92e60e1475" # Prefect Github Stats
    ]
]

…and a task to manage their execution:

@task
def run_airbyte_sync(connection: AirbyteConnection) -> AirbyteSyncResult:
    job_run = connection.trigger()
    job_run.wait_for_completion()
    return job_run.fetch_result()

💡 Note: All Prefect blocks that implement the JobBlock interface (like AirbyteConnection) have a trigger method that returns a JobRun , which has wait_for_completion and fetch_result methods.

Run dbt models

To trigger dbt run, you can leverage the predefined Prefect task from the dbt collection called trigger_dbt_cli_command. You can pass the command to run and the path to your dbt project as follows:

dbt_result = trigger_dbt_cli_command( # run dbt models
    command="dbt run",
    project_dir="github_common_contributors",
    wait_for=airbyte_results
)

Notice the use of wait_for here. This tells Prefect that we want to make sure the airbyte tasks have successfully dropped raw data in Snowflake before we try to transform it with dbt.

💡 Note: If you’re using dbt Cloud, you can use the functionality defined in prefect_dbt.cloud instead.

Run some queries against the transformed data

Assuming we now have transformed data, we need to tell Prefect how to connect to Snowflake and which queries to run. The first part can be accomplished by creating a SnowflakeConnector block:

💡 Note: A SnowflakeConnector needs a set of SnowflakeCredentials to authenticate with Snowflake. As with any block, you can create this in the UI or programmatically.

Now we can use our SnowflakeConnector to connect to the right place in Snowflake and snowflake_multiquery (another pre-built task) to ask questions about the transformed data that lives there:

common_authors, common_issue_submitters = snowflake_multiquery(
    queries=[
        "select login from commit_authors",
        "select login from issue_submitters"
    ],
    snowflake_connector=SnowflakeConnector.load("github-contributors"),
    wait_for=[dbt_result]
)

In the same way as before, we use wait_for to make sure the dbt models have successfully run before we try to query the transformed data.

The whole enchilada

Pretty nifty, huh? Let’s put it all together into a @flow:

@flow(log_prints=True)
def my_elt_flow():

    # run airbyte syncs
    airbyte_results = run_airbyte_sync.map(airbyte_connections)

    # run dbt models
    dbt_result = trigger_dbt_cli_command(
        command="dbt run",
        project_dir="github_common_contributors",
        wait_for=airbyte_results
    )

    # query our fancy new tables
    common_authors, common_issue_submitters = snowflake_multiquery(
        queries=[
            "select login from commit_authors",
            "select login from issue_submitters"
        ],
        snowflake_connector=SnowflakeConnector.load("github-contributors"),
        wait_for=dbt_result
    )
    
    # log the common contributors among our repos
    print(
        f"Common authors: {common_authors} "
        f"Common issue submitters: {common_issue_submitters}"
    )

if __name__ == "__main__":
    my_elt_flow()

Running our flow

To run our flow locally, we can simply run the script:

python elt_orchestrator.py

…and watch it go 🚀

Turns out dependabot is the only common contributor since last August — laaaaaaaame.

… but oh well — we have a nice and sturdy ELT pattern now.

Let’s create a deployment for our flow so we can run it when and where we want.

To do this, we need to make two decisions:

• where do we want to run our flow?
• where do we want to store our flow?

To keep things simple, I’ll run our code locally and keep it in GitHub (where it is, anyway!).

In terms of prefect deployments, that means I’m using the default Process infrastructure and a GitHub filesystem block.

The command to create the deployment is:

prefect deployment build flows/elt_orchestrator.py:my_elt_flow \
	-n "elt-pipeline" \
	-sb github/demo-repo \
	-a

💡 Note: If we wanted to use a different runtime infrastructure (like Kubernetes), we could specify that in the prefect deployment build command by passing -ib kubernetes/my-k8s-job-template. -a tells prefect to go ahead and apply the deployment.

With our new deployment created, we can hop over to the UI and run it manually or attach an interval/cron schedule.

… and that’s it!

We now have an intelligent Prefect workflow that will not waste time/resources running downstream tasks if upstream tasks fail.

When the flow fails (as it probably will at some point), it’s easy to identify and triage issues by observing the failing flow run in the Prefect UI!

If we ever want to change how or when our flow executes, we can do so without ever changing our code!

• change how often it runs → update the deployment schedule in the UI
• change where we store our code → add new storage block, rebuild deployment with prefect deployment build ... -sb block-type/name (same goes for infrastructure block determining where it runs)

Extending this example

While this pattern is sophisticated enough to handle many teams’ ELT workflows, there are many ways we could extend it:

• parameterize my_elt_flow to accept an arbitrary number of Airbyte Connections or dbt commands to run
• create flow runs from this deployment via API calls when you know Airbyte has new data to sync instead of on a schedule
• use the Airbyte sync row count in AirbyteSyncResult to decide how many threads dbt should run!
• do basically anything you want in Python at any point.

Next steps

If you haven’t already, you can create a Prefect Cloud account for free to serve as your launch pad into a world of more effective data engineering!

Questions? Feel free to ask them in our Slack Community!

Happy Engineering!

Coordinate ELT in 2023 with Airbyte, dbt and Prefect was originally published in The Prefect Blog on Medium, where people are continuing the conversation by highlighting and responding to this story.