diff --git a/Dockerfile b/Dockerfile index e54c6a3..93868a4 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,5 +1,7 @@ FROM australia-southeast1-docker.pkg.dev/cpg-common/images/cpg_flow:1.3.1 +# CI's get_version.py reads this line to derive the image tag (VERSION-); +# it must stay a literal semver, so the VERSION build-arg CI passes is unused. ENV VERSION=0.1.0 # Set the working directory diff --git a/README.md b/README.md index 8c047a5..943ad3a 100644 --- a/README.md +++ b/README.md @@ -8,12 +8,34 @@ This pipeline is designed to automate the conversion of dense, single-sample GTC ## Pipeline Architecture The pipeline is composed of several sequential stages, orchestrated by `cpg-flow`. Each stage is responsible for a specific part of the data processing workflow. -![Pipeline DAG](pipeline_dag.png) +```mermaid +flowchart TB + subgraph phase1 ["Phase 1 — per plate cohort"] + GtcToBcfs --> BafRegress + GtcToBcfs --> CohortBcfToPlink + end + subgraph phase2 ["Phase 2 — super cohort"] + MergeCohortPlink --> ExportCohortDatasets + MergeCohortPlink --> KingIbdseg + ExportCohortDatasets --> Plink2Qc + ExportCohortDatasets --> SnpQcReport + Plink2Qc --> QcReport + KingIbdseg --> QcReport + end + prev["Previous aggregate
(array_aggregate_pgen)"] -. Metamist .-> MergeCohortPlink + CohortBcfToPlink -. "Metamist (array_cohort_bed)" .-> MergeCohortPlink + BafRegress -. "Metamist (array_bafregress)" .-> QcReport +``` + +The dashed edges are not stage dependencies: phase 2 discovers phase-1 outputs (and the +previous aggregate) by querying registered Metamist analyses, which is what lets the two +phases run as separate submissions with the manual super-cohort creation in between. ### Stages - **GtcToBcfs**: Converts raw GTC files into two BCF formats: a "Heavy" BCF containing full intensity data and a "Light" BCF containing only genotype calls (GT) and quality scores (GQ). - **BafRegress**: Estimates sample contamination by analyzing B-Allele Frequencies (BAF) against a population reference. If no reference is provided, it will estimate AF from the cohort. Output is written to durable, version-independent storage and registered as an `array_bafregress` Metamist analysis (one per plate cohort). - **CohortBcfToPlink**: Converts the Light BCF into PLINK 1.9 binary format (`.bed`, `.bim`, `.fam`), preparing it for merging. Output is written to durable, version-independent storage and registered as an `array_cohort_bed` Metamist analysis (one per plate cohort). See [Per-plate outputs are immutable](#per-plate-outputs-are-immutable). +- **SubmitPhase2**: The final phase-1 stage. Creates the super cohort in Metamist (previous aggregate SGs ∪ this run's plate SGs) and submits `second_workflow` against it via analysis-runner. See [Rolling aggregate & two-phase run](#rolling-aggregate--two-phase-run). - **MergeCohortPlink**: Merges PLINK files from multiple cohorts into a single, unified dataset. This stage also supports a "rolling aggregate" workflow, where new samples are added to a previously generated aggregate. See [Rolling aggregate & two-phase run](#rolling-aggregate--two-phase-run). - **ExportCohortDatasets**: Converts the merged PLINK 1.9 dataset into PLINK2 (`.pgen`) format for long-term storage and analysis, and `.bcf` format in temporary storage for ancestry analysis. - **Plink2Qc**: Performs a standard suite of quality control checks on the final PLINK2 dataset, including sample/variant missingness, allele frequency, HWE, heterozygosity, and kinship. @@ -36,21 +58,52 @@ new cohort rather than overwriting an existing one. ### Rolling aggregate & two-phase run Aggregate datasets are registered against a **super cohort** (previous aggregate SGs + new plate SGs) so downstream consumers (e.g. the genomic atlas) can query array data by cohort. -cpg-flow cannot create or validate a cohort mid-run, so the pipeline runs in two phases: - -1. **Phase 1** — run against the **new plate cohorts** (`input_cohorts=[new plates]`). Produces - and registers the per-plate `array_cohort_bed` and `array_bafregress` outputs. -2. **Create the super cohort** manually in Swagger (previous aggregate SGs ∪ new plate SGs). -3. **Phase 2** — run against the **super cohort** (`input_cohorts=[super]`). Rolls the previous - aggregate forward and merges only the new plates. +cpg-flow cannot create or validate a cohort mid-run — a cohort must exist before the DAG is +built — so the pipeline runs as two chained analysis-runner runs with separate entry points: + +1. **Phase 1 (`first_workflow`)** — run against the **new plate cohorts** + (`input_cohorts=[new plates]`, `config_phase1.toml`). Produces and registers the per-plate + `array_cohort_bed` and `array_bafregress` outputs. The final `SubmitPhase2` stage then, in + a batch job: creates the super cohort (previous aggregate SGs ∪ this run's plate SGs, + reusing an existing cohort with identical membership rather than duplicating it) and + submits phase 2 against it. The hand-off works because the cohort is created at batch + runtime, before the phase-2 driver builds its DAG. The job POSTs to the analysis-runner + server directly and fails loudly on any HTTP error (the `run_analysis_runner` helper + swallows them). +2. **Phase 2 (`second_workflow`)** — runs against the **super cohort** + (`input_cohorts=[super]`, set automatically by `SubmitPhase2`; `config_phase2.toml` for a + manual run). Rolls the previous aggregate forward and merges only the new plates. + +Every stage is a `CohortStage`, so nothing in the stage graph itself separates the phases: a +phase-2 submission would otherwise also run the per-plate stages on the super cohort, and a +phase-1 submission would run the aggregate stages once per plate. The split entry points are +what pin each submission to its phase's stages. Each entry point additionally rejects a +`workflow.only_stages` selection naming stages outside its phase (cpg-flow skips stages by +exact name match, so a typo or other-phase name would be silently skipped), and +`second_workflow` refuses to run against anything other than exactly one cohort (the super +cohort) — a mismatched config fails at submission (in the driver job's log), before any job +is queued. + +To accumulate plates across several phase-1 runs before a single aggregation, set +`workflow.last_stages = ['BafRegress', 'CohortBcfToPlink']` on the early runs so only the +final one hands off to phase 2. Phase 2 can also be launched manually against a hand-made +super cohort. + +The hand-off is guarded on both sides: cohort creation fails if the created cohort is +missing any requested SG (rather than shipping a quietly smaller cohort), and the +submission record is written *before* the phase-2 submission so a re-run of phase 1 can +never submit phase 2 twice (the job also refuses at runtime if the record already exists). +If the submission itself fails, delete the sentinel TOML at +`/popgen_genotyping/SubmitPhase2//_phase2_submitted.toml` +and re-run phase 1 — the created cohort is found by membership and reused. Keep +`check_expected_outputs = true`: the existing sentinel is what skips the stage on re-run. The previous aggregate is selected explicitly by **cohort ID** (`previous_aggregate_cohort_id`). The new plates are **not listed** in phase-2 config — they are **derived** (`NEW = super − previous aggregate`), and each new SG is resolved to its plate via the registered `array_cohort_bed` analysis. The resolved plan (the contributing plate cohorts and their new-SG -counts) is printed to the driver log at submission: **confirm the plates match what you ran in -phase 1** before letting the run proceed. Use `scripts/list_aggregates.py` to pick the previous -aggregate cohort. +counts) is printed to the phase-2 driver log: check the plates match what you ran in phase 1. +Use `scripts/list_aggregates.py` to pick the previous aggregate cohort. **Why derive the new plates instead of reusing the phase-1 plate list?** It might look simpler to just carry the plate cohorts forward from phase 1, but deriving from the super cohort's membership @@ -62,17 +115,23 @@ keeps that cohort the single source of truth and is robust to things a hand-carr - **Custom / partial selection** — the super cohort can be any hand-picked SG set (a subset of a plate, or spanning plates); per-SG resolution handles this, whereas a plate list cannot. - **No config drift** — the merged output cannot disagree with the cohort it is registered against; - a `super_cohort ⊆ merged .psam` assert (PR 3b) plus the per-SG coverage check catch any plate + the post-`--keep` kept-sample-count assert plus the per-SG coverage check catch any plate never run through phase 1, failing loudly instead of producing a silently short aggregate. -The plan printout is what makes the derivation trustworthy: you get to eyeball the derived plate -set against your phase-1 runs rather than trusting it blind. +With the automatic hand-off the plan printout is a post-hoc audit rather than a gate; the hard +guarantee is the merge-time membership assert below, which fails the run outright on any +disagreement. -*(Deferred to PR 3a / PR 3b: the two-phase stage conversion, the Metamist-query-based -resolution of plate/BafRegress inputs, the cohort-ID selector, the final `--keep` to super-cohort -membership, and the `super_cohort ⊆ merged .psam` reconciliation assert. The `array_cohort_bed` -registration added now is currently write-only; consumption remains via cpg-flow stage-wiring -until then.)* +The final `plink --keep` trims the merged fileset to super-cohort membership (`merged ⊆ super`) and +asserts the kept-sample count equals the super cohort (`super ⊆ merged`) before the aggregate is +registered, so the released dataset cannot silently disagree with the cohort it registers against. + +All phase-2 output filenames embed the super-cohort ID (e.g. `_merged.bed`, +`.pgen`). Every rolling aggregate gets a new super cohort, so successive +aggregates at the same `workflow.version` land on distinct paths — cpg-flow's skip-if-exists +can therefore never reuse a previous super cohort's merge for a new one. The filenames carry +no datestamp: paths are stable across days, so an interrupted phase 2 can be resumed (or a +single stage re-run with `only_stages`) later without recomputing everything upstream. ## Prerequisites Before running the pipeline, ensure you have the following tools installed and configured: @@ -81,12 +140,19 @@ Before running the pipeline, ensure you have the following tools installed and c - **Docker**: Required for running the local reproduction scripts. ## Configuration -The pipeline is configured using a TOML file (e.g., `config.toml`). A template is provided in `src/popgen_genotyping/config_template.toml`. +The pipeline is configured using a TOML file, one per phase: start from +`src/popgen_genotyping/config_phase1.toml` (per-plate processing) or +`src/popgen_genotyping/config_phase2.toml` (aggregation against the super cohort). ### Key Parameters - `[workflow]`: - `dataset`: The analysis dataset for the output. - - `input_cohorts`: A list of cohort IDs to include in the run. + - `input_cohorts`: A list of cohort IDs to include in the run — the new plate cohorts in + phase 1, exactly the super cohort in phase 2. + - `only_stages` (optional): A within-phase subset to re-run (e.g. just `QcReport`). + The entry point pins the phase's stage list; a selection naming stages outside the + phase is rejected (see + [Rolling aggregate & two-phase run](#rolling-aggregate--two-phase-run)). - `sequencing_type`: Must be set to `array`. - `driver_image`: The Docker image for the main `cpg-flow` driver. - `bcftools_image`, `plink_image`, `king_image`: Docker images for the respective tools. @@ -96,19 +162,45 @@ The pipeline is configured using a TOML file (e.g., `config.toml`). A template i - `egt_cluster_path`: Path to the Illumina EGT cluster file. - `af_ref_path` (optional): Path to a VCF containing population allele frequencies for `BafRegress`. - `[popgen_genotyping.merge_cohort_plink]`: - - `previous_aggregate_cohort_id` (optional): The Metamist **cohort ID** of a previous aggregate to roll forward. Omit for a from-scratch (bootstrap) build. Use `scripts/list_aggregates.py` to list registered aggregate cohorts and pick one. See [Rolling aggregate & two-phase run](#rolling-aggregate--two-phase-run). *(Deferred to PR 3b: the current code still selects the previous aggregate by analysis ID via `merge_cohort_plink.previous_analysis_id`; the cohort-ID switch lands with the phase-2 conversion.)* + - `previous_aggregate_cohort_id` (required): The Metamist **cohort ID** of a previous aggregate to roll forward, or the literal `'bootstrap'` to declare a from-scratch build. There is no default: a forgotten entry fails the run rather than silently building a new-plates-only aggregate. Used by `SubmitPhase2` (super-cohort membership) and `MergeCohortPlink` (carried aggregate), so the two phases cannot drift. Use `scripts/list_aggregates.py` to list registered aggregate cohorts and pick one. See [Rolling aggregate & two-phase run](#rolling-aggregate--two-phase-run). +- `[popgen_genotyping.submit_phase2]`: + - `super_cohort_name` (required for phase 1): Name for the super cohort `SubmitPhase2` creates. Must not collide with an existing cohort name; ignored when a cohort with identical membership already exists (it is reused). ## Execution -To run the pipeline, use the `analysis-runner` command. You will need to specify the path to your configuration file, the output directory, and the script to execute. +Launch phase 1 with the `analysis-runner` command against this repo's image (the phase-2 run +is submitted automatically): + +From the repo root: ```bash -analysis-runner - --dataset - --output-dir - --config config.toml - run_workflow.py +analysis-runner \ + --skip-repo-checkout \ + --image australia-southeast1-docker.pkg.dev/cpg-common/images/popgen_genotyping:0.1.0-28 \ + --dataset \ + --access-level full \ + --output-dir \ + --config src/popgen_genotyping/config_phase1.toml \ + --description 'popgen genotyping phase 1' \ + first_workflow ``` +The image must match `workflow.driver_image` in the config; nothing validates this, so be +precise about which value governs what: the phase-1 driver runs in the CLI `--image`, while +the `SubmitPhase2` job and the whole of phase 2 use `workflow.driver_image` from the config — +if they drift, the two phases run different code. Pin an exact tag, never `:latest`: phase 2 +resolves the image string at its own start, so a floating tag can also run the two phases on +different code. CI builds a new `-` tag on every merge to main; list them with: + +```bash +gcloud artifacts docker tags list australia-southeast1-docker.pkg.dev/cpg-common/images/popgen_genotyping +``` + +To run phase 2 manually against an existing super cohort, swap in `config_phase2.toml` with +`workflow.input_cohorts = []` (and a matching description) and substitute +`second_workflow` above. Note the submission-time checks and the merge-plan printout run on +the **driver job**, after `analysis-runner` has already returned — check the driver batch's +log for the plan (phase 2) or for the ValueError if the config was rejected. + ## Local Development & Testing This repository includes scripts for local development and testing. diff --git a/pipeline_dag.mmd b/pipeline_dag.mmd index 7332a83..e93b740 100644 --- a/pipeline_dag.mmd +++ b/pipeline_dag.mmd @@ -6,18 +6,29 @@ graph TD L3[Conditional Job]:::conditional end - %% DAG Structure - GTC[Raw GTC Files]:::data --> GtoB[GtcToBcfs
GTC to Heavy/Light BCF]:::cohort - GtoB --> BAF[BafRegress
Contamination Estimation]:::cohort - GtoB --> BtoP[CohortBcfToPlink
BCF to PLINK 1.9]:::cohort - BtoP --> MCP[MergeCohortPlink
Merge Samples + Rolling Aggregate]:::multicohort - P2P1[Plink2ToPlink1
PLINK2 to PLINK1.9 if needed]:::conditional --> MCP - MCP --> ECD[ExportCohortDatasets
Export to PLINK2/BCF]:::multicohort - MCP --> KING[KingIbdseg
KING --ibdseg Relatedness]:::multicohort - ECD --> QC[Plink2Qc
Per-sample QC]:::multicohort - ECD --> SNP[SnpQcReport
EGT + Call-Rate + HWE Filter]:::multicohort - QC --> QR[QcReport
Combined QC CSV]:::multicohort - BAF --> QR + %% Phase 1: per plate cohort (first_workflow) + subgraph Phase1 [Phase 1 - first_workflow, per plate cohort] + GTC[Raw GTC Files]:::data --> GtoB[GtcToBcfs
GTC to Heavy/Light BCF]:::cohort + GtoB --> BAF[BafRegress
Contamination Estimation]:::cohort + GtoB --> BtoP[CohortBcfToPlink
BCF to PLINK 1.9]:::cohort + BAF --> SUB[SubmitPhase2
Create Super Cohort + Submit Phase 2]:::multicohort + BtoP --> SUB + end + + %% Phase 2 runs as a separate analysis-runner submission against the super cohort; + %% MergeCohortPlink resolves its inputs from Metamist, not stage wiring. + SUB -. analysis-runner submission .-> MCP + + subgraph Phase2 [Phase 2 - second_workflow, super cohort] + MCP[MergeCohortPlink
Merge Plates + Rolling Aggregate]:::cohort + P2P1[Plink2ToPlink1
PLINK2 to PLINK1.9 if needed]:::conditional --> MCP + MCP --> ECD[ExportCohortDatasets
Export to PLINK2/BCF]:::cohort + MCP --> KING[KingIbdseg
KING --ibdseg Relatedness]:::cohort + ECD --> QC[Plink2Qc
Per-sample QC]:::cohort + ECD --> SNP[SnpQcReport
EGT + Call-Rate + HWE Filter]:::cohort + QC --> QR[QcReport
Combined QC CSV]:::cohort + KING --> QR + end %% Component styling classDef cohort fill:#d1ecf1,stroke:#007bff,stroke-width:2px; @@ -26,6 +37,6 @@ graph TD classDef data fill:#fff,stroke:#333,stroke-dasharray: 5 5; %% Explicit classes - class GtoB,BAF,BtoP cohort; - class MCP,ECD,QC,KING,QR,SNP multicohort; + class GtoB,BAF,BtoP,MCP,ECD,QC,KING,QR,SNP cohort; + class SUB multicohort; class P2P1 conditional; diff --git a/pipeline_dag.png b/pipeline_dag.png deleted file mode 100644 index fbf972a..0000000 Binary files a/pipeline_dag.png and /dev/null differ diff --git a/pyproject.toml b/pyproject.toml index f3d7449..830c1db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,10 +23,16 @@ classifiers=[ # TODO Audit these dependencies later, lifted from # Harper's temp_large_cohort_migration dependencies=[ + # SubmitPhase2 submits phase 2 by POSTing to the server endpoint from analysis_runner.util + 'analysis-runner>=3.3.0', 'cpg-flow>=v1.0.1', 'loguru', 'pandas>=2.0', 'pydantic>=2.12', + # used directly by the SubmitPhase2 job; declared so they cannot vanish with a + # transitive-dependency change + 'requests', + 'toml', ] [project.urls] @@ -44,8 +50,10 @@ test = [ ] [project.scripts] -# the workflow runner script - entrypoint for the pipeline -run_workflow = 'popgen_genotyping.run_workflow:cli_main' +# phase 1: per-plate stages, then super-cohort creation + phase-2 submission +first_workflow = 'popgen_genotyping.first_workflow:cli_main' +# phase 2: rolling merge, export and QC against the super cohort +second_workflow = 'popgen_genotyping.second_workflow:cli_main' [tool.hatch.build.targets.wheel] packages = ["src/popgen_genotyping"] diff --git a/src/popgen_genotyping/config_phase1.toml b/src/popgen_genotyping/config_phase1.toml new file mode 100644 index 0000000..7d815d8 --- /dev/null +++ b/src/popgen_genotyping/config_phase1.toml @@ -0,0 +1,77 @@ +# Phase 1: per-plate processing. Run against the new plate cohorts (first_workflow). +# Produces and registers the per-plate `array_cohort_bed` and `array_bafregress` +# outputs that phase 2 later discovers via Metamist, then the SubmitPhase2 stage +# creates the super cohort and submits phase 2 automatically. See "Rolling +# aggregate & two-phase run" in the README. +[workflow] +name = 'popgen_genotyping' + +# mandatory fields for cpg-flow, not populated by analysis-runner +dataset = 'ourdna' +# The new plate cohorts to process (one or more). +input_cohorts = ['COH101', 'COH102'] +sequencing_type = 'array' + +# The first_workflow entry point pins this run to the phase-1 stages; only_stages is +# only needed for a within-phase subset. To accumulate plates across several phase-1 +# runs without the automatic phase-2 hand-off, stop before SubmitPhase2 with: +# last_stages = ['BafRegress', 'CohortBcfToPlink'] + +# used to make sure we don't repeat previously completed stages +check_expected_outputs = true + +# the method to register outputs, can be missing - will not generate metamist analysis entries +status_reporter = 'metamist' + +# TODO(you): set to the current tag of this repo's own image. The SubmitPhase2 job runs +# in it and needs the popgen_genotyping package, and it doubles as the image the phase-2 +# run is submitted with (no repo checkout). Pin an exact tag, never ':latest': phase 2 +# resolves this string at its own start, so a floating tag can run the two phases on +# different code. CI builds a new tag on every merge to main; list them with: +# gcloud artifacts docker tags list australia-southeast1-docker.pkg.dev/cpg-common/images/popgen_genotyping +driver_image = 'australia-southeast1-docker.pkg.dev/cpg-common/images/popgen_genotyping:0.1.0-28' +bcftools_image = 'australia-southeast1-docker.pkg.dev/cpg-common/images/bcftools:1.23-1' +plink_image = 'australia-southeast1-docker.pkg.dev/cpg-common/images/plink:1.9-20250819-PLINK-2.0-20260228-1' +king_image = 'australia-southeast1-docker.pkg.dev/cpg-common/images/king:2.3.2-1' +version = 1 + +[popgen_genotyping.gtc_to_bcfs] +cpu = 2 +memory = 'standard' +storage = '50G' + +[popgen_genotyping.baf_regress] +cpu = 1 +memory = 'standard' +storage = '10G' + +[popgen_genotyping.cohort_bcf_to_plink] +cpu = 8 +memory = 'highmem' +storage = '50G' + +[popgen_genotyping.merge_cohort_plink] +# TODO(you): uncomment and set. Required: the Metamist cohort ID of the previous +# aggregate to roll forward, or the literal 'bootstrap' to declare a from-scratch build. +# There is no default, so a forgotten entry fails the run instead of silently building a +# new-plates-only aggregate. Use scripts/list_aggregates.py to pick a cohort ID. +# The SubmitPhase2 stage reads the same key when computing the super-cohort membership, +# so the two phases cannot drift. +# previous_aggregate_cohort_id = 'COH123' + +[popgen_genotyping.submit_phase2] +# TODO(you): uncomment and set. Required for phase 1 (first_workflow): name for the super +# cohort the SubmitPhase2 stage creates (previous aggregate SGs + this run's plate SGs). +# Must not collide with an existing cohort name; if a cohort with identical membership +# already exists it is reused and this name is ignored. +# super_cohort_name = 'array-aggregate-2026-08-10' + +[popgen_genotyping.references] +fasta_ref_path = 'gs://cpg-common-main/references/hg38/v0/Homo_sapiens_assembly38.fasta' +# Illumina chip manifest + cluster file. These match the GDA chip shipped under +# gs://cpg-common-main/references/illumina_microarray/; override when running +# against a different chip / cluster-file revision. +bpm_manifest_path = 'gs://cpg-common-main/references/illumina_microarray/GDA-8v1-0_D2.bpm' +egt_cluster_path = 'gs://cpg-common-main/references/illumina_microarray/GDA-8v1-0_D1_ClusterFile.egt' +# Optional: population allele frequency reference for BAFRegress +# af_ref_path = 'gs://...' diff --git a/src/popgen_genotyping/config_template.toml b/src/popgen_genotyping/config_phase2.toml similarity index 58% rename from src/popgen_genotyping/config_template.toml rename to src/popgen_genotyping/config_phase2.toml index fc3e353..2f59e4b 100644 --- a/src/popgen_genotyping/config_template.toml +++ b/src/popgen_genotyping/config_phase2.toml @@ -1,48 +1,47 @@ -# this is a template for the config file -# demonstrating some of the common options used in cpg-flow +# Phase 2: aggregation. Run against the super cohort (previous aggregate SGs + +# new plate SGs) — normally created and submitted automatically by phase 1's +# SubmitPhase2 stage; use this config for a manual phase-2 run (second_workflow) +# against an existing super cohort. Rolls the previous aggregate forward, merges +# the new plates, and registers the aggregate and QC outputs against the super +# cohort. See "Rolling aggregate & two-phase run" in the README. [workflow] - -# TODO(you): update this to reflect the workflow's name name = 'popgen_genotyping' # mandatory fields for cpg-flow, not populated by analysis-runner dataset = 'ourdna' -input_cohorts = [] +# Exactly one cohort: the manually created super cohort. +input_cohorts = ['COH200'] sequencing_type = 'array' +# The second_workflow entry point pins this run to the phase-2 stages; set +# only_stages only for a within-phase subset (e.g. re-running just the QC report). + # used to make sure we don't repeat previously completed stages check_expected_outputs = true # the method to register outputs, can be missing - will not generate metamist analysis entries status_reporter = 'metamist' -driver_image = 'australia-southeast1-docker.pkg.dev/cpg-common/images/cpg_flow:latest' +# TODO(you): set to the current tag of this repo's own image — the phase-2 driver runs +# in it and needs the popgen_genotyping package. An automatic hand-off forwards phase 1's +# driver_image instead of reading this file; pin an exact tag, never ':latest'. CI builds +# a new tag on every merge to main; list them with: +# gcloud artifacts docker tags list australia-southeast1-docker.pkg.dev/cpg-common/images/popgen_genotyping +driver_image = 'australia-southeast1-docker.pkg.dev/cpg-common/images/popgen_genotyping:0.1.0-28' bcftools_image = 'australia-southeast1-docker.pkg.dev/cpg-common/images/bcftools:1.23-1' plink_image = 'australia-southeast1-docker.pkg.dev/cpg-common/images/plink:1.9-20250819-PLINK-2.0-20260228-1' king_image = 'australia-southeast1-docker.pkg.dev/cpg-common/images/king:2.3.2-1' version = 1 -[popgen_genotyping.gtc_to_bcfs] -cpu = 2 -memory = 'standard' -storage = '50G' - -[popgen_genotyping.baf_regress] -cpu = 1 -memory = 'standard' -storage = '10G' - -[popgen_genotyping.cohort_bcf_to_plink] -cpu = 8 -memory = 'highmem' -storage = '50G' - [popgen_genotyping.merge_cohort_plink] cpu = 4 memory = 'highmem' storage = '100G' -# The Metamist analysis ID of the previous multi-cohort aggregate -# previous_analysis_id = 123456 +# TODO(you): uncomment and set. Required: the Metamist cohort ID of the previous +# aggregate to roll forward, or the literal 'bootstrap' to declare a from-scratch build. +# There is no default, so a forgotten entry fails the run instead of silently building a +# new-plates-only aggregate. Use scripts/list_aggregates.py to pick a cohort ID. +# previous_aggregate_cohort_id = 'COH123' [popgen_genotyping.export_cohort_datasets] cpu = 2 @@ -96,10 +95,3 @@ hwe_keep_fewhet = true cpu = 2 memory = 'standard' storage = '20G' - -[popgen_genotyping.references] -fasta_ref_path = 'gs://cpg-common-main/references/hg38/v0/Homo_sapiens_assembly38.fasta' -bpm_manifest_path = 'gs://...' -egt_cluster_path = 'gs://...' -# Optional: population allele frequency reference for BAFRegress -# af_ref_path = 'gs://...' diff --git a/src/popgen_genotyping/first_workflow.py b/src/popgen_genotyping/first_workflow.py new file mode 100755 index 0000000..0e9bcdc --- /dev/null +++ b/src/popgen_genotyping/first_workflow.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python3 + +""" +Phase-1 entry point: per-plate processing, then super-cohort creation and phase-2 submission. + +Run against the new plate cohorts (``input_cohorts = [new plates]``). Registers the durable +per-plate outputs, then the final stage creates the super cohort and submits +``second_workflow`` against it — no manual Swagger step or second launch needed. To stop +before the automatic hand-off (e.g. accumulating plates across several phase-1 runs), set +``workflow.last_stages = ['BafRegress', 'CohortBcfToPlink']``. +""" + +from argparse import ArgumentParser + +from cpg_flow.workflow import run_workflow +from cpg_utils.config import config_retrieve + +from popgen_genotyping.stages import ( + BafRegress, + CohortBcfToPlink, + GtcToBcfs, + SubmitPhase2, +) +from popgen_genotyping.utils import validate_only_stages + +# The per-plate stages this entry point submits, plus the phase-2 hand-off. The +# aggregate stages live in second_workflow.py — see the README. +PHASE_1_STAGES: list = [GtcToBcfs, BafRegress, CohortBcfToPlink, SubmitPhase2] + + +def cli_main() -> None: + """ + Command line entry point for phase 1 of the genotyping pipeline. + """ + parser = ArgumentParser(description='Genotyping microarray pipeline: phase 1 (per-plate)') + parser.add_argument('--dry_run', action='store_true', help='Dry run') + args = parser.parse_args() + + validate_only_stages( + only_stages=config_retrieve(['workflow', 'only_stages'], default=[]), + phase_stages=PHASE_1_STAGES, + entry_point='first_workflow (phase 1)', + ) + + # The workflow name is derived from the package name + workflow_name: str = __package__ or 'popgen_genotyping' + run_workflow(name=workflow_name, stages=PHASE_1_STAGES, dry_run=args.dry_run) + + +if __name__ == '__main__': + cli_main() diff --git a/src/popgen_genotyping/jobs/merge_cohort_plink_job.py b/src/popgen_genotyping/jobs/merge_cohort_plink_job.py index 817d193..bb37c0a 100644 --- a/src/popgen_genotyping/jobs/merge_cohort_plink_job.py +++ b/src/popgen_genotyping/jobs/merge_cohort_plink_job.py @@ -20,38 +20,39 @@ def run_merge_plink( cohort_plink_paths: list[dict[str, str]], output_prefix: str, + keep_samples: list[str], previous_aggregate_resource: ResourceGroup | None = None, samples_to_remove: list[str] | None = None, - keep_samples: list[str] | None = None, job_name: str = 'merge_cohort_plink', ) -> BashJob: """ Merge multiple PLINK 1.9 datasets into a single unified dataset, with rolling aggregate support. + Whole plate filesets are merged into an intermediate fileset (a plate may carry + withdrawn/excluded SGs), then a final ``plink --keep`` pass trims the merged result to + exactly ``keep_samples`` (the super-cohort membership). + Args: cohort_plink_paths (list[dict[str, str]]): List of dicts, each with 'bed', 'bim', 'fam' cloud paths. output_prefix (str): Cloud prefix for the final merged PLINK 1.9 files. + keep_samples (list[str]): SG IDs to retain in the final fileset (the super cohort). The merged + result is trimmed to exactly this membership. previous_aggregate_resource (ResourceGroup, optional): Resource group for the previous rolling aggregate. samples_to_remove (list[str], optional): List of SG IDs to remove from the previous aggregate. - keep_samples (list[str], optional): SG IDs to retain in the final fileset. Whole plate filesets are - merged (a plate may carry withdrawn/excluded SGs), so a final ``plink --keep`` pass trims the - merged result to exactly this membership (the super cohort). When None (the default) no trim is - applied and the merged fileset is written as-is — the pre-two-phase behaviour. job_name (str): Name for the Hail Batch job. Returns: Job: A Hail Batch job object. Raises: - ValueError: If ``keep_samples`` is an empty list — distinct from ``None`` (the - documented no-trim opt-out), an empty list means a caller/config bug and would - otherwise silently skip the trim and write the untrimmed merge as the aggregate. + ValueError: If ``keep_samples`` is empty — trimming the merge to the super cohort is + mandatory, so an empty membership means a caller/config bug and would otherwise + write the untrimmed merge as the aggregate. """ - # Distinguish None (opt out of trimming) from [] (a caller bug). The truthiness checks - # below would treat both the same and skip the trim — the exact outcome the trim exists - # to prevent — so reject an empty list up front, before building the batch. - if keep_samples is not None and not keep_samples: - raise ValueError('keep_samples was provided but empty — refusing to write an untrimmed fileset') + # Trimming to the super cohort is mandatory; an empty membership is a caller/config bug + # that would otherwise write the untrimmed merge as the aggregate. Reject it up front. + if not keep_samples: + raise ValueError('keep_samples is empty — refusing to write an untrimmed fileset') b = get_batch() j = register_job( @@ -104,9 +105,8 @@ def run_merge_plink( resource = b.read_input_group(bed=paths['bed'], bim=paths['bim'], fam=paths['fam']) staged_prefixes.append(str(resource)) - # 3. Define output resource groups. When keep_samples is set the merge lands in an - # intermediate fileset and a final --keep pass (step 5) trims it to super-cohort - # membership; otherwise the merge writes straight to the final output. + # 3. Define output resource groups. The merge lands in an intermediate fileset and a + # final --keep pass (step 5) trims it to super-cohort membership. j.declare_resource_group( output_plink={ 'bed': '{root}.bed', @@ -114,17 +114,14 @@ def run_merge_plink( 'fam': '{root}.fam', } ) - if keep_samples: - j.declare_resource_group( - merged_untrimmed={ - 'bed': '{root}.bed', - 'bim': '{root}.bim', - 'fam': '{root}.fam', - } - ) - merge_target = j.merged_untrimmed - else: - merge_target = j.output_plink + j.declare_resource_group( + merged_untrimmed={ + 'bed': '{root}.bed', + 'bim': '{root}.bim', + 'fam': '{root}.fam', + } + ) + merge_target = j.merged_untrimmed # 4. Construct merge list and execute # Note: PLINK 1.9 --merge-list expects prefixes of datasets to merge @@ -168,32 +165,37 @@ def run_merge_plink( # merged (a plate may carry withdrawn/excluded SGs), so this final --keep is what # guarantees merged ⊆ super cohort. Uses the same FID=0 / IID convention as the # --remove list above. - if keep_samples: - # --keep matches on set membership, so a duplicate ID writes a duplicate line but still - # yields one .fam row — counting the raw list would fail the check below for no real - # reason. Dedupe once so the written file and the expected count cannot disagree, sorted - # so the file is reproducible when the caller passes an unordered collection. - unique_keep_samples = sorted(set(keep_samples)) - keep_list_path = f'{output_prefix}_samples_to_keep.txt' - to_path(keep_list_path).write_text('\n'.join([f'0\t{s}' for s in unique_keep_samples])) - keep_samples_resource = b.read_input(keep_list_path) - - # --keep-allele-order: preserve the A1=ALT/A2=REF orientation through the trim - # (PLINK 1.9 otherwise resets A1 to the minor allele). - j.command( - f""" - set -ex - plink --bfile {merge_target} --allow-extra-chr --output-chr chrM \\ - --keep {keep_samples_resource} \\ - --keep-allele-order --make-bed --out {j.output_plink} - - kept=$(wc -l < {j.output_plink.fam}) - if [ "$kept" -ne {len(unique_keep_samples)} ]; then - echo "expected {len(unique_keep_samples)} samples after --keep, got $kept" >&2 - exit 1 - fi - """ - ) + # --keep matches on set membership, so a duplicate ID writes a duplicate line but still + # yields one .fam row — counting the raw list would fail the check below for no real + # reason. Dedupe once so the written file and the expected count cannot disagree, sorted + # so the file is reproducible when the caller passes an unordered collection. + unique_keep_samples = sorted(set(keep_samples)) + keep_list_path = f'{output_prefix}_samples_to_keep.txt' + to_path(keep_list_path).write_text('\n'.join([f'0\t{s}' for s in unique_keep_samples])) + keep_samples_resource = b.read_input(keep_list_path) + + # --keep-allele-order: preserve the A1=ALT/A2=REF orientation through the trim + # (PLINK 1.9 otherwise resets A1 to the minor allele). + # + # Membership assert: plink --keep silently drops keep-list IIDs absent from the merged + # fileset, so the trim yields merged ∩ keep_samples, not keep_samples — a plate whose + # array_cohort_bed claims an SG its .fam does not contain would vanish with exit 0. The + # trimmed set is already a subset of keep_samples, so an equal count proves equality; + # fail loudly otherwise (guards super ⊆ merged). + j.command( + f""" + set -ex + plink --bfile {merge_target} --allow-extra-chr --output-chr chrM \\ + --keep {keep_samples_resource} \\ + --keep-allele-order --make-bed --out {j.output_plink} + + kept=$(wc -l < {j.output_plink.fam}) + if [ "$kept" -ne {len(unique_keep_samples)} ]; then + echo "expected {len(unique_keep_samples)} samples after --keep, got $kept" >&2 + exit 1 + fi + """ + ) # 6. Write outputs back to cloud b.write_output(j.output_plink, output_prefix) diff --git a/src/popgen_genotyping/jobs/submit_phase2_job.py b/src/popgen_genotyping/jobs/submit_phase2_job.py new file mode 100644 index 0000000..14be7f0 --- /dev/null +++ b/src/popgen_genotyping/jobs/submit_phase2_job.py @@ -0,0 +1,200 @@ +""" +Create the phase-2 super cohort in Metamist and submit the second-phase workflow. + +Phase 2 must run against a cohort that does not exist when phase 1 builds its DAG +(cpg-flow cannot create or validate a cohort mid-run). This job sidesteps that by doing +both steps at batch runtime, after every plate cohort has registered its outputs: it +creates the super cohort, then submits a fresh analysis-runner run whose driver sees the +cohort already existing. The submission POSTs to the analysis-runner server directly +rather than going through ``run_analysis_runner``: that helper prompts interactively for +full access and, worse, catches HTTP errors and returns normally, which would leave this +job green with phase 2 never submitted. +""" + +from typing import TYPE_CHECKING + +from cpg_utils import config, hail_batch + +if TYPE_CHECKING: + from hailtop.batch.job import PythonJob + +# Run-specific workflow config keys that must not leak into the phase-2 submission: +# the ar-guid is unique per run, and any phase-1 stage selection would mask phase-2 stages. +NON_PORTABLE_WORKFLOW_KEYS = (config.AR_GUID_NAME, 'first_stages', 'last_stages', 'only_stages', 'skip_stages') + + +def create_super_cohort_and_submit( + plate_sg_ids: list[str], + plate_cohort_ids: list[str], + previous_aggregate_cohort_id: str | None, + super_cohort_name: str, + output_path: str, +) -> str: + """ + Create (or reuse) the super cohort, then submit the phase-2 workflow via analysis-runner. + + Runs inside the driver image as a PythonJob. The super cohort membership is the union + of the phase-1 plate SGs and the previous aggregate cohort's current membership. An + existing cohort with exactly this membership is reused, so a re-run cannot register a + duplicate; a different cohort already using ``super_cohort_name`` is an error. + + Args: + plate_sg_ids (list[str]): SGs of the phase-1 plate cohorts. + plate_cohort_ids (list[str]): The phase-1 plate cohort IDs, whose durable outputs + must be registered in Metamist before phase 2 is submitted. + previous_aggregate_cohort_id (str, optional): Cohort ID of the previous aggregate, + or None for a from-scratch (bootstrap) build. + super_cohort_name (str): Name for the new super cohort. + output_path (str): Cloud path for the submission-record sentinel TOML. + + Returns: + str: The super cohort ID. + + Raises: + ValueError: If ``super_cohort_name`` is taken by a cohort with different membership. + """ + import copy # noqa: PLC0415 + + import requests # noqa: PLC0415 + import toml # noqa: PLC0415 + from analysis_runner.util import get_server_endpoint # noqa: PLC0415 + from cpg_utils import to_path # noqa: PLC0415 + from cpg_utils.cloud import get_google_identity_token # noqa: PLC0415 + from loguru import logger # noqa: PLC0415 + + from popgen_genotyping import metamist_utils # noqa: PLC0415 + + config_dict = copy.deepcopy(dict(config.get_config())) + phase1_ar_guid = config_dict['workflow'].get(config.AR_GUID_NAME) + + # 0. Refuse to double-submit. The stage-level skip only gates a later run's DAG + # construction, so a forced re-run (e.g. check_expected_outputs = false) or a + # rescheduled attempt of this job would otherwise launch a second phase-2 driver + # racing the first on the same outputs. + if to_path(output_path).exists(): + raise ValueError( + f'Phase-2 submission record already exists at {output_path}; refusing to submit again. ' + f'If the recorded submission actually failed, delete the sentinel and re-run phase 1.' + ) + + # 1. Wait until every plate cohort has registered its durable outputs: cpg-flow's + # registration jobs are not dependencies of this job (see wait_for_cohort_analyses), + # and phase 2 resolves those analyses at driver startup. + metamist_utils.wait_for_cohort_analyses( + cohort_ids=plate_cohort_ids, + analysis_types=('array_cohort_bed', 'array_bafregress'), + ) + + # 2. Resolve the target membership and check for an existing identical cohort. + cohorts = metamist_utils.query_cohorts_with_analyses() + membership = metamist_utils.resolve_super_cohort_membership( + plate_sg_ids=plate_sg_ids, + previous_aggregate_cohort_id=previous_aggregate_cohort_id, + cohorts=cohorts, + ) + + existing = metamist_utils.find_cohort_by_membership(membership, cohorts=cohorts) + if existing: + cohort_id = str(existing['id']) + logger.info(f'Reusing existing cohort {cohort_id} ({existing.get("name")}) with identical membership') + else: + name_clash = next((c for c in cohorts if c.get('name') == super_cohort_name), None) + if name_clash: + raise ValueError( + f'Cohort name {super_cohort_name!r} is already used by {name_clash.get("id")} with different ' + f'membership; choose a new super_cohort_name' + ) + description = ( + f'popgen-genotyping aggregate: {len(membership)} SGs = ' + f'previous aggregate {previous_aggregate_cohort_id or "none (bootstrap)"} + ' + f'{len(plate_sg_ids)} plate SGs (phase-1 ar-guid {phase1_ar_guid or "unknown"})' + ) + cohort_id = metamist_utils.create_custom_cohort( + name=super_cohort_name, + description=description, + sg_ids=membership, + ) + logger.info(f'Created super cohort {cohort_id} ({super_cohort_name}) with {len(membership)} SGs') + + # 3. Rewrite the run config for phase 2. + for key in NON_PORTABLE_WORKFLOW_KEYS: + config_dict['workflow'].pop(key, None) + config_dict['workflow']['input_cohorts'] = [cohort_id] + + # 4. Record the hand-off BEFORE submitting: a crash between submission and record + # would otherwise let a phase-1 re-run submit phase 2 twice, and two concurrent + # phase-2 runs race on the same outputs. The inverse failure (record written, + # submission failed) is recoverable: delete this sentinel and re-run phase 1. + record = { + 'super_cohort_id': cohort_id, + 'super_cohort_size': len(membership), + 'previous_aggregate_cohort_id': previous_aggregate_cohort_id or '', + 'phase1_ar_guid': phase1_ar_guid or '', + } + with to_path(output_path).open('w') as output_file: + output_file.write(toml.dumps(record)) + + # POST to the analysis-runner server directly and raise on any HTTP error; see the + # module docstring for why run_analysis_runner is not used. No repo/commit fields + # means no repo checkout: the driver image alone carries the code. + server_endpoint = get_server_endpoint() + response = requests.post( + server_endpoint, + json={ + 'dataset': config_dict['workflow']['dataset'], + # Phase 1's output prefix, so the analysis-runner audit record carries a + # meaningful output field (the forwarded config sets the same value anyway). + 'output': config_dict['workflow']['output_prefix'], + 'accessLevel': config_dict['workflow']['access_level'], + 'script': ['second_workflow'], + 'description': f'popgen-genotyping phase 2: aggregate cohort {cohort_id}', + 'image': config_dict['workflow']['driver_image'], + 'config': config_dict, + }, + headers={'Authorization': f'Bearer {get_google_identity_token(server_endpoint)}'}, + timeout=60, + ) + response.raise_for_status() + logger.info(f'Phase-2 submission accepted: {response.text}') + + return cohort_id + + +def run_submit_phase2( + plate_sg_ids: list[str], + plate_cohort_ids: list[str], + previous_aggregate_cohort_id: str | None, + super_cohort_name: str, + output_path: str, + job_name: str = 'SubmitPhase2', +) -> 'PythonJob': + """ + Queue the super-cohort creation + phase-2 submission as a PythonJob in the driver image. + + Args: + plate_sg_ids (list[str]): SGs of the phase-1 plate cohorts. + plate_cohort_ids (list[str]): The phase-1 plate cohort IDs. + previous_aggregate_cohort_id (str, optional): Cohort ID of the previous aggregate. + super_cohort_name (str): Name for the new super cohort. + output_path (str): Cloud path for the submission-record sentinel TOML. + job_name (str): Name for the job. + + Returns: + PythonJob: The queued job. + """ + batch = hail_batch.get_batch() + j: PythonJob = batch.new_python_job(job_name) + j.image(config.config_retrieve(['workflow', 'driver_image'])) + # Not exactly-once on spot: a preemption between the submission POST and job + # completion would rerun the function (and trip the sentinel check as a red job + # even though phase 2 was submitted), so keep this job off spot instances. + j.spot(is_spot=False) + j.call( + create_super_cohort_and_submit, + plate_sg_ids, + plate_cohort_ids, + previous_aggregate_cohort_id, + super_cohort_name, + output_path, + ) + return j diff --git a/src/popgen_genotyping/metamist_utils.py b/src/popgen_genotyping/metamist_utils.py index c6c68fd..78af7b1 100644 --- a/src/popgen_genotyping/metamist_utils.py +++ b/src/popgen_genotyping/metamist_utils.py @@ -4,12 +4,16 @@ import csv import functools +import time from collections.abc import Iterable from typing import TYPE_CHECKING, Any from cpg_utils import to_path +from loguru import logger from cpg_utils.config import config_retrieve +from metamist.apis import CohortApi from metamist.graphql import gql, query +from metamist.models import BodyCreateCohortFromCriteria, CohortBody, CohortCriteria from popgen_genotyping.utils import get_sequencing_group_cohort @@ -49,22 +53,6 @@ """ ) -# GQL to retrieve previous aggregate metadata and active SGs -QUERY_PREVIOUS_AGGREGATE = gql( - """ - query PreviousAggregateQuery($id: Int!) { - analyses(id: {eq: $id}) { - outputs - project { - sequencingGroups(activeOnly: {eq: true}) { - id - } - } - } - } - """ -) - # GQL to list every cohort in a project with its membership and registered analyses. # Analyses register against the cohort (with an empty sequencing_group_ids list), not # against individual SGs, so per-SG discovery is impossible — we enumerate cohorts and @@ -93,15 +81,15 @@ ) -def query_genotyping_manifests(project: str | None = None) -> list[dict[str, Any]]: +def metamist_project(project: str | None = None) -> str: """ - Query Metamist for genotyping manifest analyses. + Resolve the namespaced Metamist project name for the current run. Args: project (str, optional): Metamist project name. Defaults to the 'dataset' from config. Returns: - list[dict[str, Any]]: List of 'outputs' dictionaries from manifest analyses. + str: The project name, with a -test suffix at test access level. """ if project is None: project = config_retrieve(['workflow', 'dataset']) @@ -110,6 +98,21 @@ def query_genotyping_manifests(project: str | None = None) -> list[dict[str, Any if config_retrieve(['workflow', 'access_level']) == 'test' and 'test' not in project: project += '-test' + return project + + +def query_genotyping_manifests(project: str | None = None) -> list[dict[str, Any]]: + """ + Query Metamist for genotyping manifest analyses. + + Args: + project (str, optional): Metamist project name. Defaults to the 'dataset' from config. + + Returns: + list[dict[str, Any]]: List of 'outputs' dictionaries from manifest analyses. + """ + project = metamist_project(project) + # Execute the query query_result: dict[str, Any] = query(QUERY_GENOTYPING_MANIFESTS, {'project': project}) @@ -263,37 +266,6 @@ def resolve_gtc_path(sequencing_group: 'SequencingGroup') -> str: return mapping[sequencing_group.id] -def query_previous_aggregate(analysis_id: int) -> tuple[dict[str, Any], list[str]]: - """ - Query Metamist for a previous aggregate analysis and its project's active samples. - - Args: - analysis_id (int): The Metamist analysis ID. - - Returns: - tuple[dict[str, Any], list[str]]: (outputs_dict, active_sg_ids) - - Raises: - ValueError: If the analysis ID is not found. - """ - query_result: dict[str, Any] = query(QUERY_PREVIOUS_AGGREGATE, {'id': analysis_id}) - - if not query_result.get('analyses'): - raise ValueError(f'Analysis with ID {analysis_id} not found in Metamist') - - analysis: dict[str, Any] = query_result['analyses'][0] - # Outputs only contains the PGEN path, we need to reconstruct PSAM and PVAR paths based on naming convention - outputs: dict[str, Any] = {'pgen': analysis.get('outputs', {}).get('path', {})} - if not outputs['pgen'] or not outputs['pgen'].endswith('.pgen'): - raise ValueError(f'Analysis with ID {analysis_id} does not have a valid PGEN output path') - outputs.update({'psam': outputs['pgen'].replace('.pgen', '.psam')}) - outputs.update({'pvar': outputs['pgen'].replace('.pgen', '.pvar')}) - project: dict[str, Any] = analysis.get('project', {}) - active_sgs: list[str] = [sg['id'] for sg in project.get('sequencingGroups', [])] - - return outputs, active_sgs - - def query_reported_sex(project: str | None = None) -> dict[str, str]: """ Query Metamist for reported sex of participants, mapped to Sequencing Group IDs. @@ -304,12 +276,7 @@ def query_reported_sex(project: str | None = None) -> dict[str, str]: Returns: dict[str, str]: Mapping of Sequencing Group ID to reported sex. """ - if project is None: - project = config_retrieve(['workflow', 'dataset']) - - # At test access level the namespaced Metamist project carries a -test suffix. - if config_retrieve(['workflow', 'access_level']) == 'test' and 'test' not in project: - project += '-test' + project = metamist_project(project) # Execute the query query_result: dict[str, Any] = query(QUERY_REPORTED_SEX, {'project': project}) @@ -336,37 +303,6 @@ def query_reported_sex(project: str | None = None) -> dict[str, str]: return dict_samples -def resolve_rolling_aggregate(prev_analysis_id: int | str) -> tuple[dict[str, str], list[str]]: - """ - Resolve paths and sample delta for a rolling multi-cohort aggregate. - - Args: - prev_analysis_id (int | str): The Metamist analysis ID of the previous aggregate. - - Returns: - tuple[dict[str, str], list[str]]: (previous_aggregate_paths, samples_to_remove) - """ - # Import here to avoid circular dependency - from popgen_genotyping.utils import parse_psam # noqa: PLC0415 - - prev_outputs, active_sg_ids = query_previous_aggregate(int(prev_analysis_id)) - - # Expecting PLINK 2.0 PGEN/PVAR/PSAM in outputs - previous_aggregate_paths: dict[str, str] = { - 'pgen': prev_outputs['pgen'], - 'pvar': prev_outputs['pvar'], - 'psam': prev_outputs['psam'], - } - - # Parse the previous .psam to find all samples that were in the aggregate - prev_samples: list[str] = parse_psam(previous_aggregate_paths['psam']) - - # Find samples that are in the previous aggregate but no longer active - samples_to_remove: list[str] = list(set(prev_samples) - set(active_sg_ids)) - - return previous_aggregate_paths, samples_to_remove - - def _extract_output_path(outputs: Any) -> str | None: """ Extract the registered file path from an analysis ``outputs`` field. @@ -400,12 +336,7 @@ def query_cohorts_with_analyses(project: str | None = None) -> list[dict[str, An list[dict[str, Any]]: Cohort dicts, each with 'id', 'name', 'sequencingGroups' and 'analyses'. Empty list if the project has no cohorts. """ - if project is None: - project = config_retrieve(['workflow', 'dataset']) - - # At test access level the namespaced Metamist project carries a -test suffix. - if config_retrieve(['workflow', 'access_level']) == 'test' and 'test' not in project: - project += '-test' + project = metamist_project(project) query_result: dict[str, Any] = query(QUERY_COHORTS_WITH_ANALYSES, {'project': project}) @@ -413,6 +344,52 @@ def query_cohorts_with_analyses(project: str | None = None) -> list[dict[str, An return project_result.get('cohorts') or [] +def wait_for_cohort_analyses( + cohort_ids: Iterable[str], + analysis_types: Iterable[str], + project: str | None = None, + timeout_seconds: float = 900, + poll_seconds: float = 30, +) -> None: + """ + Block until every cohort has a registered analysis of every requested type. + + cpg-flow's Metamist registration jobs are not stage dependencies (the status + reporter creates them but never adds them to the stage's output jobs), so a job + that consumes registered analyses can start while registration is still in flight. + Polling closes that race; the deadline keeps a genuinely failed registration loud. + + Args: + cohort_ids (Iterable[str]): Cohorts whose analyses must exist. + analysis_types (Iterable[str]): Analysis types required on every cohort. + project (str, optional): Metamist project name. Defaults to the 'dataset' from config. + timeout_seconds (float): Deadline before giving up. Defaults to 900. + poll_seconds (float): Delay between Metamist queries. Defaults to 30. + + Raises: + TimeoutError: If any (cohort, analysis type) pair is still unregistered at the deadline. + """ + ids = list(cohort_ids) + types = list(analysis_types) + deadline = time.monotonic() + timeout_seconds + while True: + by_id = {c.get('id'): c for c in query_cohorts_with_analyses(project)} + missing = [ + (cohort_id, analysis_type) + for cohort_id in ids + for analysis_type in types + if not any(a.get('type') == analysis_type for a in ((by_id.get(cohort_id) or {}).get('analyses') or [])) + ] + if not missing: + return + if time.monotonic() >= deadline: + raise TimeoutError( + f'Timed out after {timeout_seconds}s waiting for Metamist analysis registration: {missing}' + ) + logger.info(f'Waiting for {len(missing)} Metamist analysis registration(s): {missing}') + time.sleep(poll_seconds) + + def _invert_cohort_analyses(cohorts: list[dict[str, Any]], analysis_type: str) -> dict[str, dict[str, str]]: """ Invert cohort-first records into an SG -> analysis-output map for one analysis type. @@ -660,19 +637,22 @@ def format_merge_plan(resolved: dict[str, Any], previous_aggregate_cohort_id: st Render a human-readable summary of a resolved rolling-merge plan for the driver log. Lets an operator confirm at a glance that phase 2 picked up the phase-1 plates: it - lists each contributing plate cohort with its new-SG count, and prints an expected - merged total that should equal the super cohort size. + lists each contributing plate cohort with its new-SG count. The printed totals are + all derived from the super cohort's membership, so they are informational, not a + cross-check — the expected merged total equals the super cohort size by construction. + The check that the merged data actually matches is the in-job kept-sample-count + assert after the final ``--keep``. Args: resolved (dict[str, Any]): The dict returned by :func:`resolve_merge_inputs`. previous_aggregate_cohort_id (str, optional): Cohort ID rolled forward, for the header. Returns: - str: A multi-line summary suitable for ``logging.info``. + str: A multi-line summary suitable for ``logger.info``. """ - plate_merge_list: list[dict[str, Any]] = resolved.get('plate_merge_list', []) - samples_to_remove: list[str] = resolved.get('samples_to_remove', []) - super_size: int = resolved.get('super_cohort_size', 0) + plate_merge_list: list[dict[str, Any]] = resolved['plate_merge_list'] + samples_to_remove: list[str] = resolved['samples_to_remove'] + super_size: int = resolved['super_cohort_size'] new_count: int = sum(p['new_count'] for p in plate_merge_list) carried_forward: int = super_size - new_count @@ -695,3 +675,131 @@ def format_merge_plan(resolved: dict[str, Any], previous_aggregate_cohort_id: st lines.append(f' {plate["cohort_id"]}: {plate["new_count"]} new SGs') lines.append(f' expected merged total: {expected_total} SGs') return '\n'.join(lines) + + +def resolve_super_cohort_membership( + plate_sg_ids: Iterable[str], + previous_aggregate_cohort_id: str | None, + project: str | None = None, + cohorts: list[dict[str, Any]] | None = None, +) -> list[str]: + """ + Compute the target super-cohort membership for a phase-2 run. + + The super cohort is the union of this phase-1 run's plate SGs and the previous + aggregate cohort's current membership (so SGs withdrawn from the previous aggregate + cohort since it was built are excluded automatically). + + Args: + plate_sg_ids (Iterable[str]): SGs of the phase-1 plate cohorts. + previous_aggregate_cohort_id (str, optional): Cohort ID of the previous aggregate, + or None for a from-scratch (bootstrap) build. + project (str, optional): Metamist project name. Defaults to the 'dataset' from config. + cohorts (list[dict[str, Any]], optional): Pre-fetched cohort records. + + Returns: + list[str]: Sorted super-cohort SG IDs. + + Raises: + ValueError: If the previous aggregate cohort is not found in the project, or if + every plate SG is already a member of it (nothing new to aggregate). + """ + sg_ids: set[str] = set(plate_sg_ids) + + if previous_aggregate_cohort_id: + if cohorts is None: + cohorts = query_cohorts_with_analyses(project) + agg_cohort = next((c for c in cohorts if c.get('id') == previous_aggregate_cohort_id), None) + if agg_cohort is None: + raise ValueError(f'Previous aggregate cohort {previous_aggregate_cohort_id} not found in project') + agg_sg_ids = {sg['id'] for sg in (agg_cohort.get('sequencingGroups') or []) if sg.get('id')} + # A super cohort identical to the previous aggregate would send phase 2 chasing + # an empty new-SG set; fail here with the real reason instead. + if sg_ids <= agg_sg_ids: + raise ValueError( + f'All {len(sg_ids)} plate SGs are already members of previous aggregate cohort ' + f'{previous_aggregate_cohort_id}; nothing new to aggregate' + ) + sg_ids |= agg_sg_ids + + return sorted(sg_ids) + + +def find_cohort_by_membership( + sg_ids: Iterable[str], + project: str | None = None, + cohorts: list[dict[str, Any]] | None = None, +) -> dict[str, Any] | None: + """ + Find an existing cohort whose membership is exactly ``sg_ids``. + + Used to make super-cohort creation idempotent: a re-run of phase 1 (e.g. after a + failed phase-2 submission) reuses the cohort it created last time instead of + registering a duplicate. + + Args: + sg_ids (Iterable[str]): The target membership. + project (str, optional): Metamist project name. Defaults to the 'dataset' from config. + cohorts (list[dict[str, Any]], optional): Pre-fetched cohort records. + + Returns: + dict[str, Any] | None: The matching cohort record (latest by ID if several + share the membership), or None. + """ + if cohorts is None: + cohorts = query_cohorts_with_analyses(project) + + target: set[str] = set(sg_ids) + matches = [c for c in cohorts if {sg['id'] for sg in (c.get('sequencingGroups') or []) if sg.get('id')} == target] + if not matches: + return None + # Cohort IDs are 'COH' + an unpadded integer (+ check digit), so 'latest' must + # compare numerically: lexicographic max would rank COH99 above COH100. + return max(matches, key=lambda c: int(str(c.get('id', '')).removeprefix('COH'))) + + +def create_custom_cohort(name: str, description: str, sg_ids: Iterable[str], project: str | None = None) -> str: + """ + Create a custom Metamist cohort from explicit sequencing group IDs. + + Args: + name (str): Cohort name (must not collide with an existing cohort). + description (str): Cohort description (provenance: source plates, previous aggregate). + sg_ids (Iterable[str]): The cohort membership. + project (str, optional): Metamist project name. Defaults to the 'dataset' from config. + + Returns: + str: The new cohort ID. + + Raises: + ValueError: If the created cohort is missing any requested SG (the cohort would + be silently smaller than intended), or Metamist did not return a cohort ID. + """ + project = metamist_project(project) + + requested = sorted(sg_ids) + body = BodyCreateCohortFromCriteria( + cohort_spec=CohortBody(name=name, description=description), + # sg_ids_internal must be the sole criterion: the server rejects an explicit SG + # list combined with any other criterion, including projects (metamist SET-839). + cohort_criteria=CohortCriteria(sg_ids_internal=requested), + ) + result = CohortApi().create_cohort_from_criteria(project=project, body_create_cohort_from_criteria=body) + + def _field(key: str) -> Any: + return result.get(key) if isinstance(result, dict) else getattr(result, key, None) + + cohort_id = _field('cohort_id') + if not cohort_id: + raise ValueError(f'Cohort creation for {name!r} returned no cohort ID: {result}') + + # Current Metamist raises on inactive SGs, but a cohort whose membership differs + # from the request would ship a wrong aggregate, so verify the created membership + # regardless of server version. + created = set(_field('sequencing_group_ids') or []) + missing = sorted(set(requested) - created) + if missing: + raise ValueError( + f'Cohort {name!r} ({cohort_id}) is missing {len(missing)} requested sequencing group(s): {missing}' + ) + return str(cohort_id) diff --git a/src/popgen_genotyping/run_workflow.py b/src/popgen_genotyping/run_workflow.py deleted file mode 100755 index 5c6153a..0000000 --- a/src/popgen_genotyping/run_workflow.py +++ /dev/null @@ -1,49 +0,0 @@ -#!/usr/bin/env python3 - -""" -This is the main entry point for the workflow. -""" - -from argparse import ArgumentParser - -from cpg_flow.workflow import run_workflow -from popgen_genotyping.stages import ( - BafRegress, - CohortBcfToPlink, - ExportCohortDatasets, - GtcToBcfs, - KingIbdseg, - MergeCohortPlink, - Plink2Qc, - QcReport, - SnpQcReport, -) - - -def cli_main() -> None: - """ - Command line entry point for the genotyping pipeline. - """ - parser = ArgumentParser(description='Genotyping microarray pipeline') - parser.add_argument('--dry_run', action='store_true', help='Dry run') - args = parser.parse_args() - - # The workflow name is derived from the package name - workflow_name: str = __package__ or 'popgen_genotyping' - stages: list = [ - GtcToBcfs, - BafRegress, - CohortBcfToPlink, - MergeCohortPlink, - ExportCohortDatasets, - Plink2Qc, - KingIbdseg, - SnpQcReport, - QcReport, - ] - - run_workflow(name=workflow_name, stages=stages, dry_run=args.dry_run) - - -if __name__ == '__main__': - cli_main() diff --git a/src/popgen_genotyping/second_workflow.py b/src/popgen_genotyping/second_workflow.py new file mode 100755 index 0000000..e503209 --- /dev/null +++ b/src/popgen_genotyping/second_workflow.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python3 + +""" +Phase-2 entry point: rolling merge, export and QC against the super cohort. + +Run against the super cohort (``input_cohorts = [super]``). Normally submitted +automatically by phase 1's ``SubmitPhase2`` stage; can also be launched manually against a +hand-made super cohort. +""" + +from argparse import ArgumentParser + +from cpg_flow.workflow import run_workflow +from cpg_utils.config import config_retrieve + +from popgen_genotyping.stages import ( + ExportCohortDatasets, + KingIbdseg, + MergeCohortPlink, + Plink2Qc, + QcReport, + SnpQcReport, +) +from popgen_genotyping.utils import validate_only_stages + +# The aggregate stages this entry point submits. The per-plate stages live in +# first_workflow.py: the two phases run against different cohorts (new plates vs the +# super cohort) and must not share a submission — see the README. +PHASE_2_STAGES: list = [MergeCohortPlink, ExportCohortDatasets, Plink2Qc, KingIbdseg, SnpQcReport, QcReport] + + +def validate_phase2_cohorts(input_cohorts: list[str]) -> None: + """ + Require exactly one input cohort: the super cohort. + + Every phase-2 stage is a CohortStage, so each listed cohort would otherwise get its + own run, each treating its cohort as "the super cohort" and rolling forward the same + ``previous_aggregate_cohort_id`` — registering multiple aggregates that claim the + same lineage. + + Args: + input_cohorts (list[str]): The ``workflow.input_cohorts`` config value. + + Raises: + ValueError: If anything other than exactly one cohort is listed. + """ + if len(input_cohorts) != 1: + raise ValueError( + f'Phase 2 runs against exactly one cohort (the super cohort), but ' + f'workflow.input_cohorts has {len(input_cohorts)}: {input_cohorts}. ' + 'Create the super cohort first and list only its ID.' + ) + + +def cli_main() -> None: + """ + Command line entry point for phase 2 of the genotyping pipeline. + """ + parser = ArgumentParser(description='Genotyping microarray pipeline: phase 2 (aggregate)') + parser.add_argument('--dry_run', action='store_true', help='Dry run') + args = parser.parse_args() + + validate_only_stages( + only_stages=config_retrieve(['workflow', 'only_stages'], default=[]), + phase_stages=PHASE_2_STAGES, + entry_point='second_workflow (phase 2)', + ) + validate_phase2_cohorts(input_cohorts=config_retrieve(['workflow', 'input_cohorts'], default=[])) + + # The workflow name is derived from the package name + workflow_name: str = __package__ or 'popgen_genotyping' + run_workflow(name=workflow_name, stages=PHASE_2_STAGES, dry_run=args.dry_run) + + +if __name__ == '__main__': + cli_main() diff --git a/src/popgen_genotyping/stages.py b/src/popgen_genotyping/stages.py index 297905d..2606323 100644 --- a/src/popgen_genotyping/stages.py +++ b/src/popgen_genotyping/stages.py @@ -4,11 +4,11 @@ from __future__ import annotations -from datetime import datetime, timezone from typing import TYPE_CHECKING from cpg_flow.stage import CohortStage, MultiCohortStage, stage from cpg_utils.config import config_retrieve, reference_path +from loguru import logger from popgen_genotyping.jobs.baf_regress_job import run_bafregress from popgen_genotyping.jobs.cohort_bcf_to_plink_job import run_cohort_bcf_to_plink @@ -20,18 +20,21 @@ from popgen_genotyping.jobs.plink2_to_plink1_job import run_plink2_to_plink1 from popgen_genotyping.jobs.qc_report_job import run_qc_report from popgen_genotyping.jobs.snp_qc_report_job import run_snp_qc_report +from popgen_genotyping.jobs.submit_phase2_job import run_submit_phase2 from popgen_genotyping.metamist_utils import ( + format_merge_plan, query_reported_sex, + resolve_bafregress_map, resolve_cohort_gtc_mapping, - resolve_rolling_aggregate, + resolve_merge_inputs, ) -from popgen_genotyping.utils import get_output_prefix +from popgen_genotyping.utils import get_output_prefix, get_previous_aggregate_cohort_id if TYPE_CHECKING: from cpg_flow.stage import StageInput, StageOutput from cpg_flow.targets import Cohort, MultiCohort from cpg_utils import Path - from hailtop.batch.job import BashJob + from hailtop.batch.job import BashJob, PythonJob from hailtop.batch.resource import ResourceGroup @@ -94,11 +97,9 @@ class BafRegress(CohortStage): Output is written to durable, version-independent storage and registered against the cohort as an ``array_bafregress`` analysis — computed once per plate and reused across - runs. Currently the QC report reads these via stage-wiring - (``inputs.as_path_by_target``), which only covers the current run's target cohorts. - - (Deferred to PR 3b: the QC report will instead query all ``array_bafregress`` analyses - so the final table covers every constituent cohort of the aggregate, not just new plates.) + runs. The phase-2 QC report resolves these by querying all ``array_bafregress`` analyses + (``resolve_bafregress_map``), so the final table covers every constituent cohort of the + aggregate. """ def expected_outputs(self, cohort: Cohort) -> Path: @@ -137,12 +138,9 @@ class CohortBcfToPlink(CohortStage): Convert the cohort-level light BCF to PLINK 1.9 format. Output is written to durable, version-independent storage and registered against the - cohort as an ``array_cohort_bed`` analysis. Downstream stages resolve it via - ``expected_outputs``; the fileset is processed once per plate and reused across runs. - - (Deferred to PR 3b: phase 2 will discover these filesets by querying the - ``array_cohort_bed`` analyses in Metamist instead of stage-wiring; the registration - added here is currently write-only.) + cohort as an ``array_cohort_bed`` analysis. Phase 2 discovers these filesets by querying the + ``array_cohort_bed`` analyses in Metamist (``resolve_merge_inputs``), not via stage-wiring; + the fileset is processed once per plate and reused across runs. """ def expected_outputs(self, cohort: Cohort) -> dict[str, Path]: @@ -184,119 +182,167 @@ def queue_jobs(self, cohort: Cohort, inputs: StageInput) -> StageOutput: return self.make_outputs(cohort, data=outputs, jobs=[j]) -@stage(required_stages=[CohortBcfToPlink]) -class MergeCohortPlink(MultiCohortStage): +@stage(required_stages=[BafRegress, CohortBcfToPlink]) +class SubmitPhase2(MultiCohortStage): + """ + Create the super cohort in Metamist and submit the phase-2 workflow. + + Final phase-1 stage: runs once, after the plate cohorts' compute jobs. cpg-flow's + Metamist registration jobs are not stage dependencies, so the batch job first waits + for every plate cohort's ``array_bafregress`` and ``array_cohort_bed`` analyses to be + registered (phase 2 resolves them at driver startup). It then computes the + super-cohort membership (previous aggregate cohort SGs, if configured, union this run's + plate SGs), creates the cohort — reusing an existing cohort with identical membership — + and submits ``second_workflow`` against it via analysis-runner. The sentinel output + prevents a duplicate submission on a re-run. + """ + + def expected_outputs(self, multicohort: MultiCohort) -> Path: + """ + Define the submission-record sentinel TOML path. + """ + prefix: Path = get_output_prefix(dataset=multicohort.analysis_dataset, stage_name=self.name) + return prefix / f'{multicohort.name}_phase2_submitted.toml' + + def queue_jobs(self, multicohort: MultiCohort, _inputs: StageInput) -> StageOutput: + """ + Queue the super-cohort creation + phase-2 submission job. + """ + outputs: Path = self.expected_outputs(multicohort) + + # Required: the name for the super cohort this run will create. + super_cohort_name: str = config_retrieve(['popgen_genotyping', 'submit_phase2', 'super_cohort_name']) + # Same key MergeCohortPlink reads in phase 2, so the two phases cannot drift. + previous_aggregate_cohort_id: str | None = get_previous_aggregate_cohort_id() + + j: PythonJob = run_submit_phase2( + plate_sg_ids=multicohort.get_sequencing_group_ids(), + plate_cohort_ids=[cohort.id for cohort in multicohort.get_cohorts()], + previous_aggregate_cohort_id=previous_aggregate_cohort_id, + super_cohort_name=super_cohort_name, + output_path=str(outputs), + ) + + return self.make_outputs(multicohort, data=outputs, jobs=[j]) + + +@stage +class MergeCohortPlink(CohortStage): """ - Merge all cohort PLINK 1.9 datasets into a single unified dataset, with rolling aggregate. - Output is stored in tmp. + Merge the phase-1 per-plate PLINK 1.9 datasets into the super cohort, with rolling aggregate. + + Runs in phase 2 against the manually-created super cohort. Inputs are resolved from Metamist + (``resolve_merge_inputs``), not stage-wiring — hence no ``required_stages``: the super cohort's + membership is the source of truth, the previous aggregate is carried forward from the configured + ``previous_aggregate_cohort_id``, and the contributing plates are derived as + ``NEW = super - previous aggregate`` and mapped to their registered ``array_cohort_bed`` filesets. + The resolved plan is logged for operator confirmation. Whole plate filesets are merged and a final + ``--keep`` trims the result to super-cohort membership. Output is stored in tmp. """ - def expected_outputs(self, multicohort: MultiCohort) -> dict[str, Path]: + def expected_outputs(self, cohort: Cohort) -> dict[str, Path]: """ - Define the expected multi-cohort PLINK 1.9 fileset in temporary storage. + Define the expected merged PLINK 1.9 fileset in temporary storage. """ - # Store in tmp per requirement - prefix: Path = get_output_prefix(dataset=multicohort.analysis_dataset, stage_name=self.name, tmp=True) + # Store in tmp per requirement. Keyed by the super-cohort ID: every rolling aggregate + # gets a new super cohort, so two runs at the same workflow.version land on distinct + # paths — otherwise cpg-flow's skip-if-exists would silently reuse the previous super + # cohort's merge (and the in-job --keep count assert would never fire). + prefix: Path = get_output_prefix(dataset=cohort.dataset, stage_name=self.name, tmp=True) return { - 'bed': prefix / 'merged_cohorts.bed', - 'bim': prefix / 'merged_cohorts.bim', - 'fam': prefix / 'merged_cohorts.fam', + 'bed': prefix / f'{cohort.id}_merged.bed', + 'bim': prefix / f'{cohort.id}_merged.bim', + 'fam': prefix / f'{cohort.id}_merged.fam', } - def queue_jobs(self, multicohort: MultiCohort, inputs: StageInput) -> StageOutput: + def queue_jobs(self, cohort: Cohort, _inputs: StageInput) -> StageOutput: """ - Queue the multi-cohort PLINK 1.9 merge job, incorporating rolling aggregates if configured. + Queue the PLINK 1.9 merge job for the super cohort, resolving inputs from Metamist. """ - outputs: dict[str, Path] = self.expected_outputs(multicohort) - - # 1. Gather cohort outputs - all_cohort_outputs: dict[str, dict[str, Path]] = inputs.as_dict_by_target(stage=CohortBcfToPlink) - cohort_plink_paths: list[dict[str, str]] = [] - for _cohort_id, cohort_outs in all_cohort_outputs.items(): - cohort_plink_paths.append( - { - 'bed': str(cohort_outs['bed']), - 'bim': str(cohort_outs['bim']), - 'fam': str(cohort_outs['fam']), - } - ) + outputs: dict[str, Path] = self.expected_outputs(cohort) - # 2. Check for rolling aggregate - prev_analysis_id: str | None = config_retrieve( - ['popgen_genotyping', 'merge_cohort_plink', 'previous_analysis_id'], default=None + # 1. Resolve the merge plan from Metamist. The super cohort is the source of truth; + # new plates are derived (NEW = super - previous aggregate), not listed in config. + previous_aggregate_cohort_id: str | None = get_previous_aggregate_cohort_id() + super_cohort_sg_ids: list[str] = cohort.get_sequencing_group_ids() + resolved: dict = resolve_merge_inputs( + super_cohort_sg_ids=super_cohort_sg_ids, + previous_aggregate_cohort_id=previous_aggregate_cohort_id, ) + # Log the plan so an operator can confirm the derived plates match the phase-1 runs. + # loguru, not stdlib logging: cpg-flow configures loguru, while unconfigured stdlib + # logging drops INFO — the plan would never reach the driver log. + logger.info(format_merge_plan(resolved, previous_aggregate_cohort_id)) + + cohort_plink_paths: list[dict[str, str]] = [ + {'bed': plate['bed'], 'bim': plate['bim'], 'fam': plate['fam']} for plate in resolved['plate_merge_list'] + ] + + # 2. Carry the previous aggregate forward, if any. It is stored as PLINK2, so convert + # it back to PLINK 1.9 before the merge. previous_aggregate_plink1_resource: ResourceGroup | None = None - samples_to_remove: list[str] | None = None merge_job_dependencies: list[BashJob] = [] - if prev_analysis_id: - # A previous aggregate exists in PLINK2 format, so we need to convert it to PLINK1.9 - - previous_aggregate_plink2_paths, samples_to_remove = resolve_rolling_aggregate( - prev_analysis_id=prev_analysis_id - ) - - # Define an output prefix for the converted PLINK1.9 files in tmp storage - plink1_prefix = get_output_prefix( - dataset=multicohort.analysis_dataset, stage_name='Plink2ToPlink1', tmp=True - ) + if resolved['previous_aggregate_paths']: + plink1_prefix = get_output_prefix(dataset=cohort.dataset, stage_name='Plink2ToPlink1', tmp=True) conversion_job, converted_plink1_resource = run_plink2_to_plink1( - pfile_prefix=previous_aggregate_plink2_paths, - output_prefix=str(plink1_prefix), + pfile_prefix=resolved['previous_aggregate_paths'], + output_prefix=str(plink1_prefix / f'{cohort.id}_plink1'), job_name='Plink2ToPlink1', ) merge_job_dependencies.append(conversion_job) - - # The converted PLINK1.9 resource group becomes the previous aggregate for the merge job previous_aggregate_plink1_resource = converted_plink1_resource - # 3. Call merge job + # 3. Call merge job. keep_samples trims the whole-plate merge to super-cohort membership. j: BashJob = run_merge_plink( cohort_plink_paths=cohort_plink_paths, output_prefix=str(outputs['bed']).replace('.bed', ''), + keep_samples=super_cohort_sg_ids, previous_aggregate_resource=previous_aggregate_plink1_resource, - samples_to_remove=samples_to_remove, + samples_to_remove=resolved['samples_to_remove'], job_name='MergeCohortPlink', ) if merge_job_dependencies: j.depends_on(*merge_job_dependencies) - return self.make_outputs(multicohort, data=outputs, jobs=[j]) + return self.make_outputs(cohort, data=outputs, jobs=[j]) @stage(required_stages=[MergeCohortPlink], analysis_type='array_aggregate_pgen', analysis_keys=['pgen']) -class ExportCohortDatasets(MultiCohortStage): +class ExportCohortDatasets(CohortStage): """ Export the merged cohort to PLINK2 format for long-term storage. BCF output goes to tmp for analysis, as it is too large for long term storage. """ - def expected_outputs(self, multicohort: MultiCohort) -> dict[str, Path]: + def expected_outputs(self, cohort: Cohort) -> dict[str, Path]: """ Define the expected PLINK2 outputs to long-term storage. BCF output goes to tmp for analysis, as it is too large for long term storage. """ - prefix: Path = get_output_prefix(dataset=multicohort.analysis_dataset, stage_name=self.name) - tmp_bcf_prefix: Path = get_output_prefix(dataset=multicohort.analysis_dataset, stage_name=self.name, tmp=True) - datestamp: str = datetime.now(tz=timezone.utc).strftime('%Y%m%d') + # No datestamp: the cohort ID (plus the versioned prefix) already distinguishes + # aggregates, and paths must be stable across days for skip-if-exists and + # single-stage reruns to find upstream outputs. Same for all phase-2 stages. + prefix: Path = get_output_prefix(dataset=cohort.dataset, stage_name=self.name) + tmp_bcf_prefix: Path = get_output_prefix(dataset=cohort.dataset, stage_name=self.name, tmp=True) return { - 'pgen': prefix / f'{datestamp}_cohort.pgen', - 'pvar': prefix / f'{datestamp}_cohort.pvar', - 'psam': prefix / f'{datestamp}_cohort.psam', - 'bcf': tmp_bcf_prefix / f'{datestamp}_cohort.bcf', + 'pgen': prefix / f'{cohort.id}.pgen', + 'pvar': prefix / f'{cohort.id}.pvar', + 'psam': prefix / f'{cohort.id}.psam', + 'bcf': tmp_bcf_prefix / f'{cohort.id}.bcf', } - def queue_jobs(self, multicohort: MultiCohort, inputs: StageInput) -> StageOutput: + def queue_jobs(self, cohort: Cohort, inputs: StageInput) -> StageOutput: """ Queue the dataset export job using PLINK2. """ - outputs: dict[str, Path] = self.expected_outputs(multicohort=multicohort) + outputs: dict[str, Path] = self.expected_outputs(cohort=cohort) # 1. Pull input from MergeCohortPlink - input_plink: dict[str, Path] = inputs.as_dict(target=multicohort, stage=MergeCohortPlink) + input_plink: dict[str, Path] = inputs.as_dict(target=cohort, stage=MergeCohortPlink) # 2. Call export job j: BashJob = run_export_cohort_datasets( @@ -310,24 +356,23 @@ def queue_jobs(self, multicohort: MultiCohort, inputs: StageInput) -> StageOutpu job_name='ExportCohortDatasets', ) - return self.make_outputs(multicohort, data=outputs, jobs=[j]) + return self.make_outputs(cohort, data=outputs, jobs=[j]) @stage(required_stages=[ExportCohortDatasets], analysis_type='array_qc_raw', analysis_keys=['log']) -class Plink2Qc(MultiCohortStage): +class Plink2Qc(CohortStage): """ Per-sample PLINK2 QC on the merged pgen/pvar/psam: missingness, inbreeding, sex-check, plus per-variant allele frequencies. Per-variant missingness and Hardy-Weinberg are computed inside ``SnpQcReport``. """ - def expected_outputs(self, multicohort: MultiCohort) -> dict[str, Path]: + def expected_outputs(self, cohort: Cohort) -> dict[str, Path]: """ - Define the expected PLINK2 QC output files for the multi-cohort. + Define the expected PLINK2 QC output files for the cohort. """ - prefix: Path = get_output_prefix(dataset=multicohort.analysis_dataset, stage_name=self.name) - datestamp: str = datetime.now(tz=timezone.utc).strftime('%Y%m%d') - output_base_name = f'{datestamp}_qc' + prefix: Path = get_output_prefix(dataset=cohort.dataset, stage_name=self.name) + output_base_name = f'{cohort.id}_qc' return { 'smiss': prefix / f'{output_base_name}.smiss', 'afreq': prefix / f'{output_base_name}.afreq', @@ -336,14 +381,14 @@ def expected_outputs(self, multicohort: MultiCohort) -> dict[str, Path]: 'log': prefix / f'{output_base_name}.log', } - def queue_jobs(self, multicohort: MultiCohort, inputs: StageInput) -> StageOutput: + def queue_jobs(self, cohort: Cohort, inputs: StageInput) -> StageOutput: """ - Queue the PLINK2 QC job for the multi-cohort. + Queue the PLINK2 QC job for the cohort. """ - outputs: dict[str, Path] = self.expected_outputs(multicohort=multicohort) + outputs: dict[str, Path] = self.expected_outputs(cohort=cohort) # Get the input PGEN file path from the ExportCohortDatasets stage - input_plink_pgen: Path = inputs.as_path(target=multicohort, stage=ExportCohortDatasets, key='pgen') + input_plink_pgen: Path = inputs.as_path(target=cohort, stage=ExportCohortDatasets, key='pgen') # The outputs_path for the run_plink2_qc job is the base prefix for all QC files. output_plink2_prefix = str(outputs['smiss']).removesuffix('.smiss') @@ -352,11 +397,11 @@ def queue_jobs(self, multicohort: MultiCohort, inputs: StageInput) -> StageOutpu j: BashJob = run_plink2_qc( pgen_path=str(input_plink_pgen), outputs_path=output_plink2_prefix, - job_name=f'Plink2Qc_{multicohort.name}', + job_name=f'Plink2Qc_{cohort.name}', ) # Return the expected outputs of this stage, referencing the outputs generated by the job - return self.make_outputs(multicohort, data=outputs, jobs=[j]) + return self.make_outputs(cohort, data=outputs, jobs=[j]) @stage( @@ -364,23 +409,22 @@ def queue_jobs(self, multicohort: MultiCohort, inputs: StageInput) -> StageOutpu analysis_type='array_relatedness_ibdseg', analysis_keys=['seg', 'seg_x'], ) -class KingIbdseg(MultiCohortStage): +class KingIbdseg(CohortStage): """ Infer pairwise IBD segments across the merged cohort with KING `--ibdseg`. """ - def expected_outputs(self, multicohort: MultiCohort) -> dict[str, Path]: + def expected_outputs(self, cohort: Cohort) -> dict[str, Path]: """ - Define the expected KING `--ibdseg` outputs for the multi-cohort. + Define the expected KING `--ibdseg` outputs for the cohort. KING 2.3.2 produces an X-chromosome companion (`{prefix}X.seg` / `{prefix}X.segments.gz`, no dot before the `X`) whenever the merged PLINK fileset has chrX SNPs. The job backfills header-only placeholders when the input has no chrX, so these outputs are always materialised. """ - prefix: Path = get_output_prefix(dataset=multicohort.analysis_dataset, stage_name=self.name) - datestamp: str = datetime.now(tz=timezone.utc).strftime('%Y%m%d') - output_base_name = f'{datestamp}_king' + prefix: Path = get_output_prefix(dataset=cohort.dataset, stage_name=self.name) + output_base_name = f'{cohort.id}_king' return { 'seg': prefix / f'{output_base_name}.seg', 'segments': prefix / f'{output_base_name}.segments.gz', @@ -389,13 +433,13 @@ def expected_outputs(self, multicohort: MultiCohort) -> dict[str, Path]: 'log': prefix / f'{output_base_name}.log', } - def queue_jobs(self, multicohort: MultiCohort, inputs: StageInput) -> StageOutput: + def queue_jobs(self, cohort: Cohort, inputs: StageInput) -> StageOutput: """ Queue the KING `--ibdseg` job against the merged PLINK 1.9 dataset. """ - outputs: dict[str, Path] = self.expected_outputs(multicohort=multicohort) + outputs: dict[str, Path] = self.expected_outputs(cohort=cohort) - merged_plink: dict[str, Path] = inputs.as_dict(target=multicohort, stage=MergeCohortPlink) + merged_plink: dict[str, Path] = inputs.as_dict(target=cohort, stage=MergeCohortPlink) jobs: list[BashJob] = run_king_ibdseg( bed_path=str(merged_plink['bed']), @@ -406,10 +450,10 @@ def queue_jobs(self, multicohort: MultiCohort, inputs: StageInput) -> StageOutpu output_seg_x_path=str(outputs['seg_x']), output_segments_x_path=str(outputs['segments_x']), output_log_path=str(outputs['log']), - job_name=f'KingIbdseg_{multicohort.name}', + job_name=f'KingIbdseg_{cohort.name}', ) - return self.make_outputs(multicohort, data=outputs, jobs=jobs) + return self.make_outputs(cohort, data=outputs, jobs=jobs) @stage( @@ -417,7 +461,7 @@ def queue_jobs(self, multicohort: MultiCohort, inputs: StageInput) -> StageOutpu analysis_type='array_snp_qc', analysis_keys=['inclusion_list'], ) -class SnpQcReport(MultiCohortStage): +class SnpQcReport(CohortStage): """ Per-SNP QC: vendor cluster scores + call rate + Hardy-Weinberg. @@ -431,26 +475,25 @@ class SnpQcReport(MultiCohortStage): not modified. """ - def expected_outputs(self, multicohort: MultiCohort) -> dict[str, Path]: + def expected_outputs(self, cohort: Cohort) -> dict[str, Path]: """ Define the audit TSV, inclusion list, and summary TSV outputs. """ - prefix: Path = get_output_prefix(dataset=multicohort.analysis_dataset, stage_name=self.name) - datestamp: str = datetime.now(tz=timezone.utc).strftime('%Y%m%d') + prefix: Path = get_output_prefix(dataset=cohort.dataset, stage_name=self.name) return { - 'audit_tsv': prefix / f'{datestamp}_snp_qc.audit.tsv.gz', - 'inclusion_list': prefix / f'{datestamp}_snp_qc.include.snplist', - 'summary_tsv': prefix / f'{datestamp}_snp_qc.summary.tsv', + 'audit_tsv': prefix / f'{cohort.id}_snp_qc.audit.tsv.gz', + 'inclusion_list': prefix / f'{cohort.id}_snp_qc.include.snplist', + 'summary_tsv': prefix / f'{cohort.id}_snp_qc.summary.tsv', } - def queue_jobs(self, multicohort: MultiCohort, inputs: StageInput) -> StageOutput: + def queue_jobs(self, cohort: Cohort, inputs: StageInput) -> StageOutput: """ Queue the EGT-INFO extract, merged-set variant-metrics, and filter jobs. """ - outputs: dict[str, Path] = self.expected_outputs(multicohort=multicohort) - merged_pgen_path: Path = inputs.as_path(target=multicohort, stage=ExportCohortDatasets, key='pgen') - merged_pvar_path: Path = inputs.as_path(target=multicohort, stage=ExportCohortDatasets, key='pvar') - merged_psam_path: Path = inputs.as_path(target=multicohort, stage=ExportCohortDatasets, key='psam') + outputs: dict[str, Path] = self.expected_outputs(cohort=cohort) + merged_pgen_path: Path = inputs.as_path(target=cohort, stage=ExportCohortDatasets, key='pgen') + merged_pvar_path: Path = inputs.as_path(target=cohort, stage=ExportCohortDatasets, key='pvar') + merged_psam_path: Path = inputs.as_path(target=cohort, stage=ExportCohortDatasets, key='psam') stage_path: list[str] = ['popgen_genotyping', 'snp_qc_report'] thresholds_path: list[str] = [*stage_path, 'thresholds'] @@ -487,42 +530,51 @@ def queue_jobs(self, multicohort: MultiCohort, inputs: StageInput) -> StageOutpu output_audit_tsv_path=str(outputs['audit_tsv']), output_inclusion_list_path=str(outputs['inclusion_list']), output_summary_tsv_path=str(outputs['summary_tsv']), - job_name=f'SnpQcReport_{multicohort.name}', + job_name=f'SnpQcReport_{cohort.name}', ) - return self.make_outputs(multicohort, data=outputs, jobs=jobs) + return self.make_outputs(cohort, data=outputs, jobs=jobs) -@stage(required_stages=[Plink2Qc, KingIbdseg, BafRegress], analysis_type='array_qc_report') -class QcReport(MultiCohortStage): +@stage(required_stages=[Plink2Qc, KingIbdseg], analysis_type='array_qc_report') +class QcReport(CohortStage): """ - Create the QC report for an input object. + Create the QC report for the super cohort. + + ``BafRegress`` is a phase-1 (per-plate) stage that does not run in phase 2, so it is not a + ``required_stages`` dependency. Its per-plate contamination outputs are resolved by querying + all registered ``array_bafregress`` analyses in Metamist (``resolve_bafregress_map``) and + selecting the super cohort's full membership — covering every constituent plate, not just the + new ones. """ - def expected_outputs(self, multicohort: MultiCohort) -> Path: + def expected_outputs(self, cohort: Cohort) -> Path: """ - Define the expected QC report output file for the multi-cohort. + Define the expected QC report output file for the cohort. """ - prefix: Path = get_output_prefix(dataset=multicohort.analysis_dataset, stage_name=self.name) - datestamp: str = datetime.now(tz=timezone.utc).strftime('%Y%m%d') - return prefix / f'{datestamp}_qc_report.csv' + prefix: Path = get_output_prefix(dataset=cohort.dataset, stage_name=self.name) + return prefix / f'{cohort.id}_qc_report.csv' - def queue_jobs(self, multicohort: MultiCohort, inputs: StageInput) -> StageOutput: + def queue_jobs(self, cohort: Cohort, inputs: StageInput) -> StageOutput: """ - Queue the QC report generation job for the multi-cohort. + Queue the QC report generation job for the cohort. """ - outputs: Path = self.expected_outputs(multicohort=multicohort) + outputs: Path = self.expected_outputs(cohort=cohort) # Get the plink2_qc_prefix from the Plink2Qc stage's 'smiss' output - plink_qc_smiss_path: Path = inputs.as_path(target=multicohort, stage=Plink2Qc, key='smiss') + plink_qc_smiss_path: Path = inputs.as_path(target=cohort, stage=Plink2Qc, key='smiss') plink_qc_prefix = str(plink_qc_smiss_path).removesuffix('.smiss') # Autosomal KING --ibdseg pairwise summary (REL_ID:KINSHIP:INFTYPE) - king_seg_path: Path = inputs.as_path(target=multicohort, stage=KingIbdseg, key='seg') + king_seg_path: Path = inputs.as_path(target=cohort, stage=KingIbdseg, key='seg') - # Get all bafregress output paths from all cohorts - bafregress_outputs: dict[str, Path] = inputs.as_path_by_target(stage=BafRegress) - bafregress_paths: list[str] = [str(baf_out) for baf_out in bafregress_outputs.values()] + # Resolve the BafRegress output for every super-cohort SG from Metamist (full membership, + # not just the new plates), since BafRegress does not run as a phase-2 stage. + # The map is sg_id -> plate-level file (one file per plate cohort), so dedupe before + # passing it on: repeating a plate file once per SG would multiply that plate's rows + # in the report's IID merge. Sorted for a reproducible job command. + bafregress_map: dict[str, str] = resolve_bafregress_map(sg_ids=cohort.get_sequencing_group_ids()) + bafregress_paths: list[str] = sorted(set(bafregress_map.values())) # Call the Hail Batch job function j: BashJob = run_qc_report( @@ -530,8 +582,8 @@ def queue_jobs(self, multicohort: MultiCohort, inputs: StageInput) -> StageOutpu king_seg_path=str(king_seg_path), bafregress_paths=bafregress_paths, output_path=str(outputs), - job_name=f'QcReport_{multicohort.name}', + job_name=f'QcReport_{cohort.name}', ) # Return the expected outputs of this stage - return self.make_outputs(multicohort, data=outputs, jobs=[j]) + return self.make_outputs(cohort, data=outputs, jobs=[j]) diff --git a/src/popgen_genotyping/utils.py b/src/popgen_genotyping/utils.py index 26b29e1..7f4eebd 100644 --- a/src/popgen_genotyping/utils.py +++ b/src/popgen_genotyping/utils.py @@ -47,6 +47,59 @@ def get_output_prefix(dataset: Dataset, stage_name: str, tmp: bool = False, vers return stage_prefix / str(version) +def get_previous_aggregate_cohort_id() -> str | None: + """ + Read the previous aggregate cohort ID from config, with bootstrap made explicit. + + The key is required so a forgotten config entry cannot silently build a + new-plates-only aggregate: a from-scratch build must be declared with the literal + 'bootstrap' (returns None); any other value must be a Metamist cohort ID. + + Returns: + str | None: The cohort ID to roll forward, or None for a declared bootstrap. + + Raises: + ConfigError: If the key is missing from config. + ValueError: If the value is neither a cohort ID (COH...) nor 'bootstrap'. + """ + value: str = config_retrieve(['popgen_genotyping', 'merge_cohort_plink', 'previous_aggregate_cohort_id']) + if value == 'bootstrap': + return None + if not isinstance(value, str) or not value.startswith('COH'): + raise ValueError( + f"previous_aggregate_cohort_id must be a cohort ID (COH...) or the literal 'bootstrap', got {value!r}" + ) + return value + + +def validate_only_stages(only_stages: list[str], phase_stages: list, entry_point: str) -> None: + """ + Reject a ``workflow.only_stages`` selection naming stages this entry point does not run. + + Each phase has its own entry point with a fixed stage list, so ``only_stages`` is only + ever a within-phase subset (e.g. re-running just the QC report). cpg-flow skips stages + by exact name match, so a typo, a wrong-cased name, or a stage from the other phase + would otherwise be silently skipped while the rest of the selection runs. + + Args: + only_stages (list[str]): The ``workflow.only_stages`` config value; empty runs + every stage of the phase. + phase_stages (list): The stage classes this entry point submits. + entry_point (str): The entry-point name, for the error message. + + Raises: + ValueError: If ``only_stages`` names a stage outside this entry point's phase. + """ + known = {cls.__name__ for cls in phase_stages} + unknown = set(only_stages) - known + if unknown: + raise ValueError( + f'workflow.only_stages names stages {sorted(unknown)} that {entry_point} does not run; ' + f'its stages are {sorted(known)} (exact case). The phases run against different cohorts ' + '(new plates vs the super cohort) and each has its own entry point — see the README.' + ) + + def get_sequencing_group_cohort(sequencing_group: SequencingGroup) -> Cohort: """ Resolve the cohort a sequencing group belongs to by searching the multi-cohort. diff --git a/test/test_king_ibdseg.py b/test/test_king_ibdseg.py index 2b06794..5263242 100644 --- a/test/test_king_ibdseg.py +++ b/test/test_king_ibdseg.py @@ -5,7 +5,6 @@ import gzip import re import subprocess -from datetime import datetime, timezone from pathlib import Path from unittest.mock import MagicMock, patch @@ -195,29 +194,25 @@ class TestKingIbdsegExpectedOutputs: """Path layout that Metamist `analysis_keys=['seg', 'seg_x']` depends on.""" def test_path_layout_and_naming(self) -> None: - """Five outputs under the standard prefix, date-stamped, no dot before X.""" + """Five outputs under the standard prefix, cohort-keyed, no dot before X.""" prefix = Path('/cohort/king_ibdseg') - mock_multicohort = MagicMock() + mock_cohort = MagicMock() + mock_cohort.id = 'COH123' mock_self = MagicMock() mock_self.name = 'KingIbdseg' - fixed_now = datetime(2026, 1, 15, tzinfo=timezone.utc) - with ( - patch('popgen_genotyping.stages.get_output_prefix', return_value=prefix) as mock_prefix, - patch('popgen_genotyping.stages.datetime') as mock_datetime, - ): - mock_datetime.now.return_value = fixed_now - result = KingIbdseg.expected_outputs(mock_self, mock_multicohort) + with patch('popgen_genotyping.stages.get_output_prefix', return_value=prefix) as mock_prefix: + result = KingIbdseg.expected_outputs(mock_self, mock_cohort) assert result == { - 'seg': prefix / '20260115_king.seg', - 'segments': prefix / '20260115_king.segments.gz', - 'seg_x': prefix / '20260115_kingX.seg', - 'segments_x': prefix / '20260115_kingX.segments.gz', - 'log': prefix / '20260115_king.log', + 'seg': prefix / 'COH123_king.seg', + 'segments': prefix / 'COH123_king.segments.gz', + 'seg_x': prefix / 'COH123_kingX.seg', + 'segments_x': prefix / 'COH123_kingX.segments.gz', + 'log': prefix / 'COH123_king.log', } mock_prefix.assert_called_once_with( - dataset=mock_multicohort.analysis_dataset, + dataset=mock_cohort.dataset, stage_name='KingIbdseg', ) @@ -230,8 +225,8 @@ class TestKingIbdsegQueueJobs: def test_passes_merged_plink_and_outputs_to_job(self) -> None: """Bed/bim/fam map to bed/bim/fam; the five outputs map to the five job kwargs.""" - mock_multicohort = MagicMock() - mock_multicohort.name = 'my_multicohort' + mock_cohort = MagicMock() + mock_cohort.name = 'my_cohort' merged_plink: dict[str, Path] = { 'bed': Path('/merged/cohort.bed'), @@ -252,10 +247,10 @@ def test_passes_merged_plink_and_outputs_to_job(self) -> None: mock_self.expected_outputs.return_value = expected_outputs with patch('popgen_genotyping.stages.run_king_ibdseg') as mock_run: - KingIbdseg.queue_jobs(mock_self, mock_multicohort, mock_inputs) + KingIbdseg.queue_jobs(mock_self, mock_cohort, mock_inputs) # Inputs are pulled from the right upstream stage. - mock_inputs.as_dict.assert_called_once_with(target=mock_multicohort, stage=MergeCohortPlink) + mock_inputs.as_dict.assert_called_once_with(target=mock_cohort, stage=MergeCohortPlink) # The job factory receives plink inputs and outputs in their named slots # -- this is the seam where #24-style cross-wiring would manifest. @@ -268,13 +263,13 @@ def test_passes_merged_plink_and_outputs_to_job(self) -> None: output_seg_x_path='/out/20260115_kingX.seg', output_segments_x_path='/out/20260115_kingX.segments.gz', output_log_path='/out/20260115_king.log', - job_name='KingIbdseg_my_multicohort', + job_name='KingIbdseg_my_cohort', ) # The same outputs dict flows through to make_outputs, and the job list # returned by run_king_ibdseg (recode + KING) is passed through verbatim. mock_self.make_outputs.assert_called_once_with( - mock_multicohort, + mock_cohort, data=expected_outputs, jobs=mock_run.return_value, ) diff --git a/test/test_merge_cohort_plink.py b/test/test_merge_cohort_plink.py index 523c8d6..711dba0 100644 --- a/test/test_merge_cohort_plink.py +++ b/test/test_merge_cohort_plink.py @@ -24,7 +24,8 @@ def _run_merge_plink( cohort_plink_paths: New per-plate filesets to merge (defaults to a single cohort). previous_aggregate_resource: Optional rolling-aggregate resource group. samples_to_remove: Optional withdrawn SGs to drop from the previous aggregate. - keep_samples: Optional super-cohort membership to trim the merged fileset to. + keep_samples: Super-cohort membership to trim to (defaults to a two-SG cohort; pass an + explicit list — including [] — to exercise the trim/guard behaviour). Returns: tuple[list[str], MagicMock]: (bash strings passed to j.command() in queue order, to_path mock). @@ -42,9 +43,9 @@ def _run_merge_plink( run_merge_plink( cohort_plink_paths=cohort_plink_paths if cohort_plink_paths is not None else _DEFAULT_COHORT_PATHS, output_prefix='gs://o/out', + keep_samples=keep_samples if keep_samples is not None else ['SG1', 'SG2'], previous_aggregate_resource=previous_aggregate_resource, samples_to_remove=samples_to_remove, - keep_samples=keep_samples, ) commands = [call.args[0] for call in mock_job.command.call_args_list] @@ -76,25 +77,13 @@ def test_rolling_aggregate_path_preserves_allele_order_throughout() -> None: assert '--keep-allele-order' in commands[1] -def test_no_keep_samples_adds_no_trim_step() -> None: - """keep_samples=None (default) must not queue a --keep pass — the pre-two-phase behaviour.""" - commands, mock_to_path = _run_merge_plink() - - # A single input fileset with no aggregate collapses to one --make-bed command. - assert len(commands) == 1 - assert '--merge-list' not in commands[0] - # '--keep ' (trailing space) is the sample-filter flag; '--keep-allele-order' is unrelated. - assert not any('--keep ' in c for c in commands) - mock_to_path.assert_not_called() - - def test_empty_keep_samples_raises() -> None: - """An empty keep list is a caller/config bug — distinct from None — and must fail fast. + """An empty keep list is a caller/config bug and must fail fast. - Both trim guards are truthiness checks, so [] would silently skip the trim and write the + Trimming to the super cohort is mandatory, so an empty membership would otherwise write the untrimmed merge as the aggregate: the exact outcome the trim exists to prevent. """ - with pytest.raises(ValueError, match='keep_samples was provided but empty'): + with pytest.raises(ValueError, match='keep_samples is empty'): _run_merge_plink(keep_samples=[]) @@ -110,6 +99,9 @@ def test_keep_samples_appends_final_keep_trim() -> None: assert '--keep ' in trim, 'final step must be the super-cohort --keep trim' assert '--keep-allele-order' in trim assert '--output-chr chrM' in trim + # Membership assert: the trimmed .fam must contain exactly len(keep_samples) rows, else + # plink --keep silently dropped a claimed SG (super ⊆ merged goes unverified otherwise). + assert 'wc -l' in trim and '-ne 2' in trim, 'final step must assert the kept-sample count' # --keep silently ignores IDs absent from the input, so the pass yields the intersection of the # merged set and keep_samples, not equality. The in-job count is what turns a shortfall into a diff --git a/test/test_metamist_utils.py b/test/test_metamist_utils.py index 5efba1b..5689085 100644 --- a/test/test_metamist_utils.py +++ b/test/test_metamist_utils.py @@ -6,18 +6,21 @@ from unittest.mock import MagicMock, patch import pytest +from metamist.model.new_cohort import NewCohort from popgen_genotyping.metamist_utils import ( + create_custom_cohort, + find_cohort_by_membership, format_merge_plan, parse_genotyping_manifest, query_cohorts_with_analyses, query_genotyping_manifests, - query_previous_aggregate, resolve_bafregress_map, resolve_cohort_bed_map, resolve_gtc_path, resolve_merge_inputs, - resolve_rolling_aggregate, + resolve_super_cohort_membership, + wait_for_cohort_analyses, ) @@ -208,111 +211,6 @@ def test_query_genotyping_manifests_no_results(mock_config, mock_query): assert manifests == [] -@patch('popgen_genotyping.metamist_utils.query') -def test_query_previous_aggregate(mock_query): - mock_query.return_value = { - 'analyses': [ - { - 'outputs': {'path': 'gs://path/merged.pgen'}, - 'project': {'sequencingGroups': [{'id': 'CPG001'}, {'id': 'CPG002'}]}, - } - ] - } - - outputs, active_sgs = query_previous_aggregate(123) - - assert outputs == { - 'pgen': 'gs://path/merged.pgen', - 'psam': 'gs://path/merged.psam', - 'pvar': 'gs://path/merged.pvar', - } - assert active_sgs == ['CPG001', 'CPG002'] - mock_query.assert_called_once() - - -@patch('popgen_genotyping.metamist_utils.query') -def test_query_previous_aggregate_active_only(mock_query): - """ - Test that query_previous_aggregate returns only active samples. - """ - # Mocking a response where one sample (CPG003) is inactive and thus excluded from the response - mock_query.return_value = { - 'analyses': [ - { - 'outputs': {'path': 'gs://path/merged.pgen'}, - 'project': { - 'sequencingGroups': [{'id': 'CPG001'}, {'id': 'CPG002'}] - # CPG003 is missing because it's inactive - }, - } - ] - } - - _outputs, active_sgs = query_previous_aggregate(123) - - expected_count = 2 - assert len(active_sgs) == expected_count - assert 'CPG001' in active_sgs - assert 'CPG002' in active_sgs - assert 'CPG003' not in active_sgs - - -@patch('popgen_genotyping.metamist_utils.query') -def test_query_previous_aggregate_missing_path_raises(mock_query): - """A previous-aggregate analysis whose outputs lack a `path` field is invalid.""" - mock_query.return_value = { - 'analyses': [ - { - 'outputs': {}, - 'project': {'sequencingGroups': []}, - } - ] - } - - with pytest.raises(ValueError, match='valid PGEN output path'): - query_previous_aggregate(123) - - -@patch('popgen_genotyping.metamist_utils.query') -def test_query_previous_aggregate_non_pgen_path_raises(mock_query): - """A previous-aggregate path that doesn't end in `.pgen` is rejected.""" - mock_query.return_value = { - 'analyses': [ - { - 'outputs': {'path': 'gs://path/merged.bed'}, - 'project': {'sequencingGroups': []}, - } - ] - } - - with pytest.raises(ValueError, match='valid PGEN output path'): - query_previous_aggregate(123) - - -@patch('popgen_genotyping.metamist_utils.query_previous_aggregate') -@patch('popgen_genotyping.utils.parse_psam') -def test_resolve_rolling_aggregate_with_removed_sample(mock_parse_psam, mock_query_prev): - """ - Test that resolve_rolling_aggregate correctly identifies samples to remove. - """ - # Mock previous aggregate outputs and current active samples - # Current active samples: CPG001 and CPG002 - mock_query_prev.return_value = ( - {'pgen': 'gs://path/merged.pgen', 'pvar': 'gs://path/merged.pvar', 'psam': 'gs://path/merged.psam'}, - ['CPG001', 'CPG002'], - ) - - # Mock previous aggregate FAM file content: CPG001, CPG002, and CPG003 - mock_parse_psam.return_value = ['CPG001', 'CPG002', 'CPG003'] - - # Execute - paths, to_remove = resolve_rolling_aggregate(123) - - # Verify - assert paths == {'pgen': 'gs://path/merged.pgen', 'pvar': 'gs://path/merged.pvar', 'psam': 'gs://path/merged.psam'} - assert to_remove == ['CPG003'] - - # --------------------------------------------------------------------------- # query_cohorts_with_analyses # --------------------------------------------------------------------------- @@ -538,3 +436,168 @@ def test_format_merge_plan_bootstrap(): assert 'bootstrap' in text assert 'carried forward: 0 SGs' in text assert 'expected merged total: 2 SGs' in text + + +def test_resolve_super_cohort_membership_rolling(): + cohorts = [_cohort('COH_AGG', ['CPG1', 'CPG2', 'CPG3'])] + + membership = resolve_super_cohort_membership( + plate_sg_ids=['CPG4', 'CPG5'], + previous_aggregate_cohort_id='COH_AGG', + cohorts=cohorts, + ) + + assert membership == ['CPG1', 'CPG2', 'CPG3', 'CPG4', 'CPG5'] + + +def test_resolve_super_cohort_membership_bootstrap(): + membership = resolve_super_cohort_membership( + plate_sg_ids=['CPG2', 'CPG1'], + previous_aggregate_cohort_id=None, + ) + + assert membership == ['CPG1', 'CPG2'] + + +def test_resolve_super_cohort_membership_nothing_new_raises(): + cohorts = [_cohort('COH_AGG', ['CPG1', 'CPG2', 'CPG3'])] + + with pytest.raises(ValueError, match='nothing new to aggregate'): + resolve_super_cohort_membership( + plate_sg_ids=['CPG1', 'CPG2'], + previous_aggregate_cohort_id='COH_AGG', + cohorts=cohorts, + ) + + +def test_resolve_super_cohort_membership_missing_aggregate_raises(): + with pytest.raises(ValueError, match='COH_MISSING'): + resolve_super_cohort_membership( + plate_sg_ids=['CPG1'], + previous_aggregate_cohort_id='COH_MISSING', + cohorts=[_cohort('COH_OTHER', ['CPG9'])], + ) + + +def test_find_cohort_by_membership(): + cohorts = [ + _cohort('COH1', ['CPG1', 'CPG2']), + _cohort('COH2', ['CPG1', 'CPG2', 'CPG3']), + ] + + match = find_cohort_by_membership(['CPG3', 'CPG2', 'CPG1'], cohorts=cohorts) + + assert match is not None + assert match['id'] == 'COH2' + + +def test_find_cohort_by_membership_no_match(): + cohorts = [_cohort('COH1', ['CPG1', 'CPG2'])] + + assert find_cohort_by_membership(['CPG1'], cohorts=cohorts) is None + + +def test_find_cohort_by_membership_latest_wins(): + # COH99 vs COH100 crosses a digit-length boundary, where lexicographic + # comparison would wrongly pick COH99. + cohorts = [ + _cohort('COH100', ['CPG1', 'CPG2']), + _cohort('COH99', ['CPG1', 'CPG2']), + _cohort('COH9', ['CPG1', 'CPG2']), + ] + + match = find_cohort_by_membership(['CPG1', 'CPG2'], cohorts=cohorts) + + assert match is not None + assert match['id'] == 'COH100' + + +@patch('popgen_genotyping.metamist_utils.CohortApi') +@patch('popgen_genotyping.metamist_utils.metamist_project') +def test_create_custom_cohort(mock_project, mock_api): + mock_project.return_value = 'ourdna' + mock_api.return_value.create_cohort_from_criteria.return_value = NewCohort( + cohort_id='COH42', + sequencing_group_ids=['CPG1', 'CPG2'], + dry_run=False, + ) + + cohort_id = create_custom_cohort( + name='array-aggregate-2026-08-10', + description='test cohort', + sg_ids=['CPG2', 'CPG1'], + ) + + assert cohort_id == 'COH42' + _, kwargs = mock_api.return_value.create_cohort_from_criteria.call_args + assert kwargs['project'] == 'ourdna' + body = kwargs['body_create_cohort_from_criteria'] + assert body.cohort_criteria.sg_ids_internal == ['CPG1', 'CPG2'] + # The server rejects an SG list combined with any other criterion (SET-839). + assert getattr(body.cohort_criteria, 'projects', None) in (None, []) + assert body.cohort_spec.name == 'array-aggregate-2026-08-10' + + +@patch('popgen_genotyping.metamist_utils.CohortApi') +@patch('popgen_genotyping.metamist_utils.metamist_project') +def test_create_custom_cohort_no_id_raises(mock_project, mock_api): + mock_project.return_value = 'ourdna' + # A malformed response without a cohort ID cannot be expressed as a NewCohort + # (cohort_id is required there), so a bare dict stands in for it. + mock_api.return_value.create_cohort_from_criteria.return_value = {} + + with pytest.raises(ValueError, match='no cohort ID'): + create_custom_cohort(name='x', description='y', sg_ids=['CPG1']) + + +@patch('popgen_genotyping.metamist_utils.CohortApi') +@patch('popgen_genotyping.metamist_utils.metamist_project') +def test_create_custom_cohort_missing_sgs_raises(mock_project, mock_api): + mock_project.return_value = 'ourdna' + mock_api.return_value.create_cohort_from_criteria.return_value = NewCohort( + cohort_id='COH42', + sequencing_group_ids=['CPG1'], + dry_run=False, + ) + + with pytest.raises(ValueError, match=r'missing 1 requested sequencing group.*CPG9'): + create_custom_cohort(name='x', description='y', sg_ids=['CPG1', 'CPG9']) + + +@patch('popgen_genotyping.metamist_utils.time') +@patch('popgen_genotyping.metamist_utils.query_cohorts_with_analyses') +def test_wait_for_cohort_analyses_returns_when_registered(mock_query, mock_time): + mock_time.monotonic.return_value = 0 + mock_query.return_value = [ + _cohort('COH1', ['CPG1'], analyses=[_analysis('array_cohort_bed', 'x'), _analysis('array_bafregress', 'y')]), + ] + + wait_for_cohort_analyses(cohort_ids=['COH1'], analysis_types=['array_cohort_bed', 'array_bafregress']) + + mock_time.sleep.assert_not_called() + + +@patch('popgen_genotyping.metamist_utils.time') +@patch('popgen_genotyping.metamist_utils.query_cohorts_with_analyses') +def test_wait_for_cohort_analyses_polls_until_registered(mock_query, mock_time): + mock_time.monotonic.side_effect = [0, 10] + mock_query.side_effect = [ + [_cohort('COH1', ['CPG1'], analyses=[_analysis('array_cohort_bed', 'x')])], + [_cohort('COH1', ['CPG1'], analyses=[_analysis('array_cohort_bed', 'x'), _analysis('array_bafregress', 'y')])], + ] + + wait_for_cohort_analyses(cohort_ids=['COH1'], analysis_types=['array_cohort_bed', 'array_bafregress']) + + mock_time.sleep.assert_called_once() + + +@patch('popgen_genotyping.metamist_utils.time') +@patch('popgen_genotyping.metamist_utils.query_cohorts_with_analyses') +def test_wait_for_cohort_analyses_times_out(mock_query, mock_time): + mock_time.monotonic.side_effect = [0, 1000] + mock_query.return_value = [_cohort('COH1', ['CPG1'], analyses=[_analysis('array_cohort_bed', 'x')])] + + with pytest.raises(TimeoutError, match=r'COH1.*array_bafregress'): + wait_for_cohort_analyses(cohort_ids=['COH1'], analysis_types=['array_cohort_bed', 'array_bafregress']) + + mock_time.sleep.assert_not_called() diff --git a/test/test_stages_phase2.py b/test/test_stages_phase2.py new file mode 100644 index 0000000..6ce66b3 --- /dev/null +++ b/test/test_stages_phase2.py @@ -0,0 +1,208 @@ +""" +Stage-level tests for the two phase-2 stages rewritten for the two-phase run: +MergeCohortPlink (Metamist-resolved merge inputs, mandatory --keep trim) and +QcReport (BafRegress resolved for the full super-cohort membership). +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any +from unittest.mock import MagicMock, patch + +from popgen_genotyping.stages import KingIbdseg, MergeCohortPlink, Plink2Qc, QcReport + +# -- Helpers ------------------------------------------------------------------ + +SUPER_SG_IDS = ['CPG1', 'CPG2', 'CPG3'] + +MERGE_OUTPUTS = { + 'bed': Path('/merge/COH999_merged.bed'), + 'bim': Path('/merge/COH999_merged.bim'), + 'fam': Path('/merge/COH999_merged.fam'), +} + + +def _mock_merge_cohort() -> MagicMock: + cohort = MagicMock() + cohort.id = 'COH999' + cohort.get_sequencing_group_ids.return_value = SUPER_SG_IDS + return cohort + + +def _resolved(previous_aggregate_paths: str | None) -> dict: + # Shape returned by resolve_merge_inputs; plate entries carry more keys than + # the bed/bim/fam triple the merge job consumes. + return { + 'plate_merge_list': [ + { + 'cohort_id': 'COHP1', + 'bed': 'gs://p1.bed', + 'bim': 'gs://p1.bim', + 'fam': 'gs://p1.fam', + 'new_count': len(SUPER_SG_IDS), + }, + ], + 'previous_aggregate_paths': previous_aggregate_paths, + 'samples_to_remove': ['CPGX'], + 'super_cohort_size': len(SUPER_SG_IDS), + } + + +# -- Tests: MergeCohortPlink.queue_jobs ---------------------------------------- + + +class TestMergeCohortPlinkQueueJobs: + """Wire-up between the Metamist-resolved plan and the merge job.""" + + def test_rolling_aggregate_run(self) -> None: + """With a previous aggregate: convert it to PLINK1, merge with it, trim to super.""" + mock_cohort = _mock_merge_cohort() + mock_self = MagicMock() + mock_self.expected_outputs.return_value = MERGE_OUTPUTS + + conversion_job = MagicMock(name='conversion_job') + converted_resource = MagicMock(name='converted_resource') + plink1_prefix = Path('/merge/plink2_to_plink1') + + with ( + patch('popgen_genotyping.stages.get_previous_aggregate_cohort_id', return_value='COH123') as mock_previous, + patch( + 'popgen_genotyping.stages.resolve_merge_inputs', + return_value=_resolved(previous_aggregate_paths='gs://prev/agg'), + ) as mock_resolve, + patch('popgen_genotyping.stages.get_output_prefix', return_value=plink1_prefix), + patch( + 'popgen_genotyping.stages.run_plink2_to_plink1', + return_value=(conversion_job, converted_resource), + ) as mock_convert, + patch('popgen_genotyping.stages.run_merge_plink') as mock_merge, + ): + MergeCohortPlink.queue_jobs(mock_self, mock_cohort, MagicMock()) + + # The previous aggregate comes from the required config key (shared with the + # SubmitPhase2 stage); the plan is resolved from the super cohort's membership + # (the cohort is the source of truth, plates are derived). + mock_previous.assert_called_once_with() + mock_resolve.assert_called_once_with( + super_cohort_sg_ids=SUPER_SG_IDS, + previous_aggregate_cohort_id='COH123', + ) + + # The PLINK2 aggregate is converted back to PLINK 1.9 at a cohort-keyed prefix. + mock_convert.assert_called_once_with( + pfile_prefix='gs://prev/agg', + output_prefix=str(plink1_prefix / 'COH999_plink1'), + job_name='Plink2ToPlink1', + ) + + # The merge receives only bed/bim/fam per plate, the converted previous aggregate, + # and keep_samples = the full super-cohort membership (the mandatory trim). + mock_merge.assert_called_once_with( + cohort_plink_paths=[{'bed': 'gs://p1.bed', 'bim': 'gs://p1.bim', 'fam': 'gs://p1.fam'}], + output_prefix='/merge/COH999_merged', + keep_samples=SUPER_SG_IDS, + previous_aggregate_resource=converted_resource, + samples_to_remove=['CPGX'], + job_name='MergeCohortPlink', + ) + + # The merge waits for the conversion, and the merge job is what the stage returns. + mock_merge.return_value.depends_on.assert_called_once_with(conversion_job) + mock_self.make_outputs.assert_called_once_with(mock_cohort, data=MERGE_OUTPUTS, jobs=[mock_merge.return_value]) + + def test_bootstrap_run_has_no_previous_aggregate(self) -> None: + """Without a previous aggregate: no PLINK2 conversion, no merge dependency.""" + mock_cohort = _mock_merge_cohort() + mock_self = MagicMock() + mock_self.expected_outputs.return_value = MERGE_OUTPUTS + + with ( + patch('popgen_genotyping.stages.get_previous_aggregate_cohort_id', return_value=None), + patch( + 'popgen_genotyping.stages.resolve_merge_inputs', + return_value=_resolved(previous_aggregate_paths=None), + ), + patch('popgen_genotyping.stages.run_plink2_to_plink1') as mock_convert, + patch('popgen_genotyping.stages.run_merge_plink') as mock_merge, + ): + MergeCohortPlink.queue_jobs(mock_self, mock_cohort, MagicMock()) + + mock_convert.assert_not_called() + assert mock_merge.call_args.kwargs['previous_aggregate_resource'] is None + mock_merge.return_value.depends_on.assert_not_called() + + def test_merge_plan_is_logged_for_operator_confirmation(self) -> None: + """The resolved plan reaches the driver log — the operator's pre-flight check.""" + mock_cohort = _mock_merge_cohort() + mock_self = MagicMock() + mock_self.expected_outputs.return_value = MERGE_OUTPUTS + + with ( + patch('popgen_genotyping.stages.get_previous_aggregate_cohort_id', return_value=None), + patch( + 'popgen_genotyping.stages.resolve_merge_inputs', + return_value=_resolved(previous_aggregate_paths=None), + ), + patch('popgen_genotyping.stages.format_merge_plan', return_value='THE MERGE PLAN') as mock_plan, + patch('popgen_genotyping.stages.logger') as mock_logger, + patch('popgen_genotyping.stages.run_merge_plink'), + ): + MergeCohortPlink.queue_jobs(mock_self, mock_cohort, MagicMock()) + + mock_plan.assert_called_once() + mock_logger.info.assert_called_once_with('THE MERGE PLAN') + + +# -- Tests: QcReport.queue_jobs ------------------------------------------------ + + +class TestQcReportQueueJobs: + """Wire-up between upstream QC outputs, the BafRegress map, and the report job.""" + + def test_passes_inputs_and_full_membership_bafregress(self) -> None: + """BafRegress is resolved for every super-cohort SG, not just the new plates. + + The resolved map is sg_id -> plate-level file, so SGs sharing a plate map to the + same file (CPG1/CPG2 below); the job must receive each plate file once — a repeat + would multiply that plate's rows in the report's IID merge. + """ + mock_cohort = MagicMock() + mock_cohort.name = 'super_cohort' + mock_cohort.get_sequencing_group_ids.return_value = SUPER_SG_IDS + + report_path = Path('/out/COH999_20260115_qc_report.csv') + mock_self = MagicMock() + mock_self.expected_outputs.return_value = report_path + + def as_path(target: object, stage: Any, key: str) -> Path: + del target + return { + (Plink2Qc, 'smiss'): Path('/qc/COH999_20260115_qc.smiss'), + (KingIbdseg, 'seg'): Path('/king/COH999_20260115_king.seg'), + }[(stage, key)] + + mock_inputs = MagicMock() + mock_inputs.as_path.side_effect = as_path + + bafregress_map = { + 'CPG1': 'gs://baf/COHP1.BAFRegress.txt', + 'CPG2': 'gs://baf/COHP1.BAFRegress.txt', + 'CPG3': 'gs://baf/COHP2.BAFRegress.txt', + } + + with ( + patch('popgen_genotyping.stages.resolve_bafregress_map', return_value=bafregress_map) as mock_baf, + patch('popgen_genotyping.stages.run_qc_report') as mock_run, + ): + QcReport.queue_jobs(mock_self, mock_cohort, mock_inputs) + + mock_baf.assert_called_once_with(sg_ids=SUPER_SG_IDS) + mock_run.assert_called_once_with( + plink_qc_prefix='/qc/COH999_20260115_qc', + king_seg_path='/king/COH999_20260115_king.seg', + bafregress_paths=['gs://baf/COHP1.BAFRegress.txt', 'gs://baf/COHP2.BAFRegress.txt'], + output_path=str(report_path), + job_name='QcReport_super_cohort', + ) + mock_self.make_outputs.assert_called_once_with(mock_cohort, data=report_path, jobs=[mock_run.return_value]) diff --git a/test/test_utils.py b/test/test_utils.py index 8d66555..a731331 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -4,7 +4,7 @@ from unittest.mock import MagicMock, patch import pytest -from popgen_genotyping.utils import get_output_prefix, parse_psam +from popgen_genotyping.utils import get_output_prefix, get_previous_aggregate_cohort_id, parse_psam @pytest.fixture @@ -82,3 +82,33 @@ def test_get_output_prefix_tmp(mock_get_workflow, mock_config_retrieve, tmp_path dataset.prefix.assert_called_once_with(category='tmp') assert result == tmp_path / 'wf' / 'MyStage' / '1' + + +@patch('popgen_genotyping.utils.config_retrieve') +def test_get_previous_aggregate_cohort_id(mock_config_retrieve): + mock_config_retrieve.return_value = 'COH123' + + assert get_previous_aggregate_cohort_id() == 'COH123' + + +@patch('popgen_genotyping.utils.config_retrieve') +def test_get_previous_aggregate_cohort_id_bootstrap(mock_config_retrieve): + mock_config_retrieve.return_value = 'bootstrap' + + assert get_previous_aggregate_cohort_id() is None + + +@patch('popgen_genotyping.utils.config_retrieve') +def test_get_previous_aggregate_cohort_id_invalid_raises(mock_config_retrieve): + mock_config_retrieve.return_value = 'true' + + with pytest.raises(ValueError, match=r"cohort ID .* or the literal 'bootstrap'"): + get_previous_aggregate_cohort_id() + + +@patch('popgen_genotyping.utils.config_retrieve') +def test_get_previous_aggregate_cohort_id_empty_raises(mock_config_retrieve): + mock_config_retrieve.return_value = '' + + with pytest.raises(ValueError): + get_previous_aggregate_cohort_id() diff --git a/test/test_workflow_entry_points.py b/test/test_workflow_entry_points.py new file mode 100644 index 0000000..a17293b --- /dev/null +++ b/test/test_workflow_entry_points.py @@ -0,0 +1,81 @@ +""" +Tests for the submission-time checks in the phase entry points +(first_workflow / second_workflow). +""" + +import pytest + +from popgen_genotyping import stages as stages_module +from popgen_genotyping.first_workflow import PHASE_1_STAGES +from popgen_genotyping.second_workflow import PHASE_2_STAGES, validate_phase2_cohorts +from popgen_genotyping.stages import CohortStage, MultiCohortStage +from popgen_genotyping.utils import validate_only_stages + +PHASE_2_NAMES = ['MergeCohortPlink', 'ExportCohortDatasets', 'Plink2Qc', 'KingIbdseg', 'SnpQcReport', 'QcReport'] + + +class TestValidateOnlyStages: + """only_stages may only name stages of the submitting entry point's phase.""" + + def test_full_phase_selection_passes(self) -> None: + validate_only_stages(only_stages=PHASE_2_NAMES, phase_stages=PHASE_2_STAGES, entry_point='second_workflow') + + def test_empty_selection_passes(self) -> None: + """No only_stages runs the whole phase — the entry point pins the stage list.""" + validate_only_stages(only_stages=[], phase_stages=PHASE_1_STAGES, entry_point='first_workflow') + + def test_single_stage_rerun_passes(self) -> None: + """Re-running one stage (e.g. just the QC report) is a valid within-phase subset.""" + validate_only_stages(only_stages=['QcReport'], phase_stages=PHASE_2_STAGES, entry_point='second_workflow') + + def test_other_phase_stage_raises(self) -> None: + """A phase-1 stage in a phase-2 submission would be silently skipped by cpg-flow.""" + with pytest.raises(ValueError, match='does not run'): + validate_only_stages( + only_stages=['CohortBcfToPlink', 'MergeCohortPlink'], + phase_stages=PHASE_2_STAGES, + entry_point='second_workflow', + ) + + def test_unknown_stage_raises(self) -> None: + with pytest.raises(ValueError, match='does not run'): + validate_only_stages(only_stages=['NotAStage'], phase_stages=PHASE_1_STAGES, entry_point='first_workflow') + + def test_wrong_case_raises(self) -> None: + """cpg-flow accepts wrong-cased only_stages names but skips stages by exact + match, so in a mixed-case list the wrong-cased stage is silently skipped while + the rest run (an all-lowercase list at least dies with 'No stages to run'). + Rejecting on exact case closes the silent-partial-skip case.""" + with pytest.raises(ValueError, match='does not run'): + validate_only_stages(only_stages=['gtctobcfs'], phase_stages=PHASE_1_STAGES, entry_point='first_workflow') + + +class TestPhase2CohortCardinality: + """Phase 2 runs against exactly one cohort: the super cohort.""" + + def test_single_cohort_passes(self) -> None: + validate_phase2_cohorts(input_cohorts=['COH200']) + + def test_multiple_cohorts_raises(self) -> None: + """Two cohorts would each be treated as 'the super cohort', registering two + aggregates that both claim the same previous-aggregate lineage.""" + with pytest.raises(ValueError, match='exactly one cohort'): + validate_phase2_cohorts(input_cohorts=['COH200', 'COH201']) + + def test_no_cohorts_raises(self) -> None: + with pytest.raises(ValueError, match='exactly one cohort'): + validate_phase2_cohorts(input_cohorts=[]) + + +def test_phase_lists_cover_all_stages() -> None: + """The two entry points must cover every stage the module defines, so a new + stage cannot be added without assigning it to a phase.""" + # @stage wraps each class in a decorator function; the class survives as __wrapped__. + defined = { + name + for name, obj in vars(stages_module).items() + if isinstance(getattr(obj, '__wrapped__', None), type) + and issubclass(obj.__wrapped__, (CohortStage, MultiCohortStage)) + } + split = {cls.__name__ for cls in PHASE_1_STAGES + PHASE_2_STAGES} + assert split == defined