gh-75876: Run bigmem tests in a subprocess

A test which really allocates the memory it asks for (that is, run with -M)
now runs in a subprocess, so that the memory it uses and the address space
it fragments are released when it ends.  A dummy run stays in the process.

The parent process watches the memory usage of the subprocess while waiting
for it, so the separate watchdog process is no longer needed.

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: serhiy-storchaka

Lib/test/_isolated_sample.py

@@ -10,6 +10,7 @@
 import sys
 import time
 import unittest
+from test import support
 from test.support import isolation
 
 # DurationSample sleeps this long in the subprocess; a parent-reported duration
@@ -178,3 +179,12 @@ class TimeoutSample(unittest.TestCase):
     @isolation.runInSubprocess(timeout=TIMEOUT)
     def test_hang(self):
         time.sleep(TIMEOUT_HANG)
+
+
+class BigmemSample(unittest.TestCase):
+
+    @support.bigmemtest(size=1024, memuse=1)
+    def test_where_it_runs(self, size):
+        # A real run is isolated by bigmemtest() itself, a dummy run is not.
+        self.assertEqual(isolation.runningInSubprocess,
+                         bool(support.real_max_memuse))

Lib/test/memory_watchdog.py

@@ -1257,26 +1257,17 @@ def set_memlimit(limit: str) -> None:
     max_memuse = memlimit
 
 
-class _MemoryWatchdog:
-    """An object which periodically watches the process' memory consumption
-    and prints it out.
-    """
-
-    def __init__(self):
-        self.started = False
+def _memory_watchdog(pid):
+    """Return a function printing the memory usage of process *pid*."""
+    # Imported here: test.support does not depend on test.libregrtest.
+    from test.libregrtest.utils import get_process_memory_usage
 
-    def start(self):
-        import subprocess
-        watchdog_script = findfile("memory_watchdog.py")
-        cmd = [sys.executable, watchdog_script, str(os.getpid())]
-        self.mem_watchdog = subprocess.Popen(cmd)
-        self.started = True
-
-    def stop(self):
-        if not self.started:
-            return
-        self.mem_watchdog.terminate()
-        self.mem_watchdog.wait()
+    def watch():
+        mem = get_process_memory_usage(pid)
+        if mem is not None:
+            print(f" ... process data size: {mem / (1024 ** 3):.1f} GiB",
+                  flush=True)
+    return watch
 
 
 def bigmemtest(size, memuse, dry_run=True):
@@ -1291,8 +1282,14 @@ def bigmemtest(size, memuse, dry_run=True):
     extra argument. If 'dry_run' is true, the value passed to the test method
     may be less than the requested value. If 'dry_run' is false, it means the
     test doesn't support dummy runs when -M is not specified.
+
+    A test that actually allocates the requested memory (that is, one run with
+    -M) runs in a subprocess, so that the memory it uses and the address space
+    it fragments are released when it ends.  A dummy run stays in the process.
     """
     def decorator(f):
+        from test.support import isolation
+
         @functools.wraps(f)
         def wrapper(self):
             size = wrapper.size
@@ -1308,20 +1305,25 @@ def wrapper(self):
                     "not enough memory: %.1fG minimum needed"
                     % (size * memuse / (1024 ** 3)))
 
-            if real_max_memuse and verbose:
+            if (real_max_memuse and verbose
+                    and not isolation.runningInSubprocess):
                 print()
                 peak = (size * memuse) / (1024 ** 3)
-                print(f" ... expected peak memory use: {peak:.1f} GiB")
-                watchdog = _MemoryWatchdog()
-                watchdog.start()
-            else:
-                watchdog = None
+                # Flushed, so that it precedes the memory usage below.
+                print(f" ... expected peak memory use: {peak:.1f} GiB",
+                      flush=True)
+
+            if (real_max_memuse and has_subprocess_support
+                    and not isolation.runningInSubprocess):
+                # Watch it from here: the output of the subprocess is captured.
+                cls = type(self)
+                qualname = f'{cls.__qualname__}.{f.__name__}'
+                proc = isolation._start_test(cls.__module__, qualname)
+                watchdog = _memory_watchdog(proc.pid) if verbose else None
+                isolation._replay_test(self, *proc.wait(tick=watchdog))
+                return
 
-            try:
-                return f(self, maxsize)
-            finally:
-                if watchdog:
-                    watchdog.stop()
+            return f(self, maxsize)
 
         wrapper.size = size
         wrapper.memuse = memuse

Lib/test/support/__init__.py

@@ -108,12 +108,75 @@ def _child_environ(env):
     return environ
 
 
-def _run_in_subprocess(module, qualname, options, env, timeout):
-    """Run module.qualname (a test method or class) in a fresh subprocess.
+class _SubprocessTest:
+    """A test running in a subprocess, started by _start_test().
 
-    Return ``(payload, output, returncode)``, where *payload* is the decoded
-    ``{'outcomes': ..., 'durations': ...}`` mapping from the subprocess, or
-    ``None`` if it did not run to completion (crash, import error, ...).
+    The parent can watch the subprocess (its pid) while the test runs, and
+    must wait() for it.
+    """
+
+    def __init__(self, proc, result_path):
+        self._proc = proc
+        self._result_path = result_path
+
+    @property
+    def pid(self):
+        return self._proc.pid
+
+    def wait(self, timeout=None, tick=None, interval=1.0):
+        """Wait for the test to finish, calling *tick* every *interval* seconds.
+
+        Return ``(payload, output, returncode)``, where *payload* is the
+        decoded ``{'outcomes': ..., 'durations': ...}`` mapping from the
+        subprocess, or ``None`` if it did not run to completion (crash,
+        import error, ...).
+        """
+        import marshal
+        import subprocess
+        import time
+        deadline = None if timeout is None else time.monotonic() + timeout
+        try:
+            while True:
+                step = None if deadline is None else max(
+                    0.0, deadline - time.monotonic())
+                # Wake up for the next tick, unless the timeout comes first.
+                ticking = tick is not None and (step is None or step > interval)
+                try:
+                    # communicate(), not wait(): a test writing more than a
+                    # pipe buffer would block.  Retrying keeps what it read.
+                    stdout, stderr = self._proc.communicate(
+                        timeout=interval if ticking else step)
+                    break
+                except subprocess.TimeoutExpired:
+                    if ticking:
+                        tick()
+                        continue
+                    # Report the hang rather than leaving the runner stuck.
+                    self._proc.kill()
+                    stdout, stderr = self._proc.communicate()
+                    raise _SubprocessTestError(
+                        f'test did not complete in a subprocess '
+                        f'within {timeout} seconds'
+                    ) from _remote(_decode(stdout) + _decode(stderr))
+            try:
+                with open(self._result_path, 'rb') as f:
+                    payload = marshal.load(f)
+            except (OSError, EOFError, ValueError):
+                payload = None
+            output = _decode(stdout) + _decode(stderr)
+            return payload, output, self._proc.returncode
+        finally:
+            try:
+                os.unlink(self._result_path)
+            except OSError:
+                pass
+
+
+def _start_test(module, qualname, options=(), env=None):
+    """Start module.qualname (a test method or class) in a fresh subprocess.
+
+    Return a _SubprocessTest.  Its wait() is what removes the temporary file
+    the subprocess writes its result to.
     """
     import marshal
     import subprocess
@@ -129,26 +192,16 @@ def _run_in_subprocess(module, qualname, options, env, timeout):
         cmd = [sys.executable, *options, '-m', 'test.support.subprocess_runner',
                module, qualname, result_path,
                marshal.dumps(_child_config()).hex()]
-        try:
-            proc = subprocess.run(cmd, capture_output=True,
-                                  env=_child_environ(env), timeout=timeout)
-        except subprocess.TimeoutExpired as exc:
-            # Report the hang rather than leaving the test runner stuck.
-            output = _decode(exc.stdout) + _decode(exc.stderr)
-            raise _SubprocessTestError(
-                f'test did not complete in a subprocess '
-                f'within {timeout} seconds') from _remote(output)
-        try:
-            with open(result_path, 'rb') as f:
-                payload = marshal.load(f)
-        except (OSError, EOFError, ValueError):
-            payload = None
-    finally:
+        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE,
+                                stderr=subprocess.PIPE, env=_child_environ(env))
+    except BaseException:
         try:
             os.unlink(result_path)
         except OSError:
             pass
-    return payload, _decode(proc.stdout) + _decode(proc.stderr), proc.returncode
+        raise
+    return _SubprocessTest(proc, result_path)
+
 
 
 def _replay_outcome(test, outcome):
@@ -200,6 +253,19 @@ def _check_returncode(returncode, output, what):
         raise exc from _remote(output)
 
 
+def _replay_test(test, payload, output, returncode):
+    """Reproduce in *test* the result that _SubprocessTest.wait() returned."""
+    if payload is None:
+        exc = _SubprocessTestError(
+            f'test did not complete in a subprocess (exit code {returncode})')
+        raise exc from _remote(output)
+    # The parent measures the test method's own duration (the real cost of the
+    # isolated run, subprocess startup included), so nothing to forward here.
+    # Replay the outcomes first: a failure of the test itself is more useful.
+    _replay_outcomes(test, payload['outcomes'])
+    _check_returncode(returncode, output, 'test')
+
+
 def _isolate_method(func, options, env, timeout):
     @functools.wraps(func)
     def wrapper(self, /, *args, **kwargs):
@@ -209,18 +275,8 @@ def wrapper(self, /, *args, **kwargs):
         _check_subprocess_support()
         cls = type(self)
         qualname = f'{cls.__qualname__}.{func.__name__}'
-        payload, output, returncode = _run_in_subprocess(cls.__module__,
-                                                         qualname, options,
-                                                         env, timeout)
-        if payload is None:
-            exc = _SubprocessTestError(
-                f'test did not complete in a subprocess (exit code {returncode})')
-            raise exc from _remote(output)
-        # The parent measures this method's own duration (the real cost of the
-        # isolated run, subprocess startup included), so nothing to forward here.
-        # Replay the outcomes first: a failure of the test itself is more useful.
-        _replay_outcomes(self, payload['outcomes'])
-        _check_returncode(returncode, output, 'test')
+        proc = _start_test(cls.__module__, qualname, options, env)
+        _replay_test(self, *proc.wait(timeout))
     return wrapper
 
 
@@ -244,9 +300,8 @@ def setUpClass(cls):
         _check_subprocess_support()
         # Run the whole class in a single subprocess and stash the outcomes
         # for the test methods to replay.
-        payload, output, returncode = _run_in_subprocess(cls.__module__,
-                                                         cls.__qualname__,
-                                                         options, env, timeout)
+        proc = _start_test(cls.__module__, cls.__qualname__, options, env)
+        payload, output, returncode = proc.wait(timeout)
         if payload is None:
             exc = _SubprocessTestError(
                 f'class did not complete in a subprocess (exit code {returncode})')

Lib/test/support/isolation.py

@@ -1231,17 +1231,28 @@ def test_timeout_reported_as_error(self):
         self.assertEqual(len(result.errors), 1)
         self.assertIn(f'within {TIMEOUT} seconds', result.errors[0][1])
 
+    @support.requires_subprocess()
+    def test_bigmemtest_isolates_a_real_run(self):
+        # A dummy run (no -M) stays in this process, a real run does not.
+        for memlimit in (0, support._1G):
+            with self.subTest(real_max_memuse=memlimit):
+                with support.swap_attr(support, 'real_max_memuse', memlimit):
+                    result = self._run('BigmemSample')
+                self.assertEqual(result.testsRun, 1)
+                self.assertEqual(self._names(result.failures), [])
+                self.assertEqual(self._names(result.errors), [])
+
     def test_skipped_without_subprocess_support(self):
         # On a platform without subprocess support the test is skipped in the
         # parent, before any subprocess is spawned.
         calls = []
-        orig = isolation._run_in_subprocess
+        orig = isolation._start_test
         with support.swap_attr(support, 'has_subprocess_support', False):
-            isolation._run_in_subprocess = lambda *a, **k: calls.append(a)
+            isolation._start_test = lambda *a, **k: calls.append(a)
             try:
                 result = self._run('MethodSample.test_pass')
             finally:
-                isolation._run_in_subprocess = orig
+                isolation._start_test = orig
         self.assertEqual(result.testsRun, 1)
         self.assertEqual(len(result.skipped), 1)
         self.assertEqual(calls, [])