ashb commented on code in PR #63932:
URL: https://github.com/apache/airflow/pull/63932#discussion_r3116206245
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -1863,8 +1863,10 @@ def main():
global SUPERVISOR_COMMS
SUPERVISOR_COMMS = CommsDecoder[ToTask, ToSupervisor](log=log)
- stats_factory = stats_utils.get_stats_factory(Stats)
- Stats.initialize(factory=stats_factory)
+ stats.initialize(
+ factory=stats_utils.get_stats_factory(),
Review Comment:
Do we ever pass anything but `factory=stats_utils.get_stats_factory()` for
this arg? Any reason we shouldn't put this inside `stats.initialize` directly?
##########
providers/openlineage/src/airflow/providers/openlineage/plugins/adapter.py:
##########
@@ -162,18 +158,18 @@ def emit(self, event: RunEvent):
transport_type = f"{self._client.transport.kind}".lower()
try:
- with ExitStack() as stack:
- if DualStatsManager is not None:
- stack.enter_context(
- DualStatsManager.timer(
- "ol.emit.attempts",
- extra_tags={"event_type": event_type,
"transport_type": transport_type},
- )
- )
- else:
-
stack.enter_context(Stats.timer(f"ol.emit.attempts.{event_type}.{transport_type}"))
- stack.enter_context(Stats.timer("ol.emit.attempts"))
-
+ ctx: AbstractContextManager[object]
+ if AIRFLOW_V_3_2_1_PLUS:
+ ctx = Stats.timer(
+ "ol.emit.attempts",
+ legacy_name_tags={"event_type": event_type,
"transport_type": transport_type},
Review Comment:
I'm still confused by this parameter name. Isn't aren't these tags the _non_
legacy version?
##########
task-sdk/src/airflow/sdk/observability/metrics/stats_utils.py:
##########
Review Comment:
Ditto -- duplicating shared code?
##########
airflow-core/src/airflow/cli/commands/daemon_utils.py:
##########
@@ -75,9 +75,7 @@ def run_command_with_daemon_option(
with ctx:
# in daemon context stats client needs to be reinitialized.
- from airflow._shared.observability.metrics.stats import Stats
-
- Stats.instance = None
+ # The stats module auto-detects fork via PID and
re-initializes lazily.
Review Comment:
Comment only makes sense in context of the diff, after that we shouldn't
keep the comment
##########
providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py:
##########
@@ -204,7 +205,13 @@ def on_running():
if not doc:
doc, doc_type = get_dag_documentation(dag)
- with Stats.timer(f"ol.extract.{event_type}.{operator_name}"):
+ if AIRFLOW_V_3_2_1_PLUS:
+ ctx = Stats.timer(
+ "ol.extract", legacy_name_tags={"event_type": event_type,
"operator_name": operator_name}
+ )
+ else:
+ ctx = Stats.timer(f"ol.extract.{event_type}.{operator_name}")
+ with ctx:
Review Comment:
Could/should this difference be handled once somewhere in
`airflow.providers.common.compat.sdk` perhaps?
##########
task-sdk/src/airflow/sdk/observability/metrics/datadog_logger.py:
##########
Review Comment:
I wonder why this file even exists -- isn't this a copy of what is in there
shared lib? (Can be addressed in a different PR)
##########
shared/observability/src/airflow_shared/observability/metrics/stats.py:
##########
@@ -62,71 +71,235 @@ def normalize_name_for_stats(name: str, log_warning: bool
= True) -> str:
return normalized
-class _Stats(type):
- factory: Callable[[], StatsLogger | NoStatsLogger] | None = None
- instance: StatsLogger | NoStatsLogger | None = None
-
- def __getattr__(cls, name: str) -> str:
- factory = type.__getattribute__(cls, "factory")
- instance = type.__getattribute__(cls, "instance")
-
- # When using OpenTelemetry, some subprocesses are short-lived and
- # often exit before flushing any metrics.
- #
- # The solution is to register a hook that performs a force flush at
exit.
- # The atexit hook is registered when initializing the instance.
- #
- # The instance gets initialized once per process. In case a process is
forked, then
- # the new subprocess, will inherit the already initialized instance of
the parent process.
- #
- # Store the instance pid so that it can be compared with the current
pid
- # to decide whether to initialize the instance again or not.
- #
- # So far, all forks are resetting their state to remove anything
inherited by the parent.
- # But in the future that might not always be true.
- current_pid = os.getpid()
- if cls.instance and cls._instance_pid != current_pid:
- log.info(
- "Stats instance was created in PID %s but accessed in PID %s.
Re-initializing.",
- cls._instance_pid,
- current_pid,
- )
- # Setting the instance to None, will force re-initialization.
- cls.instance = None
- cls._instance_pid = None
+def initialize(
+ *,
+ factory: Callable[[], StatsLogger | NoStatsLogger],
+ export_legacy_names: bool,
+) -> None:
+ """Initialize the stats module with a backend factory and legacy name
configuration."""
+ global _factory, _backend, _instance_pid, _export_legacy_names
+ _factory = factory
+ _backend = None
+ _instance_pid = None
+ _export_legacy_names = export_legacy_names
- if instance is None:
- if factory is None:
- factory = NoStatsLogger
- type.__setattr__(cls, "factory", factory)
- try:
- instance = factory()
- cls._instance_pid = current_pid
- except (socket.gaierror, ImportError) as e:
- log.error("Could not configure StatsClient: %s, using
NoStatsLogger instead.", e)
- instance = NoStatsLogger()
- cls._instance_pid = current_pid
+def _get_backend() -> StatsLogger | NoStatsLogger:
+ """Return the current backend, re-initializing if the process has been
forked."""
+ global _backend, _instance_pid
+ current_pid = os.getpid()
- type.__setattr__(cls, "instance", instance)
+ if _backend is not None and _instance_pid != current_pid:
+ log.info(
+ "Stats backend was created in PID %s but accessed in PID %s.
Re-initializing.",
+ _instance_pid,
+ current_pid,
+ )
+ _backend = None
+ _instance_pid = None
- return getattr(instance, name)
+ if _backend is None:
+ factory = _factory if _factory is not None else NoStatsLogger
+ try:
+ _backend = factory()
+ _instance_pid = current_pid
+ except (socket.gaierror, ImportError) as e:
+ log.error("Could not configure StatsClient: %s, using
NoStatsLogger instead.", e)
+ _backend = NoStatsLogger()
+ _instance_pid = current_pid
- def initialize(cls, *, factory: Callable) -> None:
- type.__setattr__(cls, "instance", None)
- type.__setattr__(cls, "factory", factory)
+ return _backend
- @classmethod
- def get_constant_tags(cls, *, tags_in_string: str | None) -> list[str]:
- """Get constant DataDog tags to add to all stats."""
- if not tags_in_string:
- return []
- return tags_in_string.split(",")
+def _get_legacy_stat(stat: str, variables: dict[str, Any]) -> str | None:
+ """
+ Look up and format the legacy name for a metric from the registry.
-if TYPE_CHECKING:
- Stats: StatsLogger
-else:
+ Returns the formatted legacy name, or None if the metric has no legacy
name.
+ Raises ValueError if the metric is not in the registry or required
variables are missing.
+ """
+ stat_from_registry = _registry.get(name=stat)
+
+ if not stat_from_registry:
+ raise ValueError(
+ f"Metric '{stat}' not found in the registry. Add the metric to the
YAML file before using it."
+ )
+
+ legacy_name = stat_from_registry.get("legacy_name", "-")
+
+ if legacy_name == "-":
+ return None
+
+ required_vars = stat_from_registry.get("name_variables", [])
+ missing_vars = set(required_vars) - set(variables.keys())
+ if missing_vars:
+ raise ValueError(
+ f"Missing required variables for metric '{stat}':
{sorted(missing_vars)}. "
+ f"Required variables found in the registry: {required_vars}. "
+ f"Provided variables: {sorted(variables.keys())}. "
+ f"Provide all required variables."
+ )
+
+ legacy_vars = {k: variables[k] for k in required_vars if k in variables}
+ return legacy_name.format(**legacy_vars)
+
+
+def _defined(**kwargs: Any) -> dict[str, Any]:
+ """Return only kwargs with a meaningful (non-None, non-empty) value."""
+ result = {}
+ for k, v in kwargs.items():
+ if v is None:
+ continue
+ try:
+ if len(v) == 0:
+ continue
+ except TypeError:
+ pass # Numerics and bools don't have len, always include
+ result[k] = v
+ return result
+
+
+def _merged(
+ tags: dict[str, Any] | None,
+ legacy_name_tags: dict[str, Any] | None,
+) -> dict[str, Any] | None:
+ """Merge tags and legacy_name_tags into a single dict, returning None if
both are empty."""
+ if not tags and not legacy_name_tags:
+ return None
+ merged: dict[str, Any] = {}
+ if tags:
+ merged.update(tags)
+ if legacy_name_tags:
+ merged.update(legacy_name_tags)
+ return merged or None
+
+
+def _emit_legacy(stat: str, legacy_name_tags: dict[str, Any]) -> str:
+ """Return the formatted legacy stat name, raising ValueError if it has
none."""
+ legacy_stat = _get_legacy_stat(stat, legacy_name_tags)
+ if legacy_stat is None:
+ raise ValueError(f"Stat '{stat}' doesn't have a legacy name registered
in the YAML file.")
+ return legacy_stat
+
+
+def incr(
+ stat: str,
+ count: int | None = None,
+ rate: int | float | None = None,
+ *,
+ tags: dict[str, Any] | None = None,
+ legacy_name_tags: dict[str, Any] | None = None,
+) -> None:
+ """Increment a counter metric."""
+ if _export_legacy_names and legacy_name_tags is not None:
+ _get_backend().incr(
+ _emit_legacy(stat, legacy_name_tags), **_defined(count=count,
rate=rate, tags=tags)
+ )
+ _get_backend().incr(stat, **_defined(count=count, rate=rate,
tags=_merged(tags, legacy_name_tags)))
+
+
+def decr(
+ stat: str,
+ count: int | None = None,
+ rate: int | float | None = None,
+ *,
+ tags: dict[str, Any] | None = None,
+ legacy_name_tags: dict[str, Any] | None = None,
+) -> None:
+ """Decrement a counter metric."""
+ if _export_legacy_names and legacy_name_tags is not None:
+ _get_backend().decr(
+ _emit_legacy(stat, legacy_name_tags), **_defined(count=count,
rate=rate, tags=tags)
+ )
+ _get_backend().decr(stat, **_defined(count=count, rate=rate,
tags=_merged(tags, legacy_name_tags)))
+
+
+def gauge(
+ stat: str,
+ value: float,
+ rate: int | float | None = None,
+ delta: bool | None = None,
+ *,
+ tags: dict[str, Any] | None = None,
+ legacy_name_tags: dict[str, Any] | None = None,
+) -> None:
+ """Set a gauge metric."""
+ if _export_legacy_names and legacy_name_tags is not None:
+ _get_backend().gauge(
+ _emit_legacy(stat, legacy_name_tags), value, **_defined(rate=rate,
delta=delta, tags=tags)
+ )
+ _get_backend().gauge(
+ stat, value, **_defined(rate=rate, delta=delta, tags=_merged(tags,
legacy_name_tags))
+ )
+
+
+def timing(
+ stat: str,
+ dt: DeltaType,
+ *,
+ tags: dict[str, Any] | None = None,
+ legacy_name_tags: dict[str, Any] | None = None,
+) -> None:
+ """Record a timing metric."""
+ if _export_legacy_names and legacy_name_tags is not None:
+ _get_backend().timing(_emit_legacy(stat, legacy_name_tags), dt,
**_defined(tags=tags))
+ _get_backend().timing(stat, dt, **_defined(tags=_merged(tags,
legacy_name_tags)))
+
+
+def timer(
+ stat: str | None = None,
+ tags: dict[str, Any] | None = None,
+ legacy_name_tags: dict[str, Any] | None = None,
+ **kwargs,
+) -> Timer:
+ """
+ Context manager that times a block and emits a timer metric.
+
+ When ``stat`` is None, returns the backend timer directly for stopwatch
+ use (duration is available but no metric is emitted).
+
+ When ``legacy_name_tags`` is provided and legacy name export is enabled,
both the
+ legacy metric name and the regular name are timed simultaneously via an
+ ExitStack. Otherwise, the backend timer is returned directly, preserving
the
+ ability to access ``timer.duration`` on the returned object.
+ """
+ if stat is None:
+ return _get_backend().timer()
+
+ regular_kw: dict[str, Any] = {**kwargs}
+ merged = _merged(tags, legacy_name_tags)
+ if merged:
+ regular_kw["tags"] = merged
+
+ if _export_legacy_names and legacy_name_tags is not None:
+ legacy_kw: dict[str, Any] = {**kwargs}
+ if tags is not None:
+ legacy_kw["tags"] = tags
+
+ stack = ExitStack()
+ stack.enter_context(
+ cast(
+ "AbstractContextManager[Any]",
+ _get_backend().timer(_emit_legacy(stat, legacy_name_tags),
**legacy_kw),
+ )
+ )
+ stack.enter_context(_get_backend().timer(stat, **regular_kw))
+ return cast("Timer", stack)
Review Comment:
Also this -- exit stack is not a timer, and the return should be a Timer.
##########
shared/observability/src/airflow_shared/observability/metrics/stats.py:
##########
@@ -22,18 +21,28 @@
import re
import socket
from collections.abc import Callable
-from typing import TYPE_CHECKING
+from contextlib import AbstractContextManager, ExitStack
+from typing import TYPE_CHECKING, Any, cast
from .base_stats_logger import NoStatsLogger
+from .metrics_registry import MetricsRegistry
if TYPE_CHECKING:
from .base_stats_logger import StatsLogger
+ from .protocols import DeltaType, Timer
log = logging.getLogger(__name__)
_VALID_STAT_NAME_CHARS_RE = re.compile(r"^[a-zA-Z0-9_.-]+$")
_INVALID_STAT_NAME_CHARS_RE = re.compile(r"[^a-zA-Z0-9_.-]")
+# Module-level singleton state.
+_factory: Callable[[], StatsLogger | NoStatsLogger] | None = None
+_backend: StatsLogger | NoStatsLogger | None = None
+_instance_pid: int | None = None
+_export_legacy_names: bool = True
+_registry: MetricsRegistry = MetricsRegistry()
+
Review Comment:
@xBis7 I think this one is worth addressing.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
