Upgrading to Airflow 2.0: Massive Performance Wins and Lessons Learned
By: Christine Betadam
Tags: Data Engineer, AWS, Airflow

At RealScout, we use Apache Airflow to orchestrate our crucial workflows such as data ingestions, health checks, and reconciliation processes. Our team recently upgraded our stack from Airflow 1.10.10 to the long awaited 2.0. We are now about two weeks into the new Airflow deployment in production and are already seeing some very big wins, which we’ll detail in later posts in this series. This upgrade was unique from previous Airflow upgrades due to both volume and type of changes included in 2.0. In this post, we highlight our experiences, challenges encountered, and advice for others who may be planning their own upgrade.

Airflow on AWS

Our existing Airflow stack is deployed on AWS, using the great Turbine project as a starting point for our CloudFormation template. For those unfamiliar, a CloudFormation template defines all resources and policies for a set of AWS infrastructure. Our production stack includes 1 Scheduler instance, 1 Web Server instance, 10 Celery Worker instances and 1 RDS PostgreSQL instance.

Aside from the instances, other important resources include SQS queues (used for Airflow’s Celery setup with SQS as a backend), an Auto Scaling Group for maintaining the desired number of worker instances, security groups to limit access to instances and the database, etc. One important consideration is that all of our Airflow instances (scheduler, web server, and workers) share a common /efs folder mount for airflow code and logs, which makes deployments and logging simpler.

Simplified version of our Airflow architecture at RealScout

Requirements and Usage

We have about 1000 active DAGs in total, which has grown over the years as we’ve become more accustomed to (and trusting of) Airflow. The DAGs run on various schedules; however, most of them run every 15 minutes or less. Some maintenance or more intensive DAGs run once a day or once a week. Perhaps unique to our setup is that almost all of our DAGs are 1–2 tasks, the bulk of which just orchestrate and monitor ECS Tasks. This lets us utilize very low-powered EC2 Airflow worker instances, which is great for cost saving. This is probably different from many other teams’ typical setups, who may have DAGs containing hundreds or even thousands of tasks in some cases. We’ll discuss the implications of these ‘shallow’ and ‘deep’ DAGs (as we call them) later, but know that there are important Airflow Scheduler throughput considerations depending on your reality.

Upgrade plan

At RealScout, we aim for predictable, ‘boring’’ migrations that minimize downtime. Taking as much as we could into consideration, we came up with the following plan of action to upgrade Airflow. While our plan has many specific references to our AWS setup, the spirit of them should apply to any Airflow upgrade.

Step 1: Test in Staging

Weeks before our planned migration date, we created a cloned stack in AWS (same Airflow 1.10.10 setup) for testing and validation. As mentioned earlier, we use CloudFormation, so spinning up a new stack based on our existing infrastructure configurations was trivial. Create a snapshot of your production Airflow database in RDS; we will use it to test the ‘airflow db upgrade’ command in staging. We recommend using the exact same instance sizes (CPU, memory, etc) for the best comparison. Stop the staging scheduler, web server, and worker instances. Since we use an auto scaling group to manage our worker instances, we simply scaled its max capacity to 0. Because this is our staging environment and a fresh Airflow database, no DAGs are active. ‘Restore’ your production RDS snapshot into your staging RDS instance. Again, your staging scheduler is already disabled, so no DAG runs will be triggered. Check out your branch with your new dependencies, and ‘pip install’ new packages on the scheduler and web server instances. If your setup utilizes continuous deployment, be careful of merging the upgrade branch into your main one, as you might unintentionally roll your code out to production! We do use CD, and so we disable that process in production when rolling out upgrades like these. Scale up your worker autoscaling group to 1 instance and ensure it comes online correctly with the new dependencies. Perform the Airflow upgrade and database migrations using the new CLI command `airflow db upgrade`. We execute these manually on our Scheduler instance, as it’s beefier than our web interface instance. For us, the database migration portion took about 30 minutes on an RDS instance with an HDD (for reference, our dag_run table had about 1.9 million records, while the task_instance table had about 2.4 million records). We also upgraded the instances to SSD after this, which we definitely recommend and would likely speed up your migration time. We recommend using the airflow-maintenance-dags to regularly trim your dag_runs and task_instances tables, as well as any on-disk log files, both for storage savings but also to minimize the runtime of database migrations. Restart the scheduler and web server Airflow processes. You may need to disable certain DAGs ahead of time if they should not be executed in your staging environment before turning the staging scheduler on! More on this later in our learnings/failures.

We suggest doing the above steps at least a week in advance so that you can monitor the staging stack and any enabled DAGs on it to ensure things execute as expected.

Step 2: Upgrade Production

If the downtime (the time it took to spin your tasks down, install Airflow 2.0, and run the database migrations) above is acceptable, you’re ready to proceed! Stop the scheduler, web server, let all in-flight tasks complete, and spool down your worker autoscaling group. Create a backup of the RDS database once everything is stopped.

Install new packages on the scheduler and web interface (i.e, run pip install). Perform the upgrade / database migration using the updated CLI command `airflow db upgrade`. Again, we ran these on our beefier scheduler instance instead of the web server. Restart the scheduler and web server processes/daemons. Spool up your worker autoscaling group. Re-enable any CI/CD processes you may have paused. Monitor critical metrics and ‘airflow.log’ files on relevant instances.

The upgrade steps, themselves, were relatively straightforward. We took the same steps in our previous upgrades for several minor airflow revisions pre-2.0, always updating our internal checklist with learnings and documenting any issues in team-written postmortems.

Issues Encountered

After completing the upgrade, we noticed some issues almost immediately in production that we had not encountered in staging:

Some DAG runs would intermittently get picked up by the scheduler, while many other DAGs would remain unscheduled. We would also get a heartbeat error in the Airflow UI: “The scheduler does not appear to be running. Last heartbeat was received X minutes ago. The DAGs list may not update, and new tasks will not be scheduled.”

DAGs were also taking a long time to get marked as successfully completed, even when all of their tasks completed successfully. Some DAG runs took over 15 minutes to get marked as completed. However, when checking the scheduler logs, we did not see any errors we would typically expect to see, such as “max active runs reached”, “no slots available in pool”, etc. The logs just did not seem as active or noisy as they normally are.

Example DAG waiting to get marked as completed

Scrambling, we checked several things. The CPU usage was fine on the scheduler instance; our job queues and worker instances remained fairly idle. The database also seemed fine; there were no locks or anything that could have caused conflicts.

We dug into the airflow.cfg file and attempted adjusting the “scheduler_heartbeat_sec” parameter. Still nothing. Then we tried pausing about 75% of our DAGs and noticed that the unpaused DAGs were immediately getting picked up by the scheduler again in a timely manner. DAG runs were also getting immediately marked as successfully completed again once their tasks were completed.

It felt like something about the scheduler was misconfigured, so we began diving into the scheduler internals on GitHub, specifically searching for the text “marking run” and “successful.” This led us to the DAG run code, which further led us to scheduler_job.py, where we found some important details around the variable max_dagruns_per_loop_to_schedule. The code comments helped us understand the DAG scheduling better and why we were seeing a long delay in DAGs progressing from “running” to “success.”

Details around the Scheduler Looping

We realized that we were using the default of only 20 DAG runs per scheduler loop, since we did not explicitly set this variable in the Airflow.cfg. We hypothesized that this number was too low for our case, as we have much shallower DAGs. We tested increasing this number to 200 and saw the scheduler roar back to life. We eventually settled on a value of 500 in our production environment and are seeing sub-second latency on scheduling tasks. A huge win!

We have also set the related variable, “max_dagruns_to_create_per_loop”, to 500, mostly as a preventative measure.

Learnings

Airflow 2.0 was a huge upgrade that touched virtually all of the underlying infrastructure (scheduler, web server, workers). Several challenges were encountered for this particular upgrade, which we compiled into an internal postmortem to identify where our upgrade plan went awry and how we will improve for future upgrades.

What went well:

Metastore/database table pruning in preparation of our upgrade: The documentation for Airflow 2.0 does a good job at calling this action item out for this upgrade. In our previous upgrade from 1.10.4 to 1.10.10, we did not prune the tables. The migration ended up taking several hours, as certain tables were especially large (e.g, Task Instance, DAG Runs, Log, etc. tables). Each contained tens of millions of rows and we had a rather painful outage as a result. Since then, we’ve implemented the aforementioned DB and Log cleanup DAGs shared by Team Clairvoyant. These DAGs run daily and purge data older than 30 days. Now that we’re running more and more DAGs every minute, we will likely increase the frequency of these DAGs.

Our production stack and database: We used our CloudFormation template to spin up a clone of our production environment. We also created a snapshot of our production database and pointed the cloned database to the snapshot. This allowed us to test our upgrade, as closely as possible. This also allowed us to get a good sense of the total downtime, especially during the database migration. Although this allowed us to do some thorough testing of the actual upgrade steps, we did not have a great way to fully test all of our production DAGs. We’ll discuss this more later.

Code updates: although more code updates were required for breaking changes, these were fairly trivial.

What didn’t go well:

Package dependencies: A few of our Python packages were deployed using versions that caused conflicts after our upgrade. This caused errors or expected behavior in our upgrade. One example was that the Celery package should be pinned to version 4.4.7 (at the time of this writing) as there have been reported issues with the latest 5.0.2 version.

For our Airflow implementation, we have about ~160 Python packages. It becomes very tricky to keep all the dependencies and versions working nicely with each other. Thankfully, starting with 1.10.10, the authors of Airflow have thoughtfully created a set of constraint files for each Python version (major and minor) along with their “known to be working” set of dependencies. In the future, we will more strictly pin to Airflow’s versions for the packages that we use. Here is an example constraint file.

Staging vs Production throughput: As previously mentioned, we have a 1000 or so active DAGs in our production environment. Most of these DAGs interface with third party data sources in one way or another; almost all of which have strict concurrency limitations. We use Airflow’s Pool feature to limit concurrency, but obviously that does not work across different Airflow deployments. We couldn’t afford to disable our incremental data jobs in production while testing our staging environment, so we never hit the “critical mass” of scheduler work in order to surface the chronic slowdown that we eventually fixed by upping the max_dagruns_per_loop_to_schedule variable. We have some neat ideas about how to account for this in the future; stay tuned! In the meantime, however, you should strongly consider running as many tasks as possible in your staging clusters to better simulate any configuration changes or upgrades.

One idea is to use an Airflow Variable, something like ‘skip_third_party_connections’, that all of your DAGs or tasks can check, and short-circuit exit. This way, you can still test the throughput of the scheduler and your worker setup. In a perfect world, we would purge dag_runs, task_instances, etc, and measure the total time to schedule and execute a known time range of DAGs to better test configuration changes in isolation.

Airflow Configuration Variables: It can be a bit daunting to understand all of the Airflow config variables. Estimating the right values for some settings may take a bit of trial and error. Initially missing and not fully understanding the new “max_dagruns_per_loop_to_schedule” variable took us quite some time to figure out. Although our test environment mirrored our actual production environment, we could not fully run all of the test DAGs to simulate what would happen in production. This caused us to miss the scheduling bottleneck we eventually encountered.

CLI Command Updates: In our last upgrade, we did not need to make any specific changes to our worker servers. But in this upgrade, there was a change to many of the CLI commands. In this case, “airflow celery worker” replaced the old “airflow worker”. As this was an additional manual step, we missed this change in our CloudFormation template during the worker instance provisioning steps. So once the ECS instances scaled up (via our auto scaling group policy), the actual celery workers on those instances never successfully started. This caused an extra half hour or so of extra downtime until we found the issue.

Big Wins

The biggest win of Airflow 2.0 is, hands down, the scheduler improvements. We saw immediate timeliness improvements in our DAG runs. The new scheduler has also enabled us to schedule DAGs on intervals as frequent as once-per-minute with very little latency. This was not possible prior to Airflow 2.0. Check out this great write up by the team at Astronomer if you are interested in the granular details on the improved scheduler.

Overall, we are extremely pleased with Airflow 2.0. The Airflow authors and community have thoughtfully added many improvements and stability features. There were many learnings gained from this upgrade and in our continued usage in the past couple of weeks. We expect more learnings to come over the upcoming weeks and months, as we look into more performance tuning and considerations such as multi-scheduler-process setups. We will also investigate ways to better simulate running all DAGs in a test environment.

Stay tuned for future posts as we share more learnings and advice on these topics! If you have questions or comments, please feel free to reach out here. We would love to hear from other users in the Airflow community!

Last, but certainly not least, I would like to recognize my colleague, Anthony Sosso, for his invaluable inputs and contributions throughout the Airflow upgrade and this article!