diff --git a/plexe/checkpointing.py b/plexe/checkpointing.py index 1b0dcb7f..161ad7cb 100644 --- a/plexe/checkpointing.py +++ b/plexe/checkpointing.py @@ -14,6 +14,7 @@ import base64 import json import logging +import tempfile from datetime import datetime, timezone from pathlib import Path @@ -93,6 +94,7 @@ def save_checkpoint( """ logger.info(f"Saving checkpoint for phase: {phase_name}") + temporary_path = None try: # Extract user feedback from context scratch space (if provided) user_feedback = context.scratch.get("_user_feedback") if hasattr(context, "scratch") else None @@ -115,8 +117,14 @@ def save_checkpoint( checkpoint_dir.mkdir(parents=True, exist_ok=True) checkpoint_path = checkpoint_dir / f"{phase_name}.json" - with open(checkpoint_path, "w") as f: + # Publish only a complete checkpoint. A failed write must not truncate + # the previous phase state, which may be the only resumable copy. + with tempfile.NamedTemporaryFile( + mode="w", encoding="utf-8", dir=checkpoint_dir, prefix=f".{phase_name}.", suffix=".tmp", delete=False + ) as f: + temporary_path = Path(f.name) json.dump(checkpoint_data, f, indent=2, default=str) + temporary_path.replace(checkpoint_path) logger.info(f"✓ Checkpoint saved locally: {checkpoint_path}") if user_feedback: @@ -126,6 +134,12 @@ def save_checkpoint( except Exception as e: logger.error(f"Failed to save checkpoint locally: {e}", exc_info=True) return None + finally: + if temporary_path is not None: + try: + temporary_path.unlink(missing_ok=True) + except OSError: + logger.warning("Could not remove temporary checkpoint %s", temporary_path, exc_info=True) def load_checkpoint( diff --git a/pyproject.toml b/pyproject.toml index 0b3a01bb..a893d90d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "plexe" -version = "1.4.4" +version = "1.4.5" description = "An agentic framework for building ML models from natural language" authors = [ "Marcello De Bernardi ", diff --git a/tests/unit/test_checkpoint_atomicity.py b/tests/unit/test_checkpoint_atomicity.py new file mode 100644 index 00000000..b26667ba --- /dev/null +++ b/tests/unit/test_checkpoint_atomicity.py @@ -0,0 +1,73 @@ +"""Checkpoint publication must preserve the last resumable state on failure.""" + +import json +from pathlib import Path + +import pytest + +from plexe import checkpointing +from plexe.models import BuildContext + + +@pytest.fixture +def context(tmp_path): + """Build real workflow state without running a model or a provider.""" + return BuildContext("user", "experiment", "data.parquet", tmp_path, "Predict a value") + + +@pytest.mark.parametrize("existing", [False, True]) +@pytest.mark.parametrize("failure", [OSError("disk full"), ValueError("serialization failed"), KeyboardInterrupt()]) +def test_failed_save_preserves_previous_checkpoint(tmp_path, context, monkeypatch, existing, failure): + """A partial write must never replace a valid checkpoint or appear as one.""" + checkpoint_dir = tmp_path / "checkpoints" + checkpoint_dir.mkdir() + destination = checkpoint_dir / "prepare_data.json" + previous = b'{"phase": "prepare_data", "context": {"intent": "previous"}}' + if existing: + destination.write_bytes(previous) + + def failing_dump(data, stream, **kwargs): + stream.write('{"context": ') + raise failure + + monkeypatch.setattr(checkpointing.json, "dump", failing_dump) + if isinstance(failure, KeyboardInterrupt): + with pytest.raises(KeyboardInterrupt): + checkpointing.save_checkpoint("experiment", "prepare_data", context, tmp_path) + else: + assert checkpointing.save_checkpoint("experiment", "prepare_data", context, tmp_path) is None + + assert list(checkpoint_dir.iterdir()) == ([destination] if existing else []) + if existing: + assert destination.read_bytes() == previous + assert checkpointing.load_checkpoint("prepare_data", tmp_path)["context"]["intent"] == "previous" + + +def test_successful_save_replaces_previous_checkpoint(tmp_path, context): + """The public path and persisted feedback remain compatible across saves.""" + first = checkpointing.save_checkpoint("experiment", "prepare_data", context, tmp_path) + assert first == tmp_path / "checkpoints" / "prepare_data.json" + context.intent = "Updated intent" + context.scratch["_user_feedback"] = {"message": "保留 • feedback"} + second = checkpointing.save_checkpoint("experiment", "prepare_data", context, tmp_path) + assert second == first + result = checkpointing.load_checkpoint("prepare_data", tmp_path) + assert result["context"]["intent"] == "Updated intent" + assert result["user_feedback"] == context.scratch["_user_feedback"] + assert json.loads(second.read_text())["version"] == "plexe-v1" + assert list(second.parent.iterdir()) == [second] + + +def test_failed_replace_preserves_previous_checkpoint(tmp_path, context, monkeypatch): + """An OS replacement error must clean the temporary file and preserve state.""" + destination = checkpointing.save_checkpoint("experiment", "prepare_data", context, tmp_path) + previous = destination.read_bytes() + + def fail_replace(self, target): + raise PermissionError("destination temporarily unavailable") + + monkeypatch.setattr(Path, "replace", fail_replace) + context.intent = "Not published" + assert checkpointing.save_checkpoint("experiment", "prepare_data", context, tmp_path) is None + assert destination.read_bytes() == previous + assert list(destination.parent.iterdir()) == [destination]