Accept run_id in upload_to_staging_hf, unblock build-and-test (#794)

* Accept run_id in upload_to_staging_hf, scope staging prefix by run

Main's build-and-test has been failing on every real-content push with
`TypeError: upload_to_staging_hf() got an unexpected keyword argument
'run_id'`. The four callers (`storage/upload_completed_datasets.py`,
`calibration/promote_local_h5s.py`, `modal_app/local_area.py`,
`modal_app/pipeline.py`) all pass `run_id=`, but the function's
signature was never updated. Existing tests mocked the function with
`**kwargs`, which hid the divergence.

Add `run_id: str = ""` to the signature and, when set, stage files under
`staging/{run_id}/{rel_path}` — same convention the surrounding helpers
(`_download_staged_dataset_artifacts`, `_promote_staging_to_hf`) already
use. Fall back to bare `staging/{rel_path}` when `run_id` is empty so
behavior is unchanged for the no-run_id call path.

New unit tests in `tests/unit/utils/test_data_upload.py` exercise the
real function (not a mock) with and without `run_id` and assert on the
resulting `path_in_repo` values, so the next caller drift will surface
in unit CI rather than on a 3-hour integration run.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Rename changelog fragment to PR number

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: MaxGhenis

changelog.d/794.fixed.md

@@ -0,0 +1,8 @@
+Add missing `run_id` parameter to `upload_to_staging_hf` in
+`policyengine_us_data.utils.data_upload`, matching the kwarg its
+callers have been passing. Files now land under
+`staging/{run_id}/{rel_path}` when `run_id` is set and under
+`staging/{rel_path}` otherwise, mirroring the prefix convention used
+elsewhere in the package. Unblocks the main `build-and-test` workflow,
+which had been failing with `TypeError: upload_to_staging_hf() got an
+unexpected keyword argument 'run_id'` on every real-content push.

policyengine_us_data/utils/data_upload.py

@@ -751,6 +751,7 @@ def upload_to_staging_hf(
     hf_repo_name: str = "policyengine/policyengine-us-data",
     hf_repo_type: str = "model",
     batch_size: int = 50,
+    run_id: str = "",
 ) -> int:
     """
     Upload files to staging/ paths in HuggingFace.
@@ -762,12 +763,16 @@ def upload_to_staging_hf(
         hf_repo_name: HuggingFace repository name
         hf_repo_type: Repository type
         batch_size: Number of files per commit batch
+        run_id: Optional per-run scope. When set, files land under
+            ``staging/{run_id}/{rel_path}`` so concurrent runs do not
+            collide; otherwise they land under ``staging/{rel_path}``.
 
     Returns:
         Number of files uploaded
     """
     token = os.environ.get("HUGGING_FACE_TOKEN")
     api = HfApi()
+    staging_prefix = f"staging/{run_id}" if run_id else "staging"
 
     total_uploaded = 0
     for i in range(0, len(files_with_paths), batch_size):
@@ -780,7 +785,7 @@ def upload_to_staging_hf(
                 continue
             operations.append(
                 CommitOperationAdd(
-                    path_in_repo=f"staging/{rel_path}",
+                    path_in_repo=f"{staging_prefix}/{rel_path}",
                     path_or_fileobj=str(local_path),
                 )
             )

tests/unit/utils/test_data_upload.py

@@ -0,0 +1,67 @@
+from pathlib import Path
+
+from policyengine_us_data.utils import data_upload
+
+
+class _FakeHfApi:
+    def __init__(self):
+        self.commits = []
+
+
+def _install_fake_hf(monkeypatch, tmp_path):
+    fake = _FakeHfApi()
+    monkeypatch.setattr(data_upload, "HfApi", lambda: fake)
+
+    captured_ops = []
+
+    def fake_commit(api, operations, repo_id, repo_type, token, commit_message):
+        captured_ops.extend(operations)
+
+    monkeypatch.setattr(data_upload, "hf_create_commit_with_retry", fake_commit)
+    return captured_ops
+
+
+def _make_files(tmp_path, rel_paths):
+    files = []
+    for rel in rel_paths:
+        local = tmp_path / Path(rel).name
+        local.write_text("stub")
+        files.append((local, rel))
+    return files
+
+
+def test_upload_to_staging_hf_accepts_run_id_kwarg(monkeypatch, tmp_path):
+    captured_ops = _install_fake_hf(monkeypatch, tmp_path)
+    files = _make_files(tmp_path, ["states/AL.h5"])
+
+    n = data_upload.upload_to_staging_hf(
+        files,
+        version="1.73.0",
+        run_id="abc123",
+    )
+
+    assert n == 1
+    assert len(captured_ops) == 1
+
+
+def test_upload_to_staging_hf_run_id_scopes_staging_prefix(monkeypatch, tmp_path):
+    captured_ops = _install_fake_hf(monkeypatch, tmp_path)
+    files = _make_files(tmp_path, ["states/AL.h5", "states/CA.h5"])
+
+    data_upload.upload_to_staging_hf(files, version="1.73.0", run_id="abc123")
+
+    assert [op.path_in_repo for op in captured_ops] == [
+        "staging/abc123/states/AL.h5",
+        "staging/abc123/states/CA.h5",
+    ]
+
+
+def test_upload_to_staging_hf_without_run_id_uses_bare_staging_prefix(
+    monkeypatch, tmp_path
+):
+    captured_ops = _install_fake_hf(monkeypatch, tmp_path)
+    files = _make_files(tmp_path, ["states/AL.h5"])
+
+    data_upload.upload_to_staging_hf(files, version="1.73.0")
+
+    assert [op.path_in_repo for op in captured_ops] == ["staging/states/AL.h5"]