diff --git a/providers/openlineage/docs/troubleshooting.rst b/providers/openlineage/docs/troubleshooting.rst index f679e6ff015d9..9c358a6273ee4 100644 --- a/providers/openlineage/docs/troubleshooting.rst +++ b/providers/openlineage/docs/troubleshooting.rst @@ -43,6 +43,22 @@ as well as the `task_success_overtime OpenLineageAdapter: + """ + Return the per-process ``OpenLineageAdapter`` used inside pool worker processes. + + Each ``ProcessPoolExecutor`` worker keeps exactly one adapter — and therefore one + ``OpenLineageClient`` with one set of transports — for its whole lifetime. + """ + return OpenLineageAdapter() + + +def _run_adapter_method(method_name: str, /, *args, **kwargs): + """ + Run the named ``OpenLineageAdapter`` method on the per-process adapter. + + Module-level so it is picklable across the ProcessPoolExecutor boundary. Bound adapter + methods must not be submitted to the pool directly: pickling them serializes the whole + adapter, so the worker unpickles a fresh adapter per event and builds a new + ``OpenLineageClient`` (with new transports) on every emit. Transports that start + background worker threads (e.g. the ``datadog`` transport, which always starts an async + HTTP worker thread) are never closed, so this leaks one thread per event and steadily + consumes scheduler CPU and memory until restart. + """ + return getattr(_get_process_adapter(), method_name)(*args, **kwargs) + + +def _emit_manual_state_change_event(adapter_method_name: str, stats_key: str, **kwargs): """ - Emit an OL event via the given adapter method and record its serialized size. + Emit an OL event via the named adapter method and record its serialized size. Module-level so it is picklable across the ProcessPoolExecutor boundary used by `_on_task_instance_manual_state_change` for scheduler-side "task state changed - externally" emissions. + externally" emissions. The method is resolved on the per-process adapter so the + pool worker reuses one client across events (see ``_run_adapter_method``). """ - event = adapter_method(**kwargs) + event = getattr(_get_process_adapter(), adapter_method_name)(**kwargs) Stats.gauge(stats_key, len(Serde.to_json(event).encode("utf-8"))) return event @@ -782,10 +809,10 @@ def _on_task_instance_manual_state_change( return if ti_state == TaskInstanceState.FAILED: - adapter_method = self.adapter.fail_task + adapter_method_name = "fail_task" event_type = RunState.FAIL.value.lower() elif ti_state in (TaskInstanceState.SUCCESS, TaskInstanceState.SKIPPED): - adapter_method = self.adapter.complete_task + adapter_method_name = "complete_task" event_type = RunState.COMPLETE.value.lower() else: raise ValueError(f"Unsupported ti_state: `{ti_state}`.") @@ -867,7 +894,7 @@ def _on_task_instance_manual_state_change( operator_name = (ti.operator or "unknown").lower() self.submit_callable( _emit_manual_state_change_event, - adapter_method, + adapter_method_name, f"ol.event.size.{event_type}.{operator_name}", **adapter_kwargs, ) @@ -979,7 +1006,8 @@ def on_dag_run_running(self, dag_run: DagRun, msg: str) -> None: doc, doc_type = get_dag_documentation(dag_run.dag) self.submit_callable( - self.adapter.dag_started, + _run_adapter_method, + "dag_started", dag_id=dag_run.dag_id, run_id=dag_run.run_id, logical_date=date, @@ -1031,7 +1059,8 @@ def on_dag_run_success(self, dag_run: DagRun, msg: str) -> None: doc, doc_type = get_dag_documentation(dag_run.dag) self.submit_callable( - self.adapter.dag_success, + _run_adapter_method, + "dag_success", dag_id=dag_run.dag_id, run_id=dag_run.run_id, end_date=dag_run.end_date, @@ -1082,7 +1111,8 @@ def on_dag_run_failed(self, dag_run: DagRun, msg: str) -> None: doc, doc_type = get_dag_documentation(dag_run.dag) self.submit_callable( - self.adapter.dag_failed, + _run_adapter_method, + "dag_failed", dag_id=dag_run.dag_id, run_id=dag_run.run_id, end_date=dag_run.end_date, diff --git a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py index f3132746cd660..f7f3b398d4bfc 100644 --- a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py +++ b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py @@ -92,15 +92,23 @@ def direct_submit_call(self, callable, *args, **kwargs): Bypasses the ``ProcessPoolExecutor`` so tests can assert against mocked adapter methods without hitting pickling of ``unittest.mock.Mock``. - When the submitted callable is ``_emit_manual_state_change_event``, skip - its ``Stats.gauge`` side effect (which would try to ``Serde.to_json`` a - ``MagicMock`` return value) and invoke the adapter method directly. + The module-level pool wrappers pass adapter method *names* and resolve them + on the per-process adapter; here we resolve them on this listener's adapter + instead, so assertions against mocked adapter methods keep working. For + ``_emit_manual_state_change_event`` this also skips its ``Stats.gauge`` + side effect (which would try to ``Serde.to_json`` a ``MagicMock`` return). """ - from airflow.providers.openlineage.plugins.listener import _emit_manual_state_change_event + from airflow.providers.openlineage.plugins.listener import ( + _emit_manual_state_change_event, + _run_adapter_method, + ) if callable is _emit_manual_state_change_event: - adapter_method, _stats_key, *_ = args - return adapter_method(**kwargs) + adapter_method_name, _stats_key, *_ = args + return getattr(self.adapter, adapter_method_name)(**kwargs) + if callable is _run_adapter_method: + adapter_method_name, *rest = args + return getattr(self.adapter, adapter_method_name)(*rest, **kwargs) return callable(*args, **kwargs) @@ -123,6 +131,34 @@ def shutdown(self, *args, **kwargs): print("Shutting down") +def _probe_process_adapter(): + """Return (pid, adapter id) from inside a pool worker; module-level so it is picklable.""" + import os + + from airflow.providers.openlineage.plugins.listener import _get_process_adapter + + return os.getpid(), id(_get_process_adapter()) + + +class TestProcessAdapterReuse: + def test_process_adapter_reused_across_pool_submissions(self): + """ + A pool worker must reuse one adapter (hence one client/transport set) across events. + + Regression test: submitting bound adapter methods pickled a fresh adapter per event, + making the worker build a new OpenLineageClient (and transport worker threads that are + never closed) for every DAG-run state change, leaking threads in the scheduler. + """ + from concurrent.futures import ProcessPoolExecutor + + with ProcessPoolExecutor(max_workers=1) as pool: + pid_first, adapter_id_first = pool.submit(_probe_process_adapter).result(timeout=60) + pid_second, adapter_id_second = pool.submit(_probe_process_adapter).result(timeout=60) + + assert pid_first == pid_second + assert adapter_id_first == adapter_id_second + + class TestExecutorInitializer: """Tests for _executor_initializer function."""