test: Check the memory estimation under every multiprocessing start method

The shared-memory regression test for `get_memory_info()` ran its
children under an explicit `fork` context only. The start methods
differ in how much memory the children share with the rest of the
process tree, and the default changed from `fork` to `forkserver` in
Python 3.14, so the test is now parametrized over `fork`, `forkserver`
and `spawn`.

The children are module-level functions so that every start method can
pickle them. Start methods that launch long-lived helper processes on
first use (the fork server, the resource tracker) get a warm-up before
the measurements, since those processes would otherwise land in the
first baseline reading and inflate it. The `fork` variant intentionally
forks the multi-threaded test runner, so the deprecation warning that
Python 3.12+ emits for it is filtered for this test.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TZuPLMf9Yf1NKZMqtGDS9J

Author: claude

tests/unit/_utils/test_system.py

@@ -5,7 +5,7 @@
 from multiprocessing import get_context, synchronize
 from multiprocessing.shared_memory import SharedMemory
 from types import SimpleNamespace
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, cast
 from unittest.mock import Mock
 
 import proclimits
@@ -19,6 +19,10 @@
 
 if TYPE_CHECKING:
     from collections.abc import Callable
+    from multiprocessing.context import ForkContext, ForkServerContext, SpawnContext
+    from multiprocessing.sharedctypes import Synchronized
+
+    _Context = ForkContext | ForkServerContext | SpawnContext
 
 HOST_TOTAL_BYTES = 8 * 1024**3
 HOST_AVAILABLE_BYTES = 3 * 1024**3
@@ -424,11 +428,159 @@ def test_log_resource_limits_lets_a_failing_sensor_surface(monkeypatch: pytest.M
     snapshot.assert_called_once()
 
 
+_EXTRA_MEMORY_SIZE = 1024 * 1024 * 100  # 100 MB
+"""The memory each child of the estimation test below allocates on top of its baseline."""
+
+_START_METHODS = ['fork', 'forkserver', 'spawn']
+"""Every start method `multiprocessing` offers on Linux.
+
+The default differs between interpreters - `fork` up to Python 3.13 and `forkserver` from 3.14 on - and the methods
+differ in how much memory the children share with the rest of the process tree, which is exactly what the estimation
+has to account for. The children below are module-level functions so that every method can pickle them.
+"""
+
+
+def _get_context(start_method: str) -> _Context:
+    """Get the context of a start method, typed as the concrete context so that its process type is known."""
+    return cast('_Context', get_context(start_method))
+
+
+def _idle_child(ready: synchronize.Barrier, measured: synchronize.Barrier) -> None:
+    """A child without additional memory, used to estimate the baseline memory usage of any child process."""
+    ready.wait()
+    measured.wait()
+
+
+def _own_memory_child(ready: synchronize.Barrier, measured: synchronize.Barrier) -> None:
+    """A child allocating additional memory of its own, shared with no other process."""
+    memory = SharedMemory(size=_EXTRA_MEMORY_SIZE, create=True)
+    assert memory.buf is not None
+    fill_buffer(memory.buf, _EXTRA_MEMORY_SIZE)
+    print(f'Using the memory... {memory.buf[-1]}')
+    ready.wait()
+    measured.wait()
+    memory.close()
+    memory.unlink()
+
+
+def _shared_memory_child(ready: synchronize.Barrier, measured: synchronize.Barrier, memory: SharedMemory) -> None:
+    """A child mapping additional memory that the parent and its siblings map as well."""
+    assert memory.buf is not None
+    # Fault every page in: untouched pages never enter the RSS (hiding the overcount this test guards against) and
+    # are reclaimed first under memory pressure, which drops them from every mapper's PSS.
+    page_sum = sum(memory.buf[::4096])
+    print(f'Using the memory... {page_sum}')
+    ready.wait()
+    measured.wait()
+
+
+def _get_additional_memory_estimation_while_running_processes(
+    ctx: _Context, *, target: Callable, count: int, use_shared_memory: bool = False
+) -> float:
+    """Estimate how much memory one child of the given kind adds to the process tree, in MB."""
+    processes = []
+    ready = ctx.Barrier(parties=count + 1)
+    measured = ctx.Barrier(parties=count + 1)
+    shared_memory: SharedMemory | None = None
+    memory_before = get_memory_info().current_size
+
+    if use_shared_memory:
+        shared_memory = SharedMemory(size=_EXTRA_MEMORY_SIZE, create=True)
+        assert shared_memory.buf is not None
+        fill_buffer(shared_memory.buf, _EXTRA_MEMORY_SIZE)
+        extra_args = [shared_memory]
+    else:
+        extra_args = []
+
+    for _ in range(count):
+        p = ctx.Process(target=target, args=[ready, measured, *extra_args])
+        p.start()
+        processes.append(p)
+
+    ready.wait()
+    memory_during = get_memory_info().current_size
+    measured.wait()
+
+    for p in processes:
+        p.join()
+
+    if shared_memory:
+        shared_memory.close()
+        shared_memory.unlink()
+
+    return (memory_during - memory_before).to_mb() / count
+
+
+def _estimate_memory_of_children(start_method: str, expectation_met: Synchronized[Any]) -> None:
+    """Run the estimation rounds in a process of its own, reporting through `expectation_met`.
+
+    Measuring from a fresh process keeps the readings clear of whatever else the test runner has going on.
+    """
+    ctx = _get_context(start_method)
+    children_count = 4
+    # Memory calculation is not exact, so allow for some tolerance.
+    test_tolerance = 0.3
+    measurement_rounds = 3
+
+    # Some start methods launch long-lived helper processes (the fork server, the resource tracker) on first use.
+    # They are part of the process tree and so part of the estimate, and started inside a round they would inflate
+    # the baseline for every child in it. Start them ahead of the measurements.
+    # The barriers are kept referenced here: `start()` drops its own reference to the arguments once they are sent,
+    # and a collected barrier takes its semaphore with it before the child gets to rebuild it.
+    ready = ctx.Barrier(parties=1)
+    measured = ctx.Barrier(parties=1)
+    warm_up = ctx.Process(target=_idle_child, args=[ready, measured])
+    warm_up.start()
+    warm_up.join()
+    assert warm_up.exitcode == 0
+
+    # Under memory pressure the kernel reclaims cold pages, which silently leave the PSS readings and skew a round's
+    # deltas, so a distorted round is re-measured. A genuine overcount of shared memory misses the expectation
+    # several times over in every round, so the retries cannot mask it.
+    for _ in range(measurement_rounds):
+        additional_memory_simple_child = _get_additional_memory_estimation_while_running_processes(
+            ctx, target=_idle_child, count=children_count
+        )
+        additional_memory_extra_memory_child = (
+            _get_additional_memory_estimation_while_running_processes(
+                ctx, target=_own_memory_child, count=children_count
+            )
+            - additional_memory_simple_child
+        )
+        additional_memory_shared_extra_memory_child = (
+            _get_additional_memory_estimation_while_running_processes(
+                ctx, target=_shared_memory_child, count=children_count, use_shared_memory=True
+            )
+            - additional_memory_simple_child
+        )
+
+        memory_estimation_difference_ratio = (
+            abs((additional_memory_shared_extra_memory_child * children_count) - additional_memory_extra_memory_child)
+            / additional_memory_extra_memory_child
+        )
+
+        if memory_estimation_difference_ratio < test_tolerance:
+            expectation_met.value = True
+            break
+
+        print(
+            f'{additional_memory_shared_extra_memory_child=}\n'
+            f'{children_count=}\n'
+            f'{additional_memory_extra_memory_child=}\n'
+            f'{memory_estimation_difference_ratio=}'
+        )
+
+
 # The estimation is asserted on absolute memory readings, which hold only as long as nothing else on the machine makes
 # the kernel reclaim the pages allocated below. Running alongside the other test workers is enough to break that.
 @pytest.mark.run_alone
 @pytest.mark.skipif(sys.platform != 'linux', reason='Improved estimation available only on Linux')
-def test_memory_estimation_does_not_overestimate_due_to_shared_memory() -> None:
+# Forking the multi-threaded test runner is deprecated since Python 3.12. It is deliberate here - `fork` is what
+# production runs under up to Python 3.13 - and the forked process uses no threads, so the deadlock it warns about
+# cannot happen.
+@pytest.mark.filterwarnings('ignore:This process .* is multi-threaded, use of fork:DeprecationWarning')
+@pytest.mark.parametrize('start_method', _START_METHODS)
+def test_memory_estimation_does_not_overestimate_due_to_shared_memory(start_method: str) -> None:
     """Test that memory usage estimation is not overestimating memory usage by counting shared memory multiple times.
 
     In this test, the parent process is started and its memory usage is measured in situations where it is running
@@ -438,121 +590,18 @@ def test_memory_estimation_does_not_overestimate_due_to_shared_memory() -> None:
     additional_memory_size_estimate_per_shared_memory_child * number_of_sharing_children_processes is approximately
     equal to additional_memory_size_estimate_per_unshared_memory_child where the additional shared memory is exactly
     the same as the unshared memory.
-    """
-
-    ctx = get_context('fork')
-    estimated_memory_expectation = ctx.Value('b', False)  # noqa: FBT003  # Common usage pattern for multiprocessing.Value
-
-    def parent_process() -> None:
-        extra_memory_size = 1024 * 1024 * 100  # 100 MB
-        children_count = 4
-        # Memory calculation is not exact, so allow for some tolerance.
-        test_tolerance = 0.3
-        measurement_rounds = 3
-
-        def no_extra_memory_child(ready: synchronize.Barrier, measured: synchronize.Barrier) -> None:
-            ready.wait()
-            measured.wait()
-
-        def extra_memory_child(ready: synchronize.Barrier, measured: synchronize.Barrier) -> None:
-            memory = SharedMemory(size=extra_memory_size, create=True)
-            assert memory.buf is not None
-            fill_buffer(memory.buf, extra_memory_size)
-            print(f'Using the memory... {memory.buf[-1]}')
-            ready.wait()
-            measured.wait()
-            memory.close()
-            memory.unlink()
-
-        def shared_extra_memory_child(
-            ready: synchronize.Barrier, measured: synchronize.Barrier, memory: SharedMemory
-        ) -> None:
-            assert memory.buf is not None
-            # Fault every page in: untouched pages never enter the RSS (hiding the overcount this test guards
-            # against) and are reclaimed first under memory pressure, which drops them from every mapper's PSS.
-            page_sum = sum(memory.buf[::4096])
-            print(f'Using the memory... {page_sum}')
-            ready.wait()
-            measured.wait()
-
-        def get_additional_memory_estimation_while_running_processes(
-            *, target: Callable, count: int = 1, use_shared_memory: bool = False
-        ) -> float:
-            processes = []
-            ready = ctx.Barrier(parties=count + 1)
-            measured = ctx.Barrier(parties=count + 1)
-            shared_memory: SharedMemory | None = None
-            memory_before = get_memory_info().current_size
-
-            if use_shared_memory:
-                shared_memory = SharedMemory(size=extra_memory_size, create=True)
-                assert shared_memory.buf is not None
-                fill_buffer(shared_memory.buf, extra_memory_size)
-                extra_args = [shared_memory]
-            else:
-                extra_args = []
-
-            for _ in range(count):
-                p = ctx.Process(target=target, args=[ready, measured, *extra_args])
-                p.start()
-                processes.append(p)
-
-            ready.wait()
-            memory_during = get_memory_info().current_size
-            measured.wait()
-
-            for p in processes:
-                p.join()
-
-            if shared_memory:
-                shared_memory.close()
-                shared_memory.unlink()
-
-            return (memory_during - memory_before).to_mb() / count
-
-        # Under memory pressure the kernel reclaims cold pages, which silently leave the PSS readings and skew a
-        # round's deltas, so a distorted round is re-measured. A genuine overcount of shared memory misses the
-        # expectation several times over in every round, so the retries cannot mask it.
-        for _ in range(measurement_rounds):
-            additional_memory_simple_child = get_additional_memory_estimation_while_running_processes(
-                target=no_extra_memory_child, count=children_count
-            )
-            additional_memory_extra_memory_child = (
-                get_additional_memory_estimation_while_running_processes(
-                    target=extra_memory_child, count=children_count
-                )
-                - additional_memory_simple_child
-            )
-            additional_memory_shared_extra_memory_child = (
-                get_additional_memory_estimation_while_running_processes(
-                    target=shared_extra_memory_child, count=children_count, use_shared_memory=True
-                )
-                - additional_memory_simple_child
-            )
 
-            memory_estimation_difference_ratio = (
-                abs(
-                    (additional_memory_shared_extra_memory_child * children_count)
-                    - additional_memory_extra_memory_child
-                )
-                / additional_memory_extra_memory_child
-            )
-
-            if memory_estimation_difference_ratio < test_tolerance:
-                estimated_memory_expectation.value = True
-                break
-
-            print(
-                f'{additional_memory_shared_extra_memory_child=}\n'
-                f'{children_count=}\n'
-                f'{additional_memory_extra_memory_child=}\n'
-                f'{memory_estimation_difference_ratio=}'
-            )
+    The children share a different amount of memory with the parent under each start method - copy-on-write pages of
+    the whole interpreter under `fork`, pages of the fork server under `forkserver`, only the shared libraries and
+    the explicitly shared block under `spawn` - so the estimation is checked under every one of them.
+    """
+    ctx = _get_context(start_method)
+    expectation_met = ctx.Value('b', False)  # noqa: FBT003  # Common usage pattern for multiprocessing.Value
 
-    process = ctx.Process(target=parent_process)
+    process = ctx.Process(target=_estimate_memory_of_children, args=[start_method, expectation_met])
     process.start()
     process.join()
 
-    assert estimated_memory_expectation.value, (
-        'Estimated memory usage for process with shared memory does not meet the expectation.'
+    assert expectation_met.value, (
+        f'Estimated memory usage for process with shared memory does not meet the expectation under {start_method}.'
     )