Fix empty checkpoints for custom JAX mappings - #1199
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #1199 +/- ##
==========================================
+ Coverage 94.24% 94.34% +0.09%
==========================================
Files 21 21
Lines 1686 1714 +28
==========================================
+ Hits 1589 1617 +28
Misses 97 97 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
|
@MilesCranmerBot it feels like we are loading up this PySRRegressor with too much stuff In particular there are a bunch of new functions here which are not really "core" operations to the regressor. I believe there is a nice way to refactor this. Please do so. |
|
Refactored in 2e59bfd: moved the checkpoint/pickle/export-column helper logic into Checks run:
|
Custom `extra_sympy_mappings` or `extra_jax_mappings` can hold objects that pickle cannot store, such as `sympy.Function` classes created at runtime. `_checkpoint` truncated `checkpoint.pkl` before pickling, so a failed pickle left an empty file which `from_file` could not read. - write checkpoints through a temporary file and `os.replace`, leaving any existing checkpoint intact when the pickle fails - clear `extra_jax_mappings` alongside the sympy and torch mappings - omit `equations_` from the checkpoint when it cannot be pickled; it is recreated by `refresh()` once the mappings are passed again - fall back to the CSV backups when a checkpoint is empty or corrupt Fixes astroautomata#1198 Co-Authored-By: Miles Cranmer <miles.cranmer@gmail.com>
aea7cde to
8f63afa
Compare
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8f63afa3f9
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| pysr_logger.info(f"Attempting to load model from {pkl_filename}...") | ||
| try: | ||
| with open(pkl_filename, "rb") as f: | ||
| return cast("PySRRegressor", pkl.load(f)) |
There was a problem hiding this comment.
When a checkpoint references a module or class that is no longer importable—for example after an environment change or because corruption alters a GLOBAL opcode—pickle.load raises ModuleNotFoundError, ImportError, or AttributeError, none of which are caught here. Consequently from_file propagates the exception instead of using the CSV recovery path introduced by this change, even when valid hall-of-fame backups and reconstruction arguments are available. Handle these deserialization failures as unusable checkpoints as well.
Useful? React with 👍 / 👎.
| model = _load_checkpoint(pkl_filename) if pkl_filename.exists() else None | ||
| if model is not None: |
There was a problem hiding this comment.
When an empty or corrupt checkpoint belongs to a multi-output run, returning None here routes loading into a CSV preflight that checks only hall_of_fame.csv and .bak. Multi-output searches instead write hall_of_fame_output1.csv, hall_of_fame_output2.csv, and so on, as get_equation_file(i) confirms, so recovery raises FileNotFoundError despite all required backups being present. The fallback should validate the output-specific files when nout > 1.
Useful? React with 👍 / 👎.
| except Exception as e: | ||
| pysr_logger.debug(f"Error checkpointing model: {e}") | ||
| if tmp_filename is not None: | ||
| tmp_filename.unlink(missing_ok=True) |
There was a problem hiding this comment.
If a later fit successfully updates the hall-of-fame CSV but serialization then fails—for example because a custom plugin or user-attached attribute is unpicklable—this handler deliberately preserves the previous valid checkpoint. A subsequent from_file prefers that old checkpoint and, because it already contains non-None equations, never refreshes from the newer CSV, silently returning results from the earlier run. Record that the write failed or make loading detect that the CSV is newer so preserving the file does not make stale experimental results appear current.
Useful? React with 👍 / 👎.
| with tempfile.NamedTemporaryFile( | ||
| mode="wb", | ||
| dir=pkl_filename.parent, | ||
| prefix="checkpoint.", | ||
| suffix=".pkl.tmp", | ||
| delete=False, |
There was a problem hiding this comment.
On POSIX systems, NamedTemporaryFile creates the temporary checkpoint with mode 0600, independent of the normal umask, and os.replace carries that mode to the destination. Previously, creating checkpoint.pkl with open(..., "wb") normally produced a group/world-readable file such as 0644, so this change prevents collaborators or service accounts from loading checkpoints in shared output directories. Apply the intended destination mode to the temporary file, or preserve the existing checkpoint's mode, before replacing it.
Useful? React with 👍 / 👎.
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
- treat missing modules or classes as an unusable checkpoint - look for the output-specific hall of fame files when `nout > 1` - keep the checkpoint's usual permissions across the atomic replace - warn when a preserved checkpoint may be older than the CSV files Co-Authored-By: Miles Cranmer <miles.cranmer@gmail.com>
2 participants
Summary
Fixes #1198
Tests
uv run --python /usr/bin/python3.12 --with pytest --with 'jax[cpu]>=0.4,<0.6' --with 'juliacall>=0.9.28,<0.9.29' --with 'sympy>=1,<2' --with 'pandas>=0.21,<4' --with 'numpy>=1.13,<3' --with 'scikit-learn>=1,<2' --with 'click>=7,<9' --with 'typing-extensions>=4,<5' python -m pytest pysr/test/test_jax.py -k "checkpoint or empty"\n-python3 -m compileall -q pysr/sr.py pysr/test/test_jax.py