diff --git a/.github/workflows/baseline-tests.yaml b/.github/workflows/baseline-tests.yaml index b0ab85dcdd9..2bf3a9c8b5c 100644 --- a/.github/workflows/baseline-tests.yaml +++ b/.github/workflows/baseline-tests.yaml @@ -11,12 +11,20 @@ jobs: env: PYTHONUNBUFFERED: 1 PYTEST_ADDOPTS: "--durations=25 -p no:unraisableexception -p no:threadexception" + # Two batch subprocesses per runner while its memory allows (see + # test_batched.py): 32% fewer runner-minutes, same cases, no shutdowns. + POLICYENGINE_BATCH_WORKERS: 2 strategy: fail-fast: false + # At most 6 Baseline groups run at once; GitHub starts the next group on a + # fresh runner as each finishes, so the lanes balance from real durations. + max-parallel: 6 matrix: include: - # One state per subprocess; retain four runners and one worker each. - # The former NY/OH/OK process peaked at 15.1 GB (run 36930199045). + # One state per batch subprocess, over four runners. The job's + # POLICYENGINE_BATCH_WORKERS overrides --workers 1, and runner memory + # decides whether two run at once. The former NY/OH/OK process peaked + # at 15.1 GB (run 36930199045). - group: states-shard-1 cmd: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 1/4 - group: states-shard-2 diff --git a/.github/workflows/contrib-tests.yaml b/.github/workflows/contrib-tests.yaml index cd1c80fcd32..19ee053f607 100644 --- a/.github/workflows/contrib-tests.yaml +++ b/.github/workflows/contrib-tests.yaml @@ -11,8 +11,14 @@ jobs: env: PYTHONUNBUFFERED: 1 PYTEST_ADDOPTS: "--durations=25 -p no:unraisableexception -p no:threadexception" + # Two batch subprocesses per runner while its memory allows (see + # test_batched.py): 32% fewer runner-minutes, same cases, no shutdowns. + POLICYENGINE_BATCH_WORKERS: 2 strategy: fail-fast: false + # Contrib starts after Baseline finishes, so its 6 lanes take the place of + # Baseline's 6. + max-parallel: 6 matrix: include: - group: states-shard-1 diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index a002c0abfcd..a793e99fcb2 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -273,9 +273,13 @@ jobs: run: bash ./update_itemization.sh - name: Run Household API Partners YAML tests run: uv run --no-sync make test-yaml-no-structural-other-partners + # Contrib starts once Baseline passes, on the same 6 lanes, so most PRs see + # their Baseline result first. A Baseline failure skips Contrib, and + # re-running the failed Baseline group then runs Contrib for the first time: + # no finished Contrib work is repeated. Contrib: name: Contrib tests - needs: ReleaseLock + needs: [ReleaseLock, Baseline] uses: ./.github/workflows/contrib-tests.yaml PythonTests: name: Python tests diff --git a/.github/workflows/push.yaml b/.github/workflows/push.yaml index 4917f0024b3..595efb8d499 100644 --- a/.github/workflows/push.yaml +++ b/.github/workflows/push.yaml @@ -175,12 +175,20 @@ jobs: env: PYTHONUNBUFFERED: 1 PYTEST_ADDOPTS: "--durations=25 -p no:unraisableexception -p no:threadexception" + # Two batch subprocesses per runner while its memory allows (see + # test_batched.py): 32% fewer runner-minutes, same cases, no shutdowns. + POLICYENGINE_BATCH_WORKERS: 2 strategy: fail-fast: false + # At most 6 Baseline groups run at once; GitHub starts the next group on a + # fresh runner as each finishes, so the lanes balance from real durations. + max-parallel: 6 matrix: include: - # One state per subprocess; retain four runners and one worker each. - # The former NY/OH/OK process peaked at 15.1 GB (run 36930199045). + # One state per batch subprocess, over four runners. The job's + # POLICYENGINE_BATCH_WORKERS overrides --workers 1, and runner memory + # decides whether two run at once. The former NY/OH/OK process peaked + # at 15.1 GB (run 36930199045). - group: states-shard-1 cmd: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 1/4 - group: states-shard-2 @@ -255,7 +263,8 @@ jobs: - name: Run Household API Partners YAML tests run: uv run --no-sync make test-yaml-no-structural-other-partners Contrib: - needs: ReleaseLock + # Waits for Baseline and skips if it fails, as in pr.yaml. + needs: [ReleaseLock, Baseline] name: Full Suite - Contrib (${{ matrix.group }}) runs-on: ubuntu-latest timeout-minutes: 60 @@ -265,8 +274,14 @@ jobs: env: PYTHONUNBUFFERED: 1 PYTEST_ADDOPTS: "--durations=25 -p no:unraisableexception -p no:threadexception" + # Two batch subprocesses per runner while its memory allows (see + # test_batched.py): 32% fewer runner-minutes, same cases, no shutdowns. + POLICYENGINE_BATCH_WORKERS: 2 strategy: fail-fast: false + # Contrib starts after Baseline finishes, so its 6 lanes take the place of + # Baseline's 6. + max-parallel: 6 matrix: include: - group: states-shard-1 @@ -380,8 +395,12 @@ jobs: # Historical group times: 32-46 minutes plus setup. # This runner also retains the legacy summary check, keeping PRs at 27 runners. timeout-minutes: 90 + # Runs after a shard-1 failure, but not when shard 1 was skipped (a failed + # release lock) or the run is cancelled. always() kept this runner alive + # after cancellation. if: | - always() + !cancelled() + && (needs.RestPython.result != 'skipped') && (github.repository == 'PolicyEngine/policyengine-us') && (github.event.head_commit.message == 'Update PolicyEngine US') env: diff --git a/.github/workflows/python-tests.yaml b/.github/workflows/python-tests.yaml index 7d22abd7b6d..82db686ed99 100644 --- a/.github/workflows/python-tests.yaml +++ b/.github/workflows/python-tests.yaml @@ -66,7 +66,9 @@ jobs: rest-variables-resources.txt Rest: needs: [RestPython] - if: ${{ always() }} + # Runs after a shard-1 failure but stops when the run is cancelled. + # always() kept this runner alive after cancellation. + if: ${{ !cancelled() }} name: Full Suite - Rest (Python + variables) runs-on: ubuntu-latest # Historical group times: 32-46 minutes plus setup. diff --git a/Makefile b/Makefile index 842b3d70fb2..a3dc2fcbe44 100644 --- a/Makefile +++ b/Makefile @@ -13,6 +13,9 @@ # test growth, and --workers 1 everywhere except the partners target # (its largest batch measured 2.5 GB, so two-wide stays trivially safe). BATCH := python policyengine_us/tests/test_batched.py +# CI's Baseline and Contrib jobs set POLICYENGINE_BATCH_WORKERS=2, which +# overrides --workers below: test_batched.py then runs two batch subprocesses +# per runner while its memory allows. Local runs keep the counts given here. TESTS := policyengine_us/tests # Run SPM construction/isolation tests in their own process on the existing # Microsimulation CI runner, before the dataset tests. Rest excludes these @@ -197,8 +200,9 @@ test-yaml-variables: $(BATCH) $(TESTS)/variables --batches 1 test-yaml-no-structural-states: # Release model/parameter caches between states. The former NY/OH/OK - # group peaked at 15.1 GB; keep one subprocess at a time on each of - # the existing four CI runners. Root-level state tests remain included. + # group peaked at 15.1 GB; keep one subprocess at a time here. CI's four + # runners run up to two while memory allows (POLICYENGINE_BATCH_WORKERS). + # Root-level state tests remain included. $(BATCH) $(TESTS)/policy/baseline/gov/states --mode per-subdir --workers 1 test-yaml-no-structural-other: $(BATCH) $(TESTS)/policy/baseline --batches 2 --exclude states @@ -212,8 +216,9 @@ test-yaml-no-structural-other-irs: $(BATCH) $(TESTS)/policy/baseline/gov/irs --mode per-subdir --workers 1 test-yaml-no-structural-other-household: # 4 batches (was 2), --workers 1: a 2-batch half of this folder (incl. - # the weights/MTR files) peaked 11.8 GB on CI run 28698452678. Quarter - # batches keep each subprocess <= ~6 GB. + # the weights/MTR files) peaked 11.8 GB on CI run 28698452678, and by + # October 2026 the last quarter batch reached 15.1 GB. test_batched.py's + # memory budget now continues a batch in a fresh subprocess past ~8 GB. $(BATCH) $(TESTS)/policy/baseline/household --batches 4 --workers 1 test-yaml-no-structural-other-contrib: # ubi_center is microsim-heavy → per-file, one ~4 GB peak at a time. @@ -234,8 +239,9 @@ test-yaml-reform: # ~8 GB file runs solo with ~8 GB spare); new reform files auto-route. $(BATCH) $(TESTS)/policy/reform --mode per-file --workers 1 test-yaml-no-structural-other-hhs: - # hhs (~2.7 GB peak) rides along with the baseline-contrib runner, - # which has spare headroom. + # hhs rides along with the baseline-contrib runner. As one process it + # peaked at 15.1 GB (CI run 38656284185); test_batched.py's memory budget + # continues it in a fresh subprocess past ~8 GB. $(BATCH) $(TESTS)/policy/baseline/gov/hhs --batches 1 # baseline contrib + gov/hhs share one CI runner. policy/reform no longer # rides along: its per-file peaks (~5.5-8 GB) plus this pair filled a job diff --git a/changelog.d/ci-sequential-suites-household-memory.fixed.md b/changelog.d/ci-sequential-suites-household-memory.fixed.md new file mode 100644 index 00000000000..4445a3eb2d9 --- /dev/null +++ b/changelog.d/ci-sequential-suites-household-memory.fixed.md @@ -0,0 +1 @@ +Stop the second Rest shard when a CI run is cancelled; run at most 6 Baseline groups at once and start Contrib (also at most 6) once Baseline passes, so the Pull request workflow runs at most 11 jobs at a time instead of about 25; run two YAML batch subprocesses per Baseline and Contrib runner while its free memory allows, cutting their runner time by about a third; and continue YAML test batches in a fresh subprocess once one passes an 8 GB memory budget, so batches no longer need re-splitting by hand as tests grow. diff --git a/policyengine_us/tests/code_health/test_batched_yaml_collection.py b/policyengine_us/tests/code_health/test_batched_yaml_collection.py index d2e6d20234c..f1a0517727e 100644 --- a/policyengine_us/tests/code_health/test_batched_yaml_collection.py +++ b/policyengine_us/tests/code_health/test_batched_yaml_collection.py @@ -1,18 +1,27 @@ -"""YAML batches must not rerun the separately routed Python suites. +"""YAML batches must not rerun the separately routed Python suites, must +continue in a fresh subprocess once one passes the memory budget, and share a +runner only while its memory allows. Exercise the real Core CLI and pytest collector with a one-variable country -package, so this runner regression never constructs the full US model. +package, so these runner regressions never construct the full US model. +Admission and eviction run against a scripted memory reading and stand-in +processes. """ import importlib.util +import io import os +import re import shlex import subprocess import sys +import threading +import time import xml.etree.ElementTree as ET from pathlib import Path import pytest +import yaml BATCHER = Path(__file__).resolve().parents[1] / "test_batched.py" SPEC = importlib.util.spec_from_file_location("test_batched", BATCHER) @@ -107,3 +116,413 @@ def test_python_only_directory_is_not_a_passing_yaml_batch(small_country): def test_python_only_subdir_does_not_create_a_yaml_batch(small_country): assert batched.split_into_batches(small_country.parent, 2, mode="per-subdir") == [] + + +def _write_age_cases(path, count, failing=0): + """`count` cases for one person's age; the first `failing` of them fail.""" + path.write_text( + yaml.safe_dump( + [ + { + "name": f"{path.stem} {i}", + "period": 2026, + "input": {"age": 30}, + "output": {"age": 31 if i < failing else 30}, + } + for i in range(count) + ] + ) + ) + + +def _part_summaries(output): + return re.findall(r"=+ ((?:\d+ failed, )?\d+ passed)", output) + + +def test_memory_budget_continues_unstarted_files_in_fresh_subprocesses( + small_country, monkeypatch +): + for name in ("a", "b", "c"): + _write_age_cases(small_country / f"{name}.yaml", 2) + # Every process is past a zero budget, so each file gets its own part. + monkeypatch.setenv(batched.MEMORY_BUDGET_ENV, "0") + report = small_country.parent / "report.xml" + monkeypatch.setenv("PYTEST_ADDOPTS", f"--junitxml={report}") + + result = batched.run_batch([str(small_country)], "budget", stream=False) + + assert result["status"] == "passed", result["output"] + assert result["parts"] == 3, result["output"] + # Each file ran exactly once, both of its cases, with none repeated. + assert _part_summaries(result["output"]) == ["2 passed"] * 3, result["output"] + # Every part rewrites the inherited report; the batch merges them back. + assert len(ET.parse(report).findall(".//testcase")) == 6 + + +def test_memory_budget_leaves_a_batch_under_it_in_one_subprocess( + small_country, monkeypatch +): + for name in ("a", "b", "c"): + _write_age_cases(small_country / f"{name}.yaml", 2) + monkeypatch.delenv(batched.MEMORY_BUDGET_ENV, raising=False) + + result = batched.run_batch([str(small_country)], "budget", stream=False) + + assert result["status"] == "passed", result["output"] + assert result["parts"] == 1, result["output"] + assert _part_summaries(result["output"]) == ["6 passed"], result["output"] + + +def test_memory_budget_never_splits_a_file(small_country, monkeypatch): + _write_age_cases(small_country / "a.yaml", 3) + monkeypatch.setenv(batched.MEMORY_BUDGET_ENV, "0") + + result = batched.run_batch([str(small_country)], "budget", stream=False) + + assert result["parts"] == 1, result["output"] + assert _part_summaries(result["output"]) == ["3 passed"], result["output"] + + +def test_a_failure_before_the_budget_stop_fails_the_batch_but_runs_the_rest( + small_country, monkeypatch +): + _write_age_cases(small_country / "a.yaml", 2, failing=1) + for name in ("b", "c"): + _write_age_cases(small_country / f"{name}.yaml", 2) + monkeypatch.setenv(batched.MEMORY_BUDGET_ENV, "0") + + result = batched.run_batch([str(small_country)], "budget", stream=False) + + assert result["status"] == "failed", result["output"] + assert result["returncode"] != 0 + assert _part_summaries(result["output"]) == [ + "1 failed, 1 passed", + "2 passed", + "2 passed", + ], result["output"] + + +WAIT = 5 + + +@pytest.fixture +def memory(monkeypatch): + """The runner's available memory, set by each test.""" + reading = {"mb": 16000.0} + monkeypatch.setattr(batched, "available_mb", lambda: reading["mb"]) + monkeypatch.setattr(batched, "ADMIT_STAGGER_SECONDS", 0) + return reading + + +@pytest.fixture +def admission(): + gate = batched.Admission() + yield gate + gate.stopped = True + + +class Process: + def __init__(self): + self.killed = threading.Event() + + def kill(self): + self.killed.set() + + +class Worker: + """A thread holding one admission slot, rerunning alone if stopped. + + As in run_batch, a stopped part's own thread asks for the exclusive slot. + """ + + def __init__(self, admission, process=None): + self.admitted = threading.Event() + self.go = threading.Event() + self.released = threading.Event() + self.rerun_start = threading.Event() + self.rerun_admitted = threading.Event() + self.rerun_go = threading.Event() + self.evicted = None + + def run(): + admission.acquire(False) + if process is not None: + admission.started(process) + self.admitted.set() + self.go.wait(WAIT * 4) + self.evicted = admission.release() + self.released.set() + if self.evicted: + self.rerun_start.wait(WAIT * 4) + admission.acquire(True) + self.rerun_admitted.set() + self.rerun_go.wait(WAIT * 4) + admission.release() + + threading.Thread(target=run, daemon=True).start() + + +def test_parent_reads_the_runners_available_memory(monkeypatch): + def reading(text): + monkeypatch.setattr( + batched, "open", lambda *args, **kwargs: io.StringIO(text), raising=False + ) + return batched.available_mb() + + assert reading("MemTotal: 16384000 kB\nMemAvailable: 8192000 kB\n") == 8000.0 + assert reading("MemTotal: 16384000 kB\n") is None + assert reading("MemAvailable: junk kB\n") is None + + def missing(*args, **kwargs): + raise OSError + + monkeypatch.setattr(batched, "open", missing, raising=False) + assert batched.available_mb() is None + + +def test_one_subprocess_at_a_time_when_memory_cannot_be_read(admission, memory): + memory["mb"] = None + admission.acquire(False) + second = Worker(admission) + assert not second.admitted.wait(1) + admission.release() + assert second.admitted.wait(WAIT) + second.go.set() + + +def test_another_subprocess_starts_only_with_memory_to_spare(admission, memory): + admission.acquire(False) + memory["mb"] = batched.ADMIT_MB - 1 + second = Worker(admission) + assert not second.admitted.wait(1.5) + memory["mb"] = batched.ADMIT_MB + assert second.admitted.wait(WAIT) + assert admission.waits == 1 + second.go.set() + + +def test_subprocesses_start_one_stagger_apart(admission, memory, monkeypatch): + monkeypatch.setattr(batched, "ADMIT_STAGGER_SECONDS", 1.5) + admission.acquire(False) + second = Worker(admission) + assert not second.admitted.wait(0.75) + assert second.admitted.wait(WAIT) + second.go.set() + + +def test_low_memory_stops_the_newest_and_admits_nothing_until_its_rerun( + admission, memory +): + older_process, newer_process = Process(), Process() + older = Worker(admission, process=older_process) + assert older.admitted.wait(WAIT) + time.sleep(0.05) + newer = Worker(admission, process=newer_process) + assert newer.admitted.wait(WAIT) + + memory["mb"] = batched.EVICT_MB - 1 + assert newer_process.killed.wait(WAIT) + assert not older_process.killed.is_set() + memory["mb"] = 16000.0 + newer.go.set() + assert newer.released.wait(WAIT) + assert newer.evicted is True + + # Nothing new starts between the stop and the rerun, nor during the rerun. + other = Worker(admission) + assert not other.admitted.wait(1) + newer.rerun_start.set() + assert not newer.rerun_admitted.wait(1) # the older one is still running + older.go.set() + assert newer.rerun_admitted.wait(WAIT) + assert not other.admitted.wait(1) + newer.rerun_go.set() + assert other.admitted.wait(WAIT) + other.go.set() + assert older.evicted is False + assert admission.evictions == 1 + + +def test_a_lone_subprocess_is_never_stopped(admission, memory): + process = Process() + admission.acquire(False) + admission.started(process) + memory["mb"] = 1.0 + assert not process.killed.wait(1) + assert admission.release() is False + + +def test_a_subprocess_that_has_reported_is_not_stopped(admission, memory): + older, newer = Process(), Process() + admission.acquire(False) + admission.started(older) + + def report(): + admission.acquire(False) + admission.started(newer) + admission.finished() + + thread = threading.Thread(target=report, daemon=True) + thread.start() + thread.join(WAIT) + memory["mb"] = 1.0 + assert not newer.killed.wait(1) + assert not older.killed.is_set() + + +def test_split_into_one_batch_per_worker_keeps_whole_files(tmp_path): + folder = tmp_path / "suite" + (folder / "nested").mkdir(parents=True) + files = [folder / "a.yaml", folder / "b.yaml", folder / "nested" / "c.yml"] + for file in files: + file.write_text("[]\n") + batches = batched.fill_workers([[str(folder)]], 2) + assert sorted(len(batch) for batch in batches) == [1, 2] + assert sorted(path for batch in batches for path in batch) == sorted( + str(file) for file in files + ) + two = [str(files[0]), str(files[1])] + assert sorted(len(batch) for batch in batched.fill_workers([two], 2)) == [1, 1] + assert batched.fill_workers([[str(files[0])]], 2) == [[str(files[0])]] + + +def test_split_keeps_no_batches_and_a_folder_without_yaml(tmp_path): + assert batched.fill_workers([], 2) == [] + empty = tmp_path / "python-only" + empty.mkdir() + (empty / "test_x.py").write_text("") + # Kept as the folder, so the subprocess still exits 5 (no tests collected) + # and fails the batch, as it does with one worker. + assert batched.fill_workers([[str(empty)]], 2) == [[str(empty)]] + + +class Gate: + """Stands in for Admission in run_batch; release() reports evictions.""" + + def __init__(self, evictions=()): + self.calls = [] + self.processes = [] + self.evictions = list(evictions) + self.forgotten = 0 + + def acquire(self, exclusive): + self.calls.append(exclusive) + + def started(self, process): + self.processes.append(process) + + def finished(self): + pass + + def forget(self): + self.forgotten = self.forgotten + 1 + + def release(self): + return self.evictions.pop(0) if self.evictions else False + + +def _scripted(flag, first_run_reports): + """A part that is stopped (or reports) on its first run, passes after.""" + first = ( + f" print({batched.BATCH_EXIT_MARKER!r} + ' 0', flush=True)\n" + if first_run_reports + else " os.kill(os.getpid(), signal.SIGKILL)\n" + ) + return [ + sys.executable, + "-c", + "import os, pathlib, signal, time\n" + f"flag = pathlib.Path({str(flag)!r})\n" + "if not flag.exists():\n" + " flag.touch()\n" + " print('==== 1 passed in 0.10s ====', flush=True)\n" + + first + + " time.sleep(30)\n" + "print('==== 2 passed in 0.10s ====')\n" + f"print({batched.BATCH_EXIT_MARKER!r} + ' 0', flush=True)\n", + ] + + +def test_an_evicted_part_reruns_alone_and_counts_once(small_country, monkeypatch): + gate = Gate(evictions=[True]) + monkeypatch.setattr(batched, "ADMISSION", gate) + command = _scripted(small_country / "first-run", first_run_reports=False) + result = batched.run_batch(["a.yaml"], "evicted", stream=False, command=command) + assert gate.calls == [False, True] + assert len(gate.processes) == 2 + assert all(hasattr(process, "kill") for process in gate.processes) + assert result["status"] == "passed", result["output"] + assert result["parts"] == 1 + assert "rerunning its 1 files alone" in result["output"] + + +def test_a_part_stopped_after_it_reported_keeps_its_result(small_country, monkeypatch): + # The watchdog can pick a part just as it reports; its result stands and + # its rerun reservation is dropped, so admission is not held forever. + gate = Gate(evictions=[True]) + monkeypatch.setattr(batched, "ADMISSION", gate) + command = _scripted(small_country / "first-run", first_run_reports=True) + result = batched.run_batch(["a.yaml"], "reported", stream=False, command=command) + assert gate.calls == [False] + assert gate.forgotten == 1 + assert result["parts"] == 1 + assert "rerunning" not in result["output"] + + +def _run_batcher(small_country, *args, workers=None): + # The runner checks it is started from the repository root. + (small_country.parent / "policyengine_us").mkdir(exist_ok=True) + env = dict(os.environ) + if workers is not None: + env[batched.WORKERS_ENV] = workers + run = subprocess.run( + [sys.executable, str(BATCHER), str(small_country), *args], + capture_output=True, + text=True, + timeout=600, + env=env, + cwd=small_country.parent, + ) + return run.returncode, run.stdout + run.stderr + + +def test_workers_variable_runs_every_case_and_keeps_failures(small_country): + for name in ("a", "b", "c"): + _write_age_cases( + small_country / f"{name}.yaml", 2, failing=1 if name == "b" else 0 + ) + code, output = _run_batcher(small_country, "--batches", "1", workers="2") + assert code == 1, output + assert "Workers: 2" in output + assert "Split into 2 batches for 2 workers" in output + assert "[memory] admission waits:" in output + passed = sum(int(n) for n in re.findall(r"(\d+) passed", output)) + failed = sum(int(n) for n in re.findall(r"(\d+) failed", output)) + assert (passed, failed) == (5, 1), output + + +def test_plain_workers_flag_keeps_running_its_batches_at_once(small_country): + for name in ("a", "b"): + (small_country / name).mkdir() + _write_age_cases(small_country / name / "x.yaml", 1) + code, output = _run_batcher(small_country, "--mode", "per-subdir", "--workers", "2") + assert code == 0, output + assert "Running 2 batches with 2 workers" in output + assert "[memory]" not in output + assert "Split into" not in output + + +@pytest.mark.skipif( + not Path("/proc/meminfo").exists(), reason="reads the runner's MemAvailable" +) +def test_low_runner_memory_hands_off_at_file_boundaries(small_country, monkeypatch): + for name in ("a", "b", "c"): + _write_age_cases(small_country / f"{name}.yaml", 2) + monkeypatch.setattr(batched, "ADMISSION", Gate()) + monkeypatch.setattr(batched, "FLOOR_MB", 10**9) + files = sorted(str(f) for f in small_country.glob("*.yaml")) + result = batched.run_batch(files, "floor", stream=False) + assert result["status"] == "passed", result["output"] + assert result["parts"] == 3 + assert "runner available" in result["output"] + assert "or runner floor" in result["output"] diff --git a/policyengine_us/tests/code_health/test_pr_runner_budget.py b/policyengine_us/tests/code_health/test_pr_runner_budget.py index 48fcb490cbb..1378ebf57b0 100644 --- a/policyengine_us/tests/code_health/test_pr_runner_budget.py +++ b/policyengine_us/tests/code_health/test_pr_runner_budget.py @@ -13,6 +13,7 @@ if "uses" in job } PYTHON_JOBS = GROUP_WORKFLOWS["PythonTests"]["jobs"] +PUSH_JOBS = yaml.safe_load((ROOT / ".github/workflows/push.yaml").read_text())["jobs"] KEEP_RUNNING = "${{ !cancelled() }}" @@ -30,6 +31,13 @@ def calls(job, name): ] +def runners(job): + matrix = job.get("strategy", {}).get("matrix", {}) + if "include" in matrix: + return len(matrix["include"]) + return math.prod(len(values) for values in matrix.values()) + + def test_pr_runner_budget_keeps_the_heavy_suites_separate(): counts = {} runner_jobs = {name: job for name, job in JOBS.items() if "uses" not in job} @@ -38,11 +46,7 @@ def test_pr_runner_budget_keeps_the_heavy_suites_separate(): assert name not in runner_jobs runner_jobs[name] = job for name, job in runner_jobs.items(): - matrix = job.get("strategy", {}).get("matrix", {}) - if "include" in matrix: - counts[name] = len(matrix["include"]) - else: - counts[name] = math.prod(len(values) for values in matrix.values()) + counts[name] = runners(job) # Rest hit the 90-minute timeout on #10024, #9913, #10139 and #10029 # in October 2026. One extra runner buys about 45 minutes of headroom: # two tested shards share variables and the legacy summary check. @@ -59,6 +63,87 @@ def test_pr_runner_budget_keeps_the_heavy_suites_separate(): } +def needs(job): + value = job.get("needs", []) + return [value] if isinstance(value, str) else value + + +def concurrent(job): + return min(job.get("strategy", {}).get("max-parallel", math.inf), runners(job)) + + +def test_lanes_cap_a_pr_at_11_runners_at_once(): + # Within pr.yaml, validation runs alone; everything that needs only it starts + # together, Rest waits for RestPython, and Contrib waits for Baseline. + # Baseline and Contrib each keep one job per group but run at most + # max-parallel of them at once: GitHub starts the next group as one + # finishes, so the lanes rebalance from real durations on every run, with + # nothing to re-pair as groups grow. + assert ( + GROUP_WORKFLOWS["Baseline"]["jobs"]["Baseline"]["strategy"]["max-parallel"] == 6 + ) + assert ( + GROUP_WORKFLOWS["Contrib"]["jobs"]["Contrib"]["strategy"]["max-parallel"] == 6 + ) + first_wave = {} + for name, job in JOBS.items(): + if needs(job) != ["ReleaseLock"]: + continue + if "uses" not in job: + first_wave[name] = concurrent(job) + continue + for inner_name, inner in GROUP_WORKFLOWS[name]["jobs"].items(): + if not needs(inner): + first_wave[inner_name] = concurrent(inner) + assert "Rest" not in first_wave and "Contrib" not in first_wave + assert sum(first_wave.values()) == 11 + # Contrib starts only after Baseline has finished, so at most Contrib's + # lanes plus every other first-wave job can be running then. + single = {name: n for name, n in first_wave.items() if name != "Baseline"} + contrib = concurrent(GROUP_WORKFLOWS["Contrib"]["jobs"]["Contrib"]) + assert contrib + sum(single.values()) <= 11 + + +def test_contrib_waits_for_baseline_and_skips_when_it_fails(): + # A plain needs (no if, or an if without a status function) skips Contrib + # when Baseline fails or is cancelled. Re-running the failed Baseline group + # then runs Contrib for the first time (fork run 38687143315); an if that + # let Contrib run anyway would rerun all of its finished groups. + assert needs(JOBS["Contrib"]) == ["ReleaseLock", "Baseline"] + assert "if" not in JOBS["Contrib"] + assert needs(PUSH_JOBS["Contrib"]) == ["ReleaseLock", "Baseline"] + for status in ("always()", "cancelled()", "failure()", "success()"): + assert status not in PUSH_JOBS["Contrib"]["if"] + + +def test_yaml_suites_run_two_memory_aware_workers_per_runner(): + for name in ("Baseline", "Contrib"): + for job in (GROUP_WORKFLOWS[name]["jobs"][name], PUSH_JOBS[name]): + assert job["env"]["POLICYENGINE_BATCH_WORKERS"] == 2 + + +def test_release_ci_runs_the_same_baseline_and_contrib_lanes(): + for name in ("Baseline", "Contrib"): + pr_job = GROUP_WORKFLOWS[name]["jobs"][name] + assert PUSH_JOBS[name]["strategy"] == pr_job["strategy"] + assert PUSH_JOBS[name]["steps"] == pr_job["steps"] + assert PUSH_JOBS[name]["env"] == pr_job["env"] + + +def test_no_test_job_keeps_running_after_cancellation(): + # A job-level always() also holds after a manual or cancel-in-progress + # cancellation, so the job keeps running or even starts after it. A job + # that waits for another uses !cancelled() to run after its failure. + workflows = {"pr.yaml": JOBS, "push.yaml": PUSH_JOBS} | { + Path(job["uses"]).name: GROUP_WORKFLOWS[name]["jobs"] + for name, job in JOBS.items() + if "uses" in job + } + for file, jobs in workflows.items(): + for name, job in jobs.items(): + assert "always()" not in str(job.get("if", "")), (file, name) + + def test_suite_groups_have_no_extra_runners_or_triggers(): groups = { "Baseline": ("Baseline tests", "baseline-tests.yaml"), @@ -69,7 +154,9 @@ def test_suite_groups_have_no_extra_runners_or_triggers(): for name, (label, file) in groups.items(): assert JOBS[name] == { "name": label, - "needs": "ReleaseLock", + "needs": ["ReleaseLock", "Baseline"] + if name == "Contrib" + else "ReleaseLock", "uses": f"./.github/workflows/{file}", } # PyYAML parses the bare `on` key as True. No push/PR trigger here: @@ -88,7 +175,8 @@ def test_suite_groups_have_no_extra_runners_or_triggers(): assert not job.get("continue-on-error") # The second tested runner also checks the first runner's result. There # is no third runner just for the summary, or a variables-only runner. - assert PYTHON_JOBS["Rest"]["if"] == "${{ always() }}" + # !cancelled(), not always(): always() kept it running after cancellation. + assert PYTHON_JOBS["Rest"]["if"] == KEEP_RUNNING assert PYTHON_JOBS["Rest"]["name"] == "Full Suite - Rest (Python + variables)" diff --git a/policyengine_us/tests/code_health/test_rest_python_groups.py b/policyengine_us/tests/code_health/test_rest_python_groups.py index f76b4e2818b..73e52e9a801 100644 --- a/policyengine_us/tests/code_health/test_rest_python_groups.py +++ b/policyengine_us/tests/code_health/test_rest_python_groups.py @@ -385,19 +385,22 @@ def test_original_rest_check_aggregates_all_shards_without_hiding_failure(workfl summary = jobs["Rest"] assert summary["name"] == "Full Suite - Rest (Python + variables)" assert summary["needs"] == ["RestPython"] - assert "always()" in summary["if"] + # Shard 2 still runs after a shard-1 failure, but a cancelled run stops + # it: a job-level always() kept it running after cancellation. + assert "always()" not in summary["if"] assert "strategy" not in summary assert summary["timeout-minutes"] == 90 assert not summary.get("continue-on-error", False) assert not jobs["RestPython"].get("continue-on-error", False) if workflow.name == "push.yaml": assert " ".join(summary["if"].split()) == ( - "always() && (github.repository == 'PolicyEngine/policyengine-us') " + "!cancelled() && (needs.RestPython.result != 'skipped') " + "&& (github.repository == 'PolicyEngine/policyengine-us') " "&& (github.event.head_commit.message == 'Update PolicyEngine US')" ) assert {"RestPython", "Rest"} <= set(jobs["Publish"]["needs"]) else: - assert summary["if"] in ("always()", "${{ always() }}") + assert summary["if"] == "${{ !cancelled() }}" (step,) = [ candidate for candidate in summary["steps"] diff --git a/policyengine_us/tests/test_batched.py b/policyengine_us/tests/test_batched.py index 73f9e510640..099c5867cfb 100644 --- a/policyengine_us/tests/test_batched.py +++ b/policyengine_us/tests/test_batched.py @@ -7,16 +7,19 @@ import subprocess import sys import os +import tempfile import gc import time import argparse import re import select +import shlex import threading from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor from concurrent.futures import wait as futures_wait from pathlib import Path from typing import FrozenSet, List, Dict, Optional, Tuple +from xml.etree import ElementTree import yaml @@ -489,7 +492,9 @@ def collect_gov_paths(folder_names): PYTEST_NO_TESTS_COLLECTED = 5 # How long a child gets to die after being asked to, before it is killed. TERMINATE_GRACE_SECONDS = 5 -# Overall budget for one batch, enforced inside the monitoring loop below. +# Budget for one subprocess, enforced inside the monitoring loop below. A +# batch the memory budget continues in fresh subprocesses gets a fresh budget +# for each part (see run_batch). BATCH_TIMEOUT_SECONDS = 1800 # How long the runner may take to hand back a status after printing its # summary. That gap is pytest's own session-finish hooks and unconfigure, not @@ -514,6 +519,148 @@ def collect_gov_paths(folder_names): # is still terminated exactly as promptly as before. BATCH_EXIT_MARKER = "__POLICYENGINE_BATCH_EXIT__" +# Memory accumulates across the YAML files one subprocess runs: household's 15 +# root files peaked at 13.3 GB together (local peak footprint, 2026-10-14, within +# 1% of CI's VmHWM on the batches compared), yet none passed 6.2 GB alone, and +# CI's household and gov/hhs batches each reached 15.1 GB on the 16 GB runner. +# Any fixed batch layout therefore outgrows the runner as tests are added. Instead, after each file the child checks its own peak memory and, +# once past this budget, hands the files it has not started to a fresh +# subprocess (run_batch continues them). The next file starts below the budget, +# so a subprocess peaks at about the budget plus one file's growth. A file that +# alone exceeds the budget still runs whole. Set the variable to override it, +# e.g. 0 for one file per subprocess. +MEMORY_BUDGET_MB = 8000 +MEMORY_BUDGET_ENV = "POLICYENGINE_BATCH_MEMORY_BUDGET_MB" +REMAINDER_ENV = "POLICYENGINE_BATCH_REMAINDER" + +# POLICYENGINE_BATCH_WORKERS lets several batch subprocesses share a runner +# while its memory allows; it overrides --workers. (A plain --workers N, as +# the partners target uses for its small batches, still runs N at once.) A +# 16 GB, 4-CPU CI runner running one subprocess is about 25% busy, but two +# fixed subprocesses ran it out of memory (fork runs 38687143268 and +# 38689387981: a runner shutdown in the crfb group in each, and 77-296 MB +# left in the contrib states groups). So with the variable above 1, memory +# decides how many run: +# - another subprocess starts only while the runner has ADMIT_MB available, +# and at most one starts per ADMIT_STAGGER_SECONDS, since a new one takes +# that long to load the model; one subprocess may always run; +# - a subprocess hands off its unstarted files at a file boundary (as for the +# memory budget) once the runner has less than FLOOR_MB available, and the +# remainder waits for admission; +# - if the runner falls below EVICT_MB mid-file, the newest of two or more +# subprocesses is stopped and its files rerun alone once the others finish; +# - batches are split, whole files only, until there is one per worker. +# Where available memory cannot be read (off Linux), one subprocess runs at a +# time. With two workers the 20 Baseline and Contrib groups averaged 345 +# runner-minutes per fork run against 508 with one (32% fewer over three runs; +# 33% over 38689387981 and 38710504296, which ran this eviction protocol), +# passed the same 34,794 cases, and never ran a runner out of memory. Three or +# four workers were no faster, because memory, not CPU, runs out first. +WORKERS_ENV = "POLICYENGINE_BATCH_WORKERS" +FLOOR_ENV = "POLICYENGINE_BATCH_MEMORY_FLOOR_MB" +ADMIT_MB = 6000 +FLOOR_MB = 4000 +EVICT_MB = 2000 +ADMIT_STAGGER_SECONDS = 20 + + +def available_mb() -> Optional[float]: + """The runner's MemAvailable in MB, or None off Linux.""" + try: + with open("/proc/meminfo") as meminfo: + for line in meminfo: + if line.startswith("MemAvailable:"): + return int(line.split()[1]) / 1024 + except (OSError, ValueError, IndexError): + return None + return None + + +# Child-side half of the memory budget. policyengine-core calls +# pytest.main(..., plugins=[...]) through the pytest module, so wrapping it +# adds this plugin without changing Core. +MEMORY_GUARD = """ +budget_mb = os.environ.get("POLICYENGINE_BATCH_MEMORY_BUDGET_MB") +remainder_path = os.environ.get("POLICYENGINE_BATCH_REMAINDER") +floor_mb = os.environ.get("POLICYENGINE_BATCH_MEMORY_FLOOR_MB") +if budget_mb and remainder_path: + import pytest + + def runner_available_mb(): + try: + with open("/proc/meminfo") as meminfo: + for line in meminfo: + if line.startswith("MemAvailable:"): + return int(line.split()[1]) / 1024 + except OSError: + pass + return None + + def peak_mb(): + try: + with open("/proc/self/status") as status: + for line in status: + if line.startswith("VmHWM:"): + return int(line.split()[1]) / 1024 + except OSError: + pass + if sys.platform == "darwin": + # ru_maxrss leaves out compressed memory, which reads 2-3x low on + # these suites. RUSAGE_INFO_V4 word 30 is the lifetime peak physical + # footprint that /usr/bin/time -l reports. + try: + import ctypes + + info = (ctypes.c_uint64 * 64)() + libc = ctypes.CDLL("/usr/lib/libSystem.B.dylib") + if libc.proc_pid_rusage(os.getpid(), 4, ctypes.byref(info)) == 0: + return info[30] / 2**20 + except (OSError, AttributeError): + pass + import resource + + peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + # ru_maxrss is bytes on macOS and kilobytes on Linux. + return peak / 2**20 if sys.platform == "darwin" else peak / 1024 + + class MemoryBudget: + @pytest.hookimpl(wrapper=True) + def pytest_runtest_protocol(self, item, nextitem): + result = yield + # Only stop between files, so every file runs whole in one process. + if nextitem is None or nextitem.path == item.path: + return result + peak = peak_mb() + low = runner_available_mb() if floor_mb else None + if low is not None and low >= float(floor_mb): + low = None + if peak < float(budget_mb) and low is None: + return result + items = item.session.items + remaining = list( + dict.fromkeys(str(later.path) for later in items[items.index(nextitem) :]) + ) + with open(remainder_path, "w") as handle: + handle.write("\\n".join(remaining)) + reason = ( + "peak %.0f MB" % peak + if low is None + else "runner available %.0f MB (own peak %.0f MB)" % (low, peak) + ) + pytest.exit( + "memory budget: %s after %s; %d files continue in a " + "fresh subprocess" % (reason, item.path.name, len(remaining)), + returncode=1 if item.session.testsfailed else 0, + ) + + run_pytest = pytest.main + + def main(args=None, plugins=None, **kwargs): + return run_pytest(args, [*(plugins or []), MemoryBudget()], **kwargs) + + pytest.main = main +""" + RUNNER_SHIM = f""" import os import runpy @@ -524,7 +671,7 @@ def collect_gov_paths(folder_names): os.environ["PYTEST_ADDOPTS"] = ( os.environ.get("PYTEST_ADDOPTS", "") + " --ignore-glob=*.py" ) - +{MEMORY_GUARD} code = 0 try: runpy.run_module( @@ -593,11 +740,267 @@ def batch_status(returncode: Optional[int]) -> str: return "passed" if returncode == 0 else "failed" +def junit_report_path() -> Optional[str]: + """The --junitxml file PYTEST_ADDOPTS asks every part to write, if any.""" + tokens = shlex.split(os.environ.get("PYTEST_ADDOPTS", "")) + path = None + for i, token in enumerate(tokens): + for flag in ("--junitxml", "--junit-xml"): + if token.startswith(flag + "="): + path = token.split("=", 1)[1] + elif token == flag and i + 1 < len(tokens): + path = tokens[i + 1] + return path + + +def report_mtime(path: Optional[str]) -> Optional[int]: + try: + return os.stat(path).st_mtime_ns if path else None + except FileNotFoundError: + return None + + +class Admission: + """Starts batch subprocesses only while the runner has memory for them. + + One subprocess may always run, so a run never stalls. An evicted part + reruns exclusively: no new subprocess starts while it waits or runs. + """ + + def __init__(self, admit_mb: float = ADMIT_MB, evict_mb: float = EVICT_MB): + self.admit_mb = admit_mb + self.evict_mb = evict_mb + self.cond = threading.Condition() + self.running: Dict[int, list] = {} + self.evicted: set = set() + self.pending_exclusive: set = set() + self.exclusive_waiting = 0 + self.exclusive_holder: Optional[int] = None + self.last_admit = 0.0 + self.waits = 0 + self.evictions = 0 + self.stopped = False + threading.Thread(target=self.watch, daemon=True).start() + + def may_admit(self) -> bool: + if ( + self.exclusive_holder is not None + or self.exclusive_waiting + or self.pending_exclusive + ): + return False + if not self.running: + return True + if time.time() - self.last_admit < ADMIT_STAGGER_SECONDS: + return False + avail = available_mb() + return avail is not None and avail >= self.admit_mb + + def acquire(self, exclusive: bool) -> None: + ident = threading.get_ident() + with self.cond: + if exclusive: + self.exclusive_waiting = self.exclusive_waiting + 1 + self.pending_exclusive.discard(ident) + while self.running or self.exclusive_holder is not None: + self.cond.wait(1) + self.exclusive_waiting = self.exclusive_waiting - 1 + self.exclusive_holder = ident + elif not self.may_admit(): + self.waits = self.waits + 1 + while not self.may_admit(): + self.cond.wait(1) + self.running[ident] = [None, time.time()] + self.last_admit = time.time() + + def started(self, process) -> None: + with self.cond: + entry = self.running.get(threading.get_ident()) + if entry is not None: + entry[0] = process + + def finished(self) -> None: + """The part reported or exited: stop treating it as live.""" + self.started(None) + + def forget(self) -> None: + """Drop this thread's rerun reservation, if it will not rerun.""" + with self.cond: + self.pending_exclusive.discard(threading.get_ident()) + self.cond.notify_all() + + def release(self) -> bool: + ident = threading.get_ident() + with self.cond: + self.running.pop(ident, None) + evicted = ident in self.evicted + self.evicted.discard(ident) + if self.exclusive_holder == ident: + self.exclusive_holder = None + self.cond.notify_all() + return evicted + + def watch(self) -> None: + while not self.stopped: + time.sleep(0.2) + avail = available_mb() + if avail is None or avail >= self.evict_mb: + continue + with self.cond: + live = [ + (entry[1], ident, entry[0]) + for ident, entry in self.running.items() + if entry[0] is not None and ident not in self.evicted + ] + if len(live) < 2: + continue + _, ident, process = max(live, key=lambda item: item[0]) + self.evicted.add(ident) + self.pending_exclusive.add(ident) + self.evictions = self.evictions + 1 + print( + f"[memory] runner available {avail:.0f} MB < {self.evict_mb:.0f} MB " + f"with {len(live)} subprocesses; stopping the newest and " + "rerunning its files alone", + flush=True, + ) + try: + process.kill() + except OSError: + pass + time.sleep(2) + + +ADMISSION: Optional[Admission] = None + + +def fill_workers(batches: List[List[str]], workers: int) -> List[List[str]]: + """Split batches, whole files only, until there is one per worker.""" + if not batches: + return batches + expanded = [] + for batch in batches: + files = [ + str(f) + for p in batch + for f in (yaml_files(Path(p)) if Path(p).is_dir() else [Path(p)]) + ] + expanded.append(files or list(batch)) + splittable = [ + len(b) > 1 and all(p.endswith(YAML_SUFFIXES) for p in b) for b in expanded + ] + while len(expanded) < workers and any(splittable): + largest = max( + (i for i in range(len(expanded)) if splittable[i]), + key=lambda i: len(expanded[i]), + ) + files = expanded[largest] + half = len(files) // 2 + expanded[largest : largest + 1] = [files[:half], files[half:]] + splittable[largest : largest + 1] = [half > 1, len(files) - half > 1] + return expanded + + def run_batch( test_paths: List[str], batch_name: str, stream: bool = True, command: Optional[List[str]] = None, +) -> Dict: + """Run a batch, continuing past the memory budget in fresh subprocesses. + + Each part runs through run_subprocess. A part that reaches + MEMORY_BUDGET_MB between files (or, when subprocesses share the runner, + finds less than FLOOR_MB available) writes the files it has not started to + a remainder file and exits; they continue in the next part. A part stopped + mid-file for runner memory is discarded and rerun alone. The batch passes + only if every part passes. Each part rewrites an inherited --junitxml file, + so the parts' suites are merged back into it. + """ + budget = os.environ.get(MEMORY_BUDGET_ENV, str(MEMORY_BUDGET_MB)) + report = junit_report_path() + suites: List[ElementTree.Element] = [] + parts: List[Dict] = [] + paths = test_paths + exclusive = False + discarded = "" + while paths: + handle, remainder = tempfile.mkstemp(prefix="batch-remainder-") + os.close(handle) + os.unlink(remainder) + name = batch_name if not parts else f"{batch_name} (part {len(parts) + 1})" + env = dict(os.environ, **{MEMORY_BUDGET_ENV: budget, REMAINDER_ENV: remainder}) + if ADMISSION is not None: + env[FLOOR_ENV] = str(FLOOR_MB) + written = report_mtime(report) + if ADMISSION is not None: + ADMISSION.acquire(exclusive) + try: + part = run_subprocess(paths, name, stream, command, env) + finally: + evicted = ADMISSION.release() if ADMISSION is not None else False + if evicted and not part.get("reported"): + note = ( + f" [memory] {name} stopped mid-file for runner memory; " + f"rerunning its {batch_cost(paths)} files alone" + ) + if stream: + print(note, flush=True) + discarded = discarded + part["output"] + note + "\n" + if os.path.exists(remainder): + os.unlink(remainder) + exclusive = True + continue + if evicted: + # Stopped only after it reported: keep its result. + ADMISSION.forget() + exclusive = False + parts.append(part) + command = None + if report_mtime(report) not in (None, written): + suites.extend(ElementTree.parse(report).getroot().iter("testsuite")) + try: + with open(remainder) as handle_in: + paths = [line for line in handle_in.read().splitlines() if line] + os.unlink(remainder) + except FileNotFoundError: + paths = [] + if paths: + limit = f"Memory budget ({budget} MB)" + if ADMISSION is not None: + limit = f"{limit} or runner floor ({FLOOR_MB} MB available)" + note = ( + f" {limit} reached; continuing " + f"{len(paths)} files in a fresh subprocess" + ) + if stream: + print(note, flush=True) + else: + parts[-1]["output"] = parts[-1]["output"] + note + "\n" + if len(parts) > 1 and suites: + merged = ElementTree.Element("testsuites") + merged.extend(suites) + ElementTree.ElementTree(merged).write( + report, encoding="utf-8", xml_declaration=True + ) + failed = [part for part in parts if part["status"] != "passed"] + peaks = [part["peak_rss_mb"] for part in parts if part["peak_rss_mb"] is not None] + return { + "elapsed": sum(part["elapsed"] for part in parts), + "status": failed[0]["status"] if failed else "passed", + "returncode": failed[0]["returncode"] if failed else 0, + "peak_rss_mb": max(peaks) if peaks else None, + "parts": len(parts), + "output": discarded + "".join(part["output"] for part in parts), + } + + +def run_subprocess( + test_paths: List[str], + batch_name: str, + stream: bool = True, + command: Optional[List[str]] = None, + env: Optional[Dict[str, str]] = None, ) -> Dict: """Run a batch of tests in an isolated subprocess. @@ -639,7 +1042,10 @@ def emit(text: str, end: str = "\n") -> None: stdout=subprocess.PIPE, stderr=subprocess.STDOUT, bufsize=0, + env=env, ) + if ADMISSION is not None: + ADMISSION.started(process) def sample_rss() -> None: # VmHWM is monotonic, so keeping only the latest reading preserves @@ -697,6 +1103,8 @@ def overdue() -> Optional[str]: poll_result = process.poll() if poll_result is not None: # Process terminated + if ADMISSION is not None: + ADMISSION.finished() break give_up = overdue() if give_up is not None: @@ -768,6 +1176,8 @@ def overdue() -> Optional[str]: marker_match = BATCH_EXIT_PATTERN.search(output_text) if marker_match: reported_returncode = int(marker_match.group(1)) + if ADMISSION is not None: + ADMISSION.finished() emit( "\n Runner reported " f"{describe_exit_status(reported_returncode)}; " @@ -782,9 +1192,9 @@ def overdue() -> Optional[str]: # If the runner never reported, wait for the process with a timeout if reported_returncode is None and process.poll() is None: try: - # Wait up to 30 minutes total + # Wait out the rest of this subprocess's budget elapsed = time.time() - start_time - remaining_timeout = max(1800 - elapsed, 1) + remaining_timeout = max(BATCH_TIMEOUT_SECONDS - elapsed, 1) process.wait(timeout=remaining_timeout) except subprocess.TimeoutExpired: emit(f"\n ⏱️ Timeout - terminating process...") @@ -827,6 +1237,7 @@ def overdue() -> Optional[str]: } status = batch_status(returncode) + reported = reported_returncode is not None if test_completed and summary_reported_pass and status != "passed": # The case that used to be reported as a pass: pytest printed no # failure count (e.g. "1 passed, 1 error") yet finished nonzero. @@ -846,6 +1257,7 @@ def overdue() -> Optional[str]: "returncode": returncode, "peak_rss_mb": peak_rss_mb, "output": "".join(buf), + "reported": reported, } except Exception as e: @@ -899,9 +1311,18 @@ def run_one(idx: int, batch_paths: List[str]) -> Dict: name = f"Batch {idx}" with in_flight_lock: in_flight[name] = time.time() + # Output is buffered until the batch ends, so record where each batch + # starts: a job timeout or lost runner otherwise leaves no trace of it. + print( + f"[Batch {idx}/{total}] started: {batch_cost(batch_paths)} files " + f"from {batch_paths[0]}", + flush=True, + ) try: return run_batch(batch_paths, name, stream=False) finally: + if ADMISSION is not None: + ADMISSION.forget() with in_flight_lock: in_flight.pop(name, None) @@ -1009,6 +1430,9 @@ def main(): ) args = parser.parse_args() + scheduled = bool(os.environ.get(WORKERS_ENV)) + if scheduled: + args.workers = int(os.environ[WORKERS_ENV]) exclude_list = [x.strip() for x in args.exclude.split(",") if x.strip()] shard_idx = shard_count = None @@ -1065,6 +1489,13 @@ def main(): f"({len(batches)} of {all_batch_count} batches)" ) + global ADMISSION + if scheduled and args.workers > 1: + if len(batches) < args.workers: + batches = fill_workers(batches, args.workers) + print(f"Split into {len(batches)} batches for {args.workers} workers") + ADMISSION = Admission() + if len(batches) != args.batches: print(f"Actual batches: {len(batches)} (optimized for {total_tests} files)") else: @@ -1107,6 +1538,12 @@ def main(): f"{result['status']}" ) print(f"Workers: {args.workers}") + if ADMISSION is not None: + ADMISSION.stopped = True + print( + f"[memory] admission waits: {ADMISSION.waits}, " + f"evictions: {ADMISSION.evictions}" + ) print(f"Total time: {total_elapsed:.1f}s") if args.workers > 1: print(f"Wall time: {time.time() - run_start:.1f}s")