feat(airflow): Airflow 3.x migration — Phase 1 scaffold + Phase 2 ephys (runnable)#96
Draft
tabedzki wants to merge 3 commits into
Draft
feat(airflow): Airflow 3.x migration — Phase 1 scaffold + Phase 2 ephys (runnable)#96tabedzki wants to merge 3 commits into
tabedzki wants to merge 3 commits into
Conversation
Add a parse-only scaffold for migrating the U19 orchestration layer to Apache Airflow 3.x (assessment + plan in #95). - airflow/dags: nightly_populate (unified Python+MATLAB nightly + alerts), ephys_processing, imaging_processing — TaskFlow DAGs with stubbed task bodies. nightly_populate encodes the FK ordering plus the implicit edges found in the make() audit (TowersSessionPsychTask after TowersBlock; a barrier before the subject-cumulative psych tables). - airflow/plugins/u19: stub hooks/helpers that wrap existing code (slurm_creator, clusters_paths_and_transfers, slack_utils, dual-write of recording_process status, matlab -batch). No logic reimplemented. - CI: .github/workflows/airflow-dag-validation.yml runs a DagBag import-error check via `uv sync --only-group airflow` + `uv run`. - pyproject.toml: add locked `airflow` dependency group (apache-airflow 3.2.2 + ssh/standard providers); uv.lock updated. All three DAGs parse with zero import errors against Airflow 3.2.2. Task bodies are TODO stubs; deployment/infra questions tracked in #95. Assisted-by: ClaudeCode:claude-opus-4.8
The ephys DAG task docstrings referenced STATUS_* names in their dual_write_status TODOs but never defined them. Add them as named constants mirroring recording_process_status_dict in params_config, so each task's status dual-write is explicit and the file is self-coherent. Implementation should import these from params_config rather than redeclare, to stay in sync. Assisted-by: ClaudeCode:claude-opus-4.8
Turn the ephys recording-process scaffold into a working, tested DAG that reuses (does not reimplement) the existing u19_pipeline code. Plugins (thin wrappers over u19_pipeline.automatic_job, imported lazily so DAG parsing stays light): - status.py: dual_write_status — advances recording_process.Processing status + appends LogStatus in one transaction (the issue #95 dual-write contract). - jobs.py: get_active_job_ids / get_job_row for dynamic task mapping. - transfers.py / slurm.py: wrap globus + slurm_creator with a dry-run switch. - slurm_sensor.py: deferrable SlurmJobTrigger + sensor so multi-hour Kilosort jobs don't pin a worker slot. - params.py: per-modality program_selection_params. DAG: ephys_processing.py now discovers active jobs and dynamically maps one task-group per job through transfer -> sbatch -> wait -> transfer-back -> populate_element, dual-writing status after each step. Tests (airflow/tests): unit tests (dry-run, no DB) + an opt-in integration test that exercises dual_write_status against the docker-compose MariaDB. The integration test caught a real bug (DataJoint update1 vs _update). 11 pass. Tooling: add `airflow-test` dep group + pytest config; CI now runs the unit tests after DAG validation. uv.lock updated. Assisted-by: ClaudeCode:claude-opus-4.8
tabedzki
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
🤖 AI text below 🤖
Airflow migration — Phase 1 (scaffold) + Phase 2 (ephys, runnable)
Builds on the assessment and plan in #95. This PR now carries two increments on top of
master, with no behavior change in production (the existing cron pipeline is untouched — Airflow shadows it until each phase is proven).Targets Apache Airflow 3.x (3.2.2, TaskFlow API). Dependencies and CI are driven by
uv syncagainst a lockedairflowdependency group.Phase 1 — orchestration scaffold (parse-only)
A self-contained
airflow/tree encoding the DAG structure from the #95 subgraph +make()-body analysis:airflow/dags/nightly_populate.py— unified Python + MATLAB nightly + alerts DAG. Encodes FK ordering plus the implicit edges found by auditing the MATLABmake()bodies:TowersSessionPsychTaskruns afterTowersBlock(not justTowersSession), and a barrier (towers_session_batch_complete) gates the subject-cumulative psych tables since they read all prior sessions.airflow/dags/imaging_processing.py— theAcquiredTiffseam forking into the Python element chain and the MATLABSyncImagingBehavior(still scaffold; Phase 3)..github/workflows/airflow-dag-validation.yml:DagBagimport-error check viauv sync --only-group airflow+uv run.pyproject.toml— lockedairflowdependency group (apache-airflow 3.2.2 + ssh/standard providers);uv.lockupdated.Phase 2 — ephys processing DAG, real bodies + tests (runnable)
ephys_processing.pyis no longer a stub. It discovers activerecording_processjobs and dynamic-task-maps one task-group per job through the full chain, dual-writing status at every step so the GUI / MATLAB / Slack keep working:Plugins wrap
u19_pipeline— no populate/transfer/SLURM logic is reimplemented. Task functions import the package lazily at runtime, so DAG parsing/CI stays light while execution uses the full real package:status.py—dual_write_status: advancesrecording_process.Processing.status_processing_idand appends aLogStatusaudit row in one DataJoint transaction (the dual-write contract from Assessment: migrate orchestration layer to Apache Airflow #95).jobs.py— active-job discovery + row fetch for the dynamic mapping.transfers.py/slurm.py— thin wrappers overclusters_paths_and_transfers/slurm_creator, with aU19_AIRFLOW_DRY_RUNswitch for CI.slurm_sensor.py— a deferrable trigger + sensor (SlurmJobTrigger/SlurmJobSensor) so multi-hour Kilosort jobs free the worker slot while waiting instead of pinning it.params.py— per-modality program-selection params.Tests (11 passing)
test_plugins_unit.py, 10 tests) — dry-run, no DB, CI-runnable: transfer/SLURM wrappers, deferrable trigger serialize + fire, sensor failure path, params.test_status_integration.py, opt-in viaU19_RUN_INTEGRATION=1) — proves the dual-write round-trip against a Dockerised MariaDB (u19_test_prefix, port 3307 — never productiondatajoint00). This test earned its keep: it caught a real bug (_update→update1) that parsing and mocked unit tests could not.uv run --only-group airflow-test).Testability split (inherent, not a gap): DB-touching tasks (
get_active_jobs,dual_write_status,populate_element) run for real against Docker MariaDB; transfer/SLURM tasks run their real code only on the PNI host with cluster creds, and via the dry-run switch in tests/CI (you can'tsbatchto a container).Not in this PR (tracked in #95)
Phase 3 imaging bodies + the
AcquiredTiffowner decision, Phase 4 cleanup (resubmit_failed_jobs.m→ Airflow retry trigger), and deployment / open infra questions (PNI hosting, MATLAB host, metadata DB, NetID SSO).Refs #95.