Skip to content

Assessment: migrate orchestration layer to Apache Airflow #95

Description

@tabedzki

🤖 AI text below 🤖

Summary

Feasibility assessment for migrating the U19-pipeline orchestration layer (the cron-driven automation) to Apache Airflow, across both U19-pipeline-python and U19-pipeline-matlab. This covers what maps cleanly, the real challenges, a recommended target architecture, and a phased path. No code changes — this is a planning/discussion issue. Scope is the scheduling/automation layer, not the DataJoint schema definitions or the make()/.populate() logic itself (those stay as-is, in whichever language they live).

1. What runs today (the thing being migrated)

The orchestration is a hand-rolled state machine driven by a 30-minute cron job, not a workflow engine.

Cron entry points (three independent subsystems):

  • Automatic job (*/30 * * * *, flock): call_cronjob_automatic_job.sh (conda + git pull --autostash) → cronjob_automatic_job.pyRecordingHandler.pipeline_handler_main()RecProcessHandler.pipeline_handler_main().
  • Alert system (3 AM daily): alert_system/cronjob_alert.pymain_alert_system.py (live-session monitor, schedule check, water-weigh, tech alert, locked-tables, log cleanup, rig maintenance) + a Streamlit water-weigh GUI.
  • MATLAB nightly (call_u19_night_cronjob.sh): pulls ViRMEn + U19-pipeline-matlab, then matlab -batch "startup_virtual_machine.m; populate_tables.m". See the dedicated MATLAB section below — this is an ordered ~20-step batch, not one opaque job.
  • Plus DB backup, pupillometry, and missing-syncbehavior crons.

The migrate-to-airflow branch exists but has no DAGs yet — greenfield.

MATLAB workflow scope (in scope, confirmed by exploring U19-pipeline-matlab)

This migration covers both the Python and MATLAB orchestration — but as scheduling, not a logic rewrite. The MATLAB .populate()/make() logic stays in MATLAB; Airflow just replaces its cron.

  • populate_tables.m is an ordered batch with real internal dependencies, not a single task. Sequence: reset_reweight_subjectspopulate_schedule_for_tomorrow / populate_technician_schedule → protocol-level/training-profile updates → behavior populates (Session, SessionBlock, TowersSession, TowersBlock, SpatialTimeBlobs) → optogenetics ingest + OptogeneticSession → subtasks → pupillometry (PupillometrySession, …SyncBehavior, model ingest) → posture-tracking → imaging_pipeline.SyncImagingBehavior → psychometrics (TowersSessionPsych, TowersSubjectCumulativePsych, …). Psychometrics/sync depend on the behavior tables populating first → worth modeling as a small DAG with explicit task ordering, not a BashOperator blob.
  • Three-repo + lib dependency: startup_virtual_machine.m requires ViRMEn, U19-pipeline-matlab, and datajoint_matlab_libs (mym, datajoint-matlab, GHToolbox) on the MATLAB path. The Airflow MATLAB host (likely SSH'd to) must provision all of these. Each step also does its own git pull today.
  • MATLAB is a second writer to the shared job-status tables. resubmit_failed_jobs.m writes recording_process.Processing.status_processing_id and recording_process.LogStatus — the same tables the Python state machine owns. So MATLAB is not an independent silo; it shares the orchestration state store. This sharpens the dual-state challenge (Python schemas/docker setup/readme. #2): whatever owns status_processing_id must stay consistent across Python, MATLAB, and the GUI.
  • Each modality has mirrored Python/MATLAB table defs (README), so the schema is shared; only the populate/orchestration entry points differ.

Unifying principle: orchestrate the shared DataJoint state, not the language

Both Python and MATLAB interact with the same MariaDB (11.4) behind the same DataJoint interface — the schema and the recording_process.Processing / status tables are the real shared substrate. The goal is one DAG set, manageable from a single Airflow instance, whose task dependencies are expressed against that shared DataJoint state regardless of whether a task runs Python or matlab -batch. Language is an implementation detail of each task; DataJoint tables are the dependency graph.

Python ↔ MATLAB overlap audit (requested)

Result: mostly disjoint by domain, with one real shared seam at imaging ingest — and the two languages are already interleaved within a single workflow, not merely parallel.

Domain / tables MATLAB nightly Python real-time Verdict
imaging_pipeline.AcquiredTiff populate() in populate_tables.m populate() in modality_preingestion and invokes MATLAB populate_Imaging_AcquiredTiff() via ingest_scaninfo_shell.sh (matlab -r) Direct overlap + synchronous cross-language call
imaging_pipeline.SyncImagingBehavior populate() in nightly referenced in real-time imaging flow Overlap
behavior (TowersSession/TowersBlock/SpatialTimeBlobs/psychometrics), optogenetics, pupillometry, posture-tracking, scheduling, reweight, protocol/training-profile updates MATLAB only MATLAB-owned
ephys element (EphysRecording, BehaviorSync, PreCluster, LFP, Clustering, Curation, CuratedClustering) Python only Python-owned
imaging element (Preprocess, Processing, MotionCorrection, Segmentation, Fluorescence, Activity) Python only Python-owned

Key takeaways for the DAG:

  1. AcquiredTiff is populated from both sides — the unified DAG must pick a single owner for that node (or make it idempotent + guarded) so the nightly MATLAB run and the real-time Python flow don't race.
  2. There is already a cross-language dependency edge: Python's imaging modality_preingestion calls a MATLAB script synchronously (ingest_scaninfo_shell.shpopulate_Imaging_AcquiredTiff). In Airflow this becomes an explicit edge: a MATLAB task (TiffSplit/AcquiredTiff ingest) that the downstream Python ScanInfo/element tasks depend on. This is the concrete proof that a unified DAG is the right model, not two siloed ones.
  3. Everything else is cleanly owned by one language, so the unified DAG is mostly "two domain subgraphs sharing a DataJoint state store + one shared imaging seam."

Subgraphs (traced edge-by-edge from DataJoint foreign keys)

Edges below are real FK declarations (-> table in each table's DataJoint header), so they are the populate ordering. [M] = MATLAB-owned task, [P] = Python-owned, [shared] = the cross-language node.

Roots (shared, ingested upstream): subject.Subject, task.Taskacquisition.SessionStartedacquisition.Session (also -> task.TaskLevelParameterSet).

A. Behavior / psychometrics subgraph [M] (nightly, all rooted on acquisition.Session):

acquisition.Session
 ├─ acquisition.SessionBlock ──┐
 ├─ behavior.TowersSession ────┼─→ behavior.TowersBlock  (-> SessionBlock, TowersSession; make() also WRITES TowersBlockTrial)
 │        │                    └─→ behavior.SpatialTimeBlobs   (-> Session, TowersSession)
 │        ├─→ behavior.TowersSessionPsych                  (reads TowersSession only)
 │        ├─→ behavior.TowersSubjectCumulativePsych        (cumulative: reads ALL prior sessions → batch barrier)
 │        └─→ behavior.TowersSubjectCumulativePsychLevel   (cumulative, per level → batch barrier)
 └─ behavior.TowersBlock ──→ behavior.TowersSessionPsychTask  (IMPLICIT: reads TowersBlock + TowersBlockTrial, not just TowersSession)

Correction after reading the make() bodies (FK headers were misleading here — see "Implicit dependencies" below): the four Towers…Psych* tables do not all share one prerequisite.

  • TowersSessionPsych and TowersSubjectCumulativePsych truly read only TowersSession/Session → can run once TowersSession is done.
  • TowersSessionPsychTask actually reads TowersBlock and TowersBlockTrial inside makeTuples → must run after TowersBlock, despite its FK saying only -> TowersSession.
  • TowersSubjectCumulativePsych / …Level are subject-cumulative: makeTuples queries all prior sessions (session_start_time <= … over Session & TowersSession). They are not row-local — they must run after the whole nightly TowersSession batch, not be mapped per-session in isolation, or cumulative curves will be computed on partial history.

B. Optogenetics subgraph [M]:

acquisition.Session
 └─ acquisition.SessionManipulation (-> Session, lab.ManipulationType; Manual, ingested via ingest_previous_optogenetic_sessions)
     └─ optogenetics.OptogeneticSession (-> Session, SessionManipulation, OptogeneticProtocol, OptogeneticSoftwareParameter)

C. Pupillometry & posture subgraphs [M] (independent, parallel):

acquisition.Session ─→ pupillometry.PupillometrySession ─→ pupillometry.PupillometrySyncBehavior
acquisition.Session ─→ posture_tracking.PostureTrackingSession ─→ posture_tracking.PostureTrackingSyncBehavior

D. Scheduling / housekeeping [M] (no FK into the above; run independently each night):
reset_reweight_subjects, populate_schedule_for_tomorrow, populate_technician_schedule, updateProtocolLevelTable, updateTrainingProfileProtocol, ingest_subtasks (behavior_subtask).

E. The imaging seam (cross-language) [shared] — the only place the two graphs touch:

u19_recording.recording  [P, real-time state machine produces this]
 └─ imaging_pipeline.ImagingPipelineSession        [P populate]
     └─ imaging_pipeline.AcquiredTiff   [shared]  ← Python AcquiredTiff.make() shells out to MATLAB
        │                                            populate_Imaging_AcquiredTiff() via ingest_scaninfo_shell.sh;
        │                                            also populated directly by the MATLAB nightly
        └─ imaging_pipeline.TiffSplit   [P writes it inside AcquiredTiff.make, allow_direct_insert]
            ├─→ imaging_pipeline.scan_element.ScanInfo → imaging element chain  [P]
            │      (Preprocess → Processing → MotionCorrection → Segmentation → Fluorescence → Activity)
            └─→ imaging_pipeline.SyncImagingBehavior   [M nightly; FK -> TiffSplit, but make ALSO needs
                                                        recording.RecordingBehaviorSession + behavior file on disk]

So AcquiredTiff is the single contended node (and both the Python and MATLAB AcquiredTiff.make write TiffSplit/TiffSplitFile — pick one owner per key). Downstream forks into a Python branch (element processing) and a MATLAB branch (SyncImagingBehavior) — but SyncImagingBehavior is not a clean TiffSplit-only consumer; it also depends on the Python recording chain and the raw behavior file (see Implicit deps #4).

Implicit dependencies found by reading the make() bodies (not in the FK headers)

The declared FKs govern most ordering, but the make bodies contain reads/writes that the FK graph does not show. These are the edges most likely to cause a unified DAG to mis-fire:

  1. TowersSessionPsychTaskTowersBlock + TowersBlockTrial (implicit). FK says -> TowersSession only; the make reads block/trial tables. Treat TowersBlock as a hard upstream.
  2. TowersBlock.make writes TowersBlockTrial as a side-effect (no separate populate). So TowersBlockTrial rows appear when TowersBlock runs — downstream consumers (add first two schemas (lightsheet and microscope) #1) depend on TowersBlock, not a separate task.
  3. TowersSubjectCumulativePsych / …Level are subject-cumulative — read all prior sessions for the subject. Must run after the full TowersSession batch (a "barrier" task), not per-session.
  4. SyncImagingBehavior is cross-modality + cross-language + reads disk. FK says -> TiffSplit only; the make/keySource also require recording.RecordingBehaviorSession, recording.Recording, acquisition.SessionStarted (behavior file path), TiffSplitFile, and it reads the raw behavior file from the mounted FS. Real upstreams: the Python recording chain and TiffSplit and the behavior file on disk.
  5. PupillometrySyncBehavior / PostureTrackingSyncBehavior read acquisition.SessionVideo + behavior file + the video file on disk (FK says -> …Session only). Implicit upstream: SessionVideo ingested and the video present on the mount.
  6. MATLAB AcquiredTiff.make also writes TiffSplit/TiffSplitFile (lines ~621–704) — identical product to Python's AcquiredTiff.make. Confirms the contended node: whichever owner runs, it must produce TiffSplit the same way. Don't let both run for the same key.
  7. Imported behavior makes read files off disk and self-mutate quality flags. acquisition.Session, behavior.TowersSession, behavior.TowersBlock each fetch1(SessionStarted,'new_remote_path_behavior_file'), read the VirMEn behavior file, and on failure update(... 'is_bad_session'/'invalid_session', 1). → These tasks have an external-data precondition (file must exist on the mounted FS) that no FK captures — model as an Airflow sensor/precondition, and note they write back upstream quality flags (so a "populate" task is not side-effect-free).

Net: the FK graph is correct as a floor on ordering, but the unified DAG needs ~4 extra edges (#1, #4, #5) and two structural notes (#3 barrier, #7 external-file sensors). None of these break the unified design — they make it more accurate.

F. Ephys subgraph [P] (real-time, no MATLAB edges):

u19_recording.recording
 └─ EphysPipelineSession → ephys_element.EphysRecording → BehaviorSync
     └─ PreCluster → (LFP if 1.0 probe) → Clustering → Curation → CuratedClustering

(Heavy step = Clustering/Kilosort on SLURM; this is the chain that flows through the integer state machine + transfers.)

The state machine (recording_handler.py, recording_process_handler.py, params_config.py):

  • Each job carries an integer status_*_id. Every cron tick advances each active job by exactly one status step.
  • The recording-process chain (params_config.py lines 86–186) is a linear DAG: RAW_FILE_TRANSFER_REQUEST → transfer-check → JOB_QUEUE (sbatch) → JOB_FINISHED (poll sacct) → PROC_FILE_TRANSFER_REQUEST → transfer-check → populate-element → delete-cluster-data.
  • State/history/errors live in recording_process.Processing + recording_process.LogStatus (slurm_id, task_copy_id_*, error message/exception). Notifications via Slack webhooks (lab.SlackWebhooks).

External execution:

  • SLURM on spock (spockmk2-loginvm.pni.princeton.edu) / della (della-gpu.princeton.edu): slurm_creator.py generates a .slurm, SSHes via paramiko (or runs locally if is_this_spock()), sbatch, polls sacct.
  • Globus + SCP transfers PNI↔cluster; auth via fair-research-login.
  • Heavy jobs: Kilosort2 (GPU, hours), Suite2p, CatGT. Results ingested into element-array-ephys / element-calcium-imaging via *_element_populate.py.

Config/deploy: DataJoint config (dj_local_conf.json / ~/.datajoint_config.json, custom block) for DB creds + paths; Slack webhooks in lab.SlackWebhooks. No .env secrets framework, no CI. Python 3.12+, uv-managed. Prod host u19prod@pu.win.princeton.edu; deploy = git pull on cron tick.

2. What maps cleanly to Airflow (easy wins)

  • The integer status chain → an explicit Airflow DAG. The DAG is currently encoded as an integer; each ProcessFunction becomes a task. Biggest structural win.
  • Polling steps → Sensors / deferrable operators. transfer_check (Globus) and slurm_job_check (sacct) are textbook deferrable-sensor cases — no more "check once per 30 min."
  • Retries / backoff / error capture → built in (replacing the exception_handler + manual LogStatus writes; on_failure_callback for Slack).
  • SLURM submission → SSHOperator / SLURM providers (or keep paramiko inside a PythonOperator initially).
  • Idempotency already exists — DataJoint .populate() is safe to re-run.
  • Nightly .populate() + alert crons → scheduled DAGs (even just BashOperator initially) — the lowest-risk first target.

3. The real challenges

  1. Dynamic, data-driven job set. The unit of work is "every active recording_process row," discovered at runtime. Needs dynamic task mapping or a controller DAG + TriggerDagRunOperator. Central design decision.
  2. Dual state / source of truth — and MATLAB is a second writer. DB status columns (recording_process.Processing.status_processing_id, LogStatus) are read by the RecordingProcessJobGUI + Slack and written by both the Python state machine and MATLAB (resubmit_failed_jobs.m). Airflow wants to own state → must dual-write the DB status initially, and the MATLAB resubmit path must be reconciled with Airflow (e.g. become an Airflow "clear/retry" trigger) so two systems don't fight over the same rows.
  3. Long-running SLURM jobs (Kilosort hours-long). Must use deferrable operators to avoid starving worker slots; define ownership of sbatch lifecycle on failure/cancel.
  4. Cross-cluster auth & networking. SSH keys to spock/della, Globus refresh tokens, PNI mounts; is_this_spock() special-casing must move into Connections/Hooks. PNI infra constraints (where a daemon may run, firewall, GPU queue policy) are a gating factor.
  5. MATLAB runtime + 3-repo provisioning. The nightly populate needs a licensed MATLAB plus ViRMEn, U19-pipeline-matlab, and datajoint_matlab_libs on the path (startup_virtual_machine.m). Decide whether MATLAB stays on a dedicated host invoked via SSH/matlab -batch, and where the per-run git pull of all three repos happens (in the DAG vs. baked into the host).
  6. Alert subsystem is a second migration (water-weigh, schedule, locked-tables, rig maintenance + Streamlit GUI) — easy but doubles surface area; scope explicitly.
  7. Secrets/config migration from DJ config + lab.SlackWebhooks → Airflow Connections/Variables (or keep reading DJ config inside tasks).
  8. Operational lift at PNI — Airflow scheduler/webserver/metadata-DB/executor is a real service to run, monitor, back up. And no CI exists to lean on during cutover.
  9. Concurrency/idempotency. Today flock serializes everything; Airflow runs concurrently → guard against double-submitting SLURM for the same job_id (pools, max_active_tis_per_dag, DB guards).

4. Recommended target architecture

One Airflow instance orchestrating a unified, DataJoint-state-centric DAG set for both Python and MATLAB. Keep the SSH + Globus execution model; replace cron + the status state-machine with Airflow running as a standalone daemon on the existing prod host (or a dedicated PNI VM).

  • Single management plane: all orchestration — Python real-time state machine, MATLAB nightly, and alerts — lives as DAGs in one Airflow instance/UI. The project is managed from that one place, which is the stated goal.
  • DataJoint state is the dependency graph. Tasks express dependencies via shared DataJoint tables (recording_process.Processing, modality element tables), not via language. A task is "run this populate / this make / this transfer," implemented in whichever language owns that node. Python tasks call DJ-Python; MATLAB tasks SSH to a MATLAB host and run matlab -batch.
  • Model the cross-language imaging seam explicitly: the MATLAB AcquiredTiff/TiffSplit ingest is an upstream task; the Python ScanInfo/imaging-element tasks depend on it (this edge exists today as a synchronous shell-out — Airflow makes it a first-class dependency). Pick a single owner for AcquiredTiff to avoid the nightly-vs-realtime race.
  • Executor: LocalExecutor (→ CeleryExecutor later) — not Kubernetes. PNI has no obvious K8s; KubernetesPodOperator is the biggest lift for the least near-term payoff. Revisit containers after the logic is proven.
  • DAG shape: lightweight controller DAG (replaces the 30-min cron) that queries active jobs and fans out via dynamic task mapping (or TriggerDagRunOperator) one task-group per active recording_process: request-transfer → (sensor) → submit-slurm → (deferrable sensor) → transfer-back → (sensor) → populate-element.
  • Long jobs: deferrable SLURM sensors.
  • State: dual-write DB status_* + LogStatus so the GUI/Slack/MATLAB consumers keep working; DB status stays the cross-system contract.
  • Reuse, don't rewrite: wrap existing slurm_creator, clusters_paths_and_transfers, *_element_populate in PythonOperators first; refactor into Hooks later.
  • Secrets: SSH/Globus → Airflow Connections; keep DJ config for the DB.
  • First target: the nightly .populate() + alert DAGs (lowest risk).

5. Suggested phased path

  1. Stand up Airflow on the prod host (LocalExecutor + own metadata DB), no DAGs.
  2. Phase 1 — nightly populates + alerts (covers MATLAB): port the Python nightly, the alert system, and the MATLAB populate_tables.m to scheduled DAGs. The MATLAB one becomes a small DAG that SSHes to a MATLAB host and runs the populate steps in dependency order (behavior → sync/psychometrics), rather than one BashOperator. Validates infra/logging/Slack. No state-machine risk.
  3. Phase 2 — ephys state machine: model the electrophysiology recording_process chain with dynamic mapping + deferrable SLURM sensor, dual-writing DB status. Run in shadow alongside the existing cron for one modality.
  4. Phase 3 — imaging + recording handler; retire the corresponding cron paths.
  5. Phase 4 — cleanup: decide whether Airflow becomes canonical state (requires GUI changes) or DB status stays the contract.

6. Open questions / decisions needed

  • Where can a long-lived Airflow daemon run at PNI, and what are the SSH/Globus/GPU-queue policy constraints? (gates executor choice)
  • Should the DB status_* remain the permanent cross-system contract, or do we eventually port the RecordingProcessJobGUI to read Airflow state?
  • Is the alert subsystem in scope for this migration or tracked separately?
  • Is there a MATLAB host Airflow can SSH to with all three repos + DataJoint libs provisioned, or does that host need to be created?
  • Should resubmit_failed_jobs.m be reimplemented as an Airflow clear/retry trigger so MATLAB stops writing status_processing_id directly?
  • Who owns imaging_pipeline.AcquiredTiff in the unified DAG (it's populated by both the MATLAB nightly and the Python real-time imaging flow)? Pick one owner or make it idempotent + guarded.

Decisions already made

  • MariaDB is 11.4 (confirmed; the earlier 10.2 in Infrastructure.md is stale).
  • Scope = both repos, one unified DataJoint-state-centric DAG set, single Airflow instance — the project is to be managed from one place.

Key files

Python (U19-pipeline-python): u19_pipeline/automatic_job/: recording_handler.py, recording_process_handler.py, params_config.py (status DAG, lines 86–186), cronjob_automatic_job.py, call_cronjob_automatic_job.sh, crontab_example, slurm_creator.py, clusters_paths_and_transfers.py, ephys_element_populate.py, imaging_element_populate.py · u19_pipeline/recording.py, recording_process.py · u19_pipeline/alert_system/ · scripts/conf_file_finding.py

MATLAB (U19-pipeline-matlab): scripts/call_u19_night_cronjob.sh, scripts/startup_virtual_machine.m, scripts/populate_tables.m, scripts/populate_schedule_for_tomorrow.m, scripts/populate_technician_schedule.m, scripts/resubmit_failed_jobs.m (shared-state writer) · Infrastructure.md (DB/host details)

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions