amoghrajesh commented on code in PR #68875:
URL: https://github.com/apache/airflow/pull/68875#discussion_r3459327606
##########
airflow-core/tests/unit/utils/test_process_utils.py:
##########
@@ -226,3 +228,76 @@ def test_session_leader(self, mock_set_pid):
mock_get_sid.return_value = pid
set_new_process_group()
assert mock_set_pid.call_count == 0
+
+
+ALL_MP_METHODS = ["fork", "spawn", "forkserver"]
+
+
+class TestResolveMpStartMethod:
+ def test_unset_returns_none(self):
+ assert process_utils.resolve_mp_start_method("scheduler") is None
+
+ @conf_vars({("core", "mp_start_method"): "forkserver"})
+ def test_core_value_used(self):
+ assert process_utils.resolve_mp_start_method("scheduler") ==
"forkserver"
+ assert process_utils.resolve_mp_start_method("triggerer") ==
"forkserver"
+
+ @conf_vars({("core", "mp_start_method"): "forkserver", ("scheduler",
"mp_start_method"): "fork"})
+ def test_component_overrides_core(self):
+ assert process_utils.resolve_mp_start_method("scheduler") == "fork"
+ assert process_utils.resolve_mp_start_method("triggerer") ==
"forkserver"
+
+ @conf_vars({("core", "mp_start_method"): " "})
+ def test_blank_returns_none(self):
+ assert process_utils.resolve_mp_start_method("scheduler") is None
+
+
+class TestSetComponentMpStartMethod:
+ @mock.patch("airflow.utils.process_utils.multiprocessing")
+ def test_noop_when_unset(self, mock_mp):
Review Comment:
Please add a `spec=multiprocessing` or use `autospec` please here and
everywhere else
##########
airflow-core/src/airflow/utils/process_utils.py:
##########
@@ -388,3 +389,67 @@ def set_new_process_group() -> None:
return
os.setpgid(0, 0)
+
+
+_MP_START_METHOD_OPTION = "mp_start_method"
+_MP_FORKSERVER_PRELOAD_OPTION = "mp_forkserver_preload"
+
+
+def resolve_mp_start_method(component: str) -> str | None:
+ """
+ Return the configured ``multiprocessing`` start method for *component*, or
``None``.
+
+ Resolution order: ``[<component>] mp_start_method`` then ``[core]
mp_start_method`` then
+ ``None`` (keep the platform default).
+ """
+ method = conf.get(component, _MP_START_METHOD_OPTION, fallback=None) or
conf.get(
+ "core", _MP_START_METHOD_OPTION, fallback=None
+ )
+ return method.strip() if method and method.strip() else None
+
+
+def set_component_mp_start_method(component: str) -> None:
+ """
+ Apply the configured ``multiprocessing`` start method for *component*, if
any.
+
+ Python 3.14 changed the Unix (non-macOS) default from ``fork`` to
``forkserver``;
+ ``forkserver``/``spawn`` children re-import Airflow instead of sharing it
copy-on-write, which
+ inflates the memory of components that spawn workers (the scheduler's
``LocalExecutor`` and the
+ triggerer). This restores control over that choice.
+
+ Must be called before the component creates any ``multiprocessing`` object
(e.g. before the
+ executor or job runner starts). A no-op when unset or when the method is
unavailable on the
+ current platform, so it is always safe to call.
+ """
+ method = resolve_mp_start_method(component)
+ if not method:
+ return
+
+ available = multiprocessing.get_all_start_methods()
+ if method not in available:
+ log.warning(
+ "Configured %s=%r for %s is not available on this platform
(available: %s); "
+ "leaving the platform default start method in place.",
+ _MP_START_METHOD_OPTION,
+ method,
+ component,
+ available,
+ )
+ return
+
+ try:
+ multiprocessing.set_start_method(method, force=True)
+ except (RuntimeError, ValueError) as exc:
+ log.warning("Could not set multiprocessing start method to %r for %s:
%s", method, component, exc)
Review Comment:
Worth adding test coverage for this path?
--
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]
