diff --git a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml index 536e439f876e2..4afcf7561d2f2 100644 --- a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml +++ b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml @@ -86,6 +86,14 @@ metrics: legacy_name: "-" name_variables: [] + - name: "resumable_job.reconnect_failure" + description: "Number of times a ResumableJobMixin operator found a stored run ID + on retry but could not determine the prior job's status because the status check + raised, so the reconnect could not proceed. Metric with operator tagging." + type: "counter" + legacy_name: "-" + name_variables: [] + - name: "ti_failures" description: "Overall task instances failures. Metric with dag_id and task_id tagging." type: "counter" diff --git a/task-sdk/src/airflow/sdk/bases/resumablejobmixin.py b/task-sdk/src/airflow/sdk/bases/resumablejobmixin.py index ef1629da97503..2788930ec6340 100644 --- a/task-sdk/src/airflow/sdk/bases/resumablejobmixin.py +++ b/task-sdk/src/airflow/sdk/bases/resumablejobmixin.py @@ -149,14 +149,29 @@ def execute_resumable(self, context: Context) -> Any: if external_id: stats.incr("resumable_job.reconnect_attempt", tags=stats_tags) - status = self.get_job_status(external_id, context) - - span.set_attribute("resumable.external_id", str(external_id)) - span.set_attribute("resumable.prior_status", status) - - if self.is_job_active(status): + try: + status = self.get_job_status(external_id, context) + + span.set_attribute("resumable.external_id", str(external_id)) + span.set_attribute("resumable.prior_status", status) + + if self.is_job_active(status): + decision = "reconnect" + elif self.is_job_succeeded(status): + decision = "already_succeeded" + else: + decision = "terminal_resubmit" + except Exception: + # Prior job status could not be determined, so the reconnect could not be + # decided; count it so reconnect_attempt reconciles with its outcomes, then + # let the error propagate unchanged. + span.set_attribute("resumable.decision", "reconnect_failure") + stats.incr("resumable_job.reconnect_failure", tags=stats_tags) + raise + + span.set_attribute("resumable.decision", decision) + if decision == "reconnect": # Job is still running, skip submission and reconnect to it. - span.set_attribute("resumable.decision", "reconnect") stats.incr("resumable_job.reconnect_success", tags=stats_tags) self.log.info( "Reconnecting to existing job", @@ -165,9 +180,8 @@ def execute_resumable(self, context: Context) -> Any: status=status, ) reconnect_to = external_id - elif self.is_job_succeeded(status): + elif decision == "already_succeeded": # Job already finished successfully, skip polling and return result directly. - span.set_attribute("resumable.decision", "already_succeeded") stats.incr("resumable_job.already_succeeded", tags=stats_tags) self.log.info( "Job already completed successfully, skipping resubmission", @@ -177,7 +191,6 @@ def execute_resumable(self, context: Context) -> Any: already_succeeded_id = external_id else: # Job is in a terminal failed state, fall through and submit a new job. - span.set_attribute("resumable.decision", "terminal_resubmit") stats.incr("resumable_job.terminal_resubmit", tags=stats_tags) self.log.warning( "Prior job in terminal state, resubmitting fresh", diff --git a/task-sdk/tests/task_sdk/bases/test_resumablejobmixin.py b/task-sdk/tests/task_sdk/bases/test_resumablejobmixin.py index 1f84baa2d89e4..49f35adf69556 100644 --- a/task-sdk/tests/task_sdk/bases/test_resumablejobmixin.py +++ b/task-sdk/tests/task_sdk/bases/test_resumablejobmixin.py @@ -296,6 +296,48 @@ def test_terminal_resubmit_fires_when_job_failed(self): assert "resumable_job.reconnect_success" not in called_names assert "resumable_job.fresh_submit" not in called_names + @pytest.mark.parametrize("failing_method", ["get_job_status", "is_job_active", "is_job_succeeded"]) + def test_reconnect_failure_fires_when_reconnect_decision_raises(self, failing_method): + def boom(*args, **kwargs): + raise ConnectionError("external system unreachable") + + op = ConcreteResumableOperator(task_id="test_task") + op._status_map["job-001"] = "SUCCEEDED" + setattr(op, failing_method, boom) + mock_incr = MagicMock() + with patch(self._PATCH, mock_incr): + with pytest.raises(ConnectionError): + op.execute_resumable(make_context(FakeTaskState({"test_job_id": "job-001"}))) + called_names = [call.args[0] for call in mock_incr.call_args_list] + assert "resumable_job.reconnect_attempt" in called_names + assert "resumable_job.reconnect_failure" in called_names + assert "resumable_job.reconnect_success" not in called_names + assert "resumable_job.already_succeeded" not in called_names + assert "resumable_job.terminal_resubmit" not in called_names + assert "resumable_job.fresh_submit" not in called_names + + @pytest.mark.parametrize( + ("status", "outcome", "log_level"), + [ + ("RUNNING", "reconnect_success", "info"), + ("SUCCEEDED", "already_succeeded", "info"), + ("FAILED", "terminal_resubmit", "warning"), + ], + ) + def test_logging_failure_does_not_add_reconnect_failure(self, status, outcome, log_level): + op = ConcreteResumableOperator(task_id="test_task") + op._status_map["job-001"] = status + mock_incr = MagicMock() + error = ConnectionError("log sink unreachable") + with patch(self._PATCH, mock_incr), patch.object(op.log, log_level, side_effect=error): + with pytest.raises(ConnectionError, match="log sink unreachable") as exc: + op.execute_resumable(make_context(FakeTaskState({"test_job_id": "job-001"}))) + assert exc.value is error + assert mock_incr.call_args_list == [ + mock.call("resumable_job.reconnect_attempt", tags=self._TAG), + mock.call(f"resumable_job.{outcome}", tags=self._TAG), + ] + @pytest.mark.parametrize( ("team_name", "expected_tag"), [ @@ -374,6 +416,30 @@ def test_non_active_stored_job_span_attributes(self, status, expected_decision): assert span.attributes["resumable.external_id"] == "job-001" assert span.attributes["resumable.prior_status"] == status + @pytest.mark.parametrize( + ("status", "expected_decision", "log_level"), + [ + ("RUNNING", "reconnect", "info"), + ("SUCCEEDED", "already_succeeded", "info"), + ("FAILED", "terminal_resubmit", "warning"), + ], + ) + def test_logging_failure_preserves_decision(self, status, expected_decision, log_level): + op = ConcreteResumableOperator(task_id="test_task") + op._status_map["job-001"] = status + exporter, module_tracer = self._make_tracer() + with ( + mock.patch(self._MODULE_TRACER, module_tracer), + patch.object(op.log, log_level, side_effect=ConnectionError("log sink unreachable")), + ): + with pytest.raises(ConnectionError, match="log sink unreachable"): + op.execute_resumable(make_context(FakeTaskState({"test_job_id": "job-001"}))) + span = self._get_decision_span(exporter) + assert span is not None + assert span.attributes["resumable.decision"] == expected_decision + assert span.attributes["resumable.external_id"] == "job-001" + assert span.attributes["resumable.prior_status"] == status + class TestLogging: def test_warning_when_task_store_unavailable(self):