Retry a learned artifact's pickle once, after a collection

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()

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)
 

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]:

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"