import math
import shutil
import tempfile
import traceback
import typing as ty
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from fileformats.generic import File, FileSet
from fileformats.medimage import DicomCollection
from frametree.core.frameset import FrameSet
from frametree.xnat import Xnat
from tqdm import tqdm
from xnat.exceptions import XNATResponseError
from xnat_ingest.helpers.remotes import (
LocalSessionListing,
SessionListing,
SessionOnlyListing,
calculate_checksums,
compare_resource_with_xnat,
dir_older_than,
get_xnat_checksums,
get_xnat_resource,
get_xnat_session,
iterate_s3_sessions,
list_session_dirs,
split_resource_by_modality,
)
from ..exceptions import IncompleteCheckumsException
from ..helpers.arg_types import StoreCredentials, UploadMethod
from ..helpers.logging import logger
from ..helpers.metadata import Metadata
from ..model.resource import ImagingResource
from ..model.session import ImagingSession
def has_scan_dicom(resources: ty.Iterable[ImagingResource]) -> bool:
"""Whether resources include DICOM files attached to an imaging scan."""
return any(
resource.scan is not None and isinstance(resource.fileset, DicomCollection)
for resource in resources
)
[docs]
def upload(
input_dir: str,
xnat_repo: Xnat,
always_include: ty.Sequence[str | FileSet] = (),
store_credentials: StoreCredentials | None = None,
require_manifest: bool = True,
methods: ty.Sequence[UploadMethod] = (),
wait_period: int = 0,
num_files_per_batch: int = 0,
check_checksums: bool = True,
s3_cache_dir: ty.Optional[Path] = None,
raise_errors: bool = False,
dry_run: bool = False,
max_workers: ty.Optional[int] = None,
) -> list[str]:
"""Upload sorted sessions in the given staging directory to XNAT
Parameters
----------
input_dir: Path
The directory containing the sessions to upload. Each session should be in a separate subdirectory.
xnat_repo: Xnat
The XNAT repository to upload to
always_include: Sequence[str]
A sequence of scan types or file paths to always include in the upload regardless of whether they are
explicitly specified in the frameset definition
raise_errors: bool
Whether to raise errors that occur during upload or to log them and continue with the next session
store_credentials: StoreCredentials
Whether to store credentials for accessing staging directories that require authentication (e.g. S3)
require_manifest: bool
Whether to require a manifest file in each session directory that specifies the resources to upload
and their checksums
methods: Sequence[UploadMethod]
The upload method to use for each datatype (e.g. 'tgz_file' or 'directory')
wait_period: int
The minimum age in seconds of session directories to upload (only applicable for local staging directories)
num_files_per_batch: int
The number of files to upload in each batch when uploading resources with the 'directory' method
(if 0, all files will be uploaded in a single batch)
check_checksums: bool
Whether to check checksums of uploaded resources against the checksums specified in the manifest file and
the checksums of the files in the staged resources (if available) to verify that they were
dry_run: bool
Whether to list the sessions that would be uploaded instead of actually uploading them
max_workers: int, optional
The number of threads to use to upload resources within a session concurrently.
Different resources map to different scans/catalogs on XNAT so are safe to
upload in parallel; a failure uploading one resource doesn't stop the others
from being attempted. If None, defaults to
`concurrent.futures.ThreadPoolExecutor`'s default.
"""
errors = []
# Ensure input_path is a string so we can check for s3://
input_dir = str(input_dir)
# Note that this context manager doesn't do anything if the connection is
# already open, so it's safe to use even if the connection is already open
with xnat_repo.connection:
# DROP THE CLIENT-SIDE VIEW OF XNAT BEFORE DECIDING ANYTHING.
#
# `upload --loop` holds ONE connection for the life of the process, and
# xnatpy caches project/subject/experiment listings on it. Without this,
# every pass answers "does this already exist on XNAT?" from a snapshot
# taken when the process started, so nothing an operator does in XNAT is
# ever visible to a long-running uploader.
#
# The failure that motivated it: an operator deletes a partially
# uploaded session in XNAT so the pipeline will re-upload it. The next
# pass reads the cache, still sees the session, logs "Skipping ... as all
# the resources already exist on XNAT", and skips it for ever. No error,
# no retry. Only restarting the process recovers it, and nothing tells
# the operator that. Confirmed on a live deployment: deleting alone
# changed nothing; deleting AND restarting uploaded all 383 files.
#
# The same staleness hides a NEWLY CREATED project, where
# `connection.projects[...]` raises and the caller reports
# "Project '<id>' does not exist on XNAT" about a project that is plainly
# visible in the web UI.
#
# THIS DOES NOT RECONNECT, AND THAT IS THE POINT. A per-pass
# close/reopen would re-authenticate 1440 times a day and recreate the
# session-per-minute churn that holding a single connection was
# introduced to avoid. XNATSession.clearcache() only empties local dicts
# and listing caches: it does not log out, re-authenticate, or touch the
# HTTP session. MEASURED against a live XNAT: 60 consecutive
# clearcache+re-read cycles produced exactly ONE session id, unchanged
# throughout and released cleanly on disconnect.
xnat_repo.connection.clearcache()
num_sessions: int
sessions: ty.Iterable[SessionListing]
if input_dir.startswith("s3://"):
if s3_cache_dir is None:
s3_cache_dir = Path(tempfile.mkdtemp())
logger.info(
f"Using temporary directory '{s3_cache_dir}' to cache S3 files during upload"
)
sessions = iterate_s3_sessions(
input_dir, store_credentials, s3_cache_dir, wait_period=wait_period
)
# bit of a hack: number of sessions is the first item in the iterator
num_sessions = next(sessions) # type: ignore[assignment]
else:
sessions = []
for session_dir in list_session_dirs(input_dir):
if dir_older_than(session_dir, wait_period):
if "." in session_dir.name:
sessions.append(LocalSessionListing(session_dir))
else:
sessions.append(SessionOnlyListing(session_dir))
else:
logger.info(
"Skipping '%s' session as it has been modified recently",
session_dir,
)
num_sessions = len(sessions)
logger.info(
"Found %d sessions in staging directory to stage'%s'",
num_sessions,
input_dir,
)
framesets: dict[str, FrameSet] = {}
for session_listing in tqdm(
sessions,
total=num_sessions,
desc=f"Processing staged sessions found in '{input_dir}'",
):
if dry_run:
logger.info(
"Would attempt to upload '%s' if not dry run",
session_listing.name,
)
continue
try:
if session_listing.all_uploaded(xnat_repo.connection):
logger.info(
"Skipping upload of '%s' as all the resources already exist on XNAT",
session_listing.name,
)
continue # skip as session already exists
if isinstance(session_listing, SessionOnlyListing):
xsession = session_listing.find_xnat_session(xnat_repo.connection)
if xsession is None:
raise RuntimeError(
f"No XNAT session found with label '{session_listing.session_id}'. "
"Ensure the session exists on XNAT before uploading session-only resources."
)
for resource_name in session_listing.resource_paths:
resource = ImagingResource.load(
session_listing.cache_path / resource_name,
require_manifest=require_manifest,
check_checksums=check_checksums,
)
uri = f"{xsession.uri}/resources/{resource.name}"
xnat_repo.connection.put(uri)
xnat_repo.connection.clearcache()
xresource = xnat_repo.connection.create_object(uri)
xresource.upload_dir(
session_listing.cache_path / resource.name,
method=UploadMethod.select_method(
methods, type(resource.fileset)
),
)
logger.info(
"Uploaded '%s' to session '%s'",
resource.name,
session_listing.session_id,
)
logger.info(
"Successfully uploaded all resources to '%s'",
session_listing.session_id,
)
continue
session = ImagingSession.load(
session_listing.cache_path,
require_manifest=require_manifest,
check_checksums=check_checksums,
)
# SIGNAL, DO NOT REFUSE. A staged session with no session-level
# __METADATA__.json is one that no completed `save()` produced:
# staged by an older version, assembled by hand, or left behind by
# a run that died between writing the scans and writing the
# metadata. Refusing it was considered and rejected, because
# `ImagingSession.load` treats the file as optional by design and
# upstream ships `upload` to people whose staging directories we
# have never seen. A skip that waits for ever is worse than the
# partial upload it prevents, because a partial upload is at
# least visible.
#
# So it uploads, and says so. The event is what makes the case
# countable: without it the only trace is a session that behaves
# slightly differently from every other one, for no visible reason.
if not (session_listing.cache_path / Metadata.FNAME).exists():
logger.warning(
"Staged session '%s' has no session-level %s. Uploading it "
"anyway, but nothing recorded what produced it, so its "
"project/subject/session ids come from the directory name "
"alone and cannot be cross-checked.",
session_listing.session_id,
Metadata.FNAME,
extra={
"event": "session_metadata_missing",
"session": session_listing.session_id,
},
)
# Create corresponding session on XNAT
logger.debug(
"Creating XNAT session for '%s' in project '%s'",
session.session_id,
session.project_id,
)
xproject = xnat_repo.connection.projects[session.project_id]
# Access Arcana frameset associated with project
try:
frameset = framesets[session.project_id]
except KeyError:
try:
frameset = FrameSet.load(session.project_id, xnat_repo)
except Exception as e:
if not always_include:
logger.error(
"Did not load frameset definition (%s) from %s project "
"on %s. Either '--always-include' flag must be used or "
"the frameset must be defined on XNAT using the `frametree` "
"command line tool (see https://arcanaframework.github.io/frametree/).",
e,
session.project_id,
xnat_repo.server,
)
continue
else:
frameset = None
framesets[session.project_id] = frameset
# Get the XNAT session object (creates it if it does not exist)
xsession = get_xnat_session(session, xproject)
# Anonymise DICOMs and save to directory prior to upload
if always_include:
logger.info(
f"Including {always_include} scans/files in upload from '{session.name}' to "
f"{session.path} regardless of whether they are explicitly specified"
)
selected_resources = sorted(
session.select_resources(frameset, always_include=always_include)
)
session_has_dicom = has_scan_dicom(selected_resources)
# Resolve which resources need uploading sequentially -- this can
# create new scans/resources on XNAT and mutates xsession's/xscan's
# shared caches, so isn't safe to do concurrently. The actual upload
# of each resource's files is independent (different resources map to
# different scan/resource catalogs on XNAT) so is safe to fan out.
to_upload: list[tuple[ImagingResource, ty.Any, ty.Any]] = []
incomplete_on_xnat: list[str] = []
repaired_on_xnat: list[str] = []
for resource in selected_resources:
# A DICOM resource whose series carries more than one modality is
# split into one part per modality here, uploaded to its own scan,
# so that XNAT never has to split it itself when it rebuilds the
# session from the headers (see `split_resource_by_modality`)
for part in split_resource_by_modality(resource):
try:
xresource, only_files = get_xnat_resource(part, xsession)
except IncompleteCheckumsException as e:
# Exists on XNAT but is short and cannot be repaired by
# uploading. Record it so the session does not report as
# fully uploaded.
logger.error("%s", e.msg)
incomplete_on_xnat.append(part.path)
continue
if xresource is None:
logger.info(
"Skipping '%s' resource as it is already uploaded",
part.path,
)
continue # skipping as resource already exists
if only_files is not None:
repaired_on_xnat.append(part.path)
to_upload.append((part, xresource, only_files))
def _upload_resource(
resource: ImagingResource,
xresource: ty.Any,
only_files: ty.Optional[ty.Set[str]] = None,
) -> None:
"""Upload a resource, or just the files XNAT is missing.
`only_files` is set when XNAT already holds a strict subset
of what we have. Re-sending the files it already has would
be pointless and, on a large resource, expensive.
"""
wanted_fspaths = select_files_to_upload(
list(resource.fileset.fspaths),
resource.fileset.parent,
only_files,
resource.path,
)
logger.debug(
"Uploading '%s' resource to '%s'",
resource.path,
xresource,
)
if isinstance(resource.fileset, File):
for fspath in wanted_fspaths:
logger.debug(
"Uploading '%s' to '%s' in %s",
fspath,
fspath.name,
xresource,
)
xresource.upload(str(fspath), fspath.name)
else:
# Upload the contents of the resource to XNAT
upload_method = UploadMethod.select_method(
methods, type(resource.fileset)
)
# Get the directory containing the files to upload
# and create a temporary upload directory alongside it
# to hardlink files to upload in each batch into
dir_to_upload = resource.fileset.parent
upload_dir = dir_to_upload.parent / (
"." + dir_to_upload.name + "-upload"
)
# Split the files to upload into batches and hardlink them into
# separate directories so we can use upload_dir
files_to_upload = wanted_fspaths
num_files = len(files_to_upload)
batch_size = (
num_files_per_batch
if num_files_per_batch > 0
else num_files
)
num_batches = math.ceil(num_files / batch_size)
logger.debug(
"Uploading %s files to '%s' in %s in %s batches of %s files using '%s' method",
num_files,
resource.path,
xresource,
num_batches,
batch_size,
upload_method,
)
for i in range(num_batches):
# Create a temporary directory to upload the batch from
if upload_dir.exists():
shutil.rmtree(upload_dir)
upload_dir.mkdir()
for fspath in files_to_upload[
i * batch_size : (i + 1) * batch_size
]:
dest = upload_dir / fspath.relative_to(dir_to_upload)
dest.hardlink_to(fspath)
logger.debug(
"Uploading batch %s of %s of '%s' to %s with '%s' method",
i,
num_batches,
upload_dir,
xresource,
upload_method,
)
xresource.upload_dir(upload_dir, method=upload_method)
shutil.rmtree(upload_dir)
if check_checksums:
logger.debug("retrieving checksums for %s", xresource)
remote_checksums = get_xnat_checksums(xresource)
# Names are always available, digests are not, so this
# runs either way rather than being skipped on a site
# without enableChecksums.
logger.debug("calculating checksums for %s", xresource)
calc_checksums = calculate_checksums(resource.fileset)
# Compared file by file, not as two whole dicts. A
# resource just topped up holds real digests for the
# files already there and empty ones for those just
# added, and a whole-dict `!=` calls that a mismatch.
comparison = compare_resource_with_xnat(
calc_checksums, remote_checksums
)
if not comparison.comparable:
logger.debug(
"Remote checksums were not calculated for %s "
"(requires `enableChecksums` to be set "
"site-wide), comparing file names only",
xresource,
)
if not comparison.complete:
raise RuntimeError(
"Checksums do not match after upload of "
f"'{resource.path}' resource.\n"
f"Extra keys were {sorted(comparison.extra)}\n"
f"Missing keys were {sorted(comparison.missing)}\n"
"Mismatching files were "
f"{sorted(comparison.differing)}\n"
f"Remote checksums were {remote_checksums}\n"
f"Calculated checksums were {calc_checksums}\n"
)
else:
logger.debug(
"Not checking checksums for '%s' resource as checksum "
"checking is disabled",
resource.path,
)
logger.info(f"Uploaded '{resource.path}' in '{session.name}'")
resource_errors: list[tuple[ImagingResource, BaseException]] = []
with ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = {
executor.submit(
_upload_resource, resource, xresource, only_files
): resource
for resource, xresource, only_files in to_upload
}
for future in tqdm(
as_completed(futures),
total=len(futures),
desc=f"Uploading resources found in {session.name}",
):
resource = futures[future]
try:
future.result()
except Exception as e:
logger.error(
"Failed to upload '%s' resource in '%s': %s\n%s",
resource.path,
session.name,
e,
traceback.format_exc(),
)
resource_errors.append((resource, e))
if repaired_on_xnat:
# A repair means an earlier pass left this session short,
# so it is worth a session-level line of its own.
logger.info(
"Repaired %d incomplete resource(s) on XNAT in '%s': %s",
len(repaired_on_xnat),
session.name,
sorted(repaired_on_xnat),
)
msg = session_upload_verdict(
session_name=session.name,
num_attempted=len(to_upload),
failed_paths=[r.path for r, _ in resource_errors],
incomplete_paths=incomplete_on_xnat,
)
if msg is not None:
errors.append(msg)
if raise_errors and resource_errors:
raise RuntimeError(msg) from resource_errors[0][1]
logger.error(msg)
# Success is not announced here: metadata extraction and
# pipeline triggering still follow, and there is one report at
# the end of all of it.
# Extract DICOM metadata
if session_has_dicom:
logger.info("Extracting metadata from DICOMs on XNAT..")
try:
xnat_repo.connection.put(
f"/data/experiments/{xsession.id}?pullDataFromHeaders=true"
)
except XNATResponseError as e:
logger.warning(
f"Failed to extract metadata: {e}\nResponse: "
f"{e.response.text if hasattr(e, 'response') else 'N/A'}"
)
try:
xnat_repo.connection.put(
f"/data/experiments/{xsession.id}?fixScanTypes=true"
)
except XNATResponseError as e:
logger.warning(
f"Failed to fix scan types in '{session.name}': {e}\nResponse: "
f"{e.response.text if hasattr(e, 'response') else 'N/A'}"
)
else:
logger.info(
"Skipping DICOM header pull / scan-type fixing for '%s' as it "
"contains no DICOM data",
session.name,
)
try:
xnat_repo.connection.put(
f"/data/experiments/{xsession.id}?triggerPipelines=true"
)
except XNATResponseError as e:
logger.warning(
f"Failed to trigger pipelines in '{session.name}': {e}\nResponse: "
f"{e.response.text if hasattr(e, 'response') else 'N/A'}"
)
# Guarded by the verdict: unconditional, this claimed success
# in the same pass that reported the session as failed.
if msg is None:
logger.info(f"Successfully uploaded all files in '{session.name}'")
except Exception as e:
if not raise_errors:
error_msg = [
(
f"Skipping upload of '{session_listing.name}' due to error: \"{e}\""
f"\n{traceback.format_exc()}\n\n"
)
]
logger.error("".join(error_msg))
errors.extend(error_msg)
continue
else:
raise
if errors:
logger.error("Upload completed with %s errors", len(errors))
else:
logger.info("Upload completed successfully")
return errors
def select_files_to_upload(
fspaths: ty.Sequence[Path],
parent: Path,
only_files: ty.Optional[ty.Set[str]],
resource_path: str,
) -> list[Path]:
"""Pick the staged files to send, and refuse to send none of them by accident.
`only_files` is set when XNAT already holds a strict subset of what we have,
and it names the files it is missing. None means upload everything.
The two name shapes are not guaranteed to agree, they are ENFORCED to agree
elsewhere. `parent` is FileSet.parent, which is commonpath() over the staged
files and so collapses to a subdirectory when every file sits in one, while
the manifest behind `only_files` is keyed to the resource directory. They
coincide for a flat resource directory and can differ for a nested one.
ImagingResource.load() is what normally makes them agree: it recomputes the
manifest keys with this same `relative_to` and raises on a mismatch. That
runs only under `check_checksums`, which also gates the post-upload
verification, so `--dont-check-checksums` removes both ends at once. Hence
this fails closed rather than trusting the shapes.
Raises
------
RuntimeError
when `only_files` names files to upload but nothing matched, which would
otherwise upload zero files and report the resource as uploaded
"""
if only_files is None:
return list(fspaths)
wanted = [p for p in fspaths if str(p.relative_to(parent)) in only_files]
if only_files and not wanted:
# Fail closed. only_files holds names XNAT is known to be missing, so
# matching nothing means the names and the paths disagree, not that
# there is nothing to do. With num_files_per_batch > 0 the batch loop
# runs zero times and the resource is then logged as uploaded.
raise RuntimeError(
f"Refusing to repair '{resource_path}': XNAT is missing "
f"{len(only_files)} file(s) {sorted(only_files)[:5]} but none of "
f"the {len(fspaths)} staged file(s) matched those names, so nothing "
"would be uploaded. The manifest keys and the staged paths "
"disagree. Delete the resource on XNAT to have it uploaded afresh."
)
return wanted
def session_upload_verdict(
session_name: str,
num_attempted: int,
failed_paths: ty.Sequence[str],
incomplete_paths: ty.Sequence[str],
) -> ty.Optional[str]:
"""Summarise how a session's upload went, or None if it went cleanly.
Kept separate from upload() because the SUCCESS branch is the one that
misled operators: resources skipped as "already uploaded" never entered
`to_upload`, so they could never reach `resource_errors`, so a session that
delivered a fraction of its files still logged "Successfully uploaded all
files". Returning None only when BOTH lists are empty makes that impossible
to reintroduce by accident, and makes the rule testable on its own.
Parameters
----------
session_name : str
the session being reported on
num_attempted : int
how many resources were actually attempted
failed_paths : Sequence[str]
resources whose upload raised
incomplete_paths : Sequence[str]
resources already on XNAT but missing files held in staging, which this
uploader will not repair
Returns
-------
str or None
an operator-facing message, or None if nothing was wrong
"""
if not failed_paths and not incomplete_paths:
return None
parts = []
if failed_paths:
parts.append(
f"{len(failed_paths)} of {num_attempted} resource(s) failed to upload: "
+ ", ".join(failed_paths)
)
if incomplete_paths:
parts.append(
f"{len(incomplete_paths)} resource(s) already on XNAT but incomplete, "
"and NOT repaired: "
+ ", ".join(incomplete_paths)
+ ". Delete them on XNAT to allow re-upload."
)
return f"'{session_name}' did not upload cleanly: " + "; ".join(parts)