diff --git a/kubernetes/loculus/templates/loculus-preprocessing-config.yaml b/kubernetes/loculus/templates/loculus-preprocessing-config.yaml index dabe36aa3b..7db9039b83 100644 --- a/kubernetes/loculus/templates/loculus-preprocessing-config.yaml +++ b/kubernetes/loculus/templates/loculus-preprocessing-config.yaml @@ -18,6 +18,14 @@ data: {{- if and (hasKey $organismConfig.schema "submissionDataTypes") (hasKey $organismConfig.schema.submissionDataTypes "maxSequencesPerEntry") }} max_sequences_per_entry: {{ $organismConfig.schema.submissionDataTypes.maxSequencesPerEntry }} {{- end }} + {{- if and (hasKey $organismConfig.schema "submissionDataTypes") (hasKey $organismConfig.schema.submissionDataTypes "files") }} + {{- if $organismConfig.schema.submissionDataTypes.files.enabled }} + submission_file_categories: + {{- range $organismConfig.schema.submissionDataTypes.files.categories }} + - {{ .name }} + {{- end }} + {{- end }} + {{- end }} processing_spec: {{- $args := dict "metadata" $metadata "referenceGenomes" $organismConfig.referenceGenomes }} {{- include "loculus.preprocessingSpecs" $args | nindent 6 }} diff --git a/kubernetes/loculus/values.schema.json b/kubernetes/loculus/values.schema.json index 134655edc6..aeac087ed5 100644 --- a/kubernetes/loculus/values.schema.json +++ b/kubernetes/loculus/values.schema.json @@ -814,6 +814,14 @@ "type": "string", "description": "Minimizer used for nextclade sort (if require_nextclade_sort_match or if segments are to be assigned using nextclade sort)" }, + "deacon_max_host_proportion": { + "groups": ["nextcladePipelineConfigFile"], + "docsIncludePrefix": false, + "type": "number", + "minimum": 0, + "maximum": 1, + "description": "Maximum proportion of host reads in submitted raw read files, as detected by Deacon." + }, "backend_request_timeout_seconds": { "groups": ["nextcladePipelineConfigFile"], "docsIncludePrefix": false, diff --git a/kubernetes/loculus/values.yaml b/kubernetes/loculus/values.yaml index 96a6c84f31..552c3bdb98 100644 --- a/kubernetes/loculus/values.yaml +++ b/kubernetes/loculus/values.yaml @@ -1558,6 +1558,7 @@ defaultOrganismConfig: &defaultOrganismConfig - name: singleReference genes: [] nextclade_dataset_server: https://data.clades.nextstrain.org/v3 + deacon_max_host_proportion: 0.01 ingest: &ingest image: ghcr.io/loculus-project/ingest configFile: diff --git a/kubernetes/loculus/values_preview_server.yaml b/kubernetes/loculus/values_preview_server.yaml index 7200693acb..7c262b00a6 100644 --- a/kubernetes/loculus/values_preview_server.yaml +++ b/kubernetes/loculus/values_preview_server.yaml @@ -142,7 +142,7 @@ resources: memory: "200Mi" cpu: "10m" limits: - memory: "3Gi" + memory: "8Gi" minio: requests: memory: "250Mi" @@ -156,3 +156,23 @@ defaultResources: limits: memory: "1Gi" cpu: "20m" + +# Preview only: run just the west-nile preprocessing pipeline while testing the deacon +# host-depletion step. Each enabled organism runs its own prepro pod and fetches the +# ~3.1Gi panhuman index, so disabling the rest keeps startup and memory manageable. +# Deep-merges `enabled: false` onto the base defaultOrganisms entries. +defaultOrganisms: + cchf: + enabled: false + cchf-multi-ref: + enabled: false + dummy-organism: + enabled: false + dummy-organism-with-files: + enabled: false + ebola-sudan: + enabled: false + enteroviruses: + enabled: false + not-aligned-organism: + enabled: false diff --git a/preprocessing/nextclade/environment.yml b/preprocessing/nextclade/environment.yml index 275cbca792..ef8ccccd54 100644 --- a/preprocessing/nextclade/environment.yml +++ b/preprocessing/nextclade/environment.yml @@ -24,3 +24,4 @@ dependencies: - pytest=9.1.1 - pydantic=2.13.4 - diamond=2.2.2 + - deacon=0.15.0 diff --git a/preprocessing/nextclade/src/loculus_preprocessing/backend.py b/preprocessing/nextclade/src/loculus_preprocessing/backend.py index 3c760e5f39..ba92bc49c2 100644 --- a/preprocessing/nextclade/src/loculus_preprocessing/backend.py +++ b/preprocessing/nextclade/src/loculus_preprocessing/backend.py @@ -17,7 +17,7 @@ from .config import Config from .datatypes import ( - FileIdAndName, + FileIdAndNameAndReadUrl, FileUploadInfo, ProcessedEntry, UnprocessedData, @@ -97,7 +97,10 @@ def parse_ndjson(ndjson_data: str) -> Sequence[UnprocessedEntry]: submitted_files = json_object["data"].get("files") file_mapping = ( { - category: [FileIdAndName(fileId=f["fileId"], name=f["name"]) for f in files] + category: [ + FileIdAndNameAndReadUrl(fileId=f["fileId"], name=f["name"], url=f.get("url")) + for f in files + ] for category, files in submitted_files.items() } if submitted_files @@ -245,6 +248,13 @@ def upload_embl_file_to_presigned_url( raise RuntimeError(msg) +def download_file(config: Config, url: str, save_path: str) -> None: + with requests.get(url, stream=True, timeout=config.backend_request_timeout_seconds) as response: + response.raise_for_status() + with Path(save_path).open("wb") as f: + f.writelines(response.iter_content(chunk_size=1024 * 1024)) + + def download_minimizer(config, save_path): if config.minimizer_url: url = config.minimizer_url diff --git a/preprocessing/nextclade/src/loculus_preprocessing/config.py b/preprocessing/nextclade/src/loculus_preprocessing/config.py index cb3757bf27..e0f890e39e 100644 --- a/preprocessing/nextclade/src/loculus_preprocessing/config.py +++ b/preprocessing/nextclade/src/loculus_preprocessing/config.py @@ -10,6 +10,7 @@ from pydantic import BaseModel, Field, model_validator from loculus_preprocessing.datatypes import ( + FileCategory, FunctionArgs, FunctionInputs, FunctionName, @@ -27,6 +28,7 @@ NEXTCLADE_PREFIX = "nextclade." ASSIGNED_REFERENCE_PREFIX = "ASSIGNED_REFERENCE" INTERNAL_INPUT_PREFIXES = (NEXTCLADE_PREFIX, ASSIGNED_REFERENCE_PREFIX) +DEACON_INDEX = "panhuman-1.k31w15.idx" class EmblInfoMetadataPropertyNames(BaseModel): @@ -119,7 +121,9 @@ class Config(BaseModel): require_nextclade_sort_match: bool = False minimizer_url: str | None = None diamond_dmnd_url: str | None = None + deacon_max_host_proportion: float | None = None + submission_file_categories: list[str] = Field(default_factory=list) create_embl_file: bool = False scientific_name: str = "Orthonairovirus haemorrhagiae" molecule_type: MoleculeType = MoleculeType.GENOMIC_RNA @@ -142,6 +146,14 @@ def finalize(self): if not self.backend_host: # Set here so we can use organism self.backend_host = f"http://127.0.0.1:8079/{self.organism}" + if FileCategory.RAW_READS in self.submission_file_categories: + if self.deacon_max_host_proportion is None: + msg = "deacon_max_host_proportion must be set if raw reads submissions are enabled" + raise ValueError(msg) + if not (0.0 <= self.deacon_max_host_proportion <= 1.0): + msg = "deacon_max_host_proportion must be in [0.0, 1.0]" + raise ValueError(msg) + self.processing_order = get_processing_order(self) return self diff --git a/preprocessing/nextclade/src/loculus_preprocessing/datatypes.py b/preprocessing/nextclade/src/loculus_preprocessing/datatypes.py index f957168f0f..8029b67eec 100644 --- a/preprocessing/nextclade/src/loculus_preprocessing/datatypes.py +++ b/preprocessing/nextclade/src/loculus_preprocessing/datatypes.py @@ -1,7 +1,8 @@ +import json from collections.abc import Iterable -from dataclasses import dataclass, field +from dataclasses import dataclass, field, fields from enum import StrEnum, unique -from typing import Any, Final +from typing import Any, Final, Self AccessionVersion = str GeneName = str @@ -82,9 +83,10 @@ def from_single(cls, name: str, type, message: str): @dataclass -class FileIdAndName: +class FileIdAndNameAndReadUrl: fileId: str # noqa: N815 name: str + url: str | None = None @dataclass @@ -95,7 +97,7 @@ class UnprocessedData: submissionId: str # noqa: N815 metadata: InputMetadata unalignedNucleotideSequences: dict[SequenceName, NucleotideSequence | None] # noqa: N815 - files: dict[FileCategory, list[FileIdAndName]] | None + files: dict[FileCategory, list[FileIdAndNameAndReadUrl]] | None @dataclass @@ -111,7 +113,7 @@ class UnprocessedEntry: @dataclass class UnprocessedAfterNextclade: inputMetadata: InputMetadata # noqa: N815 - files: dict[FileCategory, list[FileIdAndName]] | None + files: dict[FileCategory, list[FileIdAndNameAndReadUrl]] | None # Derived metadata produced by Nextclade nextcladeMetadata: dict[SequenceName, Any] | None # noqa: N815 unalignedNucleotideSequences: dict[SequenceName, NucleotideSequence | None] # noqa: N815 @@ -127,7 +129,7 @@ class UnprocessedAfterNextclade: @dataclass class ProcessedData: metadata: ProcessedMetadata - files: dict[FileCategory, list[FileIdAndName]] | None + files: dict[FileCategory, list[FileIdAndNameAndReadUrl]] | None unalignedNucleotideSequences: dict[SequenceName, Any] # noqa: N815 alignedNucleotideSequences: dict[SequenceName, Any] # noqa: N815 nucleotideInsertions: dict[SequenceName, Any] # noqa: N815 @@ -265,3 +267,19 @@ class MoleculeType(StrEnum): class Topology(StrEnum): LINEAR = "linear" CIRCULAR = "circular" + + +@dataclass +class DeaconSummary: + time: float + seqs_in: int + seqs_out: int + seqs_removed: int + seqs_removed_proportion: float + + @classmethod + def from_json(cls, json_path: str) -> Self: + with open(json_path, encoding="utf-8") as f: + data = json.load(f) + wanted = {f.name for f in fields(cls)} + return cls(**{k: v for k, v in data.items() if k in wanted}) diff --git a/preprocessing/nextclade/src/loculus_preprocessing/file_processing_functions.py b/preprocessing/nextclade/src/loculus_preprocessing/file_processing_functions.py index 464659992a..dba2938ff2 100644 --- a/preprocessing/nextclade/src/loculus_preprocessing/file_processing_functions.py +++ b/preprocessing/nextclade/src/loculus_preprocessing/file_processing_functions.py @@ -1,9 +1,19 @@ import logging +import os +import subprocess # noqa: S404 +from tempfile import TemporaryDirectory +from typing import cast +from uuid import uuid4 +from requests import HTTPError + +from loculus_preprocessing.backend import download_file +from loculus_preprocessing.config import DEACON_INDEX, Config from loculus_preprocessing.datatypes import ( AnnotationSourceType, + DeaconSummary, FileCategory, - FileIdAndName, + FileIdAndNameAndReadUrl, ProcessingAnnotation, ) @@ -11,7 +21,9 @@ def process_submitted_files( - file_mapping: dict[FileCategory, list[FileIdAndName]], + config: Config, + dataset_dir: str, + file_mapping: dict[FileCategory, list[FileIdAndNameAndReadUrl]], ) -> tuple[list[ProcessingAnnotation], list[ProcessingAnnotation]]: errors: list[ProcessingAnnotation] = [] warnings: list[ProcessingAnnotation] = [] @@ -22,7 +34,11 @@ def process_submitted_files( continue match category: case FileCategory.RAW_READS: - rr_errors, rr_warnings = validate_raw_reads_submission(files) + rr_errors, rr_warnings = validate_raw_reads_submission( + config, + dataset_dir, + files, + ) errors.extend(rr_errors) warnings.extend(rr_warnings) case _: @@ -42,7 +58,9 @@ def process_submitted_files( def validate_raw_reads_submission( - files: list[FileIdAndName], + config: Config, + dataset_dir: str, + files: list[FileIdAndNameAndReadUrl], ) -> tuple[list[ProcessingAnnotation], list[ProcessingAnnotation]]: errors: list[ProcessingAnnotation] = [] warnings: list[ProcessingAnnotation] = [] @@ -74,4 +92,125 @@ def validate_raw_reads_submission( ) ) + if errors: + # only do in-depth validation if basic checks pass + return errors, warnings + + deacon_index = os.path.join(dataset_dir, DEACON_INDEX) + with TemporaryDirectory() as tmp_dir: + local_files = [] + for file in files: + if (url := file.url) is None: + message = ( + f"Internal error: no URL for file '{file.name}' with id " + f"'{file.fileId}'. Please contact the administrator." + ) + errors.append( + ProcessingAnnotation.from_single( + name=file.name, type=AnnotationSourceType.FILE, message=message + ) + ) + continue + + file_name_internal = os.path.join(tmp_dir, f"{file.fileId}-{file.name}") + try: + download_file(config, url, file_name_internal) + except HTTPError as e: + message = f"Error downloading file '{file.name}' from S3: {e}" + logger.error(message) + errors.append( + ProcessingAnnotation.from_single( + name=file.name, type=AnnotationSourceType.FILE, message=message + ) + ) + continue + local_files.append(file_name_internal) + if errors: + return errors, warnings + + summary_json_path = os.path.join(tmp_dir, f"{uuid4()}_summary.json") + deacon_summary = run_deacon_filter(deacon_index, local_files, summary_json_path) + # deacon_max_host_proportion is guaranteed not None by Config.finalize() + if deacon_summary.seqs_removed_proportion > cast(float, config.deacon_max_host_proportion): + names = ", ".join(file.name for file in files) + message = ( + f"File(s) '{names}' had a host reads proportion of " + f"{deacon_summary.seqs_removed_proportion}, maximum allowed proportion " + f"is {config.deacon_max_host_proportion}" + ) + errors.append( + ProcessingAnnotation.from_fields( + input_fields=[f.name for f in files], + output_fields=[f.name for f in files], + type=AnnotationSourceType.FILE, + message=message, + ) + ) + return errors, warnings + + +def run_deacon_fetch(index_path: str) -> None: + args = ["deacon", "index", "fetch", "--output", index_path, "panhuman-1"] + logger.debug(f"Running Deacon fetch: {args}") + + exit_code = subprocess.run(args, check=False).returncode # noqa: S603 + if exit_code != 0: + message = f"Deacon fetch failed with exit code {exit_code}" + logger.error(message) + raise RuntimeError(message) + + +def run_deacon_filter(index: str, input_files: list[str], summary_json: str) -> DeaconSummary: + args = [ + "deacon", + "--use-server", + "filter", + "--deplete", + "--threads", + "1", + "--summary", + summary_json, + index, + *input_files, + ] + logger.debug(f"Running Deacon filter on '{', '.join(input_files)}': {args}") + + exit_code = subprocess.run( # noqa: S603 + args, check=False, stdout=subprocess.DEVNULL + ).returncode + if exit_code != 0: + message = f"Deacon filter failed with exit code {exit_code}" + logger.error(message) + raise Exception(message) + return DeaconSummary.from_json(summary_json) + + +def start_deacon_server() -> subprocess.Popen: + args = [ + "deacon", + "server", + "start", + ] + logger.debug("Starting Deacon server") + + return subprocess.Popen(args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) # noqa: S603 + + +def stop_deacon_server(proc: subprocess.Popen) -> None: + args = ["deacon", "--use-server", "server", "stop"] + logger.debug("Stopping Deacon server") + + subprocess.run( # noqa: S603 + args, + check=False, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + + try: + proc.wait(timeout=5) + except subprocess.TimeoutExpired: + logger.warning("Failed to stop Deacon server gracefully, sending SIGKILL") + proc.kill() + proc.wait() diff --git a/preprocessing/nextclade/src/loculus_preprocessing/nextclade.py b/preprocessing/nextclade/src/loculus_preprocessing/nextclade.py index 8545472b63..b15882d13b 100644 --- a/preprocessing/nextclade/src/loculus_preprocessing/nextclade.py +++ b/preprocessing/nextclade/src/loculus_preprocessing/nextclade.py @@ -27,7 +27,7 @@ AnnotationSourceType, FastaId, FileCategory, - FileIdAndName, + FileIdAndNameAndReadUrl, GeneName, GenericSequence, NucleotideInsertion, @@ -803,9 +803,9 @@ def enrich_with_nextclade( # noqa: PLR0914 } for entry in unprocessed } - input_files: dict[AccessionVersion, dict[FileCategory, list[FileIdAndName]] | None] = { - entry.accessionVersion: entry.data.files for entry in unprocessed - } + input_files: dict[ + AccessionVersion, dict[FileCategory, list[FileIdAndNameAndReadUrl]] | None + ] = {entry.accessionVersion: entry.data.files for entry in unprocessed} batch = assign_segment_for_alignment(unprocessed, config=config, dataset_dir=dataset_dir) unaligned_nucleotide_sequences = batch.unalignedNucleotideSequences diff --git a/preprocessing/nextclade/src/loculus_preprocessing/prepro.py b/preprocessing/nextclade/src/loculus_preprocessing/prepro.py index 6ea6a7eeb9..10a6a77108 100644 --- a/preprocessing/nextclade/src/loculus_preprocessing/prepro.py +++ b/preprocessing/nextclade/src/loculus_preprocessing/prepro.py @@ -1,11 +1,17 @@ import logging +import os import time from collections import defaultdict from collections.abc import Sequence from tempfile import TemporaryDirectory from typing import Any -from loculus_preprocessing.file_processing_functions import process_submitted_files +from loculus_preprocessing.file_processing_functions import ( + process_submitted_files, + run_deacon_fetch, + start_deacon_server, + stop_deacon_server, +) from .backend import ( download_diamond_db, @@ -17,6 +23,7 @@ ) from .config import ( ASSIGNED_REFERENCE_PREFIX, + DEACON_INDEX, METADATA_DEPENDENCY_PREFIX, NEXTCLADE_PREFIX, AlignmentRequirement, @@ -31,7 +38,7 @@ AnnotationSource, AnnotationSourceType, FileCategory, - FileIdAndName, + FileIdAndNameAndReadUrl, GeneName, InputData, InputMetadata, @@ -518,6 +525,7 @@ def unpack_annotations(config, nextclade_metadata: dict[str, Any] | None) -> dic def process_single( accession_version: AccessionVersion, unprocessed: UnprocessedAfterNextclade, + dataset_dir: str, config: Config, ) -> SubmissionData: """Process a single sequence per config""" @@ -537,7 +545,9 @@ def process_single( accession_version, unprocessed, config ) - file_errors, file_warnings = process_submitted_files(unprocessed.files or {}) + file_errors, file_warnings = process_submitted_files( + config, dataset_dir, unprocessed.files or {} + ) processed_entry = ProcessedEntry( accession=accession_from_str(accession_version), @@ -578,6 +588,7 @@ def process_single( def process_single_unaligned( accession_version: AccessionVersion, unprocessed: UnprocessedData, + dataset_dir: str, config: Config, ) -> SubmissionData: """Process a single sequence per config""" @@ -592,7 +603,9 @@ def process_single_unaligned( accession_version, unprocessed, config ) - file_errors, file_warnings = process_submitted_files(unprocessed.files or {}) + file_errors, file_warnings = process_submitted_files( + config, dataset_dir, unprocessed.files or {} + ) return processed_entry_no_alignment( accession_version=accession_version, @@ -646,7 +659,7 @@ def process_all( nextclade_results = enrich_with_nextclade(unprocessed, dataset_dir, config) for id, result in nextclade_results.items(): try: - processed_single = process_single(id, result, config) + processed_single = process_single(id, result, dataset_dir, config) except Exception as e: logger.error(f"Processing failed for {id} with error: {e}") processed_single = processed_entry_with_errors(id) @@ -655,7 +668,7 @@ def process_all( for entry in unprocessed: try: processed_single = process_single_unaligned( - entry.accessionVersion, entry.data, config + entry.accessionVersion, entry.data, dataset_dir, config ) except Exception as e: logger.error(f"Processing failed for {entry.accessionVersion} with error: {e}") @@ -680,7 +693,7 @@ def upload_flatfiles(processed: Sequence[SubmissionData], config: Config) -> Non upload_embl_file_to_presigned_url(file_content, upload_info.url, upload_info.headers) processed_files = submission_data.processed_entry.data.files or {} processed_files.setdefault(FileCategory.ANNOTATIONS, []).append( - FileIdAndName(fileId=file_id, name=file_name) + FileIdAndNameAndReadUrl(fileId=file_id, name=file_name) ) submission_data.processed_entry.data.files = processed_files except Exception as e: @@ -699,7 +712,7 @@ def upload_flatfiles(processed: Sequence[SubmissionData], config: Config) -> Non ) -def run(config: Config) -> None: # noqa: C901 +def run(config: Config) -> None: # noqa: C901, PLR0912 with TemporaryDirectory(delete=not config.keep_tmp_dir) as dataset_dir: if config.alignment_requirement != AlignmentRequirement.NONE: download_nextclade_dataset(dataset_dir, config) @@ -713,6 +726,11 @@ def run(config: Config) -> None: # noqa: C901 msg = "Diamond database URL must be provided for diamond segment classification" raise ValueError(msg) download_diamond_db(config, dataset_dir + "/diamond/diamond.dmnd") + can_submit_raw_reads = FileCategory.RAW_READS in config.submission_file_categories + if can_submit_raw_reads: + index_path = os.path.join(dataset_dir, DEACON_INDEX) + run_deacon_fetch(index_path) + total_processed = 0 etag = None last_force_refresh = time.time() @@ -731,6 +749,9 @@ def run(config: Config) -> None: # noqa: C901 # Don't use etag if we just got data # preprocessing only asks for 100 sequences to process at a time, so there might be more etag = None + deacon_server_process = None + if can_submit_raw_reads: + deacon_server_process = start_deacon_server() try: processed = process_all(unprocessed, dataset_dir, config) except Exception as e: @@ -738,6 +759,9 @@ def run(config: Config) -> None: # noqa: C901 f"Processing failed. Traceback : {e}. Unprocessed data: {unprocessed}" ) continue + finally: + if deacon_server_process is not None: + stop_deacon_server(deacon_server_process) if config.create_embl_file: upload_flatfiles(processed, config) diff --git a/preprocessing/nextclade/tests/factory_methods.py b/preprocessing/nextclade/tests/factory_methods.py index 679601e744..02e5f2159c 100644 --- a/preprocessing/nextclade/tests/factory_methods.py +++ b/preprocessing/nextclade/tests/factory_methods.py @@ -10,7 +10,7 @@ AnnotationSource, AnnotationSourceType, FileCategory, - FileIdAndName, + FileIdAndNameAndReadUrl, NucleotideSequence, ProcessedData, ProcessedEntry, @@ -79,7 +79,7 @@ def create_unprocessed_entry( accession_id: str, sequences: dict[SegmentName, NucleotideSequence | None], group_id: int = 2, - files: dict[FileCategory, list[FileIdAndName]] | None = None, + files: dict[FileCategory, list[FileIdAndNameAndReadUrl]] | None = None, ) -> UnprocessedEntry: return UnprocessedEntry( accessionVersion=f"LOC_{accession_id}.1", @@ -134,7 +134,7 @@ def create_processed_entry( errors: list[ProcessingAnnotation] | None = None, warnings: list[ProcessingAnnotation] | None = None, processed_alignment: ProcessedAlignment | None = None, - files: dict[FileCategory, list[FileIdAndName]] | None = None, + files: dict[FileCategory, list[FileIdAndNameAndReadUrl]] | None = None, ) -> ProcessedEntry: if errors is None: errors = [] @@ -169,11 +169,11 @@ def create_processed_entry( class Case: name: str input_metadata: dict[str, str | None] = field(default_factory=dict) - input_files: dict[FileCategory, list[FileIdAndName]] | None = None + input_files: dict[FileCategory, list[FileIdAndNameAndReadUrl]] | None = None input_sequence: dict[str, str | None] = field(default_factory=lambda: {"main": None}) accession_id: str = "000999" expected_metadata: dict[str, ProcessedMetadataValue] = field(default_factory=dict) - expected_files: dict[FileCategory, list[FileIdAndName]] | None = None + expected_files: dict[FileCategory, list[FileIdAndNameAndReadUrl]] | None = None expected_errors: list[ProcessingAnnotation] | None = None expected_warnings: list[ProcessingAnnotation] | None = None expected_processed_alignment: ProcessedAlignment | None = None diff --git a/preprocessing/nextclade/tests/test_file_processing.py b/preprocessing/nextclade/tests/test_file_processing.py new file mode 100644 index 0000000000..99a27b5c06 --- /dev/null +++ b/preprocessing/nextclade/tests/test_file_processing.py @@ -0,0 +1,192 @@ +# ruff: noqa: S101 + +from unittest.mock import MagicMock + +import pytest + +from loculus_preprocessing import file_processing_functions, prepro +from loculus_preprocessing.config import Config +from loculus_preprocessing.datatypes import ( + DeaconSummary, + FileCategory, + FileIdAndNameAndReadUrl, +) +from loculus_preprocessing.file_processing_functions import ( + process_submitted_files, + validate_raw_reads_submission, +) + +MAX_HOST_PROPORTION = 0.05 + + +@pytest.fixture +def config() -> Config: + return Config( + submission_file_categories=[FileCategory.RAW_READS], + deacon_max_host_proportion=MAX_HOST_PROPORTION, + ) + + +@pytest.fixture +def mock_deacon(monkeypatch): + """Mock the deacon index file download and deacon filter calls. + + Returns a setter so each test can choose the reported host-read proportion (default 0.0). + """ + monkeypatch.setattr( + file_processing_functions, "download_file", lambda config, url, save_path: None + ) + + def set_removed_proportion(removed_proportion: float = 0.0) -> None: + monkeypatch.setattr( + file_processing_functions, + "run_deacon_filter", + lambda index, input_file, summary_json: deacon_summary(removed_proportion), + ) + + set_removed_proportion(0.0) + return set_removed_proportion + + +def deacon_summary(removed_proportion: float) -> DeaconSummary: + """Build a DeaconSummary reporting the given proportion of removed host reads.""" + seqs_in = 100 + seqs_removed = round(seqs_in * removed_proportion) + return DeaconSummary( + time=0.0, + seqs_in=seqs_in, + seqs_out=seqs_in - seqs_removed, + seqs_removed=seqs_removed, + seqs_removed_proportion=removed_proportion, + ) + + +def fastq_file( + file_id: str, name: str, url: str | None = "https://example.test/reads" +) -> FileIdAndNameAndReadUrl: + return FileIdAndNameAndReadUrl(fileId=file_id, name=name, url=url) + + +def test_fetch_deacon_idx_only_when_raw_reads_enabled(config, monkeypatch): + class _ExitLoop(Exception): # noqa: N818 + """Attaches to `fetch_unprocessed_sequences` to exit run()'s + processing loop after startup. + """ + + fetch_index_mock = MagicMock(return_value=True) + monkeypatch.setattr(prepro, "run_deacon_fetch", fetch_index_mock) + monkeypatch.setattr(prepro, "fetch_unprocessed_sequences", MagicMock(side_effect=_ExitLoop)) + + # RAW_READS enabled -> the index is fetched once + with pytest.raises(_ExitLoop): + prepro.run(config) + fetch_index_mock.assert_called_once() + (index_path,) = fetch_index_mock.call_args.args + assert index_path.endswith(prepro.DEACON_INDEX) + + fetch_index_mock = MagicMock(return_value=True) + monkeypatch.setattr(prepro, "run_deacon_fetch", fetch_index_mock) + + # RAW_READS not enabled -> the index must not be fetched. + config.submission_file_categories = [] + config.deacon_max_host_proportion = None + with pytest.raises(_ExitLoop): + prepro.run(config) + fetch_index_mock.assert_not_called() + + +def test_process_submitted_files_runs_raw_reads_check(config, tmp_path, monkeypatch): + validate_mock = MagicMock(return_value=([], [])) + monkeypatch.setattr(file_processing_functions, "validate_raw_reads_submission", validate_mock) + + files = [fastq_file("f1", "reads.fastq")] + errors, warnings = process_submitted_files( + config, str(tmp_path), {FileCategory.RAW_READS: files} + ) + + validate_mock.assert_called_once_with(config, str(tmp_path), files) + assert errors == [] + assert warnings == [] + + +def test_process_submitted_files_skips_empty_category(config, tmp_path, monkeypatch): + # The backend sends an empty list for file categories without files + monkeypatch.setattr( + file_processing_functions, + "validate_raw_reads_submission", + lambda *a: pytest.fail("empty category must not be validated"), + ) + file_mapping = {FileCategory.RAW_READS: []} + errors, warnings = process_submitted_files(config, str(tmp_path), file_mapping) + assert errors == [] + assert warnings == [] + + +def test_process_submitted_files_reports_unsupported_category(config, tmp_path): + file_mapping = {FileCategory.ANNOTATIONS: [fastq_file("f1", "notes.txt")]} + errors, _ = process_submitted_files(config, str(tmp_path), file_mapping) + assert len(errors) == 1 + assert "not supported by preprocessing" in errors[0].message + + +@pytest.mark.usefixtures("mock_deacon") +def test_single_valid_fastq_passes(config, tmp_path): + files = [fastq_file("f1", "reads.fastq")] + errors, warnings = validate_raw_reads_submission(config, str(tmp_path), files) + assert errors == [] + assert warnings == [] + + +@pytest.mark.usefixtures("mock_deacon") +def test_paired_valid_fastq_passes(config, tmp_path): + files = [fastq_file("f1", "reads_R1.fastq"), fastq_file("f2", "reads_R2.fastq.gz")] + errors, warnings = validate_raw_reads_submission(config, str(tmp_path), files) + assert errors == [] + assert warnings == [] + + +@pytest.mark.usefixtures("mock_deacon") +def test_too_many_files_reports_error(config, tmp_path): + files = [ + fastq_file("f1", "r1.fastq"), + fastq_file("f2", "r2.fastq"), + fastq_file("f3", "r3.fastq"), + ] + errors, _ = validate_raw_reads_submission(config, str(tmp_path), files) + assert len(errors) == 1 + assert "Received 3" in errors[0].message + + +@pytest.mark.usefixtures("mock_deacon") +def test_unrecognized_extension_reports_error(config, tmp_path): + files = [fastq_file("f1", "reads.txt")] + errors, _ = validate_raw_reads_submission(config, str(tmp_path), files) + assert len(errors) == 1 + assert "unrecognized extension" in errors[0].message + + +def test_host_contamination_within_threshold(config, tmp_path, mock_deacon): + mock_deacon(0.01) + files = [fastq_file("f1", "reads.fastq")] + errors, _ = validate_raw_reads_submission(config, str(tmp_path), files) + assert errors == [] + + +def test_host_contamination_exceeds_threshold(config, tmp_path, mock_deacon): + mock_deacon(0.9) + files = [fastq_file("f1", "reads.fastq")] + errors, _ = validate_raw_reads_submission(config, str(tmp_path), files) + assert len(errors) == 1 + assert "host reads proportion" in errors[0].message + + +def test_missing_url_reports_error_and_skips_deacon(config, tmp_path, monkeypatch): + def fail_if_called(*args, **kwargs): + msg = "deacon must not run when the file URL is missing" + raise AssertionError(msg) + + monkeypatch.setattr(file_processing_functions, "run_deacon_filter", fail_if_called) + files = [fastq_file("f1", "reads.fastq", url=None)] + errors, _ = validate_raw_reads_submission(config, str(tmp_path), files) + assert len(errors) == 1 + assert "no URL for file" in errors[0].message diff --git a/preprocessing/nextclade/tests/test_nextclade_preprocessing.py b/preprocessing/nextclade/tests/test_nextclade_preprocessing.py index 71352cd060..e65d3b1578 100644 --- a/preprocessing/nextclade/tests/test_nextclade_preprocessing.py +++ b/preprocessing/nextclade/tests/test_nextclade_preprocessing.py @@ -26,8 +26,6 @@ ) from loculus_preprocessing.datatypes import ( AnnotationSourceType, - FileCategory, - FileIdAndName, SegmentClassificationMethod, SubmissionData, UnprocessedData, @@ -286,141 +284,6 @@ def invalid_sequence() -> str: sequenceNameToFastaId={"main": "fastaHeader"}, ), ), - Case( - name="with file", - input_metadata={}, - input_files={ - FileCategory.RAW_READS: [ - FileIdAndName(fileId="file-id-0001", name="reads_R1.fastq"), - FileIdAndName(fileId="file-id-0002", name="reads_R2.fastq"), - ] - }, - input_sequence={"fastaHeader": sequence_with_mutation("single")}, - accession_id="1", - expected_metadata={ - "completeness": 1.0, - "totalInsertedNucs": 0, - "totalSnps": 1, - "totalDeletedNucs": 0, - "length": len(consensus_sequence("single")), - "nonExistentField": "None", - "variant": True, - }, - expected_files={ - FileCategory.RAW_READS: [ - FileIdAndName(fileId="file-id-0001", name="reads_R1.fastq"), - FileIdAndName(fileId="file-id-0002", name="reads_R2.fastq"), - ] - }, - expected_errors=[], - expected_warnings=[], - expected_processed_alignment=ProcessedAlignment( - unalignedNucleotideSequences={"main": sequence_with_mutation("single")}, - alignedNucleotideSequences={"main": sequence_with_mutation("single")}, - nucleotideInsertions={}, - alignedAminoAcidSequences={ - "NPEbolaSudan": ebola_sudan_aa(sequence_with_mutation("single"), "NP"), - "VP35EbolaSudan": ebola_sudan_aa(sequence_with_mutation("single"), "VP35"), - }, - aminoAcidInsertions={}, - sequenceNameToFastaId={"main": "fastaHeader"}, - ), - ), - Case( - name="with too many raw read files", - input_metadata={}, - input_files={ - FileCategory.RAW_READS: [ - FileIdAndName(fileId="file-id-0001", name="reads_R1.fastq"), - FileIdAndName(fileId="file-id-0002", name="reads_R2.fastq"), - FileIdAndName(fileId="file-id-0003", name="reads_R3.fastq"), - ] - }, - input_sequence={"fastaHeader": sequence_with_mutation("single")}, - accession_id="1", - expected_metadata={ - "completeness": 1.0, - "totalInsertedNucs": 0, - "totalSnps": 1, - "totalDeletedNucs": 0, - "length": len(consensus_sequence("single")), - "nonExistentField": "None", - "variant": True, - }, - expected_files={ - FileCategory.RAW_READS: [ - FileIdAndName(fileId="file-id-0001", name="reads_R1.fastq"), - FileIdAndName(fileId="file-id-0002", name="reads_R2.fastq"), - FileIdAndName(fileId="file-id-0003", name="reads_R3.fastq"), - ] - }, - expected_errors=build_processing_annotations( - [ - ProcessingAnnotationHelper( - ["reads_R1.fastq", "reads_R2.fastq", "reads_R3.fastq"], - ["reads_R1.fastq", "reads_R2.fastq", "reads_R3.fastq"], - "Received 3 for raw reads upload. Please submit raw reads as one or two FASTQ files containing raw single-end or paired-end reads.", - AnnotationSourceType.FILE, - ) - ] - ), - expected_warnings=[], - expected_processed_alignment=ProcessedAlignment( - unalignedNucleotideSequences={"main": sequence_with_mutation("single")}, - alignedNucleotideSequences={"main": sequence_with_mutation("single")}, - nucleotideInsertions={}, - alignedAminoAcidSequences={ - "NPEbolaSudan": ebola_sudan_aa(sequence_with_mutation("single"), "NP"), - "VP35EbolaSudan": ebola_sudan_aa(sequence_with_mutation("single"), "VP35"), - }, - aminoAcidInsertions={}, - sequenceNameToFastaId={"main": "fastaHeader"}, - ), - ), - Case( - name="with unrecognized raw read file extension", - input_metadata={}, - input_files={ - FileCategory.RAW_READS: [FileIdAndName(fileId="file-id-0001", name="reads.txt")] - }, - input_sequence={"fastaHeader": sequence_with_mutation("single")}, - accession_id="1", - expected_metadata={ - "completeness": 1.0, - "totalInsertedNucs": 0, - "totalSnps": 1, - "totalDeletedNucs": 0, - "length": len(consensus_sequence("single")), - "nonExistentField": "None", - "variant": True, - }, - expected_files={ - FileCategory.RAW_READS: [FileIdAndName(fileId="file-id-0001", name="reads.txt")] - }, - expected_errors=build_processing_annotations( - [ - ProcessingAnnotationHelper( - ["reads.txt"], - ["reads.txt"], - "Raw reads file 'reads.txt' has unrecognized extension. " - "Allowed extensions: .fastq, .fastq.gz, .fq, .fq.gz", - AnnotationSourceType.FILE, - ) - ] - ), - expected_warnings=[], - expected_processed_alignment=ProcessedAlignment( - unalignedNucleotideSequences={"main": sequence_with_mutation("single")}, - alignedNucleotideSequences={"main": sequence_with_mutation("single")}, - nucleotideInsertions={}, - alignedAminoAcidSequences={ - "NPEbolaSudan": ebola_sudan_aa(sequence_with_mutation("single"), "NP"), - "VP35EbolaSudan": ebola_sudan_aa(sequence_with_mutation("single"), "VP35"), - }, - aminoAcidInsertions={}, - sequenceNameToFastaId={"main": "fastaHeader"}, - ), - ), ] single_segment_failed_case_definitions = [ @@ -1271,51 +1134,6 @@ def invalid_sequence() -> str: sequenceNameToFastaId={"ebola-sudan": "ebola-sudan"}, ), ), - Case( - name="validate raw read files when sequences are not aligned", - input_metadata={}, - input_sequence={ - "prefix_ebola-sudan": sequence_with_mutation("ebola-sudan"), - "other_prefix_ebola-zaire": sequence_with_mutation("ebola-zaire"), - }, - accession_id="1", - input_files={ - FileCategory.RAW_READS: [FileIdAndName(fileId="file-id-0001", name="reads.txt")] - }, - expected_files={ - FileCategory.RAW_READS: [FileIdAndName(fileId="file-id-0001", name="reads.txt")] - }, - expected_metadata={ - "length_ebola-sudan": len(consensus_sequence("ebola-sudan")), - "length_ebola-zaire": len(consensus_sequence("ebola-zaire")), - }, - expected_errors=build_processing_annotations( - [ - ProcessingAnnotationHelper( - ["reads.txt"], - ["reads.txt"], - "Raw reads file 'reads.txt' has unrecognized extension. " - "Allowed extensions: .fastq, .fastq.gz, .fq, .fq.gz", - AnnotationSourceType.FILE, - ) - ] - ), - expected_warnings=[], - expected_processed_alignment=ProcessedAlignment( - unalignedNucleotideSequences={ - "ebola-sudan": sequence_with_mutation("ebola-sudan"), - "ebola-zaire": sequence_with_mutation("ebola-zaire"), - }, - alignedNucleotideSequences={}, - nucleotideInsertions={}, - alignedAminoAcidSequences={}, - aminoAcidInsertions={}, - sequenceNameToFastaId={ - "ebola-sudan": "prefix_ebola-sudan", - "ebola-zaire": "other_prefix_ebola-zaire", - }, - ), - ), ]