From c8e9f3fbb8b52aec800046495bb4c2a017fd2043 Mon Sep 17 00:00:00 2001 From: Amber Li Date: Tue, 18 Aug 2026 17:07:58 -0400 Subject: [PATCH] Retry a learned artifact's pickle once, after a collection Saving learned NSRTs and GNN weights intermittently dies with `TypeError: cannot pickle '_abc._abc_data' object` -- the C-level cache behind an abstract base class. It has blocked three PRs on three different shards this week and it predates all of them: the same nine-test cluster reproduces on unmodified master. WHAT IS ESTABLISHED Two call sites, both reached from the CI tracebacks: nsrt_learning_approach.py :114 (`pkl.dump(self._nsrts, f)`) and gnn_approach.py:207 (`pkl.dump(info, f)`). Every failure seen so far funnels through one of them. On CI it is DETERMINISTIC for a given test set: shard 8 failed twice on the same single test, and shard 6 twice on the same three, across re-runs of untouched jobs. Which shard is hit moves with pytest-split's packing, which is why adding tests to an unrelated PR appears to "cause" it. Locally it is stochastic -- 3 of 12 identical runs, with code, test order and PYTHONHASHSEED all fixed. The likely difference is that CI containers are uniform while a developer machine is not. WHAT IS NOT ESTABLISHED The root cause. An `_abc_data` holds WEAK references, so whether dill trips over one plausibly turns on collection timing -- the one thing that still varies with everything else pinned. That is the hypothesis this mitigation is aimed at, and it is a hypothesis: it is stated here rather than dressed up as a diagnosis. If the failures stop, that is also the evidence for it. If they do not, it rules the hypothesis out, which is worth knowing too. Ruled out along the way, none of them the cause: pytest-randomly (not installed -- `-p no:randomly` is a silent no-op here), execution order (progress lines match character-for-character between a red and a green run), hash randomisation (PYTHONHASHSEED=0 is exported), and any single polluting test (a bisection appeared to find one, then the control showed removing its neighbours worked equally well -- that bisection was invalid, having used single runs to measure a 25% event). WHY IT IS SHAPED THIS WAY Serialising to bytes before writing, rather than retrying into the handle: a dump that raises part-way has already written a prefix, and appending a retry to that leaves a corrupt file which only fails at LOAD time -- much worse than the error being fixed. A test pins the file being empty after a persistent failure. Any TypeError is retried, not just the `_abc_data` one. Matching on the message would break silently when it is reworded, and an object that is genuinely unpicklable fails the second time too and raises exactly as before -- so the broad catch costs one wasted attempt and hides nothing. Two tests, both mutation-checked: removing the retry fails the first and not the second, and dumping straight into the handle fails both. Full local gates: mypy clean, pylint 10.00/10, yapf/docformatter/isort applied. The previously failing shard passes locally, though a single local run is weak evidence at a 25% rate -- CI is the real test of this, and is why it is worth landing to find out. --- predicators/approaches/gnn_approach.py | 2 +- .../approaches/nsrt_learning_approach.py | 2 +- predicators/utils.py | 41 +++++++++++++- tests/test_utils.py | 54 +++++++++++++++++++ 4 files changed, 95 insertions(+), 4 deletions(-) diff --git a/predicators/approaches/gnn_approach.py b/predicators/approaches/gnn_approach.py index 75b4ef29ab..00fffac8c2 100644 --- a/predicators/approaches/gnn_approach.py +++ b/predicators/approaches/gnn_approach.py @@ -204,7 +204,7 @@ def learn_from_offline_dataset(self, dataset: Dataset) -> None: self._add_output_specific_fields_to_save_info(info) save_path = utils.get_approach_save_path_str() with open(f"{save_path}_None.gnn", "wb") as f: - pkl.dump(info, f) + utils.pkl_dump_with_retry(info, f) def load(self, online_learning_cycle: Optional[int]) -> None: save_path = utils.get_approach_load_path_str() diff --git a/predicators/approaches/nsrt_learning_approach.py b/predicators/approaches/nsrt_learning_approach.py index b852b7d00c..24e28fd0f4 100644 --- a/predicators/approaches/nsrt_learning_approach.py +++ b/predicators/approaches/nsrt_learning_approach.py @@ -111,7 +111,7 @@ def _learn_nsrts(self, trajectories: List[LowLevelTrajectory], annotations=annotations) save_path = utils.get_approach_save_path_str() with open(f"{save_path}_{online_learning_cycle}.NSRTs", "wb") as f: - pkl.dump(self._nsrts, f) + utils.pkl_dump_with_retry(self._nsrts, f) if CFG.compute_sidelining_objective_value: self._compute_sidelining_objective_value(trajectories) diff --git a/predicators/utils.py b/predicators/utils.py index f27c36b77b..362279d1ab 100644 --- a/predicators/utils.py +++ b/predicators/utils.py @@ -27,8 +27,8 @@ from dataclasses import dataclass, field from functools import cached_property from pathlib import Path -from typing import TYPE_CHECKING, Any, Callable, ClassVar, Collection, Dict, \ - FrozenSet, Generator, Generic, Hashable, Iterable, Iterator, List, \ +from typing import IO, TYPE_CHECKING, Any, Callable, ClassVar, Collection, \ + Dict, FrozenSet, Generator, Generic, Hashable, Iterable, Iterator, List, \ Optional, Sequence, Set, Tuple from typing import Type as TypingType from typing import TypeVar, Union, cast @@ -3723,6 +3723,43 @@ def save_ground_atom_dataset(ground_atom_dataset: List[GroundAtomTrajectory], pkl.dump(ground_atom_dataset_to_pkl, f) +def pkl_dump_with_retry(obj: Any, f: IO[bytes]) -> None: + """``pkl.dump``, retried once after a collection if it raises TypeError. + + Saving a learned artifact intermittently dies with ``TypeError: cannot + pickle '_abc._abc_data' object``, which is the C-level cache behind an + abstract base class. On CI it is reproducible for a given set of tests -- + the same failures twice on the same shard, three times over -- and it has + been seen from two call sites, ``nsrt_learning_approach._learn_nsrts`` and + ``gnn_approach.learn_from_offline_dataset``. Locally it appears at roughly + one run in four with the code, test order and PYTHONHASHSEED all fixed. + + The root cause is NOT established. What is: an ``_abc_data`` holds WEAK + references, so whether dill trips over one plausibly depends on collection + timing, which is the one thing that varies run to run under everything + else being pinned. ``gc.collect()`` before retrying is aimed at exactly + that. **This is a mitigation on a hypothesis, not a fix on a diagnosis** -- + if it stops the failures it is also the evidence for the hypothesis, and + if it does not, that rules the hypothesis out. + + Serialising to bytes first rather than retrying into ``f`` matters: a dump + that raises part-way has already written a prefix, and a retry appending + to that would leave a corrupt file that only fails at load time, which is + much worse than the error being fixed here. + + Any TypeError is retried, not only the ``_abc_data`` one. Matching on the + message would break silently when it is reworded, and an object that is + genuinely unpicklable fails the second time too and raises as it always + would -- so the broad catch costs one wasted attempt and hides nothing. + """ + try: + blob = pkl.dumps(obj) + except TypeError: + gc.collect() + blob = pkl.dumps(obj) + f.write(blob) + + def merge_ground_atom_datasets( gad1: List[GroundAtomTrajectory], gad2: List[GroundAtomTrajectory]) -> List[GroundAtomTrajectory]: diff --git a/tests/test_utils.py b/tests/test_utils.py index a456cc93bc..c7a3d71efd 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -3746,3 +3746,57 @@ def test_parse_model_output_into_option_plan_strict(): strict=True) assert len(no_seed) == 1 assert no_seed[0][2] == [] + + +def test_pkl_dump_with_retry_survives_a_transient_failure( + tmp_path, monkeypatch): + """A TypeError on the first attempt must not lose the artifact. + + The real failure is ``cannot pickle '_abc._abc_data' object``, seen + intermittently while saving learned NSRTs and GNN weights. It is + faked here because it does not reproduce on demand -- which is the + whole reason the retry exists rather than a targeted fix. + """ + attempts = [] + real_dumps = utils.pkl.dumps + + def _flaky_dumps(obj, *args, **kwargs): + """Fail once with the real error, then behave.""" + attempts.append(obj) + if len(attempts) == 1: + raise TypeError("cannot pickle '_abc._abc_data' object") + return real_dumps(obj, *args, **kwargs) + + monkeypatch.setattr(utils.pkl, "dumps", _flaky_dumps) + path = tmp_path / "artifact.pkl" + with open(path, "wb") as f: + utils.pkl_dump_with_retry({"learned": [1, 2, 3]}, f) + + assert len(attempts) == 2, "the failed dump was not retried" + with open(path, "rb") as f: + assert utils.pkl.load(f) == {"learned": [1, 2, 3]} + + +def test_pkl_dump_with_retry_writes_nothing_when_it_fails( + tmp_path, monkeypatch): + """A persistent failure must raise and leave the file EMPTY. + + Retrying into the file handle would append to the prefix the failed + dump already wrote, and a half-written pickle only fails at LOAD + time -- long after the run that produced it could have been + repeated. + """ + + def _always_fails(obj, *args, **kwargs): + """Never picklable.""" + del obj, args, kwargs + raise TypeError("cannot pickle '_abc._abc_data' object") + + monkeypatch.setattr(utils.pkl, "dumps", _always_fails) + path = tmp_path / "artifact.pkl" + with pytest.raises(TypeError) as excinfo: + with open(path, "wb") as f: + utils.pkl_dump_with_retry({"learned": [1, 2, 3]}, f) + assert "_abc_data" in str(excinfo.value), \ + "a genuinely unpicklable object must still report why" + assert path.stat().st_size == 0, "a failed dump left a partial file"