diff --git a/predicators/approaches/gnn_approach.py b/predicators/approaches/gnn_approach.py index 75b4ef29a..00fffac8c 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 b852b7d00..24e28fd0f 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 f27c36b77..362279d1a 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 a456cc93b..c7a3d71ef 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"