Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
118 changes: 118 additions & 0 deletions src/exaspim_swc_processing/acquisition.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,118 @@
"""Obtain a dataset's ``acquisition`` from the registry, falling back to S3.

The transform needs the acquisition's axis orientation, and ``RegistrationPipeline``
takes it as a file path. Where that file comes from matters: a dataset's mounted bundle
may carry a stale copy, or none, while the registry holds the record the indexer
published.

The order matches :func:`~exaspim_swc_processing.sources.default_sources` -- DocDB v2,
then v1, then the object in S3 -- so the registry is authoritative and S3 answers only
for assets it has not indexed.
"""

from __future__ import annotations

import json
import logging
from collections.abc import Sequence
from pathlib import Path

from exaspim_swc_processing.parent_metadata import MetadataSource
from exaspim_swc_processing.sources import (
DOCDB_HOST,
OPEN_DATA_BUCKET,
Fetcher,
docdb_fetcher,
s3_fetcher,
)

logger = logging.getLogger(__name__)

ACQUISITION_FIELD = "acquisition"
"""Top-level field of a DocDB record holding the acquisition."""

ACQUISITION_KEY = "acquisition.json"
"""Object at the asset root holding the acquisition."""


class AcquisitionNotFoundError(LookupError):
"""Raised when no source holds an acquisition for the dataset."""


def acquisition_sources(
host: str = DOCDB_HOST, bucket: str = OPEN_DATA_BUCKET
) -> list[tuple[MetadataSource, Fetcher]]:
"""Build the resolution order for an acquisition.

Parameters
----------
host : str, optional
DocDB host, by default :data:`~exaspim_swc_processing.sources.DOCDB_HOST`.
bucket : str, optional
S3 bucket, by default :data:`~exaspim_swc_processing.sources.OPEN_DATA_BUCKET`.

Returns
-------
list[tuple[MetadataSource, Fetcher]]
Sources in priority order: DocDB v2, DocDB v1, then S3.
"""
return [
(MetadataSource.DOCDB_V2, docdb_fetcher("v2", host, field=ACQUISITION_FIELD)),
(MetadataSource.DOCDB_V1, docdb_fetcher("v1", host, field=ACQUISITION_FIELD)),
(MetadataSource.S3, s3_fetcher(bucket, filename=ACQUISITION_KEY)),
]


def resolve_acquisition(
asset_name: str,
destination: Path,
sources: Sequence[tuple[MetadataSource, Fetcher]] | None = None,
local_fallback: str = "",
) -> tuple[Path, MetadataSource | None]:
"""Write the dataset's acquisition to a file and say where it came from.

Parameters
----------
asset_name : str
Processed dataset name.
destination : Path
File to write the acquisition to when it comes from a registry or S3.
sources : Sequence[tuple[MetadataSource, Fetcher]] | None, optional
Sources in priority order. Defaults to :func:`acquisition_sources`.
local_fallback : str, optional
An acquisition already on disk, used only when no source answers.

Returns
-------
tuple[Path, MetadataSource | None]
The file to read, and the source it came from -- ``None`` when the local
fallback was used.

Raises
------
AcquisitionNotFoundError
If no source answers and there is no usable local fallback.
"""
for source, fetch in sources if sources is not None else acquisition_sources():
try:
document = fetch(asset_name)
except Exception as error: # noqa: BLE001 - an unreachable source is not fatal
logger.info("Acquisition source %s failed: %s", source.value, error)
continue
if not document:
continue
destination.parent.mkdir(parents=True, exist_ok=True)
destination.write_text(json.dumps(document, indent=2), encoding="utf-8")
logger.info("Acquisition for %s read from %s", asset_name, source.value)
return destination, source

if local_fallback and Path(local_fallback).is_file():
logger.warning(
"No acquisition for %s in DocDB or S3; using the copy staged at %s",
asset_name,
local_fallback,
)
return Path(local_fallback), None
raise AcquisitionNotFoundError(
f"No acquisition for {asset_name} in DocDB v2, DocDB v1 or S3, and no local copy"
)
75 changes: 8 additions & 67 deletions src/exaspim_swc_processing/packaging.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,12 @@
requirement that the upstream assets cannot currently satisfy, so the asset would advertise
itself as incomplete. Quality control belongs to a later stage and is not written here.

The packaging step records itself as a :class:`~aind_data_schema.core.processing.DataProcess`
appended to each cell's record. It is a real processing step — it is what produced the
directory — and it carries the provenance of where the parent metadata was read from.
Packaging does not record itself as a
:class:`~aind_data_schema.core.processing.DataProcess`. It changes no data -- the
reconstructions are hardlinked, byte for byte -- so a record of it would describe
assembling the asset rather than producing the data, in a document meant for the latter.
Where the parent metadata was read from is carried on ``DataDescription.tags``, which is
indexed and therefore searchable, unlike ``output_parameters``.
"""

from __future__ import annotations
Expand All @@ -26,8 +29,7 @@
from pathlib import Path

from aind_data_schema.components.identifiers import Code
from aind_data_schema.core.processing import DataProcess, ProcessStage
from aind_data_schema_models.process_names import ProcessName
from aind_data_schema.core.processing import DataProcess

from exaspim_swc_processing.data_description import (
DataDescriptionDerivationError,
Expand All @@ -41,12 +43,6 @@

logger = logging.getLogger(__name__)

PACKAGING_STEP_NAME = "exaspim_swc_packaging"
"""``DataProcess.name`` of the step this module performs."""

DEFAULT_EXPERIMENTERS = ("Peter Grotz",)
"""Default ``DataProcess.experimenters`` for the packaging step."""

DATA_DESCRIPTION_FILENAME = "data_description.json"
PROCESSING_FILENAME = "processing.json"

Expand Down Expand Up @@ -103,65 +99,13 @@ class PackagingResult:
skipped: tuple[SkippedCell, ...]


def build_packaging_process(
parent: ParentMetadata,
code: Code,
start_time: datetime,
end_time: datetime,
output_path: str,
cell_count: int,
experimenters: Sequence[str] = DEFAULT_EXPERIMENTERS,
) -> DataProcess:
"""Describe the packaging step itself.

Parameters
----------
parent : ParentMetadata
Resolved parent metadata, supplying the provenance of where it was read from.
code : Code
The packaging code: repository URL, name and version.
start_time : datetime
When packaging started.
end_time : datetime
When packaging finished.
output_path : str
Where the cell directories were written, relative to ``/results``.
cell_count : int
Number of cells packaged in the run.
experimenters : Sequence[str], optional
Who is responsible for the run, by default :data:`DEFAULT_EXPERIMENTERS`.

Returns
-------
DataProcess
The step record, carrying :meth:`ParentMetadata.provenance` in
``output_parameters``. Typed ``OTHER`` because regrouping files into per-cell
assets is not one of the vocabulary's defined operations; the docs direct
unmatched operations to ``ANALYSIS`` or ``OTHER``, and ``OTHER`` requires the
``name`` and ``notes`` set here.
"""
return DataProcess(
process_type=ProcessName.OTHER,
name=PACKAGING_STEP_NAME,
stage=ProcessStage.PROCESSING,
code=code,
experimenters=list(experimenters),
start_date_time=start_time,
end_date_time=end_time,
output_path=output_path,
output_parameters={**parent.provenance(), "cells_packaged": cell_count},
notes="Regrouped pipeline outputs into one derived asset per reconstruction.",
)


def package_cells(
stage_root: Path,
output_root: Path,
parent: ParentMetadata,
stage_processes: Sequence[DataProcess],
pipeline: Code,
creation_time: datetime,
packaging_process: DataProcess,
specs: Sequence[ArtifactSpec] = ARTIFACT_SPECS,
overrides: Mapping[str, object] | None = None,
) -> PackagingResult:
Expand All @@ -181,8 +125,6 @@ def package_cells(
The pipeline that produced the run.
creation_time : datetime
Creation time shared by every cell in the run.
packaging_process : DataProcess
Record of this packaging step, appended to each cell's processing record.
specs : Sequence[ArtifactSpec], optional
Artifacts to collect, by default :data:`~exaspim_swc_processing.layout.ARTIFACT_SPECS`.
overrides : Mapping[str, object] | None, optional
Expand All @@ -195,7 +137,7 @@ def package_cells(
skipped rather than written without metadata.
"""
specs = tuple(specs)
processes = [*stage_processes, packaging_process]
record = build_cell_processing(list(stage_processes), pipeline)
packaged: list[PackagedCell] = []
skipped: list[SkippedCell] = []

Expand All @@ -222,7 +164,6 @@ def package_cells(

cell_dir = output_root / description.name
_write_cell_directory(cell, cell_dir, specs)
record = build_cell_processing(processes, pipeline)
(cell_dir / DATA_DESCRIPTION_FILENAME).write_text(
description.model_dump_json(indent=2), encoding="utf-8"
)
Expand Down
31 changes: 19 additions & 12 deletions src/exaspim_swc_processing/sources.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,10 @@ def _docdb_client(version: str, host: str) -> DocDbClient:


def docdb_fetcher(
version: str, host: str = DOCDB_HOST, client: DocDbClient | None = None
version: str,
host: str = DOCDB_HOST,
client: DocDbClient | None = None,
field: str = CORE_FIELD,
) -> Fetcher:
"""Build a fetcher reading one DocDB endpoint.

Expand All @@ -112,12 +115,15 @@ def docdb_fetcher(
Host serving the index, by default :data:`DOCDB_HOST`.
client : DocDbClient | None, optional
A pre-built client. Constructed on first use when omitted.
field : str, optional
Core field to read, by default :data:`CORE_FIELD`. Any top-level field of the
record works, e.g. ``"acquisition"``.

Returns
-------
Fetcher
Returns the asset's raw ``data_description``, or ``None`` if the endpoint has no
record for it.
Returns the asset's raw document for ``field``, or ``None`` if the endpoint has
no record for it.
"""

def _fetch(asset_name: str) -> dict | None:
Expand All @@ -136,16 +142,14 @@ def _fetch(asset_name: str) -> dict | None:
resolved = client if client is not None else _docdb_client(version, host)
records = resolved.retrieve_docdb_records(
filter_query={"name": asset_name},
projection={CORE_FIELD: 1},
projection={field: 1},
limit=1,
)
if not records:
return None
document = records[0].get(CORE_FIELD)
document = records[0].get(field)
if not document:
logger.warning(
"DocDB %s has a record for %r with no %s", version, asset_name, CORE_FIELD
)
logger.warning("DocDB %s has a record for %r with no %s", version, asset_name, field)
return None
return document

Expand Down Expand Up @@ -173,6 +177,7 @@ def _s3_client() -> S3Client:
def s3_fetcher(
bucket: str = OPEN_DATA_BUCKET,
client: S3Client | None = None,
filename: str = DATA_DESCRIPTION_KEY,
) -> Fetcher:
"""Build a fetcher reading the asset's ``data_description.json`` from S3.

Expand All @@ -185,12 +190,14 @@ def s3_fetcher(
Bucket holding the asset, by default :data:`OPEN_DATA_BUCKET`.
client : S3Client | None, optional
A pre-built S3 client. Constructed on first use when omitted.
filename : str, optional
Object to read from the asset root, by default :data:`DATA_DESCRIPTION_KEY`.

Returns
-------
Fetcher
Returns the raw ``data_description``, or ``None`` if the object is absent or is
not readable as JSON.
Returns the raw document, or ``None`` if the object is absent or is not readable
as JSON.
"""

def _fetch(asset_name: str) -> dict | None:
Expand All @@ -207,12 +214,12 @@ def _fetch(asset_name: str) -> dict | None:
The raw document, or ``None`` if absent or unreadable.
"""
resolved = client if client is not None else _s3_client()
key = f"{asset_name}/{DATA_DESCRIPTION_KEY}"
key = f"{asset_name}/{filename}"
try:
response = resolved.get_object(Bucket=bucket, Key=key)
return json.loads(response["Body"].read())
except Exception as error: # noqa: BLE001 - any read failure means "not here"
logger.info("No %s in s3://%s/%s (%s)", DATA_DESCRIPTION_KEY, bucket, key, error)
logger.info("No %s in s3://%s/%s (%s)", filename, bucket, key, error)
return None

return _fetch
Expand Down
Loading
Loading