55from multiprocessing import get_context , synchronize
66from multiprocessing .shared_memory import SharedMemory
77from types import SimpleNamespace
8- from typing import TYPE_CHECKING
8+ from typing import TYPE_CHECKING , cast
99from unittest .mock import Mock
1010
1111import proclimits
1919
2020if TYPE_CHECKING :
2121 from collections .abc import Callable
22+ from multiprocessing .context import ForkContext , ForkServerContext , SpawnContext
2223
2324HOST_TOTAL_BYTES = 8 * 1024 ** 3
2425HOST_AVAILABLE_BYTES = 3 * 1024 ** 3
@@ -424,11 +425,44 @@ def test_log_resource_limits_lets_a_failing_sensor_surface(monkeypatch: pytest.M
424425 snapshot .assert_called_once ()
425426
426427
428+ _EXTRA_MEMORY_SIZE = 1024 * 1024 * 100 # 100 MB
429+
430+
431+ # The children of the estimation test below live at module level, so that every start method can pickle them.
432+ def no_extra_memory_child (ready : synchronize .Barrier , measured : synchronize .Barrier ) -> None :
433+ ready .wait ()
434+ measured .wait ()
435+
436+
437+ def extra_memory_child (ready : synchronize .Barrier , measured : synchronize .Barrier ) -> None :
438+ memory = SharedMemory (size = _EXTRA_MEMORY_SIZE , create = True )
439+ assert memory .buf is not None
440+ fill_buffer (memory .buf , _EXTRA_MEMORY_SIZE )
441+ print (f'Using the memory... { memory .buf [- 1 ]} ' )
442+ ready .wait ()
443+ measured .wait ()
444+ memory .close ()
445+ memory .unlink ()
446+
447+
448+ def shared_extra_memory_child (ready : synchronize .Barrier , measured : synchronize .Barrier , memory : SharedMemory ) -> None :
449+ assert memory .buf is not None
450+ # Fault every page in: untouched pages never enter the RSS (hiding the overcount this test guards
451+ # against) and are reclaimed first under memory pressure, which drops them from every mapper's PSS.
452+ page_sum = sum (memory .buf [::4096 ])
453+ print (f'Using the memory... { page_sum } ' )
454+ ready .wait ()
455+ measured .wait ()
456+
457+
427458# The estimation is asserted on absolute memory readings, which hold only as long as nothing else on the machine makes
428459# the kernel reclaim the pages allocated below. Running alongside the other test workers is enough to break that.
429460@pytest .mark .run_alone
430461@pytest .mark .skipif (sys .platform != 'linux' , reason = 'Improved estimation available only on Linux' )
431- def test_memory_estimation_does_not_overestimate_due_to_shared_memory () -> None :
462+ # The start methods differ in how much memory the children share with the rest of the process tree, which is what the
463+ # estimation has to account for. The default is `fork` up to Python 3.13 and `forkserver` from 3.14 on.
464+ @pytest .mark .parametrize ('start_method' , ['fork' , 'forkserver' , 'spawn' ])
465+ def test_memory_estimation_does_not_overestimate_due_to_shared_memory (start_method : str ) -> None :
432466 """Test that memory usage estimation is not overestimating memory usage by counting shared memory multiple times.
433467
434468 In this test, the parent process is started and its memory usage is measured in situations where it is running
@@ -440,41 +474,18 @@ def test_memory_estimation_does_not_overestimate_due_to_shared_memory() -> None:
440474 the same as the unshared memory.
441475 """
442476
443- ctx = get_context ('fork' )
444- estimated_memory_expectation = ctx .Value ('b' , False ) # noqa: FBT003 # Common usage pattern for multiprocessing.Value
477+ # The measuring process is a closure, which only `fork` can start. The children are started by `start_method`.
478+ fork_ctx = get_context ('fork' )
479+ estimated_memory_expectation = fork_ctx .Value ('b' , False ) # noqa: FBT003 # Common usage pattern for multiprocessing.Value
445480
446481 def parent_process () -> None :
447- extra_memory_size = 1024 * 1024 * 100 # 100 MB
482+ ctx = cast ('ForkContext | ForkServerContext | SpawnContext' , get_context (start_method ))
483+ extra_memory_size = _EXTRA_MEMORY_SIZE
448484 children_count = 4
449485 # Memory calculation is not exact, so allow for some tolerance.
450486 test_tolerance = 0.3
451487 measurement_rounds = 3
452488
453- def no_extra_memory_child (ready : synchronize .Barrier , measured : synchronize .Barrier ) -> None :
454- ready .wait ()
455- measured .wait ()
456-
457- def extra_memory_child (ready : synchronize .Barrier , measured : synchronize .Barrier ) -> None :
458- memory = SharedMemory (size = extra_memory_size , create = True )
459- assert memory .buf is not None
460- fill_buffer (memory .buf , extra_memory_size )
461- print (f'Using the memory... { memory .buf [- 1 ]} ' )
462- ready .wait ()
463- measured .wait ()
464- memory .close ()
465- memory .unlink ()
466-
467- def shared_extra_memory_child (
468- ready : synchronize .Barrier , measured : synchronize .Barrier , memory : SharedMemory
469- ) -> None :
470- assert memory .buf is not None
471- # Fault every page in: untouched pages never enter the RSS (hiding the overcount this test guards
472- # against) and are reclaimed first under memory pressure, which drops them from every mapper's PSS.
473- page_sum = sum (memory .buf [::4096 ])
474- print (f'Using the memory... { page_sum } ' )
475- ready .wait ()
476- measured .wait ()
477-
478489 def get_additional_memory_estimation_while_running_processes (
479490 * , target : Callable , count : int = 1 , use_shared_memory : bool = False
480491 ) -> float :
@@ -510,6 +521,15 @@ def get_additional_memory_estimation_while_running_processes(
510521
511522 return (memory_during - memory_before ).to_mb () / count
512523
524+ # Some start methods launch long-lived helper processes (the fork server, the resource tracker) on first use.
525+ # They belong to the process tree and so to the estimate, but started inside a round they would inflate its
526+ # baseline. Start them ahead of the measurements, keeping the barriers referenced until the child is done.
527+ ready , measured = ctx .Barrier (parties = 1 ), ctx .Barrier (parties = 1 )
528+ warm_up = ctx .Process (target = no_extra_memory_child , args = [ready , measured ])
529+ warm_up .start ()
530+ warm_up .join ()
531+ assert warm_up .exitcode == 0
532+
513533 # Under memory pressure the kernel reclaims cold pages, which silently leave the PSS readings and skew a
514534 # round's deltas, so a distorted round is re-measured. A genuine overcount of shared memory misses the
515535 # expectation several times over in every round, so the retries cannot mask it.
@@ -549,10 +569,10 @@ def get_additional_memory_estimation_while_running_processes(
549569 f'{ memory_estimation_difference_ratio = } '
550570 )
551571
552- process = ctx .Process (target = parent_process )
572+ process = fork_ctx .Process (target = parent_process )
553573 process .start ()
554574 process .join ()
555575
556576 assert estimated_memory_expectation .value , (
557- 'Estimated memory usage for process with shared memory does not meet the expectation.'
577+ f 'Estimated memory usage for process with shared memory does not meet the expectation under { start_method } .'
558578 )
0 commit comments