Explicitly initialize Task SDK Stats in API server lifespan (hardening; backport fix for 3.2.x)

[diogosilva30]

closes: #68077

Why

The API server serves the Edge Worker REST API (/edge_worker/v1/...). A worker heartbeat (PATCH /edge_worker/v1/worker/) runs set_state → set_metrics, which records edge_worker.* metrics through the Task SDK Stats singleton (resolved by the Edge provider via airflow.providers.common.compat).

On main: hardening against a fragile import side effect

The bug does not reproduce on current main. On main, task_runner.py imports serde at module level, which triggers serde._register() → Stats.initialize(...) as a side effect. That fires on any import airflow — including in the API server — so Stats gets initialized accidentally.

However, this is fragile: it would silently regress if the module-level serde import in task_runner were ever made lazy, or if the mask-secrets import chain changed. The API server doesn't actually depend on serde, so relying on this side effect is not intentional.

Adding an explicit Stats.initialize(...) in the FastAPI lifespan makes the dependency intentional and self-contained. It's a no-op on main (already initialized) and removes the accidental coupling.

On 3.2.x: real bug fix

On 3.2.x, task_runner imports serde lazily (inside functions), so the accidental Stats.initialize side effect never fires in the API server process. As a result, the Task SDK Stats singleton stays a NoStatsLogger and every Edge Worker metric is silently dropped.

This explains why api_server.* metrics still work on 3.2.x (they use the separately-initialized core stats path) while all edge_worker.* metrics vanish.

A backport of this change to the 3.2.x stable branch is the actual user-visible fix.

What

Initialize the Task SDK Stats singleton from the FastAPI lifespan (runs once per worker, post-fork), mirroring the existing explicit init in other long-running components (scheduler_job_runner.py, triggerer_job_runner.py, dag_processing/manager.py, executors/base_executor.py, task_runner.py). The call is guarded so a metrics misconfiguration can never block API server startup.

How verified

• New unit tests in airflow-core/tests/unit/api_fastapi/test_app.py: asserts lifespan initializes Task SDK Stats with the configured factory, and that an init failure is swallowed (startup not blocked).
• Verified on a live Airflow 3.2.2 + edge3 3.7.0 deployment: before the fix Stats.instance is NoStatsLogger and no edge_worker.* series export; after the fix 200+ edge_worker.* series export, correctly tagged with worker_name.
• Verified on main with breeze start-airflow --executor EdgeExecutor --integration statsd: edge_worker.* series present with and without the lifespan init (confirming main is not broken; the explicit init becomes a no-op).

Relationship to #67328

Not a duplicate. #67328 reintroduces a DualStatsManager code path in edge3's set_metrics, guarded by try: from airflow.sdk.observability.stats import DualStatsManager. DualStatsManager was removed in #63932, so on any Airflow that includes #63932 (i.e. 3.2+) that import fails and edge3 falls back to the existing Stats.gauge(..., tags={"worker_name": ...}) path. Neither branch of #67328 calls Stats.initialize(), so Edge Worker metrics on 3.2.x still depend on this PR.

Relationship to #64523

Not redundant. #64523 adds HttpMetricsMiddleware that emits REST API request metrics (api.requests, api.request.duration, api.request.errors) using the Stats singleton. If this PR is merged first, those metrics flow through on main regardless (due to the accidental init). But on 3.2.x or if the accidental import side effect ever regresses on main, #64523's metrics would also silently drop — making this an implicit prerequisite for #64523 as well.

---

Was generative AI tooling used to co-author this PR?

• Yes — GitHub Copilot (Claude Sonnet 4.6)

Generated-by: GitHub Copilot (Claude Sonnet 4.6) following the guidelines

[diogosilva30]

Some interesting follow-up findings here: #68077 (comment)

[jscheffl]

Looks good, but there is another parallel PR in #64523 which adds API server metrics in general - is this PR redundant if the other is merged? Still needed?

[diogosilva30]

Hey @jscheffl

Not redundant, I'd say it's actually complementary and something of a prerequisite for #64523.

After more investigation the situation turned out to be more nuanced than the original PR description suggested, so I've updated the PR body accordingly.

On main: the bug does not reproduce. Stats gets initialized in the API server as an unintentional side effect: task_runner.py imports serde at module level, which triggers serde._register() → Stats.initialize(...) during import airflow. Details in the issue. So the metrics introduced by #64523 work on main today, but only because of that fragile coupling. If the module-level serde import in task_runner ever becomes lazy again (as it was on 3.2.x before #63932), both Edge Worker metrics and the new API metrics from #64523 would silently regress with no obvious cause. This PR makes the initialization explicit and adds a regression test, so neither set of metrics can be silently dropped by a future import change.

On 3.2.x: the bug is real and verified. The serde import in task_runner is lazy on 3.2.x so the accidental init never fires in the API server process. Stats stays a NoStatsLogger and all edge_worker.* metrics are silently dropped. A backport to the 3.2.x stable branch would be the actual user-visible fix. Happy to open one if that's desirable.

[jscheffl]

General makes sense and looks good to me - would propose to make this as a safety net as described. The current "enablement" via the description really looks brittle and some refactoring (actually planned as well) might break it apart.

One safety question... if the reason for the bad init in Edge is the remobal of DualStatsMgr and the usage of the Stats in common-compat - would it also be a fix to add such handling in common-compat such that both core as well as edge can stay "naive" and the common compat bridges the gap? Then also a upgrade of common-compat provider would be a fix to make metrics in Airflow 3.2.x working again knowing that there will not be a next 3.2.x release.

[jscheffl]

Note: Static checks seems to be not related to this PR.

[diogosilva30]

@jscheffl Good point. A common-compat bridge could actually fix 3.2.x retroactively, which this PR can't since there won't be another 3.2.x release.

I'd still keep the init here though. Every long-running component (scheduler, triggerer, dag-processor, executors, task-runner) already calls Stats.initialize(...) itself, so the API server was just missing from that list and this brings it in line. Doing it in common-compat would mean the provider reaching into core internals for the factory and legacy_names_on config to set the singleton up correctly, which adds coupling I'd rather avoid. The lifespan init also covers consumers that don't go through compat at all, like the HttpMetricsMiddleware in #64523 that uses the Stats singleton directly.

So I see them as complementary. This is the clean fix going forward, and if we want to help existing 3.2.x users I'm happy to open a follow-up for a targeted common-compat change.

[jscheffl]

Note: If comments adjusted and no objections would propose to merge to get this into 3.3.0 - and as ^^^an improvement in common-compat can be made irrespective in a parallel PR.

[ashb]

@jscheffl This PR makes no sense to me.

It's talking about Task SDK, but it doesn't import anything from in task SDK (nor should it).

Task SDK code should not run in the APi Server. Anything relying on that is fundamentally broken and breaking boundaries. If that is the reason this change is introduced then this should be reverted. #68481

[jscheffl]

@jscheffl This PR makes no sense to me.

It's talking about Task SDK, but it doesn't import anything from in task SDK (nor should it).

Task SDK code should not run in the APi Server. Anything relying on that is fundamentally broken and breaking boundaries. If that is the reason this change is introduced then this should be reverted. #68481

Sorry, it seems I was not reading through in detail. In earlier commits that was referring to "Edge stats sent via API server not recorded" and I requested to remove Edge because this was a general shortcoming that API server did not properly initialize stats. Somehow I over-looked the "Task SDK" term.

No of course Task SDK is not running in API server and this PR does not change this. But workding is still wrong. Let me have a corrective PR. Will do this in the afternoon. Note it is about commit message, exception text and comments, yes should be corrected. But still the general stats initialization was missing.

Positively speaking you can read it "if Task SDK is sending something that should be recorded in stats via API server it was missing" (which was the case of Edge in 3.2).