Skip to content

Commit ba2746c

Browse files
Eduardo Hirata-Miyasakiclaude
andcommitted
feat(viscy-data): CSV metadata sidecar for write-protected datasets
Add viscy_data.meta_csv (build_provenance_fields, metadata_to_rows, write_meta_rows_csv, read_meta_rows_csv) and thread a csv_dir option through the preprocessing path so normalization/QC statistics can be written to and read from per-store CSV sidecars instead of the zarr .zattrs, opening the store read-only. Datasets mounted without write access can then be prepared and trained on. Wires csv_dir through VisCyTrainer.preprocess, generate_normalization_metadata, the QC metrics run + config, the dynaclr cell-index preprocessing (CLI + Nextflow module + workflow) and its parquet reader, and airtable prepare (check_preprocessed + batch discovery). qc gains a viscy-data dependency for the shared module; uv.lock updated accordingly. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent a44d5a0 commit ba2746c

25 files changed

Lines changed: 2351 additions & 1377 deletions

File tree

applications/airtable/src/airtable_utils/prepare.py

Lines changed: 66 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
from iohub import open_ome_zarr
1212
from pydantic import BaseModel, Field
1313

14+
from viscy_data.meta_csv import read_meta_rows_csv
15+
1416
logger = logging.getLogger(__name__)
1517

1618
# ---------------------------------------------------------------------------
@@ -58,6 +60,7 @@ class SlurmStageConfig(BaseModel):
5860
time: str = "06:00:00"
5961
gres: str | None = None
6062
constraint: str | None = None
63+
qos: str | None = None
6164

6265

6366
class SlurmConfig(BaseModel):
@@ -92,6 +95,8 @@ class PrepareConfig(BaseModel):
9295
nfs_root: Path = Path("/hpc/projects/intracellular_dashboard/organelle_dynamics")
9396
vast_root: Path = Path("/hpc/projects/organelle_phenotyping/datasets")
9497
workspace_dir: Path = Path("/hpc/mydata/eduardo.hirata/repos/viscy")
98+
csv_dir: Path | None = None
99+
jobs_dir: Path | None = None
95100
concatenate: ConcatenateConfig = Field(default_factory=ConcatenateConfig)
96101
qc: QCParams = Field(default_factory=QCParams)
97102
preprocess: PreprocessParams = Field(default_factory=PreprocessParams)
@@ -195,19 +200,28 @@ def check_zarr_version(zarr_path: Path) -> dict[str, int | str | None]:
195200
return result
196201

197202

198-
def check_preprocessed(zarr_path: Path) -> bool:
199-
"""Check if normalization metadata has been written to the zarr store.
203+
def check_preprocessed(zarr_path: Path, csv_dir: Path | None = None) -> bool:
204+
"""Check if normalization metadata has been written for the zarr store.
200205
201206
Parameters
202207
----------
203208
zarr_path : Path
204209
Path to the zarr store root.
210+
csv_dir : Path or None
211+
If given, check the per-store CSV sidecar under this directory
212+
(written by ``viscy preprocess --csv_dir ...``) instead of the
213+
store's ``.zattrs``/``zarr.json`` — for datasets prepared without
214+
write access.
205215
206216
Returns
207217
-------
208218
bool
209219
True if normalization stats are present.
210220
"""
221+
if csv_dir is not None:
222+
sidecar = read_meta_rows_csv(csv_dir, zarr_path)
223+
return sidecar is not None and bool((sidecar["field_name"] == "normalization").any())
224+
211225
zarr_json = zarr_path / "zarr.json"
212226
zattrs = zarr_path / ".zattrs"
213227

@@ -223,6 +237,34 @@ def check_preprocessed(zarr_path: Path) -> bool:
223237
return False
224238

225239

240+
# ---------------------------------------------------------------------------
241+
# Batch discovery (filesystem glob, no NFS/Airtable dependency)
242+
# ---------------------------------------------------------------------------
243+
244+
245+
def discover_zarr_stores(data_root: Path) -> list[Path]:
246+
"""Find dataset zarr stores under a root laid out like VAST's convention.
247+
248+
Matches ``{data_root}/{dataset_name}/{dataset_name}.zarr`` exactly (the
249+
zarr's own basename must match its containing directory's name) — this
250+
excludes sibling ``tracking.zarr`` dirs as well as any other variant
251+
zarr stores (e.g. ``{dataset_name}_chunked.zarr``,
252+
``{dataset_name}_sharded.zarr``) that may sit alongside the canonical
253+
store for rechunking/testing purposes.
254+
255+
Parameters
256+
----------
257+
data_root : Path
258+
Root directory containing one subdirectory per dataset.
259+
260+
Returns
261+
-------
262+
list[Path]
263+
Sorted paths to each dataset's zarr store.
264+
"""
265+
return sorted(p for p in Path(data_root).glob("*/*.zarr") if p.stem == p.parent.name)
266+
267+
226268
# ---------------------------------------------------------------------------
227269
# Discovery (reads NFS zarr via iohub)
228270
# ---------------------------------------------------------------------------
@@ -339,7 +381,7 @@ def generate_crop_concat_config(
339381
}
340382

341383

342-
def generate_qc_config(data_path: Path, qc_params: QCParams) -> dict:
384+
def generate_qc_config(data_path: Path, qc_params: QCParams, csv_dir: Path | None = None) -> dict:
343385
"""Build a QC config dict compatible with ``qc run -c``.
344386
345387
Parameters
@@ -348,6 +390,9 @@ def generate_qc_config(data_path: Path, qc_params: QCParams) -> dict:
348390
Path to the VAST zarr (target of QC).
349391
qc_params : QCParams
350392
Focus-slice QC parameters.
393+
csv_dir : Path or None
394+
If given, ``qc run`` writes results to a per-store CSV sidecar
395+
under this directory (read-only store) instead of ``.zattrs``.
351396
352397
Returns
353398
-------
@@ -357,6 +402,7 @@ def generate_qc_config(data_path: Path, qc_params: QCParams) -> dict:
357402
return {
358403
"data_path": str(data_path),
359404
"num_workers": qc_params.num_workers,
405+
"csv_dir": str(csv_dir) if csv_dir else None,
360406
"focus_slice": {
361407
"channel_names": qc_params.channel_names,
362408
"NA_det": qc_params.NA_det,
@@ -412,6 +458,8 @@ def _slurm_header(job_name: str, output_dir: Path, cfg: SlurmStageConfig) -> str
412458
lines.append(f"#SBATCH --gres={cfg.gres}")
413459
if cfg.constraint:
414460
lines.append(f'#SBATCH --constraint="{cfg.constraint}"')
461+
if cfg.qos:
462+
lines.append(f"#SBATCH --qos={cfg.qos}")
415463
return "\n".join(lines)
416464

417465

@@ -532,7 +580,7 @@ def generate_qc_slurm(
532580
export PYTHONNOUSERSITE=1
533581
534582
echo "=== QC: focus slice detection ==="
535-
uv run --project "{workspace_dir}" --package qc \
583+
uv run --project "{workspace_dir}" --package viscy-qc \
536584
qc run -c "{qc_config_path}"
537585
echo "QC complete."
538586
""")
@@ -546,6 +594,7 @@ def generate_preprocess_slurm(
546594
workspace_dir: Path,
547595
preprocess_params: PreprocessParams,
548596
slurm_cfg: SlurmStageConfig,
597+
csv_dir: Path | None = None,
549598
) -> str:
550599
"""Generate SLURM script for normalization preprocessing (CPU only).
551600
@@ -563,6 +612,10 @@ def generate_preprocess_slurm(
563612
Normalization preprocessing parameters.
564613
slurm_cfg : SlurmStageConfig
565614
SLURM resource parameters.
615+
csv_dir : Path or None
616+
If given, ``viscy preprocess`` writes results to a per-store CSV
617+
sidecar under this directory (read-only store) instead of
618+
``.zattrs``.
566619
567620
Returns
568621
-------
@@ -576,6 +629,7 @@ def generate_preprocess_slurm(
576629
ch_flag = f"--channel_names={ch_arg}"
577630
else:
578631
ch_flag = " ".join(f"--channel_names={c}" for c in ch_arg)
632+
csv_dir_flag = f' --csv_dir "{csv_dir}"' if csv_dir else ""
579633

580634
body = dedent(f"""\
581635
@@ -586,7 +640,7 @@ def generate_preprocess_slurm(
586640
uv run --project "{workspace_dir}" --package dynaclr \
587641
viscy preprocess --data_path "{vast_zarr_path}" \
588642
{ch_flag} --num_workers {preprocess_params.num_workers} \
589-
--block_size {preprocess_params.block_size}
643+
--block_size {preprocess_params.block_size}{csv_dir_flag}
590644
echo "Preprocess complete."
591645
""")
592646
return header + "\n" + body
@@ -597,7 +651,9 @@ def generate_preprocess_slurm(
597651
# ---------------------------------------------------------------------------
598652

599653

600-
def check_dataset_status(dataset_name: str, nfs_root: Path, vast_root: Path) -> dict[str, str]:
654+
def check_dataset_status(
655+
dataset_name: str, nfs_root: Path, vast_root: Path, csv_dir: Path | None = None
656+
) -> dict[str, str]:
601657
"""Check existence and version info for a dataset across NFS and VAST.
602658
603659
Parameters
@@ -608,6 +664,9 @@ def check_dataset_status(dataset_name: str, nfs_root: Path, vast_root: Path) ->
608664
NFS root directory.
609665
vast_root : Path
610666
VAST root directory.
667+
csv_dir : Path or None
668+
If given, check the ``preprocessed`` status via the CSV sidecar
669+
under this directory instead of ``.zattrs`` (see ``check_preprocessed``).
611670
612671
Returns
613672
-------
@@ -629,7 +688,7 @@ def check_dataset_status(dataset_name: str, nfs_root: Path, vast_root: Path) ->
629688
ver = check_zarr_version(vast["zarr"])
630689
zarr_fmt = str(ver["zarr_format"]) if ver["zarr_format"] else "?"
631690
ome_ver = str(ver["ome_version"]) if ver["ome_version"] else "?"
632-
preprocessed = "yes" if check_preprocessed(vast["zarr"]) else "no"
691+
preprocessed = "yes" if check_preprocessed(vast["zarr"], csv_dir=csv_dir) else "no"
633692

634693
return {
635694
"dataset": dataset_name,

applications/airtable/src/airtable_utils/prepare_cli.py

Lines changed: 96 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import subprocess
88

99
import click
10+
from viscy_data.meta_csv import backfill_zattrs_to_csv
1011

1112
from airtable_utils.prepare import (
1213
PrepareConfig,
@@ -15,6 +16,7 @@
1516
check_zarr_version,
1617
discover_channels,
1718
discover_wells,
19+
discover_zarr_stores,
1820
filter_raw_channels,
1921
format_status_table,
2022
generate_concatenate_script,
@@ -100,7 +102,7 @@ def run(dataset_name: str, config_path: str, dry_run: bool, force: bool) -> None
100102
ver = check_zarr_version(vast["zarr"])
101103
is_v3 = ver["zarr_format"] == 3
102104
is_ome05 = ver["ome_version"] == "0.5"
103-
is_preprocessed = check_preprocessed(vast["zarr"])
105+
is_preprocessed = check_preprocessed(vast["zarr"], csv_dir=cfg.csv_dir)
104106

105107
if is_v3 and is_ome05 and is_preprocessed:
106108
click.echo(
@@ -149,7 +151,7 @@ def run(dataset_name: str, config_path: str, dry_run: bool, force: bool) -> None
149151
click.echo(f" Wrote: {crop_concat_path}")
150152

151153
# 8. Generate qc_config.yml
152-
qc_cfg = generate_qc_config(vast["zarr"], cfg.qc)
154+
qc_cfg = generate_qc_config(vast["zarr"], cfg.qc, csv_dir=cfg.csv_dir)
153155
qc_config_path = vast["output_dir"] / "qc_config.yml"
154156
write_yaml(qc_cfg, qc_config_path)
155157
click.echo(f" Wrote: {qc_config_path}")
@@ -192,6 +194,7 @@ def run(dataset_name: str, config_path: str, dry_run: bool, force: bool) -> None
192194
workspace_dir=cfg.workspace_dir,
193195
preprocess_params=cfg.preprocess,
194196
slurm_cfg=cfg.slurm.preprocess,
197+
csv_dir=cfg.csv_dir,
195198
)
196199
preprocess_script_path = vast["output_dir"] / "03_preprocess.sh"
197200
preprocess_script_path.write_text(preprocess_script)
@@ -246,10 +249,100 @@ def status(dataset_names: tuple[str, ...], config_path: str) -> None:
246249
"""Check NFS/VAST existence and version status for one or more datasets."""
247250
cfg = _load_prepare_config(config_path)
248251

249-
rows = [check_dataset_status(name, cfg.nfs_root, cfg.vast_root) for name in dataset_names]
252+
rows = [check_dataset_status(name, cfg.nfs_root, cfg.vast_root, csv_dir=cfg.csv_dir) for name in dataset_names]
250253
click.echo(format_status_table(rows))
251254

252255

256+
@prepare.command("batch-preprocess")
257+
@click.argument("data_root", type=click.Path(exists=True, file_okay=False))
258+
@click.option(
259+
"-c",
260+
"--config",
261+
"config_path",
262+
required=True,
263+
type=click.Path(exists=True),
264+
help="Path to prepare config YAML.",
265+
)
266+
@click.option("--dry-run", is_flag=True, help="Generate configs/scripts without submitting SLURM jobs.")
267+
def batch_preprocess(data_root: str, config_path: str, dry_run: bool) -> None:
268+
"""Preprocess every not-yet-preprocessed dataset zarr under DATA_ROOT.
269+
270+
Discovers `{name}/{name}.zarr` stores under DATA_ROOT, skips any already
271+
marked preprocessed in the CSV sidecar, mirrors metadata that's already
272+
sitting in a store's own `.zattrs` (e.g. computed elsewhere with write
273+
access, then copied here) into the CSV sidecar for free, and only
274+
generates + submits QC/normalization SLURM jobs for stores with neither.
275+
Safe to re-run repeatedly as new datasets appear — anything already
276+
covered (via CSV or zattrs) is always skipped, so this never redoes work.
277+
"""
278+
cfg = _load_prepare_config(config_path)
279+
if cfg.csv_dir is None:
280+
raise click.ClickException(
281+
"batch-preprocess requires 'csv_dir' in the prepare config — datasets under "
282+
"DATA_ROOT may be owned by teammates without write access to their .zattrs."
283+
)
284+
jobs_root = cfg.jobs_dir or cfg.csv_dir
285+
286+
stores = discover_zarr_stores(data_root)
287+
click.echo(f"Discovered {len(stores)} dataset(s) under {data_root}.")
288+
289+
pending = [store for store in stores if not check_preprocessed(store, csv_dir=cfg.csv_dir)]
290+
291+
backfilled = []
292+
todo = []
293+
for store in pending:
294+
if backfill_zattrs_to_csv(store, cfg.csv_dir):
295+
backfilled.append(store)
296+
else:
297+
todo.append(store)
298+
299+
click.echo(
300+
f" {len(stores) - len(pending)} already in CSV sidecar, "
301+
f"{len(backfilled)} backfilled from existing zattrs, {len(todo)} need SLURM jobs."
302+
)
303+
304+
for store in todo:
305+
dataset_name = store.stem
306+
job_dir = jobs_root / "_jobs" / dataset_name
307+
job_dir.mkdir(parents=True, exist_ok=True)
308+
309+
qc_config_path = job_dir / "qc_config.yml"
310+
write_yaml(generate_qc_config(store, cfg.qc, csv_dir=cfg.csv_dir), qc_config_path)
311+
312+
qc_script_path = job_dir / "02_qc.sh"
313+
qc_script_path.write_text(
314+
generate_qc_slurm(dataset_name, job_dir, qc_config_path, cfg.workspace_dir, cfg.slurm.qc)
315+
)
316+
317+
preprocess_script_path = job_dir / "03_preprocess.sh"
318+
preprocess_script_path.write_text(
319+
generate_preprocess_slurm(
320+
dataset_name,
321+
job_dir,
322+
store,
323+
cfg.workspace_dir,
324+
cfg.preprocess,
325+
cfg.slurm.preprocess,
326+
csv_dir=cfg.csv_dir,
327+
)
328+
)
329+
click.echo(f" Wrote: {job_dir}")
330+
331+
if dry_run:
332+
continue
333+
334+
result_qc = subprocess.run(["sbatch", str(qc_script_path)], capture_output=True, text=True, check=True)
335+
qc_job_id = _parse_slurm_job_id(result_qc.stdout)
336+
result_pp = subprocess.run(["sbatch", str(preprocess_script_path)], capture_output=True, text=True, check=True)
337+
pp_job_id = _parse_slurm_job_id(result_pp.stdout)
338+
click.echo(f" {dataset_name}: QC job {qc_job_id}, preprocess job {pp_job_id}")
339+
340+
if dry_run:
341+
click.echo("\n--dry-run: configs and scripts generated, nothing submitted.")
342+
else:
343+
click.echo(f"\nSubmitted jobs for {len(todo)} dataset(s).")
344+
345+
253346
def main() -> None:
254347
"""Entry point for the prepare CLI."""
255348
prepare()

0 commit comments

Comments
 (0)