kacpermuda commented on code in PR #74577:
URL: https://github.com/apache/airflow/pull/74577#discussion_r4443491857
##########
providers/openlineage/src/airflow/providers/openlineage/conf.py:
##########
@@ -163,6 +163,12 @@ def dag_state_change_process_pool_size() -> int:
return conf.getint(_CONFIG_SECTION, "dag_state_change_process_pool_size",
fallback="1")
+@cache
+def api_server_max_pending_events() -> int:
Review Comment:
`api_server_max_pending_events` isn't registered in
`providers/openlineage/provider.yaml`, so there's no
description/default/version_added and it won't show in the config reference.
It's also a new user-facing knob that isn't in the PR description. More
broadly, #74544 is about forks and leaked processes, not queue growth. Can we
drop the semaphore and this option and do the bounded queue in a separate PR
where the drop-vs-block tradeoff gets its own discussion?
##########
airflow-core/src/airflow/api_fastapi/app.py:
##########
@@ -100,6 +101,13 @@ def _initialize_api_server_stats() -> None:
@asynccontextmanager
async def lifespan(app: FastAPI):
_initialize_api_server_stats()
+ os.environ["_AIRFLOW_COMPONENT"] = "api_server"
+ from airflow.listeners.listener import get_listener_manager
+
+ try:
+ get_listener_manager().hook.on_starting(component=app)
Review Comment:
This is the first time the API server fires `on_starting`/`before_stopping`,
and that's for every listener plugin, not only OpenLineage. Third-party
listeners may assume these hooks only come from a `Job` (`job.py`) and touch
`component.job_type` and similar. Worth a note in the PR description, and maybe
a significant newsfragment? The inline `get_listener_manager` import also needs
a reason (circular import?) or should move to module level.
##########
providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py:
##########
@@ -1117,41 +1146,76 @@ def _fork_execute(self, callable, callable_name: str):
os._exit(0)
@property
- def executor(self) -> ProcessPoolExecutor:
+ def executor(self) -> Executor:
+ self._check_pid()
+ if self._is_api_server:
+ if not self._thread_executor:
+ with self._lock:
+ if not self._thread_executor:
+ # The API server uses ThreadPoolExecutor with
max_workers=1 to serialize
+ # access to the module-level cached OpenLineageAdapter
(_get_process_adapter)
+ # across manual DAG run and task instance state change
emissions.
+ # Multiple threads must not access the adapter
concurrently because
+ # OpenLineageAdapter.get_or_create_openlineage_client
is not thread-safe.
+ # Scheduler uses ProcessPoolExecutor where each child
process has its
+ # own memory space and isolated adapter.
+ self._thread_executor = ThreadPoolExecutor(
+ max_workers=1,
+ thread_name_prefix="openlineage-api-executor",
+ )
+ return self._thread_executor
+
if not self._executor:
- self._executor = ProcessPoolExecutor(
- max_workers=conf.dag_state_change_process_pool_size(),
- initializer=_executor_initializer,
- )
+ with self._lock:
+ if not self._executor:
+ self._executor = ProcessPoolExecutor(
+ max_workers=conf.dag_state_change_process_pool_size(),
+ initializer=_executor_initializer,
+ )
return self._executor
@hookimpl
def on_starting(self, component) -> None:
self.log.debug("on_starting: %s", component.__class__.__name__)
+ if component.__class__.__name__ == "FastAPI" or getattr(component,
"__module__", "").startswith(
+ "fastapi"
+ ):
+ self._is_api_server = True
@hookimpl
def before_stopping(self, component) -> None:
self.log.debug("before_stopping: %s", component.__class__.__name__)
- if self._executor is None:
+ self._check_pid()
+ with self._lock:
+ thread_executor, self._thread_executor = self._thread_executor,
None
+ self._semaphore = None
+ process_executor, self._executor = self._executor, None
+
+ if thread_executor is not None:
+ try:
+ thread_executor.shutdown(wait=False, cancel_futures=True)
Review Comment:
`cancel_futures=True` on graceful shutdown silently drops queued lineage
events, whereas the scheduler path waits up to 30s. That might be the right
tradeoff here, but could we log how many were cancelled (or at least say in the
docstring/PR that they're dropped)? Otherwise it's a silent data loss that
on-call can't see.
##########
providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py:
##########
@@ -1310,21 +1374,64 @@ def on_dag_run_failed(self, dag_run: DagRun, msg: str)
-> None:
except BaseException as e:
self.log.warning("OpenLineage received exception in method
on_dag_run_failed", exc_info=e)
+ def _get_semaphore(self) -> threading.BoundedSemaphore | None:
+ if not self._is_api_server:
+ return None
+ if self._semaphore is None:
+ with self._lock:
+ if self._semaphore is None:
+ max_pending = conf.api_server_max_pending_events()
+ if max_pending > 0:
+ # Capacity includes 1 actively running thread plus
queue buffer
+ self._semaphore =
threading.BoundedSemaphore(max_pending)
+ return self._semaphore
+
def submit_callable(self, callable, *args, **kwargs):
Review Comment:
`submit_callable` is now ~55 lines with a nested try/except BaseException
and two retry branches (thread / process) that diverge in side effects: only
the thread branch resets `_semaphore` and passes `cancel_futures=True`. Can we
pull the "recreate executor if still the one that failed" step into a small
helper so the two paths are visibly the same?
##########
airflow-core/src/airflow/api_fastapi/app.py:
##########
@@ -100,6 +101,13 @@ def _initialize_api_server_stats() -> None:
@asynccontextmanager
async def lifespan(app: FastAPI):
_initialize_api_server_stats()
+ os.environ["_AIRFLOW_COMPONENT"] = "api_server"
Review Comment:
`_AIRFLOW_COMPONENT` isn't read anywhere else in the repo, so this
introduces a new process-wide env var that child processes inherit, and it's
popped unconditionally on shutdown. The listener already learns it's in the API
server via `on_starting(component=app)` right below, so can we drop the env var
and rely on that single signal? If `on_starting` raises, we log and the
listener stays in scheduler mode, which is the fork risk, but the fix for that
is a more robust detection, not a second global.
##########
providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py:
##########
@@ -1117,41 +1146,76 @@ def _fork_execute(self, callable, callable_name: str):
os._exit(0)
@property
- def executor(self) -> ProcessPoolExecutor:
+ def executor(self) -> Executor:
+ self._check_pid()
+ if self._is_api_server:
+ if not self._thread_executor:
+ with self._lock:
+ if not self._thread_executor:
+ # The API server uses ThreadPoolExecutor with
max_workers=1 to serialize
+ # access to the module-level cached OpenLineageAdapter
(_get_process_adapter)
+ # across manual DAG run and task instance state change
emissions.
+ # Multiple threads must not access the adapter
concurrently because
+ # OpenLineageAdapter.get_or_create_openlineage_client
is not thread-safe.
+ # Scheduler uses ProcessPoolExecutor where each child
process has its
+ # own memory space and isolated adapter.
+ self._thread_executor = ThreadPoolExecutor(
+ max_workers=1,
+ thread_name_prefix="openlineage-api-executor",
+ )
+ return self._thread_executor
+
if not self._executor:
- self._executor = ProcessPoolExecutor(
- max_workers=conf.dag_state_change_process_pool_size(),
- initializer=_executor_initializer,
- )
+ with self._lock:
+ if not self._executor:
+ self._executor = ProcessPoolExecutor(
+ max_workers=conf.dag_state_change_process_pool_size(),
+ initializer=_executor_initializer,
+ )
return self._executor
@hookimpl
def on_starting(self, component) -> None:
self.log.debug("on_starting: %s", component.__class__.__name__)
+ if component.__class__.__name__ == "FastAPI" or getattr(component,
"__module__", "").startswith(
Review Comment:
Matching on class name / `__module__.startswith("fastapi")` is brittle: a
`FastAPI` subclass defined elsewhere matches via the name only by luck, and
anything else with a fastapi-module class gets misclassified. Can we use
`isinstance(component, FastAPI)` (lazy import) instead? Also `_is_api_server`
as a property+setter over `_is_api_server_flag` seems heavier than a plain
attribute.
##########
providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py:
##########
@@ -1310,21 +1374,64 @@ def on_dag_run_failed(self, dag_run: DagRun, msg: str)
-> None:
except BaseException as e:
self.log.warning("OpenLineage received exception in method
on_dag_run_failed", exc_info=e)
+ def _get_semaphore(self) -> threading.BoundedSemaphore | None:
+ if not self._is_api_server:
+ return None
+ if self._semaphore is None:
+ with self._lock:
+ if self._semaphore is None:
+ max_pending = conf.api_server_max_pending_events()
+ if max_pending > 0:
+ # Capacity includes 1 actively running thread plus
queue buffer
+ self._semaphore =
threading.BoundedSemaphore(max_pending)
+ return self._semaphore
+
def submit_callable(self, callable, *args, **kwargs):
+ sem = self._get_semaphore()
+ if sem is not None and not sem.acquire(blocking=False):
+ max_pending = conf.api_server_max_pending_events()
+ self.log.warning(
+ "OpenLineage API server event queue saturated (limit=%d);
dropping event to prevent memory growth under collector latency/outage.",
Review Comment:
This logs one WARNING per dropped event, so during a collector outage it
becomes a hot-loop log line, and it re-reads conf just to format it. If this
stays, can we log once per N drops or keep a counter? Also the line is over
length and the inline `from concurrent.futures import Future` on L1397 should
be a top-level import.
##########
providers/openlineage/tests/unit/openlineage/plugins/test_listener.py:
##########
@@ -3267,3 +3269,393 @@ def _failing_emit():
listener.log.warning.assert_called_once()
assert "on_running" in listener.log.warning.call_args.args
+
+
@mock.patch("airflow.providers.openlineage.plugins.listener.ThreadPoolExecutor")
Review Comment:
Both `uses_thread_pool` tests patch out both executors, so they only prove
wiring. Can we add one test that submits through the real `ThreadPoolExecutor`
and checks the event reaches a fake transport? That's the behaviour that
matters for #74544. Also the repeated `MagicMock` dag-run setup could be a
fixture.
##########
providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py:
##########
@@ -1310,21 +1374,64 @@ def on_dag_run_failed(self, dag_run: DagRun, msg: str)
-> None:
except BaseException as e:
self.log.warning("OpenLineage received exception in method
on_dag_run_failed", exc_info=e)
+ def _get_semaphore(self) -> threading.BoundedSemaphore | None:
+ if not self._is_api_server:
+ return None
+ if self._semaphore is None:
+ with self._lock:
+ if self._semaphore is None:
+ max_pending = conf.api_server_max_pending_events()
+ if max_pending > 0:
+ # Capacity includes 1 actively running thread plus
queue buffer
Review Comment:
Capacity is tied to the semaphore object, but a retry in `submit_callable`
sets `self._semaphore = None` and a new one is built. In-flight futures release
the old semaphore while new submissions acquire the new one, so the effective
bound can briefly be up to 2x the limit. The named cap doesn't actually enforce
what it says across a recreate.
##########
providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py:
##########
@@ -166,11 +167,40 @@ class OpenLineageListener:
"""OpenLineage listener sends events on task instance and dag run starts,
completes and failures."""
def __init__(self):
- self._executor = None
+ self._executor: ProcessPoolExecutor | None = None
+ self._thread_executor: ThreadPoolExecutor | None = None
+ self._semaphore: threading.BoundedSemaphore | None = None
+ self._is_api_server_flag: bool = False
+ self._pid: int = os.getpid()
+ self._lock = threading.Lock()
self.log = logging.getLogger(__name__)
self.extractor_manager = ExtractorManager()
self.adapter = OpenLineageAdapter()
+ @property
+ def _is_api_server(self) -> bool:
+ return self._is_api_server_flag or
os.environ.get("_AIRFLOW_COMPONENT") == "api_server"
+
+ @_is_api_server.setter
+ def _is_api_server(self, value: bool) -> None:
+ self._is_api_server_flag = value
+
+ def _check_pid(self) -> None:
+ """Reset executors if the process was forked (e.g. Gunicorn master to
worker)."""
+ current_pid = os.getpid()
+ if self._pid != current_pid:
+ # Fork occurred. The inherited lock might be held by a thread that
no longer exists in this child.
+ # Replace it with a fresh lock using double-checked locking to
avoid races between child threads.
+ new_lock = threading.Lock()
Review Comment:
`new_lock` is local to each calling thread, so `with new_lock:` serializes
nothing. Two threads hitting the first post-fork call both pass the `_pid`
check and both reset, and the second can wipe an executor the first just
created. Can we use `os.register_at_fork(after_in_child=...)` to reset state
once in the child, before any thread exists? That also removes `_check_pid()`
from the hot path of `executor` and `before_stopping`.
--
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]
