Make long-run Modal retries preserve outputs (#1018)

* Make long-run Modal retries preserve outputs

* Add long-run preemption changelog

Author: MaxGhenis

changelog.d/1018.fixed.md

@@ -0,0 +1 @@
+Preserve resumable long-run Modal projection artifacts across preemption retries.

modal_app/long_term_projection.py

@@ -251,6 +251,45 @@ def _output_files(output_dir: Path) -> list[str]:
     )
 
 
+def _has_resume_artifacts(output_dir: Path) -> bool:
+    """Return whether deleting output_dir would lose resumable long-run work."""
+    if not output_dir.exists():
+        return False
+    progress_patterns = (
+        "*.h5",
+        "*.h5.metadata.json",
+        "calibration_manifest.json",
+        "long_run_production_manifest.json",
+        ".parallel_tmp/*/*.h5",
+        ".parallel_tmp/*/*.h5.metadata.json",
+        ".parallel_tmp/*/calibration_manifest.json",
+        ".parallel_logs/*.log",
+    )
+    return any(
+        next(output_dir.glob(pattern), None) is not None
+        for pattern in progress_patterns
+    )
+
+
+def _prepare_output_dir(
+    output_dir: Path,
+    *,
+    clear_output: bool,
+) -> None:
+    if clear_output and output_dir.exists():
+        if _has_resume_artifacts(output_dir):
+            print(
+                "clear_output requested, but existing long-run artifacts were "
+                f"found in {output_dir}; preserving them for resumable execution. "
+                "Use a separate explicit Modal volume removal command for an "
+                "intentional destructive restart.",
+                flush=True,
+            )
+        else:
+            shutil.rmtree(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+
 def _commit_output_volume(*, suppress_errors: bool = False) -> None:
     try:
         output_volume.commit()
@@ -286,7 +325,7 @@ def build_long_term_projection(
     upload_to_hf_staging: bool = False,
     allow_validation_failures: bool = False,
     keep_temp: bool = False,
-    clear_output: bool = True,
+    clear_output: bool = False,
     support_augmentation_profile: str = "",
     support_augmentation_target_year: int | None = None,
     support_augmentation_align_to_run_year: bool = False,
@@ -306,9 +345,10 @@ def build_long_term_projection(
 
     run_id = sanitize_run_id(run_id)
     output_dir = _OUTPUT_MOUNT / run_id
-    if clear_output and output_dir.exists():
-        shutil.rmtree(output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
+    _prepare_output_dir(
+        output_dir,
+        clear_output=clear_output,
+    )
 
     command = _build_command(
         years=years,
@@ -384,7 +424,7 @@ def main(
     upload_to_hf_staging: bool = False,
     allow_validation_failures: bool = False,
     keep_temp: bool = False,
-    clear_output: bool = True,
+    clear_output: bool = False,
     support_augmentation_profile: str = "",
     support_augmentation_target_year: int | None = None,
     support_augmentation_align_to_run_year: bool = False,

tests/unit/test_long_term_modal_projection.py

@@ -176,6 +176,7 @@ def fake_spawn(**kwargs):
     assert captured_kwargs["run_id"] == "crfb-sentinel"
     assert captured_kwargs["source_sha"] == "abc123"
     assert captured_kwargs["upload_to_hf_staging"] is False
+    assert captured_kwargs["clear_output"] is False
 
 
 def test_main_refuses_dirty_source_by_default(monkeypatch):
@@ -245,3 +246,61 @@ def test_remote_result_reports_hf_staging_prefix(monkeypatch, tmp_path):
     assert result["run_id"] == "run-123"
     assert result["hf_staging_prefix"] == "staging/abc123-run-123/long_term"
     assert Path(result["output_dir"]) == tmp_path / "run-123"
+
+
+def test_remote_preserves_resume_artifacts_when_clear_output_requested(
+    monkeypatch,
+    tmp_path,
+):
+    long_term = _load_long_term_projection_module(monkeypatch)
+    monkeypatch.setattr(long_term, "_OUTPUT_MOUNT", tmp_path)
+
+    output_dir = tmp_path / "run-123"
+    resume_dir = output_dir / ".parallel_tmp" / "2026"
+    resume_dir.mkdir(parents=True)
+    completed_h5 = resume_dir / "2026.h5"
+    completed_h5.write_text("already-complete", encoding="utf-8")
+    (resume_dir / "2026.h5.metadata.json").write_text("{}", encoding="utf-8")
+    (resume_dir / "calibration_manifest.json").write_text("{}", encoding="utf-8")
+
+    def fake_stream_command(command, env):
+        assert completed_h5.exists()
+
+    monkeypatch.setattr(long_term, "_stream_command", fake_stream_command)
+
+    result = long_term.build_long_term_projection(
+        years="2026-2100",
+        run_id="Run 123",
+        source_sha="abc123",
+        clear_output=True,
+    )
+
+    assert completed_h5.read_text(encoding="utf-8") == "already-complete"
+    assert ".parallel_tmp/2026/2026.h5" in result["files"]
+
+
+def test_remote_clear_output_can_remove_non_resumable_scratch(
+    monkeypatch,
+    tmp_path,
+):
+    long_term = _load_long_term_projection_module(monkeypatch)
+    monkeypatch.setattr(long_term, "_OUTPUT_MOUNT", tmp_path)
+
+    output_dir = tmp_path / "run-123"
+    output_dir.mkdir(parents=True)
+    scratch_file = output_dir / "scratch.txt"
+    scratch_file.write_text("scratch", encoding="utf-8")
+
+    def fake_stream_command(command, env):
+        assert not scratch_file.exists()
+
+    monkeypatch.setattr(long_term, "_stream_command", fake_stream_command)
+
+    result = long_term.build_long_term_projection(
+        years="2026-2100",
+        run_id="Run 123",
+        source_sha="abc123",
+        clear_output=True,
+    )
+
+    assert "scratch.txt" not in result["files"]