Merge pull request #819 from PolicyEngine/maria/chunk_parallelization

Pipeline environment fixes

Author: anth-volk

.github/workflows/pr.yaml

@@ -86,6 +86,26 @@ jobs:
         env:
           CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
 
+  optimized-integration-tests:
+    runs-on: ubuntu-latest
+    needs: [check-fork, lint]
+    env:
+      MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
+      MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
+      MODAL_ENVIRONMENT: main
+      MODAL_APP_NAME: policyengine-us-data-pipeline
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.14"
+      - name: Install optimized test deps
+        run: pip install modal pytest
+      - name: Deploy Modal pipeline app
+        run: modal deploy --env="${MODAL_ENVIRONMENT}" modal_app/pipeline.py
+      - name: Run optimized integration tests
+        run: python -m pytest tests/optimized/test_modal_pipeline_seams.py -v
+
   smoke-test:
     runs-on: ubuntu-latest
     needs: [check-fork, lint]
@@ -140,6 +160,7 @@ jobs:
     env:
       MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
       MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
+      MODAL_ENVIRONMENT: main
       HUGGING_FACE_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN }}
     steps:
       - uses: actions/checkout@v4
@@ -164,6 +185,6 @@ jobs:
             } >> "$GITHUB_STEP_SUMMARY"
           fi
 
-          modal run modal_app/data_build.py \
+          modal run --env="${MODAL_ENVIRONMENT}" modal_app/data_build.py \
             --branch=${{ github.head_ref || github.ref_name }} \
             ${STAGE_ARGS}

.github/workflows/push.yaml

@@ -12,29 +12,32 @@ jobs:
       - run: pip install ruff>=0.9.0
       - run: ruff format --check .
 
-  # ── Build datasets and run integration tests ────────────────
-  build-and-test:
+  # ── Build and linear integration tests ──────────────────────
+  build-and-linear-integration-tests:
+    name: Build and linear integration tests
     runs-on: ubuntu-latest
     needs: lint
     if: github.event.head_commit.message != 'Update package version'
     env:
       MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
       MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
+      MODAL_ENVIRONMENT: main
       HUGGING_FACE_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN }}
     steps:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
         with:
           python-version: "3.14"
       - run: pip install modal
-      - name: Build datasets and run integration tests on Modal
+      - name: Run linear integration tests on Modal
         run: |
-          modal run modal_app/data_build.py \
+          modal run --env="${MODAL_ENVIRONMENT}" modal_app/data_build.py \
             --upload \
             --branch=${{ github.ref_name }}
 
   # ── Documentation ──────────────────────────────────────────
   docs:
+    name: Documentation
     runs-on: ubuntu-latest
     if: github.event.head_commit.message != 'Update package version'
     permissions:
@@ -62,6 +65,7 @@ jobs:
 
   # ── Versioning (bump + changelog on non-version-bump pushes) ──
   versioning:
+    name: Versioning
     runs-on: ubuntu-latest
     if: github.event.head_commit.message != 'Update package version'
     steps:

Makefile

@@ -41,10 +41,10 @@ changelog:
 	python .github/bump_version.py
 	towncrier build --yes --version $$(python -c "import re; print(re.search(r'version = \"(.+?)\"', open('pyproject.toml').read()).group(1))")
 download:
-	python policyengine_us_data/storage/download_private_prerequisites.py
+	python -m policyengine_us_data.storage.download_private_prerequisites
 
 upload:
-	python policyengine_us_data/storage/upload_completed_datasets.py
+	python -m policyengine_us_data.storage.upload_completed_datasets
 
 docker:
 	docker buildx build --platform linux/amd64 . -t policyengine-us-data:latest
@@ -78,19 +78,19 @@ DATABASE_YEAR ?= 2024
 
 database:
 	rm -f policyengine_us_data/storage/calibration/policy_data.db
-	python policyengine_us_data/db/create_database_tables.py
-	python policyengine_us_data/db/create_initial_strata.py --year $(YEAR)
-	python policyengine_us_data/db/etl_national_targets.py --year $(YEAR)
-	python policyengine_us_data/db/etl_age.py --year $(YEAR)
-	python policyengine_us_data/db/etl_medicaid.py --year $(YEAR)
-	python policyengine_us_data/db/etl_snap.py --year $(YEAR)
-	python policyengine_us_data/db/etl_tanf.py --year $(YEAR)
-	python policyengine_us_data/db/etl_state_income_tax.py --year $(YEAR)
-	python policyengine_us_data/db/etl_irs_soi.py --year $(YEAR)
-	python policyengine_us_data/db/etl_aca_agi_state_targets.py --year $(YEAR)
-	python policyengine_us_data/db/etl_aca_marketplace.py --year $(YEAR)
-	python policyengine_us_data/db/etl_pregnancy.py --year $(YEAR)
-	python policyengine_us_data/db/validate_database.py
+	python -m policyengine_us_data.db.create_database_tables
+	python -m policyengine_us_data.db.create_initial_strata --year $(YEAR)
+	python -m policyengine_us_data.db.etl_national_targets --year $(YEAR)
+	python -m policyengine_us_data.db.etl_age --year $(YEAR)
+	python -m policyengine_us_data.db.etl_medicaid --year $(YEAR)
+	python -m policyengine_us_data.db.etl_snap --year $(YEAR)
+	python -m policyengine_us_data.db.etl_tanf --year $(YEAR)
+	python -m policyengine_us_data.db.etl_state_income_tax --year $(YEAR)
+	python -m policyengine_us_data.db.etl_irs_soi --year $(YEAR)
+	python -m policyengine_us_data.db.etl_aca_agi_state_targets --year $(YEAR)
+	python -m policyengine_us_data.db.etl_aca_marketplace --year $(YEAR)
+	python -m policyengine_us_data.db.etl_pregnancy --year $(YEAR)
+	python -m policyengine_us_data.db.validate_database
 
 database-refresh:
 	rm -f policyengine_us_data/storage/calibration/policy_data.db
@@ -105,18 +105,18 @@ promote-dataset:
 	@echo "Dataset promoted to HF."
 
 data: download database
-	python policyengine_us_data/utils/uprating.py
-	python policyengine_us_data/datasets/acs/acs.py
-	python policyengine_us_data/datasets/cps/cps.py
-	python policyengine_us_data/datasets/puf/irs_puf.py
-	python policyengine_us_data/datasets/puf/puf.py
-	python policyengine_us_data/datasets/cps/extended_cps.py
-	python policyengine_us_data/calibration/create_stratified_cps.py
-	python policyengine_us_data/calibration/create_source_imputed_cps.py
+	python -m policyengine_us_data.utils.uprating
+	python -m policyengine_us_data.datasets.acs.acs
+	python -m policyengine_us_data.datasets.cps.cps
+	python -m policyengine_us_data.datasets.puf.irs_puf
+	python -m policyengine_us_data.datasets.puf.puf
+	python -m policyengine_us_data.datasets.cps.extended_cps
+	python -m policyengine_us_data.calibration.create_stratified_cps
+	python -m policyengine_us_data.calibration.create_source_imputed_cps
 
 data-legacy: data
-	python policyengine_us_data/datasets/cps/enhanced_cps.py
-	python policyengine_us_data/datasets/cps/small_enhanced_cps.py
+	python -m policyengine_us_data.datasets.cps.enhanced_cps
+	python -m policyengine_us_data.datasets.cps.small_enhanced_cps
 
 calibrate: data
 	python -m policyengine_us_data.calibration.unified_calibration \
@@ -151,7 +151,7 @@ validate-data:
 	python -c "from policyengine_us_data.storage.upload_completed_datasets import validate_all_datasets; validate_all_datasets()"
 
 refresh-soi-targets:
-	python policyengine_us_data/storage/calibration_targets/refresh_soi_table_targets.py \
+	python -m policyengine_us_data.storage.calibration_targets.refresh_soi_table_targets \
 		--source-year $(SOI_SOURCE_YEAR) \
 		--target-year $(SOI_TARGET_YEAR)
 

changelog.d/modal-runtime-seams.fixed.md

@@ -0,0 +1 @@
+Refactor Modal runtime setup to rely on `Image.uv_sync()` and the active Python interpreter rather than manual venv wiring, and add an optimized deployed-image seam test lane to the push workflow.

changelog.d/modal-venv-activation.fixed.md

@@ -0,0 +1 @@
+Activate the uv-managed venv inside Modal pipeline containers so module-scope imports from `policyengine_us_data` (notably `pandas` via `geography/__init__.py`) resolve at container boot. `uv sync --frozen` installs dependencies into `/root/policyengine-us-data/.venv/`, but Modal boots the container with the system Python, so `pipeline.py` imports were failing with `ModuleNotFoundError: No module named 'pandas'`. The image now sets `VIRTUAL_ENV`, prepends `.venv/bin` to `PATH`, and adds the venv site-packages to `PYTHONPATH`.

modal_app/README.md

@@ -145,7 +145,7 @@ Produces `source_imputed_stratified_extended_cps_2024.h5` from raw CPS/PUF/ACS/S
 | Method | Command |
 |--------|---------|
 | **Local** | `make data` |
-| **Modal (CI)** | `modal run modal_app/data_build.py --branch=<branch>` |
+| **Modal (CI)** | `modal run --env=main modal_app/data_build.py --branch=<branch>` |
 | **GitHub Actions** | Automatic on merge to `main` via `code_changes.yaml` → `reusable_test.yaml` (with `full_suite: true`). Also triggered by `pr_code_changes.yaml` on PRs. |
 
 Notes:

modal_app/data_build.py

@@ -87,6 +87,11 @@
 ]
 
 
+def _python_cmd(*args: str) -> list[str]:
+    """Build a command that uses the current interpreter."""
+    return [sys.executable, *args]
+
+
 def setup_gcp_credentials():
     """Write GCP credentials JSON to a temp file for google.auth.default()."""
     creds_json = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS_JSON")
@@ -207,7 +212,7 @@ def run_script(
     env: Optional[dict] = None,
     log_file: IO = None,
 ) -> str:
-    """Run a script with uv and return its path for logging.
+    """Run a script with the current interpreter and return its path.
 
     Args:
         script_path: Path to the Python script to run.
@@ -220,7 +225,15 @@ def run_script(
     Raises:
         subprocess.CalledProcessError: If the script fails.
     """
-    cmd = ["uv", "run", "python", "-u", script_path]
+    script = Path(script_path)
+    if (
+        script.suffix == ".py"
+        and script.parts
+        and script.parts[0] in {"policyengine_us_data", "modal_app"}
+    ):
+        cmd = _python_cmd("-u", "-m", ".".join(script.with_suffix("").parts))
+    else:
+        cmd = _python_cmd("-u", script_path)
     if args:
         cmd.extend(args)
     run_env = env or os.environ.copy()
@@ -371,7 +384,7 @@ def run_tests_with_checkpoints(
 
         print(f"Running tests: {module}")
         result = subprocess.run(
-            ["uv", "run", "python", "-u", "-m", "pytest", module, "-v"],
+            _python_cmd("-u", "-m", "pytest", module, "-v"),
             env=env,
         )
 
@@ -467,7 +480,7 @@ def build_datasets(
     )
     # Build policy_data.db from source
     subprocess.run(
-        ["uv", "run", "make", "database"],
+        ["make", "database"],
         check=True,
         cwd="/root/policyengine-us-data",
         env=env,

modal_app/images.py

@@ -6,8 +6,10 @@
 """
 
 import subprocess
-import modal
 from pathlib import Path
+from typing import Callable
+
+import modal
 
 REPO_ROOT = Path(__file__).resolve().parent.parent
 
@@ -30,41 +32,85 @@
 except Exception:
     pass
 
-_IGNORE = [
-    ".git",
-    "__pycache__",
-    "*.egg-info",
-    ".pytest_cache",
-    "*.h5",
-    "*.npy",
-    "*.pkl",
-    "*.db",
-    "node_modules",
-    "venv",
-    ".venv",
-    "docs/_build",
-    "paper",
-    "presentations",
-]
+
+# Extra paths the Modal image must never include, beyond what .gitignore
+# already covers. `.git` holds hundreds of MB of pack data that Modal never
+# reads; `paper` and `presentations` are authoring directories.
+_MODAL_EXTRA_IGNORE = {".git", "paper", "presentations"}
+
+
+def _build_ignore_callable(repo_root: Path) -> Callable[[Path], bool]:
+    """Return an ignore predicate that mirrors .gitignore for Modal.
+
+    Modal's `add_local_dir(ignore=...)` uses dockerignore semantics, where
+    bare patterns like `*.h5` only match root-level files. Our `.gitignore`
+    mixes gitignore semantics (bare names match at any depth). Translating
+    patterns by hand is error-prone and drifts over time. Instead, we ask
+    git directly for the set of locally-ignored paths and exclude those.
+    Untracked-but-not-ignored files still ship so uncommitted edits to
+    Modal code (e.g. `modal_app/images.py` itself) make it into the image.
+    """
+    repo_root = repo_root.resolve()
+    ignored_paths: set[Path] = set()
+    try:
+        out = subprocess.check_output(
+            [
+                "git",
+                "-C",
+                str(repo_root),
+                "ls-files",
+                "--others",
+                "--ignored",
+                "--exclude-standard",
+                "--directory",
+            ],
+            text=True,
+            stderr=subprocess.DEVNULL,
+        )
+        for line in out.splitlines():
+            entry = line.strip().rstrip("/")
+            if entry:
+                ignored_paths.add((repo_root / entry).resolve())
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        pass
+
+    for name in _MODAL_EXTRA_IGNORE:
+        ignored_paths.add((repo_root / name).resolve())
+
+    def should_ignore(path: Path) -> bool:
+        try:
+            resolved = path.resolve()
+        except (OSError, ValueError):
+            return False
+        if resolved in ignored_paths:
+            return True
+        for parent in resolved.parents:
+            if parent in ignored_paths:
+                return True
+            if parent == repo_root:
+                return False
+        return False
+
+    return should_ignore
 
 
 def _base_image(extras: list[str] | None = None):
-    extra_flags = " ".join(f"--extra {e}" for e in (extras or []))
     return (
         modal.Image.debian_slim(python_version="3.14")
         .apt_install("git", "make")
-        .pip_install("uv>=0.8")
+        .uv_sync(
+            uv_project_dir=str(REPO_ROOT),
+            frozen=True,
+            extras=extras,
+        )
         .add_local_dir(
             str(REPO_ROOT),
             remote_path="/root/policyengine-us-data",
             copy=True,
-            ignore=_IGNORE,
+            ignore=_build_ignore_callable(REPO_ROOT),
         )
+        .workdir("/root/policyengine-us-data")
         .env(GIT_ENV)
-        .run_commands(
-            f"cd /root/policyengine-us-data && "
-            f"UV_HTTP_TIMEOUT=300 uv sync --frozen {extra_flags}"
-        )
     )