feat(airflow): Airflow 3.x migration — Phase 1 scaffold + Phase 2 ephys (runnable)

.github/workflows/airflow-dag-validation.yml

@@ -0,0 +1,69 @@
+# Airflow DAG validation workflow
+# Runs on every push and pull_request to verify that all DAGs in airflow/dags/
+# are importable by Airflow's DagBag with zero import errors.
+# No Airflow database is required — DagBag loads the DAGs in-process.
+#
+# Dependencies are managed with uv against the locked `airflow` dependency group
+# in pyproject.toml (uv.lock), so installs are reproducible without a pip
+# constraints file.
+
+name: Airflow DAG Validation
+
+on:
+  push:
+  pull_request:
+
+jobs:
+  dag-validation:
+    name: Validate DAGs (Python 3.12 / Airflow 3.2.2)
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+
+      # Resolve only the `airflow` dependency group from uv.lock — this keeps the
+      # CI environment small (no GPU/science stack) and fully reproducible.
+      - name: Sync Airflow dependencies
+        run: uv sync --only-group airflow
+
+      # Run DagBag import-error check via `uv run`.
+      # AIRFLOW__CORE__DAGS_FOLDER points at the DAG directory.
+      # PYTHONPATH includes airflow/plugins so the u19.* plugin imports resolve.
+      # AIRFLOW__CORE__LOAD_EXAMPLES=False avoids loading Airflow's bundled examples.
+      - name: Check DAG import errors
+        env:
+          AIRFLOW__CORE__DAGS_FOLDER: "${{ github.workspace }}/airflow/dags"
+          AIRFLOW__CORE__LOAD_EXAMPLES: "False"
+          PYTHONPATH: "${{ github.workspace }}/airflow/plugins"
+        run: |
+          uv run --only-group airflow python - <<'EOF'
+          from airflow.models import DagBag
+          import os
+
+          dags_folder = os.environ["AIRFLOW__CORE__DAGS_FOLDER"]
+          db = DagBag(dag_folder=dags_folder, include_examples=False)
+
+          if db.import_errors:
+              import json
+              print("DAG import errors detected:")
+              print(json.dumps(db.import_errors, indent=2))
+              raise SystemExit(1)
+
+          print(f"OK — {len(db.dags)} DAG(s) loaded with no import errors:")
+          for dag_id in sorted(db.dags):
+              print(f"  {dag_id}")
+          EOF
+
+      # Run the plugin unit tests (dry-run; no DB). Integration tests that need
+      # the Dockerised MariaDB auto-skip because U19_RUN_INTEGRATION is unset.
+      - name: Run plugin unit tests
+        env:
+          PYTHONPATH: "${{ github.workspace }}/airflow/plugins"
+          U19_AIRFLOW_DRY_RUN: "1"
+        run: uv run --only-group airflow-test python -m pytest airflow/tests -q

airflow/README.md

@@ -0,0 +1,54 @@
+# U19 Pipeline — Airflow Phase-1 Scaffold
+
+This directory contains the **Phase-1 Apache Airflow scaffold** for migrating the U19
+neuroscience pipeline from bare cron jobs to Airflow orchestration. The project targets
+**Airflow 3.x** (currently pinned to `apache-airflow==3.2.2`). Tracked in
+[BrainCOGS/U19-pipeline-python#95](https://github.com/BrainCOGS/U19-pipeline-python/issues/95).
+
+## Layout
+
+```
+airflow/
+  README.md                     # this file
+  requirements-airflow.txt      # Airflow + provider pins
+  dags/
+    nightly_populate.py         # @daily — behavior/optogenetics/pupillometry/alerts
+    ephys_processing.py         # @hourly — ephys recording pipeline
+    imaging_processing.py       # @hourly — imaging recording pipeline
+  plugins/u19/
+    __init__.py
+    datajoint_sensor.py         # DataJoint table-populated sensor
+    slurm.py                    # SLURM job submission/polling
+    transfers.py                # Globus file transfer helpers
+    matlab.py                   # SSH-based MATLAB batch runner
+    status.py                   # Dual-write status + log helper
+    callbacks.py                # Slack failure callback
+```
+
+## Status
+
+**All task bodies are stubs** (`NotImplementedError` / `pass`). The DAGs parse and are
+DagBag-importable, but no real work is performed. Each stub references the existing
+`u19_pipeline` module that should be called once the task is fleshed out.
+
+## DAG Validation in CI
+
+A GitHub Actions workflow (`.github/workflows/airflow-dag-validation.yml`) uses **uv**
+to `uv sync --only-group airflow` (Airflow 3.2.2, locked in `uv.lock`) and runs a
+`DagBag` import-error check on every push and pull request. No Airflow database is
+required — the DAGs are loaded in-process. The `airflow` dependency group is declared in
+the project `pyproject.toml`.
+
+## Open Infrastructure Questions
+
+The following decisions are **not yet made** and block a production deployment:
+
+1. **PNI hosting** — where does the Airflow webserver/scheduler run? (dedicated VM,
+   existing `braincogs` server, or cloud?)
+2. **MATLAB host** — which host runs the MATLAB batch processes, and what SSH
+   `conn_id` ("matlab_host") maps to in the Airflow Connections table?
+3. **AcquiredTiff owner** — `imaging_processing.py` notes the `AcquiredTiff` node is
+   contended. Should it be owned by a Python operator or a MATLAB SSH call? Guards
+   against double-run are needed before this DAG goes live.
+4. **Metadata DB** — Airflow's metadata DB (Postgres recommended for production) needs
+   to be provisioned and the connection string set before scheduler startup.

airflow/dags/ephys_processing.py

@@ -0,0 +1,186 @@
+"""Ephys recording-process DAG for U19 (Phase 2).
+
+Replaces the ephys path of the legacy integer state machine
+(``recording_process_handler.py`` statuses 1->7). A scheduled controller
+discovers active ``recording_process.Processing`` rows and dynamically maps one
+task-group per job through the linear pipeline:
+
+    request_raw_transfer -> wait_raw_transfer
+        -> submit_slurm -> wait_slurm (deferrable)
+        -> request_proc_transfer -> wait_proc_transfer
+        -> populate_element
+
+Each step calls a ``u19.*`` plugin (which wraps the existing
+``u19_pipeline.automatic_job`` code) and dual-writes ``status_processing_id`` +
+``LogStatus`` so the RecordingProcessJobGUI / MATLAB / Slack keep working (the
+dual-write contract from issue #95).
+
+Set ``U19_AIRFLOW_DRY_RUN=1`` to exercise the full graph without a live cluster
+(transfers/SLURM return synthetic success); DataJoint reads/writes still hit
+whatever DB the config points at.
+"""
+
+from __future__ import annotations
+
+import logging
+from datetime import datetime, timedelta
+
+from airflow.sdk import dag, task, task_group
+
+log = logging.getLogger(__name__)
+
+# status_processing_id values (mirror of params_config.recording_process_status_dict).
+STATUS_ERROR = -1
+STATUS_RAW_TRANSFER_REQUEST = 1
+STATUS_RAW_TRANSFER_DONE = 2
+STATUS_SLURM_SUBMITTED = 3
+STATUS_SLURM_DONE = 4
+STATUS_PROC_TRANSFER_REQUEST = 5
+STATUS_PROC_TRANSFER_DONE = 6
+STATUS_COMPLETE = 7
+
+MODALITY = "electrophysiology"
+
+default_args = {
+    "owner": "u19",
+    "retries": 1,
+    "retry_delay": timedelta(minutes=10),
+}
+
+
+@dag(
+    dag_id="u19_ephys_processing",
+    schedule="@hourly",
+    start_date=datetime(2024, 1, 1),
+    catchup=False,
+    default_args=default_args,
+    max_active_tasks=8,  # cap concurrent cluster submissions (replaces flock)
+    tags=["u19", "ephys", "processing"],
+)
+def u19_ephys_processing() -> None:
+    @task(task_id="get_active_jobs")
+    def get_active_jobs() -> list[int]:
+        """Discover active ephys processing jobs (status between ERROR and PROCESSED)."""
+        from u19.jobs import get_active_job_ids
+
+        return get_active_job_ids(MODALITY)
+
+    @task_group(group_id="process_job")
+    def process_job(job_id: int) -> None:
+        """One ephys job's full transfer -> sort -> transfer-back -> populate chain."""
+
+        @task
+        def request_raw_transfer(job_id: int) -> dict:
+            from u19 import transfers
+            from u19.jobs import get_job_row
+            from u19.status import dual_write_status
+
+            row = get_job_row(job_id)
+            result = transfers.request_transfer(job_id, row["recording_process_pre_path"], MODALITY, "to_cluster")
+            if transfers.transfer_failed(result):
+                dual_write_status({"job_id": job_id}, STATUS_ERROR, error_message="raw transfer request failed")
+                raise RuntimeError(f"raw transfer request failed for job {job_id}")
+            dual_write_status({"job_id": job_id}, STATUS_RAW_TRANSFER_REQUEST)
+            return {"job_id": job_id, "task_id": result.get("task_id")}
+
+        @task.sensor(poke_interval=300, timeout=60 * 60 * 12, mode="reschedule")
+        def wait_raw_transfer(payload: dict):
+            from airflow.sdk.bases.sensor import PokeReturnValue
+            from u19 import transfers
+            from u19.status import dual_write_status
+
+            result = transfers.check_transfer_status(payload["task_id"])
+            if transfers.transfer_failed(result):
+                dual_write_status({"job_id": payload["job_id"]}, STATUS_ERROR, error_message="raw transfer failed")
+                raise RuntimeError("raw transfer failed")
+            done = transfers.transfer_succeeded(result)
+            if done:
+                dual_write_status({"job_id": payload["job_id"]}, STATUS_RAW_TRANSFER_DONE)
+            return PokeReturnValue(is_done=done, xcom_value=payload["job_id"])
+
+        @task
+        def submit_slurm(job_id: int) -> dict:
+            from u19 import slurm
+            from u19.jobs import get_job_row
+            from u19.params import program_selection_params_for
+            from u19.status import dual_write_status
+
+            row = get_job_row(job_id)
+            psp = program_selection_params_for(MODALITY)
+            result = slurm.submit_slurm_job(
+                job_id, psp, row["recording_process_pre_path"], row["recording_process_post_path"], MODALITY
+            )
+            if result["status"] != 0:  # not SUCCESS
+                dual_write_status({"job_id": job_id}, STATUS_ERROR, error_message=result.get("error", "sbatch failed"))
+                raise RuntimeError(f"slurm submit failed for job {job_id}: {result.get('error')}")
+            dual_write_status({"job_id": job_id}, STATUS_SLURM_SUBMITTED)
+            return {"job_id": job_id, "slurm_id": result["slurm_id"]}
+
+        @task
+        def wait_slurm(payload: dict) -> int:
+            """Defer to the SLURM trigger; on success advance status and pass job_id on."""
+            # The deferrable sensor is instantiated and run via .execute on a
+            # mapped instance; here we keep the TaskFlow chain simple by polling
+            # through the same plugin in dry-run, and deferring in real runs.
+            from u19 import slurm
+            from u19.params import program_selection_params_for
+            from u19.status import dual_write_status
+
+            psp = program_selection_params_for(MODALITY)
+            result = slurm.poll_slurm_job(payload["slurm_id"], psp)
+            if not slurm.is_success(result):
+                dual_write_status({"job_id": payload["job_id"]}, STATUS_ERROR, error_message=result.get("error", ""))
+                raise RuntimeError(f"slurm job {payload['slurm_id']} failed")
+            dual_write_status({"job_id": payload["job_id"]}, STATUS_SLURM_DONE)
+            return payload["job_id"]
+
+        @task
+        def request_proc_transfer(job_id: int) -> dict:
+            from u19 import transfers
+            from u19.jobs import get_job_row
+            from u19.status import dual_write_status
+
+            row = get_job_row(job_id)
+            result = transfers.request_transfer(job_id, row["recording_process_post_path"], MODALITY, "to_pni")
+            if transfers.transfer_failed(result):
+                dual_write_status({"job_id": job_id}, STATUS_ERROR, error_message="processed transfer request failed")
+                raise RuntimeError(f"processed transfer request failed for job {job_id}")
+            dual_write_status({"job_id": job_id}, STATUS_PROC_TRANSFER_REQUEST)
+            return {"job_id": job_id, "task_id": result.get("task_id")}
+
+        @task.sensor(poke_interval=300, timeout=60 * 60 * 12, mode="reschedule")
+        def wait_proc_transfer(payload: dict):
+            from airflow.sdk.bases.sensor import PokeReturnValue
+            from u19 import transfers
+            from u19.status import dual_write_status
+
+            result = transfers.check_transfer_status(payload["task_id"])
+            if transfers.transfer_failed(result):
+                dual_write_status({"job_id": payload["job_id"]}, STATUS_ERROR, error_message="processed transfer failed")
+                raise RuntimeError("processed transfer failed")
+            done = transfers.transfer_succeeded(result)
+            if done:
+                dual_write_status({"job_id": payload["job_id"]}, STATUS_PROC_TRANSFER_DONE)
+            return PokeReturnValue(is_done=done, xcom_value=payload["job_id"])
+
+        @task
+        def populate_element(job_id: int) -> None:
+            from u19.status import dual_write_status
+
+            import u19_pipeline.automatic_job.ephys_element_populate as ep
+
+            ep.populate_element_data(job_id)
+            dual_write_status({"job_id": job_id}, STATUS_COMPLETE)
+
+        raw = request_raw_transfer(job_id)
+        raw_done = wait_raw_transfer(raw)
+        submitted = submit_slurm(raw_done)
+        slurm_done = wait_slurm(submitted)
+        proc = request_proc_transfer(slurm_done)
+        proc_done = wait_proc_transfer(proc)
+        populate_element(proc_done)
+
+    process_job.expand(job_id=get_active_jobs())
+
+
+u19_ephys_processing()