Tolerate post-H5 Modal volume reload conflicts (#988)

* Tolerate post-H5 volume reload conflicts

* Add changelog for H5 reload fix

* Stage local H5s with candidate scope

Author: MaxGhenis

changelog.d/988.fixed.md

@@ -0,0 +1 @@
+Tolerated Modal open-file conflicts when reloading the pipeline volume after local H5 child builds complete.

modal_app/local_area.py

@@ -44,7 +44,10 @@
 )
 from policyengine_us_data.pipeline_metadata import pipeline_node  # noqa: E402
 from policyengine_us_data.pipeline_schema import PipelineNode  # noqa: E402
-from policyengine_us_data.utils.run_context import resolve_run_id  # noqa: E402
+from policyengine_us_data.utils.run_context import (  # noqa: E402
+    resolve_candidate_version,
+    resolve_run_id,
+)
 
 app = modal.App(
     os.environ.get("US_DATA_LOCAL_AREA_APP_NAME") or "policyengine-us-data-local-area"
@@ -361,6 +364,17 @@ def get_version() -> str:
     return pyproject["project"]["version"]
 
 
+def get_staging_candidate_version(
+    fallback_version: str,
+    explicit_candidate_version: str = "",
+) -> str:
+    """Resolve the HF staging candidate scope for local-area artifacts."""
+    return (
+        resolve_candidate_version(explicit_candidate_version, env=os.environ)
+        or fallback_version
+    )
+
+
 @pipeline_node(
     PipelineNode(
         id="build_publishing_input_bundle",
@@ -901,7 +915,11 @@ def validate_staging(branch: str, run_id: str, version: str = "") -> Dict:
     )
 )
 def upload_to_staging(
-    branch: str, version: str, manifest: Dict, run_id: str = ""
+    branch: str,
+    version: str,
+    manifest: Dict,
+    run_id: str = "",
+    candidate_version: str = "",
 ) -> str:
     """
     Upload files to HuggingFace staging only.
@@ -912,6 +930,10 @@ def upload_to_staging(
     setup_repo(branch)
 
     manifest_json = json.dumps(manifest)
+    staging_candidate_version = get_staging_candidate_version(
+        version,
+        explicit_candidate_version=candidate_version,
+    )
 
     result = subprocess.run(
         _python_cmd(
@@ -924,6 +946,7 @@ def upload_to_staging(
 
 manifest = json.loads('''{manifest_json}''')
 version = "{version}"
+staging_candidate_version = "{staging_candidate_version}"
 run_id = "{run_id}"
 staging_dir = Path("{VOLUME_MOUNT}")
 run_dir = staging_dir / run_id
@@ -947,10 +970,14 @@ def upload_to_staging(
 
 # Upload to HuggingFace staging/
 print(f"Uploading {{len(files_with_paths)}} files to HuggingFace staging/...")
-hf_count = upload_to_staging_hf(files_with_paths, version, run_id=run_id)
+hf_count = upload_to_staging_hf(
+    files_with_paths,
+    candidate_version=staging_candidate_version,
+    run_id=run_id,
+)
 print(f"Uploaded {{hf_count}} files to HuggingFace staging/")
 
-print(f"Staged version {{version}} for promotion")
+print(f"Staged candidate {{staging_candidate_version}} for promotion")
 """,
         ),
         text=True,
@@ -961,7 +988,8 @@ def upload_to_staging(
         raise RuntimeError(f"Upload failed: {result.stderr}")
 
     return (
-        f"Staged version {version} with {len(manifest['files'])} files. "
+        f"Staged candidate {staging_candidate_version} with "
+        f"{len(manifest['files'])} files. "
         f"Run promote workflow to publish to HuggingFace production and GCS."
     )
 
@@ -1065,6 +1093,7 @@ def coordinate_publish(
     n_clones: int = 430,
     validate: bool = True,
     run_id: str = "",
+    candidate_version: str = "",
     expected_fingerprint: str = "",
     work_items_override: List[Dict] | None = None,
 ) -> Dict:
@@ -1073,6 +1102,10 @@ def coordinate_publish(
     setup_repo(branch)
 
     version = get_version()
+    staging_candidate_version = get_staging_candidate_version(
+        version,
+        explicit_candidate_version=candidate_version,
+    )
 
     run_id = run_id or resolve_run_id()
     if not run_id:
@@ -1085,6 +1118,7 @@ def coordinate_publish(
     print(f"Run ID: {run_id}")
     print("=" * 60)
     print(f"Publishing version {version} from branch {branch}")
+    print(f"Staging candidate {staging_candidate_version}")
     print(f"Using {num_workers} parallel workers")
 
     staging_dir = Path(VOLUME_MOUNT)
@@ -1320,7 +1354,11 @@ def coordinate_publish(
 
     print("\nStarting upload to staging...")
     result = upload_to_staging.remote(
-        branch=branch, version=version, manifest=manifest, run_id=run_id
+        branch=branch,
+        version=version,
+        manifest=manifest,
+        run_id=run_id,
+        candidate_version=staging_candidate_version,
     )
     print(result)
 
@@ -1382,12 +1420,17 @@ def coordinate_national_publish(
     validate: bool = True,
     run_id: str = "",
     skip_upload: bool = False,
+    candidate_version: str = "",
 ) -> Dict:
     """Build and upload a national US.h5 from national weights."""
     setup_gcp_credentials()
     setup_repo(branch)
 
     version = get_version()
+    staging_candidate_version = get_staging_candidate_version(
+        version,
+        explicit_candidate_version=candidate_version,
+    )
 
     run_id = run_id or resolve_run_id()
     if not run_id:
@@ -1400,6 +1443,7 @@ def coordinate_national_publish(
     print(f"Run ID: {run_id}")
     print("=" * 60)
     print(f"Building national H5 for version {version} from branch {branch}")
+    print(f"Staging candidate {staging_candidate_version}")
 
     staging_dir = Path(VOLUME_MOUNT)
 
@@ -1557,7 +1601,7 @@ def coordinate_national_publish(
 )
 upload_to_staging_hf(
     [("{national_h5}", "national/US.h5")],
-    "{version}",
+    candidate_version="{staging_candidate_version}",
     run_id="{run_id}",
 )
 print("Done")

modal_app/pipeline.py

@@ -140,6 +140,23 @@ def _python_cmd(*args: str) -> list[str]:
     return [sys.executable, *args]
 
 
+def _try_reload_pipeline_volume_after_h5_builds(vol) -> bool:
+    """Reload the pipeline volume unless Modal still sees child-open artifacts."""
+
+    try:
+        vol.reload()
+        return True
+    except RuntimeError as exc:
+        message = str(exc)
+        if "open files preventing the operation" not in message:
+            raise
+        print(
+            "WARNING: Skipping pipeline volume reload after H5 builds because "
+            f"Modal still reports an open artifact file: {message}"
+        )
+        return False
+
+
 def _calibration_package_parameters(
     *,
     workers: int,
@@ -1611,6 +1628,7 @@ def run_pipeline(
                 n_clones=n_clones,
                 validate=True,
                 run_id=run_id,
+                candidate_version=candidate_version,
                 expected_fingerprint=(
                     meta.regional_fingerprint or meta.fingerprint or ""
                 ),
@@ -1636,6 +1654,7 @@ def run_pipeline(
                     n_clones=n_clones,
                     validate=True,
                     run_id=run_id,
+                    candidate_version=candidate_version,
                 )
                 print(
                     f"    → coordinate_national_publish fc: {national_h5_handle.object_id}"
@@ -1672,7 +1691,7 @@ def run_pipeline(
                 )
                 print(f"  National H5: {national_msg}")
 
-            pipeline_volume.reload()
+            _try_reload_pipeline_volume_after_h5_builds(pipeline_volume)
             staging_volume.reload()
 
             if isinstance(regional_h5_result, dict) and regional_h5_result.get(

tests/support/modal_local_area.py

@@ -103,6 +103,24 @@ def decorator(func):
 
         fake_pipeline_metadata.pipeline_node = _fake_pipeline_node
         fake_pipeline_schema.PipelineNode = _FakePipelineNode
+
+        def _fake_resolve_candidate_version(
+            explicit="",
+            *,
+            env=None,
+            **kwargs,
+        ):
+            env = env or {}
+            return (
+                explicit
+                or env.get("US_DATA_CANDIDATE_SCOPE", "")
+                or env.get("US_DATA_CANDIDATE_VERSION", "")
+                or env.get("CANDIDATE_SCOPE", "")
+                or env.get("CANDIDATE_VERSION", "")
+                or env.get("US_DATA_PACKAGE_VERSION", "")
+            )
+
+        fake_run_context.resolve_candidate_version = _fake_resolve_candidate_version
         fake_run_context.resolve_run_id = lambda explicit="", **kwargs: explicit
         fake_partitioning.partition_weighted_work_items = lambda *args, **kwargs: []
 

tests/unit/test_modal_local_area.py

@@ -127,6 +127,41 @@ def fake_run(cmd, **kwargs):
     assert f'version = "{run_id}"' not in script
 
 
+def test_staging_candidate_version_prefers_pipeline_candidate(monkeypatch):
+    local_area = load_local_area_module()
+
+    monkeypatch.setenv("US_DATA_CANDIDATE_VERSION", "1.115.2-patch")
+
+    assert local_area.get_staging_candidate_version("1.115.2") == "1.115.2-patch"
+
+
+def test_upload_to_staging_uses_candidate_scope(monkeypatch):
+    local_area = load_local_area_module()
+    captured = {}
+
+    monkeypatch.setattr(local_area, "setup_repo", lambda branch: None)
+
+    def fake_run(cmd, **kwargs):
+        captured["cmd"] = cmd
+        return SimpleNamespace(returncode=0, stderr="")
+
+    monkeypatch.setattr(local_area.subprocess, "run", fake_run)
+
+    result = local_area.upload_to_staging(
+        branch="main",
+        version="1.115.2",
+        manifest={"files": {"states/NC.h5": {"sha256": "abc"}}},
+        run_id="usdata-gha123-a1",
+        candidate_version="1.115.2-patch",
+    )
+
+    script = captured["cmd"][-1]
+    assert 'version = "1.115.2"' in script
+    assert 'staging_candidate_version = "1.115.2-patch"' in script
+    assert "candidate_version=staging_candidate_version" in script
+    assert result.startswith("Staged candidate 1.115.2-patch")
+
+
 def test_build_publishing_input_bundle_preserves_traceability_inputs():
     local_area = load_local_area_module(stub_policyengine=False)
 

tests/unit/test_pipeline.py

@@ -16,6 +16,7 @@
     _new_run_metadata,
     _pipeline_error_summary,
     _run_required_promotion_subprocess,
+    _try_reload_pipeline_volume_after_h5_builds,
 )
 from modal_app.step_manifests.state import RunMetadata  # noqa: E402
 from modal_app.step_manifests.store import (  # noqa: E402
@@ -71,6 +72,29 @@ def test_national_fit_lambda_matches_national_preset():
     assert NATIONAL_FIT_LAMBDA_L0 == pytest.approx(1e-4)
 
 
+def test_try_reload_pipeline_volume_after_h5_builds_tolerates_modal_open_file():
+    class VolumeWithOpenFileConflict:
+        def reload(self):
+            raise RuntimeError(
+                "there are open files preventing the operation: "
+                "path artifacts/run/policy_data.db is open"
+            )
+
+    assert (
+        _try_reload_pipeline_volume_after_h5_builds(VolumeWithOpenFileConflict())
+        is False
+    )
+
+
+def test_try_reload_pipeline_volume_after_h5_builds_reraises_other_errors():
+    class BrokenVolume:
+        def reload(self):
+            raise RuntimeError("volume service unavailable")
+
+    with pytest.raises(RuntimeError, match="volume service unavailable"):
+        _try_reload_pipeline_volume_after_h5_builds(BrokenVolume())
+
+
 def test_pipeline_error_summary_uses_traceback_ref_when_available():
     ref = ArtifactReference(
         path="runs/run-1/errors/error.json",

tests/unit/test_pipeline_source_contracts.py

@@ -109,6 +109,31 @@ def test_run_pipeline_refreshes_diagnostics_even_when_h5_outputs_reused() -> Non
     assert "Upload validation diagnostics even when H5 outputs are reused." in source
 
 
+def test_run_pipeline_tolerates_post_h5_pipeline_volume_open_files() -> None:
+    source_text = PIPELINE_SOURCE.read_text()
+    tree = ast.parse(source_text)
+    run_pipeline = _function_def(tree, "run_pipeline")
+    reload_helper = _function_def(tree, "_try_reload_pipeline_volume_after_h5_builds")
+    run_pipeline_source = ast.get_source_segment(source_text, run_pipeline)
+    helper_source = ast.get_source_segment(source_text, reload_helper)
+
+    assert "_try_reload_pipeline_volume_after_h5_builds(pipeline_volume)" in (
+        run_pipeline_source
+    )
+    assert "pipeline_volume.reload()" not in run_pipeline_source
+    assert "open files preventing the operation" in helper_source
+    assert "return False" in helper_source
+
+
+def test_run_pipeline_passes_candidate_version_to_h5_publishers() -> None:
+    source_text = PIPELINE_SOURCE.read_text()
+    tree = ast.parse(source_text)
+    run_pipeline = _function_def(tree, "run_pipeline")
+    source = ast.get_source_segment(source_text, run_pipeline)
+
+    assert source.count("candidate_version=candidate_version") >= 2
+
+
 def test_full_release_path_combines_base_regional_and_national_outputs():
     tree = ast.parse(PIPELINE_SOURCE.read_text())
     helper = _function_def(tree, "_full_release_staging_rel_paths")