From fb9667fcfe9e9651832696f91c2c66fcf525d721 Mon Sep 17 00:00:00 2001 From: Christian Tabedzki <35670232+tabedzki@users.noreply.github.com> Date: Fri, 12 Jun 2026 12:14:26 -0400 Subject: [PATCH] feat: implement NWB export handler conversion, validation, and status flow Completes the cronjob export handler that PR #90 only reconciled to the schema. 1. Recording/scan key derivation (handler + nwb_production_utils): the job record carries only the acquisition.Session key, not recording_id/scan keys, so the old `{k: job[k] for k in recording.Recording.primary_key}` collapsed to {} and matched every recording (vacuous pass). Now resolves recording_id(s) via the new recording_ids_for_session() helper (recording.Recording.BehaviorSession). Applied to both the ephys and imaging branches; imaging fails loud (Scan<-> session linkage not yet confirmed) instead of silently passing. 2. process_nwb_conversion: no longer raises NotImplementedError. Conversion logic is extracted into the shared module u19_pipeline/nwb_export/conversion.py (build_source_data / query_metadata / run_conversion_to_file, ported from run_nwb_export.py), so the CLI and cronjob share one code path. Input paths are resolved via resolve_input_paths() (export_params or NWB_EXPORT_DATA_ROOT) with an actionable error when unresolved. 3. process_validation: removed the broken stub that inserted a validation record then raised NotImplementedError (which duplicate-key-crashed on retry). Now runs a real HDF5 integrity check + optional NWB Inspector, records truthful results, and inserts the row idempotently (replace on retry). 4. Status flow reconciled to NwbExportStatusEnum (QUEUED->DATA_VALIDATION-> PROCESSING->VALIDATION->[UPLOAD]->COMPLETED). The old handler treated status 3 as completion, conflicting with the enum (VALIDATION=3, COMPLETED=6) and run_nwb_export.py. Added process_upload for the optional DANDI stage (no-DANDI jobs go straight to COMPLETED; requested uploads use DandiUploadClient and fail loud if credentials are missing). Co-Authored-By: Claude Fable 5 --- scripts/run_nwb_export.py | 251 +++-------- .../automatic_job/nwb_export_handler.py | 420 ++++++++++++++---- u19_pipeline/nwb_export/conversion.py | 306 +++++++++++++ u19_pipeline/nwb_production_utils.py | 65 ++- 4 files changed, 748 insertions(+), 294 deletions(-) create mode 100644 u19_pipeline/nwb_export/conversion.py diff --git a/scripts/run_nwb_export.py b/scripts/run_nwb_export.py index 86ba4844..423e82c5 100644 --- a/scripts/run_nwb_export.py +++ b/scripts/run_nwb_export.py @@ -34,10 +34,8 @@ import argparse import json import logging -import os import sys import traceback -from datetime import datetime from pathlib import Path logging.basicConfig( @@ -52,6 +50,7 @@ # Helpers # ────────────────────────────────────────────────────────────────────────────── + def _connect_dj() -> None: """Ensure DataJoint is connected, trying conf_file_finding first.""" try: @@ -66,146 +65,23 @@ def _connect_dj() -> None: dj.conn() -def _find_kilosort_output(probe_dir: Path) -> Path | None: - """ - Return the highest-numbered job's Kilosort output directory under probe_dir. - - Expected layout: - /_imec/job_id_/kilosort_output/ - """ - job_dirs = sorted( - [p for p in probe_dir.glob("job_id_*") if p.name.split("_")[-1].isdigit()], - key=lambda p: int(p.name.split("_")[-1]), - ) - if not job_dirs: - return None - kilosort_outputs = list(job_dirs[-1].glob("kilosort*_output")) - return kilosort_outputs[0] if kilosort_outputs else None - - -def _build_source_data( - job: dict, - export_params: dict, - virmen_file: Path | None, - kilosort_dir: Path | None, -) -> dict: - """ - Translate the DataJoint job record + export_params into a - `source_data` dict accepted by TowersNWBConverter. - - Returns - ------- - dict - Ready to pass as ``source_data=`` to TowersNWBConverter. - - Raises - ------ - FileNotFoundError - If a required data file cannot be located. - """ - source_data: dict = {} - - # ── Behavior (always required) ──────────────────────────────────────────── - if virmen_file is None: - raise FileNotFoundError( - "No --virmen-file provided. Cannot locate the behavioral .mat file. " - "Re-run with --virmen-file /path/to/session.mat" - ) - if not Path(virmen_file).exists(): - raise FileNotFoundError(f"Virmen file not found: {virmen_file}") - source_data["VirmenData"] = {"file_path": str(virmen_file)} - - # ── Ephys ───────────────────────────────────────────────────────────────── - if export_params.get("include_ephys") and kilosort_dir is not None: - kilosort_dir = Path(kilosort_dir) - probe_dirs = sorted(kilosort_dir.glob("*_imec*")) - if not probe_dirs: - log.warning("--kilosort-dir given but no *_imec* subdirectories found – skipping ephys.") - else: - for probe_dir in probe_dirs: - probe_idx = "".join(filter(str.isdigit, probe_dir.name.split("imec")[-1])) - interface_name = f"KilosortProbe{probe_idx}" if probe_idx else "Kilosort" - ks_output = _find_kilosort_output(probe_dir) - if ks_output is None: - log.warning(f"No Kilosort output found under {probe_dir} – skipping.") - continue - source_data[interface_name] = {"folder_path": str(ks_output)} - log.info(f" {interface_name}: {ks_output}") - elif export_params.get("include_ephys"): - log.warning( - "include_ephys=True but no --kilosort-dir provided. " - "Re-run with --kilosort-dir to include ephys data." - ) - - return source_data - - -def _query_metadata(session_key: dict) -> dict: - """ - Pull experimenter, subject sex/DoB and sync timestamps from DataJoint. - - Returns a dict with keys: experimenter, subject_sex, subject_dob, - sync_timestamps (may be None if not found). - """ - import datajoint as dj - - result: dict = { - "experimenter": [], - "subject_sex": "U", - "subject_dob": None, - "sync_timestamps": None, - } - - try: - subject = dj.create_virtual_module("subject", dj.config["custom"]["database.prefix"] + "subject") - lab = dj.create_virtual_module("lab", dj.config["custom"]["database.prefix"] + "lab") - - subject_fullname = session_key["subject_fullname"] - sub_info = (subject.Subject() * lab.User() & f"subject_fullname = '{subject_fullname}'").fetch1() - - # Owner name - owner_full = sub_info.get("full_name", sub_info.get("user_id", "")) - if " " in owner_full: - parts = owner_full.rsplit(" ", 1) - result["experimenter"].append(f"{parts[-1]}, {parts[0]}") - else: - result["experimenter"].append(owner_full) - - # Sex - sex_map = {"Male": "M", "Female": "F", "Unknown": "U", "m": "M", "f": "F"} - result["subject_sex"] = sex_map.get(str(sub_info.get("sex", "U")), "U") - - # Date of birth - dob = sub_info.get("dob") - if dob is not None: - result["subject_dob"] = datetime.combine(dob, datetime.min.time()) if hasattr(dob, "year") else dob - - except Exception as exc: - log.warning(f"Could not query all metadata from DB: {exc}") - - # Sync timestamps (optional BehaviorSync table) - try: - nwb_prod = dj.create_virtual_module( - "nwb_production", dj.config["custom"]["database.prefix"] + "nwb_production" - ) - sync_rows = (nwb_prod.BehaviorSync & session_key).fetch("sync_timestamps", as_dict=True) - if sync_rows: - import numpy as np - - result["sync_timestamps"] = np.array(sync_rows[0]["sync_timestamps"]) - except Exception: - pass # BehaviorSync is optional - - return result - +# Shared conversion logic lives in u19_pipeline.nwb_export.conversion so the CLI +# and the cronjob handler share one code path. Re-exported under the historical +# private names to keep this module's internal references working. +from u19_pipeline.nwb_export.conversion import ( # noqa: E402 + run_conversion_to_file as _run_conversion_to_file, +) # ────────────────────────────────────────────────────────────────────────────── # Status helpers # ────────────────────────────────────────────────────────────────────────────── -def _transition(nwb_production, job_key: dict, new_status_id: int, dry_run: bool) -> None: - from u19_pipeline.nwb_production import update_job_status # type: ignore + +def _transition( + nwb_production, job_key: dict, new_status_id: int, dry_run: bool +) -> None: from u19_pipeline.nwb_export_enums import NwbExportStatusEnum # type: ignore + from u19_pipeline.nwb_production import update_job_status # type: ignore new_status = NwbExportStatusEnum(new_status_id) log.info(f" → {new_status.name}") @@ -214,9 +90,9 @@ def _transition(nwb_production, job_key: dict, new_status_id: int, dry_run: bool def _fail(nwb_production, job_key: dict, exc: Exception, dry_run: bool) -> None: - from u19_pipeline.nwb_production import update_job_status # type: ignore - from u19_pipeline.nwb_export_enums import NwbExportStatusEnum # type: ignore from u19_pipeline.nwb_export.error_capture import capture_exception # type: ignore + from u19_pipeline.nwb_export_enums import NwbExportStatusEnum # type: ignore + from u19_pipeline.nwb_production import update_job_status # type: ignore tb = capture_exception(exc) log.error(f"Job FAILED: {tb['error_message']}") @@ -233,15 +109,15 @@ def _fail(nwb_production, job_key: dict, exc: Exception, dry_run: bool) -> None: # Main # ────────────────────────────────────────────────────────────────────────────── + def run( job_id: int, virmen_file: Path | None, kilosort_dir: Path | None, dry_run: bool, ) -> None: - import datajoint as dj - from u19_pipeline.nwb_export_enums import NwbExportStatusEnum # type: ignore from u19_pipeline import nwb_production # type: ignore + from u19_pipeline.nwb_export_enums import NwbExportStatusEnum # type: ignore job_key = {"nwb_job_id": job_id} @@ -275,10 +151,13 @@ def run( except json.JSONDecodeError: # Streamlit stored it as a Python repr string in early versions import ast + try: export_params = ast.literal_eval(raw_params) except Exception: - log.warning(f"Could not parse export_parameters: {raw_params!r}; proceeding with empty params") + log.warning( + f"Could not parse export_parameters: {raw_params!r}; proceeding with empty params" + ) session_key = { "subject_fullname": job["subject_fullname"], @@ -288,15 +167,21 @@ def run( log.info(f" export_params: {export_params}") if dry_run: - log.info("[DRY RUN] Would transition: QUEUED → DATA_VALIDATION → PROCESSING → VALIDATION → COMPLETED") + log.info( + "[DRY RUN] Would transition: QUEUED → DATA_VALIDATION → PROCESSING → VALIDATION → COMPLETED" + ) log.info("[DRY RUN] Exiting without any DB writes or file operations.") return # ── 2. DATA_VALIDATION ──────────────────────────────────────────────────── - _transition(nwb_production, job_key, int(NwbExportStatusEnum.DATA_VALIDATION), dry_run) + _transition( + nwb_production, job_key, int(NwbExportStatusEnum.DATA_VALIDATION), dry_run + ) try: - from u19_pipeline.nwb_production_utils import validate_behavior_data_exists # type: ignore + from u19_pipeline.nwb_production_utils import ( + validate_behavior_data_exists, # type: ignore + ) ok, msg = validate_behavior_data_exists(session_key) if not ok: @@ -310,59 +195,36 @@ def run( _transition(nwb_production, job_key, int(NwbExportStatusEnum.PROCESSING), dry_run) try: - from tank_lab_to_nwb.convert_towers_task.towersnwbconverter import TowersNWBConverter # type: ignore + from tank_lab_to_nwb.convert_towers_task.towersnwbconverter import ( + TowersNWBConverter, # type: ignore + ) except ImportError: log.error( "tank-lab-to-nwb is not importable. Install it with:\n" " pip install -e /path/to/tank-lab-to-nwb-clean\n" "or add it to PYTHONPATH." ) - _fail(nwb_production, job_key, ImportError("tank_lab_to_nwb not installed"), dry_run) + _fail( + nwb_production, + job_key, + ImportError("tank_lab_to_nwb not installed"), + dry_run, + ) sys.exit(1) + output_path = job["output_filepath"] try: - source_data = _build_source_data(job, export_params, virmen_file, kilosort_dir) - log.info(f" source_data keys: {list(source_data.keys())}") - - metadata = _query_metadata(session_key) - - converter = TowersNWBConverter( - source_data=source_data, - sync_timestamps=metadata["sync_timestamps"], - ) - - raw_metadata = converter.get_metadata() - - # Inject DB-sourced metadata - raw_metadata["NWBFile"]["session_description"] = ( - f"U19 pipeline export – {session_key['subject_fullname']} " - f"{session_key['session_date']}" - ) - if metadata["experimenter"]: - raw_metadata["NWBFile"]["experimenter"] = metadata["experimenter"] - - if "Subject" not in raw_metadata: - raw_metadata["Subject"] = {} - if metadata["subject_sex"]: - raw_metadata["Subject"]["sex"] = metadata["subject_sex"] - if metadata["subject_dob"] is not None: - raw_metadata["Subject"]["date_of_birth"] = metadata["subject_dob"] - - output_path = job["output_filepath"] - Path(output_path).parent.mkdir(parents=True, exist_ok=True) - - log.info(f" Writing NWB to: {output_path}") - converter.run_conversion( - nwbfile_path=output_path, - metadata=raw_metadata, - overwrite=True, + size_gb = _run_conversion_to_file( + job=job, + export_params=export_params, + session_key=session_key, + virmen_file=virmen_file, + kilosort_dir=kilosort_dir, + output_path=output_path, ) - log.info(" ✓ Conversion complete") # Record actual file size - size_gb = Path(output_path).stat().st_size / (1024**3) nwb_production.NwbExportJob.update1({**job_key, "actual_file_size_gb": size_gb}) - log.info(f" ✓ File size: {size_gb:.3f} GB") except Exception as exc: log.error(traceback.format_exc()) @@ -392,13 +254,15 @@ def run( # CLI entry point # ────────────────────────────────────────────────────────────────────────────── + def _parse_args() -> argparse.Namespace: p = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter, ) p.add_argument( - "--job-id", "-j", + "--job-id", + "-j", type=int, default=None, help=( @@ -434,9 +298,10 @@ def _parse_args() -> argparse.Namespace: def _get_pending_job_ids() -> list[int]: """Return nwb_job_ids for all jobs not in a terminal state (COMPLETED / FAILED).""" - import datajoint as dj - from u19_pipeline import acquisition # noqa: ensure FK context - from u19_pipeline import nwb_production # type: ignore + + from u19_pipeline import ( + nwb_production, # type: ignore + ) from u19_pipeline.nwb_export_enums import NwbExportStatusEnum terminal_ids = [ @@ -444,9 +309,11 @@ def _get_pending_job_ids() -> list[int]: int(NwbExportStatusEnum.FAILED), ] restriction = " AND ".join(f"status_id != {s}" for s in terminal_ids) - job_ids = (nwb_production.NwbExportJob & restriction).fetch( - "nwb_job_id", order_by="submission_timestamp ASC" - ).tolist() + job_ids = ( + (nwb_production.NwbExportJob & restriction) + .fetch("nwb_job_id", order_by="submission_timestamp ASC") + .tolist() + ) return job_ids @@ -469,7 +336,7 @@ def _get_pending_job_ids() -> list[int]: else: log.info(f"Found {len(job_ids)} pending job(s): {job_ids}") for job_id in job_ids: - log.info(f"\n{'='*60}\nProcessing job #{job_id}\n{'='*60}") + log.info(f"\n{'=' * 60}\nProcessing job #{job_id}\n{'=' * 60}") try: run( job_id=job_id, diff --git a/u19_pipeline/automatic_job/nwb_export_handler.py b/u19_pipeline/automatic_job/nwb_export_handler.py index 5025b0e2..4e0e1dbb 100644 --- a/u19_pipeline/automatic_job/nwb_export_handler.py +++ b/u19_pipeline/automatic_job/nwb_export_handler.py @@ -14,12 +14,12 @@ import u19_pipeline.automatic_job.params_config as config import u19_pipeline.utils.slack_utils as slack_utils -from u19_pipeline import acquisition, nwb_production, recording -from u19_pipeline.imaging_pipeline import imaging_element +from u19_pipeline import acquisition, nwb_production +from u19_pipeline.nwb_export_enums import NwbExportStatusEnum from u19_pipeline.nwb_production_utils import ( + recording_ids_for_session, validate_behavior_data_exists, validate_ephys_data_exists, - validate_imaging_data_exists, ) @@ -43,6 +43,26 @@ def _parse_number_list(raw) -> list: return [] +def _parse_export_params(raw) -> dict: + """ + Parse a job's export_parameters value into a dict. + + Stored as a JSON string (or a Python-repr string in early versions); may also + already be a dict or None. Returns an empty dict on failure. + """ + if raw is None: + return {} + if isinstance(raw, dict): + return raw + try: + return dict(json.loads(raw)) + except (ValueError, TypeError): + try: + return dict(ast.literal_eval(raw)) + except (ValueError, SyntaxError, TypeError): + return {} + + class NwbExportHandler: """Handler for NWB export job processing pipeline.""" @@ -54,8 +74,14 @@ def pipeline_handler_main(): This function is called repeatedly by the cronjob to process all active NWB export jobs through their pipeline stages. """ - # Get active jobs (status < COMPLETED and not FAILED) - active_jobs = (nwb_production.NwbExportJob & "status_id >= 0 AND status_id < 3").fetch(as_dict=True) + # Active jobs = non-terminal: status in [QUEUED .. UPLOADED], i.e. + # status_id >= 0 AND status_id < COMPLETED(6) AND status_id != FAILED(-1). + completed = int(NwbExportStatusEnum.COMPLETED) + failed = int(NwbExportStatusEnum.FAILED) + restriction = ( + f"status_id >= 0 AND status_id < {completed} AND status_id != {failed}" + ) + active_jobs = (nwb_production.NwbExportJob & restriction).fetch(as_dict=True) print(f"Processing {len(active_jobs)} active NWB export jobs...") @@ -63,18 +89,46 @@ def pipeline_handler_main(): current_status = job["status_id"] try: - # Dispatch to appropriate handler based on current status - if current_status == 0: # QUEUED -> DATA_VALIDATION + # Dispatch to the appropriate handler based on current status, + # using NwbExportStatusEnum values consistently. + if current_status == int(NwbExportStatusEnum.QUEUED): success, error_info = NwbExportHandler.process_data_validation(job) - next_status = 1 if success else -1 + next_status = ( + int(NwbExportStatusEnum.DATA_VALIDATION) if success else failed + ) - elif current_status == 1: # DATA_VALIDATION -> PROCESSING + elif current_status == int(NwbExportStatusEnum.DATA_VALIDATION): success, error_info = NwbExportHandler.process_nwb_conversion(job) - next_status = 2 if success else -1 + next_status = ( + int(NwbExportStatusEnum.PROCESSING) if success else failed + ) - elif current_status == 2: # PROCESSING -> COMPLETED + elif current_status == int(NwbExportStatusEnum.PROCESSING): success, error_info = NwbExportHandler.process_validation(job) - next_status = 3 if success else -1 + next_status = ( + int(NwbExportStatusEnum.VALIDATION) if success else failed + ) + + elif current_status == int(NwbExportStatusEnum.VALIDATION): + success, error_info = NwbExportHandler.process_upload(job) + # process_upload returns the next status (UPLOAD or COMPLETED) + next_status = ( + error_info.pop("_next_status", completed) if success else failed + ) + + elif current_status == int(NwbExportStatusEnum.UPLOAD): + success, error_info = NwbExportHandler.process_upload(job) + next_status = ( + error_info.pop("_next_status", completed) if success else failed + ) + + elif current_status == int(NwbExportStatusEnum.UPLOADED): + # Upload finished previously; finalize. + success, error_info = ( + True, + {"error_message": None, "error_exception": None}, + ) + next_status = completed else: continue # Unknown status, skip @@ -82,23 +136,35 @@ def pipeline_handler_main(): # Update status if changed if next_status != current_status: NwbExportHandler.update_status_pipeline( - {"nwb_job_id": job["nwb_job_id"]}, current_status, next_status, error_info + {"nwb_job_id": job["nwb_job_id"]}, + current_status, + next_status, + error_info, ) - # Send Slack notifications for completion/failure - if next_status == 3: # COMPLETED - slack_message = f"NWB export completed: Job #{job['nwb_job_id']}" + if next_status == completed: + slack_message = ( + f"NWB export completed: Job #{job['nwb_job_id']}" + ) try: slack_utils.send_slack_update_notification( - config.slack_webhooks_dict.get("nwb_export_notification"), slack_message, job + config.slack_webhooks_dict.get( + "nwb_export_notification" + ), + slack_message, + job, ) except Exception as e: print(f"Failed to send Slack notification: {e}") - elif next_status == -1: # FAILED + elif next_status == failed: try: slack_utils.send_slack_error_notification( - config.slack_webhooks_dict.get("nwb_export_notification"), error_info, job + config.slack_webhooks_dict.get( + "nwb_export_notification" + ), + error_info, + job, ) except Exception as e: print(f"Failed to send Slack error notification: {e}") @@ -108,11 +174,14 @@ def pipeline_handler_main(): print(f"Error processing job {job['nwb_job_id']}: {e}") traceback.print_exc() - error_info = {"error_message": str(e)[:255], "error_exception": traceback.format_exc()[:4095]} + error_info = { + "error_message": str(e)[:255], + "error_exception": traceback.format_exc()[:4095], + } NwbExportHandler.update_status_pipeline( {"nwb_job_id": job["nwb_job_id"]}, current_status, - -1, # FAILED + failed, error_info, ) @@ -155,22 +224,35 @@ def process_data_validation(job: dict) -> tuple[bool, dict]: raise ValueError(f"Behavior validation failed: {error_msg}") elif modality_name == "ephys": - recording_key = { - k: job[k] for k in recording.Recording.primary_key if k in job - } + # The job record carries only the acquisition.Session key, not + # recording_id. Resolve the recording(s) for this session via + # the BehaviorSession Part table, then validate each. + recording_ids = recording_ids_for_session(session_key) + if not recording_ids: + raise ValueError( + f"Ephys validation failed: no recording linked to " + f"session {session_key}" + ) probe_numbers = _parse_number_list(modality.get("probe_numbers")) - valid, error_msg = validate_ephys_data_exists(recording_key, probe_numbers) - if not valid: - raise ValueError(f"Ephys validation failed: {error_msg}") + for rid in recording_ids: + recording_key = {"recording_id": rid} + valid, error_msg = validate_ephys_data_exists( + recording_key, probe_numbers + ) + if not valid: + raise ValueError( + f"Ephys validation failed for recording {rid}: {error_msg}" + ) elif modality_name == "imaging": - scan_key = { - k: job[k] for k in imaging_element.Scan.primary_key if k in job - } - fov_numbers = _parse_number_list(modality.get("fov_numbers")) - valid, error_msg = validate_imaging_data_exists(scan_key, fov_numbers) - if not valid: - raise ValueError(f"Imaging validation failed: {error_msg}") + # TODO: the imaging_element.Scan <-> acquisition.Session linkage + # is not reliably known. Do NOT fabricate a scan key (an empty + # restriction matches all rows and passes vacuously). Fail loud + # until the linkage is confirmed and wired here. + raise ValueError( + f"Imaging validation failed: could not resolve imaging Scan " + f"for session {session_key}; imaging export not yet wired." + ) print(f"Data validation passed for job {job['nwb_job_id']}") return True, error_info @@ -184,14 +266,12 @@ def process_data_validation(job: dict) -> tuple[bool, dict]: @staticmethod def process_nwb_conversion(job: dict) -> tuple[bool, dict]: """ - Convert data to NWB format. + Convert source data to an NWB file using the shared converter. - This is a placeholder implementation. In production, this would: - 1. Initialize NWB file with metadata - 2. Add behavior data using VirmenDataInterface - 3. Add ephys data from Kilosort outputs - 4. Add imaging data from ROI traces - 5. Write to output_filepath + Resolves input paths, runs TowersNWBConverter via the shared + u19_pipeline.nwb_export.conversion module (same code path as the CLI in + scripts/run_nwb_export.py), writes the NWB file and records the actual + file size on the job. Args: job: Job record dictionary @@ -199,41 +279,44 @@ def process_nwb_conversion(job: dict) -> tuple[bool, dict]: Returns: Tuple of (success, error_info) """ + from u19_pipeline.nwb_export.conversion import ( + resolve_input_paths, + run_conversion_to_file, + ) + error_info = {"error_message": None, "error_exception": None} try: print(f"Converting data to NWB for job {job['nwb_job_id']}...") - # Ensure output directory exists + export_params = _parse_export_params(job.get("export_parameters")) + session_key = { + "subject_fullname": job["subject_fullname"], + "session_date": str(job["session_date"]), + "session_number": int(job["session_number"]), + } + + virmen_file, kilosort_dir = resolve_input_paths(job, export_params) + output_path = Path(job["output_filepath"]) output_path.parent.mkdir(parents=True, exist_ok=True) - # TODO: Implement actual NWB conversion - # This would involve: - # - from pynwb import NWBFile, NWBHDF5IO - # - Creating NWBFile with session metadata - # - Adding behavior module with VirmenDataInterface - # - Adding ecephys module with Units tables - # - Adding ophys module with ImageSegmentation + Fluorescence - # - Writing to HDF5 file - - # For now, create a placeholder - # In production, replace this with actual NWB conversion code - raise NotImplementedError( - "NWB conversion not yet implemented. " - "This requires integration with VirmenDataInterface and NWB conversion tools." + size_gb = run_conversion_to_file( + job=job, + export_params=export_params, + session_key=session_key, + virmen_file=virmen_file, + kilosort_dir=kilosort_dir, + output_path=str(output_path), + ) + + (nwb_production.NwbExportJob & {"nwb_job_id": job["nwb_job_id"]}).update1( + {"actual_file_size_gb": size_gb} ) print(f"NWB conversion completed for job {job['nwb_job_id']}") return True, error_info - except NotImplementedError as e: - # Special handling for not-implemented - this is expected - error_info["error_message"] = "NWB conversion awaiting implementation" - error_info["error_exception"] = str(e)[:4095] - print(f"NWB conversion not implemented for job {job['nwb_job_id']}") - return False, error_info - except Exception as e: error_info["error_message"] = str(e)[:255] error_info["error_exception"] = traceback.format_exc()[:4095] @@ -243,9 +326,11 @@ def process_nwb_conversion(job: dict) -> tuple[bool, dict]: @staticmethod def process_validation(job: dict) -> tuple[bool, dict]: """ - Validate NWB file and insert validation record. + Validate the generated NWB file and record a validation result. - This would run NWB Inspector and other validation checks. + Runs an HDF5 integrity check and, if nwbinspector is importable, the NWB + Inspector. Inserts a single NwbExportValidation row (idempotent on retry) + and updates actual_file_size_gb. Args: job: Job record dictionary @@ -253,55 +338,117 @@ def process_validation(job: dict) -> tuple[bool, dict]: Returns: Tuple of (success, error_info) """ + import json + error_info = {"error_message": None, "error_exception": None} try: print(f"Validating NWB file for job {job['nwb_job_id']}...") output_path = Path(job["output_filepath"]) - - # Check file exists if not output_path.exists(): raise FileNotFoundError(f"NWB file not found: {output_path}") - # Get actual file size file_size_gb = output_path.stat().st_size / (1024**3) - # TODO: Implement actual NWB validation - # This would involve: - # - Running NWB Inspector: nwbinspector.inspect_nwbfile(filepath) - # - Checking HDF5 integrity: h5py.File(filepath, 'r') - # - Verifying required metadata present - # - Counting warnings and errors + # ── HDF5 integrity: open and read top-level keys ────────────────── + hdf5_ok = False + top_keys: list = [] + try: + import h5py + + with h5py.File(output_path, "r") as f: + top_keys = list(f.keys()) + hdf5_ok = bool(top_keys) + except Exception as exc: # noqa: BLE001 + hdf5_ok = False + error_info["error_message"] = f"HDF5 integrity check failed: {exc}"[ + :255 + ] + + # ── NWB Inspector (optional) ────────────────────────────────────── + inspector_ran = False + inspector_passed = True + warnings_count = 0 + errors_count = 0 + report_messages: list = [] + try: + import nwbinspector + + inspect_fn = getattr(nwbinspector, "inspect_nwbfile", None) + if inspect_fn is None: + raise AttributeError("nwbinspector.inspect_nwbfile not available") + for msg in inspect_fn(nwbfile_path=str(output_path)): + importance = str(getattr(msg, "importance", "")) + report_messages.append( + { + "importance": importance, + "message": str(getattr(msg, "message", "")), + "check": str(getattr(msg, "check_function_name", "")), + } + ) + if "ERROR" in importance.upper(): + errors_count += 1 + else: + warnings_count += 1 + inspector_ran = True + inspector_passed = errors_count == 0 + except ImportError: + print( + "nwbinspector not importable; skipping NWB Inspector (not a failure)." + ) + except Exception as exc: # noqa: BLE001 + # Inspector itself errored; record but do not crash validation. + print(f"NWB Inspector run failed: {exc}") + report_messages.append({"inspector_error": str(exc)}) + + validation_passed = bool( + hdf5_ok and (not inspector_ran or inspector_passed) + ) + + report_json = json.dumps( + { + "top_level_keys": top_keys, + "inspector_ran": inspector_ran, + "messages": report_messages, + } + ) - # For now, create a placeholder validation record validation_record = { "nwb_job_id": job["nwb_job_id"], "validation_timestamp": datetime.now(), - "validation_passed": False, # Set to False until implemented - "validation_report_json": {}, # Would contain Inspector output + "validation_passed": validation_passed, + "validation_report_json": report_json, "file_size_gb": file_size_gb, - "nwb_inspector_passed": False, - "hdf5_integrity_passed": False, - "metadata_complete_passed": False, - "validation_warnings_count": 0, - "validation_errors_count": 1, # Count "not implemented" as error + "nwb_inspector_passed": inspector_passed if inspector_ran else True, + "hdf5_integrity_passed": hdf5_ok, + "metadata_complete_passed": hdf5_ok, + "validation_warnings_count": warnings_count, + "validation_errors_count": errors_count, } + # Insert ONCE. On a cronjob retry a row may already exist for this + # nwb_job_id (primary key); replace it rather than crashing on a + # duplicate-key insert. + job_pk = {"nwb_job_id": job["nwb_job_id"]} + if nwb_production.NwbExportValidation & job_pk: + (nwb_production.NwbExportValidation & job_pk).delete_quick() nwb_production.NwbExportValidation.insert1(validation_record) - # Update actual file size in main job record - (nwb_production.NwbExportJob & {"nwb_job_id": job["nwb_job_id"]}).update1( + (nwb_production.NwbExportJob & job_pk).update1( {"actual_file_size_gb": file_size_gb} ) - raise NotImplementedError("NWB validation not yet implemented") + if not validation_passed: + error_info["error_message"] = ( + error_info["error_message"] + or f"NWB validation failed (hdf5_ok={hdf5_ok}, errors={errors_count})" + ) + print(f"NWB validation failed for job {job['nwb_job_id']}") + return False, error_info - except NotImplementedError as e: - error_info["error_message"] = "NWB validation awaiting implementation" - error_info["error_exception"] = str(e)[:4095] - print(f"NWB validation not implemented for job {job['nwb_job_id']}") - return False, error_info + print(f"NWB validation passed for job {job['nwb_job_id']}") + return True, error_info except Exception as e: error_info["error_message"] = str(e)[:255] @@ -310,7 +457,90 @@ def process_validation(job: dict) -> tuple[bool, dict]: return False, error_info @staticmethod - def update_status_pipeline(job_key: dict, old_status: int, new_status: int, error_info: dict): + def process_upload(job: dict) -> tuple[bool, dict]: + """ + Optional DANDI upload stage. + + If the job has no NwbExportJobDandi row, no upload was requested and the + job transitions straight to COMPLETED. If a dandiset was selected, upload + the NWB file to DANDI and record the upload status; on success transition + to COMPLETED, on failure mark FAILED with an actionable error (a requested + upload is never silently skipped). + + The status to advance to on success is returned via the ``_next_status`` + key of ``error_info`` (the dispatch loop pops it). + + Args: + job: Job record dictionary + + Returns: + Tuple of (success, error_info) + """ + error_info = {"error_message": None, "error_exception": None} + completed = int(NwbExportStatusEnum.COMPLETED) + + try: + job_pk = {"nwb_job_id": job["nwb_job_id"]} + + dandi_rows = (nwb_production.NwbExportJobDandi & job_pk).fetch(as_dict=True) + if not dandi_rows: + # No DANDI upload requested for this job -> finalize. + print( + f"No DANDI upload requested for job {job['nwb_job_id']}; finalizing." + ) + error_info["_next_status"] = completed + return True, error_info + + dandiset_id = dandi_rows[0]["dandiset_id"] + user_id = job["user_id"] + + from u19_pipeline.nwb_production import ( + can_upload_to_dandi, + get_dandi_credentials, + ) + + if not can_upload_to_dandi(user_id): + raise RuntimeError( + f"DANDI upload requested (dandiset {dandiset_id}) but user " + f"'{user_id}' has no usable DANDI credentials configured." + ) + + api_key, _default_dandiset = get_dandi_credentials(user_id) + if not api_key: + raise RuntimeError( + f"DANDI upload requested but no API key available for user '{user_id}'." + ) + + from u19_pipeline.nwb_export.dandi.upload_client import DandiUploadClient + + output_path = job["output_filepath"] + print( + f"Uploading job {job['nwb_job_id']} to DANDI dandiset {dandiset_id}..." + ) + client = DandiUploadClient(api_key=api_key, dandiset_id=dandiset_id) + client.upload(output_path) + + # Record upload success. The neuroconv upload returns organised file + # paths, not a DANDI asset id; leave dandi_asset_id NULL until a real + # asset-id lookup is wired (TODO). + (nwb_production.NwbExportJobDandi & job_pk).update1( + {**job_pk, "upload_status": 1, "upload_timestamp": datetime.now()} + ) + + print(f"DANDI upload complete for job {job['nwb_job_id']}") + error_info["_next_status"] = completed + return True, error_info + + except Exception as e: + error_info["error_message"] = str(e)[:255] + error_info["error_exception"] = traceback.format_exc()[:4095] + print(f"DANDI upload failed for job {job['nwb_job_id']}: {e}") + return False, error_info + + @staticmethod + def update_status_pipeline( + job_key: dict, old_status: int, new_status: int, error_info: dict + ): """ Update job status and log transition. @@ -320,14 +550,18 @@ def update_status_pipeline(job_key: dict, old_status: int, new_status: int, erro new_status: New status ID error_info: Dictionary with error_message and error_exception """ - print(f"Updating job {job_key['nwb_job_id']}: status {old_status} -> {new_status}") + print( + f"Updating job {job_key['nwb_job_id']}: status {old_status} -> {new_status}" + ) # Update job status (nwb_production.NwbExportJob & job_key).update1({"status_id": new_status}) # Set completion timestamp if completed if new_status == 3: # COMPLETED - (nwb_production.NwbExportJob & job_key).update1({"completion_timestamp": datetime.now()}) + (nwb_production.NwbExportJob & job_key).update1( + {"completion_timestamp": datetime.now()} + ) # Log status change log_entry = { diff --git a/u19_pipeline/nwb_export/conversion.py b/u19_pipeline/nwb_export/conversion.py new file mode 100644 index 00000000..b818ca6b --- /dev/null +++ b/u19_pipeline/nwb_export/conversion.py @@ -0,0 +1,306 @@ +""" +Shared NWB conversion logic for the U19 export pipeline. + +This module holds the importable functions that drive a real NWB conversion via +``TowersNWBConverter``. Both the CLI (``scripts/run_nwb_export.py``) and the +cronjob handler (``u19_pipeline/automatic_job/nwb_export_handler.py``) import +from here so there is a single, shared code path. + +The logic was ported from ``scripts/run_nwb_export.py`` (``_build_source_data``, +``_query_metadata`` and the PROCESSING block) so it can be reused without +duplication. +""" + +from __future__ import annotations + +import logging +import os +from datetime import datetime +from pathlib import Path + +log = logging.getLogger(__name__) + + +# ────────────────────────────────────────────────────────────────────────────── +# Input-path resolution +# ────────────────────────────────────────────────────────────────────────────── + + +def resolve_input_paths( + job: dict, + export_params: dict, +) -> tuple[Path | None, Path | None]: + """ + Resolve the behavioral Virmen file and the Kilosort base directory. + + Resolution order: + 1. Explicit paths from ``export_params`` (keys ``virmen_file`` / + ``kilosort_dir``). + 2. A data-root env var (``NWB_EXPORT_DATA_ROOT``) joined with the session + identifiers. This is intentionally a *simple* convention; we do not try + to guess a specific lab directory layout we cannot verify. + 3. If nothing resolves, raise a clear, actionable error. + + Args: + job: NwbExportJob record dict (carries subject_fullname/session_date/...). + export_params: Parsed export_parameters dict for the job. + + Returns: + ``(virmen_file, kilosort_dir)`` — ``kilosort_dir`` may be ``None`` when + the job has no ephys component. + + Raises: + FileNotFoundError: If the Virmen behavioral file cannot be resolved. + """ + virmen_file: Path | None = None + kilosort_dir: Path | None = None + + # 1. Explicit paths from export_params. + if export_params.get("virmen_file"): + virmen_file = Path(export_params["virmen_file"]) + if export_params.get("kilosort_dir"): + kilosort_dir = Path(export_params["kilosort_dir"]) + + # 2. Fall back to a data-root env var + session identifiers. + data_root = os.environ.get("NWB_EXPORT_DATA_ROOT") + if virmen_file is None and data_root: + candidate = ( + Path(data_root) + / str(job.get("subject_fullname", "")) + / f"{job.get('session_date', '')}_{job.get('session_number', '')}.mat" + ) + if candidate.exists(): + virmen_file = candidate + + # 3. Could not resolve the (required) behavioral file → fail loud. + if virmen_file is None: + raise FileNotFoundError( + "Could not resolve the Virmen behavioral .mat file for this job. " + "Provide it explicitly via export_parameters['virmen_file'] " + "(and export_parameters['kilosort_dir'] for ephys), or set the " + "NWB_EXPORT_DATA_ROOT environment variable to a data root containing " + "/_.mat." + ) + if not virmen_file.exists(): + raise FileNotFoundError(f"Virmen file not found: {virmen_file}") + + return virmen_file, kilosort_dir + + +# ────────────────────────────────────────────────────────────────────────────── +# source_data construction (ported from run_nwb_export.py) +# ────────────────────────────────────────────────────────────────────────────── + + +def _find_kilosort_output(probe_dir: Path) -> Path | None: + """ + Return the highest-numbered job's Kilosort output directory under probe_dir. + + Expected layout: + /_imec/job_id_/kilosort_output/ + """ + job_dirs = sorted( + [p for p in probe_dir.glob("job_id_*") if p.name.split("_")[-1].isdigit()], + key=lambda p: int(p.name.split("_")[-1]), + ) + if not job_dirs: + return None + kilosort_outputs = list(job_dirs[-1].glob("kilosort*_output")) + return kilosort_outputs[0] if kilosort_outputs else None + + +def build_source_data( + job: dict, + export_params: dict, + virmen_file: Path | None, + kilosort_dir: Path | None, +) -> dict: + """ + Translate the DataJoint job record + export_params into a ``source_data`` + dict accepted by ``TowersNWBConverter``. + + Raises: + FileNotFoundError: If a required data file cannot be located. + """ + source_data: dict = {} + + # ── Behavior (always required) ──────────────────────────────────────────── + if virmen_file is None: + raise FileNotFoundError( + "No Virmen .mat file provided. Cannot locate the behavioral data." + ) + if not Path(virmen_file).exists(): + raise FileNotFoundError(f"Virmen file not found: {virmen_file}") + source_data["VirmenData"] = {"file_path": str(virmen_file)} + + # ── Ephys ───────────────────────────────────────────────────────────────── + if export_params.get("include_ephys") and kilosort_dir is not None: + kilosort_dir = Path(kilosort_dir) + probe_dirs = sorted(kilosort_dir.glob("*_imec*")) + if not probe_dirs: + log.warning( + "kilosort_dir given but no *_imec* subdirectories found – skipping ephys." + ) + else: + for probe_dir in probe_dirs: + probe_idx = "".join( + filter(str.isdigit, probe_dir.name.split("imec")[-1]) + ) + interface_name = ( + f"KilosortProbe{probe_idx}" if probe_idx else "Kilosort" + ) + ks_output = _find_kilosort_output(probe_dir) + if ks_output is None: + log.warning( + f"No Kilosort output found under {probe_dir} – skipping." + ) + continue + source_data[interface_name] = {"folder_path": str(ks_output)} + log.info(f" {interface_name}: {ks_output}") + elif export_params.get("include_ephys"): + log.warning( + "include_ephys=True but no kilosort_dir provided; ephys data will not be included." + ) + + return source_data + + +# ────────────────────────────────────────────────────────────────────────────── +# DB metadata query (ported from run_nwb_export.py) +# ────────────────────────────────────────────────────────────────────────────── + + +def query_metadata(session_key: dict) -> dict: + """ + Pull experimenter, subject sex/DoB and sync timestamps from DataJoint. + + Returns a dict with keys: experimenter, subject_sex, subject_dob, + sync_timestamps (may be None if not found). + """ + import datajoint as dj + + result: dict = { + "experimenter": [], + "subject_sex": "U", + "subject_dob": None, + "sync_timestamps": None, + } + + try: + subject = dj.create_virtual_module( + "subject", dj.config["custom"]["database.prefix"] + "subject" + ) + lab = dj.create_virtual_module( + "lab", dj.config["custom"]["database.prefix"] + "lab" + ) + + subject_fullname = session_key["subject_fullname"] + sub_info = ( + subject.Subject() * lab.User() & f"subject_fullname = '{subject_fullname}'" + ).fetch1() + + owner_full = sub_info.get("full_name", sub_info.get("user_id", "")) + if " " in owner_full: + parts = owner_full.rsplit(" ", 1) + result["experimenter"].append(f"{parts[-1]}, {parts[0]}") + else: + result["experimenter"].append(owner_full) + + sex_map = {"Male": "M", "Female": "F", "Unknown": "U", "m": "M", "f": "F"} + result["subject_sex"] = sex_map.get(str(sub_info.get("sex", "U")), "U") + + dob = sub_info.get("dob") + if dob is not None: + result["subject_dob"] = ( + datetime.combine(dob, datetime.min.time()) + if hasattr(dob, "year") + else dob + ) + + except Exception as exc: + log.warning(f"Could not query all metadata from DB: {exc}") + + # Sync timestamps (optional BehaviorSync table) + try: + nwb_prod = dj.create_virtual_module( + "nwb_production", dj.config["custom"]["database.prefix"] + "nwb_production" + ) + sync_rows = (nwb_prod.BehaviorSync & session_key).fetch( + "sync_timestamps", as_dict=True + ) + if sync_rows: + import numpy as np + + result["sync_timestamps"] = np.array(sync_rows[0]["sync_timestamps"]) + except Exception: + pass # BehaviorSync is optional + + return result + + +# ────────────────────────────────────────────────────────────────────────────── +# Conversion driver (ported from run_nwb_export.py PROCESSING block) +# ────────────────────────────────────────────────────────────────────────────── + + +def run_conversion_to_file( + job: dict, + export_params: dict, + session_key: dict, + virmen_file: Path | None, + kilosort_dir: Path | None, + output_path: str, +) -> float: + """ + Run the full NWB conversion to ``output_path`` and return the size in GB. + + Builds source_data, queries DB metadata, runs ``TowersNWBConverter`` and + writes the NWB file (overwriting any existing file at the path). + + Raises: + ImportError: If ``tank_lab_to_nwb`` is not installed. + Exception: Propagated from the converter on conversion failure. + """ + from tank_lab_to_nwb.convert_towers_task.towersnwbconverter import ( + TowersNWBConverter, + ) + + source_data = build_source_data(job, export_params, virmen_file, kilosort_dir) + log.info(f" source_data keys: {list(source_data.keys())}") + + metadata = query_metadata(session_key) + + converter = TowersNWBConverter( + source_data=source_data, + sync_timestamps=metadata["sync_timestamps"], + ) + + raw_metadata = converter.get_metadata() + + raw_metadata["NWBFile"]["session_description"] = ( + f"U19 pipeline export – {session_key['subject_fullname']} " + f"{session_key['session_date']}" + ) + if metadata["experimenter"]: + raw_metadata["NWBFile"]["experimenter"] = metadata["experimenter"] + + if "Subject" not in raw_metadata: + raw_metadata["Subject"] = {} + if metadata["subject_sex"]: + raw_metadata["Subject"]["sex"] = metadata["subject_sex"] + if metadata["subject_dob"] is not None: + raw_metadata["Subject"]["date_of_birth"] = metadata["subject_dob"] + + Path(output_path).parent.mkdir(parents=True, exist_ok=True) + + log.info(f" Writing NWB to: {output_path}") + converter.run_conversion( + nwbfile_path=output_path, + metadata=raw_metadata, + overwrite=True, + ) + log.info(" ✓ Conversion complete") + + size_gb = Path(output_path).stat().st_size / (1024**3) + log.info(f" ✓ File size: {size_gb:.3f} GB") + return size_gb diff --git a/u19_pipeline/nwb_production_utils.py b/u19_pipeline/nwb_production_utils.py index 666151b8..5a61ade8 100644 --- a/u19_pipeline/nwb_production_utils.py +++ b/u19_pipeline/nwb_production_utils.py @@ -4,6 +4,35 @@ Provides functions to estimate file sizes and validate data availability. """ +import logging + +log = logging.getLogger(__name__) + + +def recording_ids_for_session(session_key: dict) -> list: + """ + Return the recording_id(s) linked to an acquisition.Session. + + The NwbExportJob record only carries the acquisition.Session primary key + (subject_fullname, session_date, session_number); it does NOT carry + recording_id. The link lives in recording.Recording.BehaviorSession, a Part + table whose secondary attribute is ``-> acquisition.Session``. This helper + resolves the session key to the set of recording_ids for that session. + + Args: + session_key: Dictionary with acquisition.Session identifiers. + + Returns: + List of recording_id ints (empty if none linked). + """ + from u19_pipeline import recording # noqa: PLC0415 + + return ( + (recording.Recording.BehaviorSession & session_key) + .fetch("recording_id") + .tolist() + ) + def estimate_behavior_size_gb(session_key: dict) -> float: """ @@ -101,30 +130,46 @@ def estimate_total_size(nwb_job_key: dict) -> float: Returns: Total estimated size in GB """ - from u19_pipeline import acquisition, nwb_production, recording - from u19_pipeline.imaging_pipeline import imaging_element + from u19_pipeline import acquisition, nwb_production total_gb = 0.0 job = (nwb_production.NwbExportJob & nwb_job_key).fetch1() + session_key = {k: job[k] for k in acquisition.Session.primary_key if k in job} modalities = (nwb_production.NwbExportModality & nwb_job_key).fetch(as_dict=True) for modality in modalities: modality_name = modality["modality_name"] if modality_name == "behavior": - session_key = {k: job[k] for k in acquisition.Session.primary_key if k in job} total_gb += estimate_behavior_size_gb(session_key) elif modality_name == "ephys": - recording_key = {k: job[k] for k in recording.Recording.primary_key if k in job} + # The job carries the session key, not recording_id. Resolve the + # recording_id(s) for the session via the BehaviorSession Part table. + recording_ids = recording_ids_for_session(session_key) probe_numbers = _parse_number_list(modality.get("probe_numbers")) - total_gb += estimate_ephys_size_gb(recording_key, probe_numbers) + if not recording_ids: + log.warning( + "estimate_total_size: no recording linked to session %s; " + "skipping ephys estimate.", + session_key, + ) + continue + for rid in recording_ids: + total_gb += estimate_ephys_size_gb({"recording_id": rid}, probe_numbers) elif modality_name == "imaging": - scan_key = {k: job[k] for k in imaging_element.Scan.primary_key if k in job} - fov_numbers = _parse_number_list(modality.get("fov_numbers")) - total_gb += estimate_imaging_size_gb(scan_key, fov_numbers) + # TODO: the imaging_element.Scan <-> acquisition.Session linkage is + # not reliably known. Until it is confirmed we cannot build a real + # scan_key; estimating against an empty restriction would be vacuous, + # so skip with a logged warning rather than fabricate a key. + log.warning( + "estimate_total_size: imaging Scan<->session linkage not yet wired " + "for session %s; skipping imaging estimate.", + session_key, + ) + continue return total_gb @@ -156,7 +201,9 @@ def validate_behavior_data_exists(session_key: dict) -> tuple[bool, str]: return False, f"Error validating behavior data: {str(e)}" -def validate_ephys_data_exists(recording_key: dict, probe_numbers: list) -> tuple[bool, str]: +def validate_ephys_data_exists( + recording_key: dict, probe_numbers: list +) -> tuple[bool, str]: """ Validate that ephys data exists for specified probes.