This is an automated email from the ASF dual-hosted git repository.
amoghrajesh pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new ead86184fb5 Emit resumable_job metrics from KubernetesPodOperator
(#75734)
ead86184fb5 is described below
commit ead86184fb5d75ae8226b8a3c13365ede3cb1675
Author: Amogh Desai <[email protected]>
AuthorDate: Tue Oct 27 16:38:13 2026 +0530
Emit resumable_job metrics from KubernetesPodOperator (#75734)
---
.../providers/cncf/kubernetes/operators/pod.py | 16 ++-
.../unit/cncf/kubernetes/operators/test_pod.py | 134 +++++++++++++++++++++
2 files changed, 149 insertions(+), 1 deletion(-)
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py
index ede4462c355..44a5d09d359 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py
@@ -87,7 +87,7 @@ from airflow.providers.cncf.kubernetes.utils.pod_manager
import (
detect_pod_terminate_early_issues,
)
from airflow.providers.cncf.kubernetes.version_compat import
AIRFLOW_V_3_1_PLUS, AIRFLOW_V_3_3_PLUS
-from airflow.providers.common.compat.sdk import XCOM_RETURN_KEY,
AirflowSkipException, conf
+from airflow.providers.common.compat.sdk import XCOM_RETURN_KEY,
AirflowSkipException, Stats, conf
if AIRFLOW_V_3_1_PLUS:
from airflow.sdk import BaseHook, BaseOperator
@@ -745,18 +745,25 @@ class KubernetesPodOperator(BaseOperator):
{"name": pod.metadata.name, "namespace": pod.metadata.namespace},
)
+ def _resume_stats_tags(self, context: Context) -> dict[str, str]:
+ ti = context.get("ti")
+ return {**(getattr(ti, "stats_tags", None) or {}), "operator":
type(self).__name__}
+
def get_or_create_pod(self, pod_request_obj: k8s.V1Pod, context: Context)
-> k8s.V1Pod:
if self.reattach_on_restart:
+ stats_tags = self._resume_stats_tags(context)
pod = self._get_pod_from_task_state_store(context)
if pod is None:
pod = self.find_pod(pod_request_obj.metadata.namespace,
context=context)
if pod is not None:
self._persist_pod_identity_to_task_state_store(context,
pod)
if pod:
+ Stats.incr("resumable_job.reconnect_attempt", tags=stats_tags)
# If pod is terminated then delete the pod an create a new as
not possible to get xcom
pod_phase = pod.status.phase if pod.status and
pod.status.phase else None
pod_reason = pod.status.reason.lower() if pod.status and
pod.status.reason else ""
if pod_phase not in (PodPhase.SUCCEEDED, PodPhase.FAILED) and
pod_reason != "evicted":
+ Stats.incr("resumable_job.reconnect_success",
tags=stats_tags)
self.log.info(
"Reusing existing pod '%s' (phase=%s, reason=%s) since
it is not terminated or evicted.",
pod.metadata.name,
@@ -765,6 +772,11 @@ class KubernetesPodOperator(BaseOperator):
)
return pod
+ if pod_phase == PodPhase.SUCCEEDED:
+ Stats.incr("resumable_job.already_succeeded",
tags=stats_tags)
+ else:
+ Stats.incr("resumable_job.terminal_resubmit",
tags=stats_tags)
+
self.log.info(
"Found terminated old matching pod %s with labels %s",
pod.metadata.name,
@@ -777,6 +789,8 @@ class KubernetesPodOperator(BaseOperator):
return pod
self.log.info("Deleted pod to handle rerun and create new
pod!")
+ else:
+ Stats.incr("resumable_job.fresh_submit", tags=stats_tags)
self.log.debug("Starting pod:\n%s",
yaml.safe_dump(pod_request_obj.to_dict()))
created_pod = self.pod_manager.create_pod(pod=pod_request_obj)
diff --git
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py
index 19a681ac97a..56ace367c95 100644
--- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py
+++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py
@@ -2998,6 +2998,140 @@ class TestKubernetesPodOperatorDurableExecution:
assert
KubernetesPodOperator._KubernetesPodOperator__supports_durable_execution is True
+class TestKubernetesPodOperatorResumableMetrics:
+ PATCH = f"{KPO_MODULE}.Stats.incr"
+
+ @pytest.fixture(autouse=True)
+ def setup_tests(self):
+ self.create_pod_patch = patch(f"{POD_MANAGER_CLASS}.create_pod")
+ self.create_mock = self.create_pod_patch.start()
+ self._default_client_patch = patch(f"{HOOK_CLASS}._get_default_client")
+ self._default_client_patch.start()
+
+ yield
+
+ patch.stopall()
+
+ def make_operator(self):
+ k = KubernetesPodOperator(
+ image="ubuntu:16.04",
+ task_id="task",
+ name="hello",
+ log_pod_spec_on_failure=False,
+ )
+ k.reattach_on_restart = True
+ return k
+
+ def make_pod(self, phase, reason=None):
+ pod = MagicMock()
+ pod.status.phase = phase
+ pod.status.reason = reason
+ return pod
+
+ def make_pod_request(self):
+ request = MagicMock()
+ request.to_dict.return_value = {"metadata": {"name": "test-pod"}}
+ return request
+
+ @pytest.mark.parametrize(
+ ("phase", "reason", "from_store", "expected"),
+ [
+ pytest.param(None, None, False, ["resumable_job.fresh_submit"],
id="fresh_submit"),
+ pytest.param(
+ PodPhase.RUNNING,
+ None,
+ True,
+ ["resumable_job.reconnect_attempt",
"resumable_job.reconnect_success"],
+ id="reconnect",
+ ),
+ pytest.param(
+ PodPhase.SUCCEEDED,
+ None,
+ True,
+ ["resumable_job.reconnect_attempt",
"resumable_job.already_succeeded"],
+ id="already_succeeded",
+ ),
+ pytest.param(
+ PodPhase.FAILED,
+ None,
+ True,
+ ["resumable_job.reconnect_attempt",
"resumable_job.terminal_resubmit"],
+ id="terminal_resubmit",
+ ),
+ pytest.param(
+ PodPhase.RUNNING,
+ "Evicted",
+ True,
+ ["resumable_job.reconnect_attempt",
"resumable_job.terminal_resubmit"],
+ id="evicted",
+ ),
+ pytest.param(
+ PodPhase.RUNNING,
+ None,
+ False,
+ ["resumable_job.reconnect_attempt",
"resumable_job.reconnect_success"],
+ id="found_by_label_search",
+ ),
+ ],
+ )
+ def test_counters_per_branch(self, phase, reason, from_store, expected):
+ k = self.make_operator()
+ context = create_context(k)
+ context["task_state_store"] = MagicMock()
+ pod = self.make_pod(phase, reason) if phase is not None else None
+
+ with (
+ patch(
+
f"{KPO_MODULE}.KubernetesPodOperator._get_pod_from_task_state_store",
+ return_value=pod if from_store else None,
+ ),
+ patch(
+ f"{KPO_MODULE}.KubernetesPodOperator.find_pod",
+ return_value=None if from_store else pod,
+ ),
+ patch(f"{KPO_MODULE}.KubernetesPodOperator.process_pod_deletion"),
+ patch(self.PATCH) as incr,
+ ):
+ k.get_or_create_pod(pod_request_obj=self.make_pod_request(),
context=context)
+
+ assert [call.args[0] for call in incr.call_args_list] == expected
+
+ def test_counters_carry_task_instance_stats_tags(self):
+ k = self.make_operator()
+ context = create_context(k)
+ context["task_state_store"] = MagicMock()
+ ti = MagicMock()
+ ti.stats_tags = {"dag_id": "dag", "task_id": "task", "run_type":
"manual"}
+ context["ti"] = ti
+
+ with (
+
patch(f"{KPO_MODULE}.KubernetesPodOperator._get_pod_from_task_state_store",
return_value=None),
+ patch(f"{KPO_MODULE}.KubernetesPodOperator.find_pod",
return_value=None),
+ patch(self.PATCH) as incr,
+ ):
+ k.get_or_create_pod(pod_request_obj=self.make_pod_request(),
context=context)
+
+ incr.assert_called_once_with(
+ "resumable_job.fresh_submit",
+ tags={
+ "dag_id": "dag",
+ "task_id": "task",
+ "run_type": "manual",
+ "operator": "KubernetesPodOperator",
+ },
+ )
+
+ def test_no_counters_when_not_durable(self):
+ k = self.make_operator()
+ k.reattach_on_restart = False
+ context = create_context(k)
+
+ with patch(self.PATCH) as incr:
+ k.get_or_create_pod(pod_request_obj=self.make_pod_request(),
context=context)
+
+ incr.assert_not_called()
+
+
class TestKubernetesPodOperatorDurableBelow3_3:
@pytest.mark.parametrize("durable_value", [True, False])
def test_durable_has_no_effect(self, durable_value):