import errno
import math
import shutil
import tempfile
import traceback
import typing as ty
from concurrent.futures import ThreadPoolExecutor, as_completed
from contextlib import ExitStack, contextmanager
from pathlib import Path
from fileformats.generic import File, FileSet
from fileformats.medimage import DicomCollection, DicomImage
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 (
DEFAULT_MAX_WORKERS,
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 readable DICOM files attached to a scan.
A full ``DicomCollection`` and a representative ``DicomImage`` can both be
consumed by XNAT's ``pullDataFromHeaders`` operation.
"""
return any(
resource.scan is not None
and isinstance(resource.fileset, (DicomCollection, DicomImage))
for resource in resources
)
@contextmanager
def _staged_upload_batch(
source_dir: Path, files: ty.Sequence[Path]
) -> ty.Iterator[Path]:
with tempfile.TemporaryDirectory(
prefix=f".{source_dir.name}-upload-", dir=source_dir.parent
) as upload_dir_str:
upload_dir = Path(upload_dir_str)
for fspath in files:
dest = upload_dir / fspath.relative_to(source_dir)
dest.parent.mkdir(parents=True, exist_ok=True)
# Put the file data in the batch, not a symbolic link. A link can
# point to a file in a different batch, or out of the resource.
# Keep the path of the link, because the manifest uses that path.
# On Python 3.11, hardlink_to() does not follow a link. So resolve it.
source = fspath.resolve(strict=True)
try:
dest.hardlink_to(source)
except OSError as error:
if error.errno != errno.EXDEV or source == fspath:
raise
# Copy a link target that is on a different file system. For
# any other file, raise the error: the batch is next to the
# resource, on the same file system.
shutil.copy2(source, dest)
yield upload_dir
[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,
s3_max_workers: int = DEFAULT_MAX_WORKERS,
xnat_max_workers: int = DEFAULT_MAX_WORKERS,
modality: ty.Optional[str] = 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 uploaded correctly
dry_run: bool
Whether to list the sessions that would be uploaded instead of actually uploading them
s3_max_workers: int, optional
The maximum number of concurrent S3 object downloads. Defaults to 10.
xnat_max_workers: int, optional
The maximum number of concurrent XNAT resource uploads. Defaults to 10.
modality: str, optional
The modality of the sessions to upload (e.g. 'MR', 'CT'). Defaults to the "Modality" metadata field if present
otherwise an error is raised.
"""
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, ExitStack() as cleanup:
# DROP THE CLIENT-SIDE VIEW OF XNAT BEFORE DECIDING ANYTHING.
#
# `upload --loop` holds ONE connection for the life of the process, and
# xnatpy caches what it reads on it. Decisions must not come from that
# cache: an operator who deletes a session in XNAT expects the next pass
# to upload it again, and a new project must be seen at once.
#
# So the check that runs on every pass, all_uploaded(), reads XNAT
# directly and builds no xnatpy objects. The cache is cleared only
# before a session is uploaded, which is the only step that uses
# xnatpy objects. Clearing it on EVERY pass made xnatpy build new
# listings on every pass. xnatpy 0.7.2 retains those listings; 0.8.1
# retains weak registry entries. Both accumulate during polling.
#
# None of this reconnects. A reconnect per pass would rebuild xnatpy's
# schema classes each time, which is a leak of its own.
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,
max_workers=s3_max_workers,
)
# Close the iterator when this pass ends, also on an error, so
# that it removes the session it downloaded.
cleanup.callback(sessions.close) # type: ignore[attr-defined]
# 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
# The steps below use xnatpy objects. Clear xnatpy's cache first
# so that they see XNAT as it is now.
xnat_repo.connection.clearcache()
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, modality=modality)
# 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.
"""
resource_dir = resource.fileset.parent
# The same walk as the checksums, which raises on broken links
fspaths = resource.fileset.all_file_paths()
if not fspaths:
raise RuntimeError(
f"Resource '{resource.path}' has no files to upload"
)
wanted_fspaths = select_files_to_upload(
fspaths,
resource_dir,
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)
)
# Keep batch directories alongside their source files so
# that hard links stay on the same filesystem.
dir_to_upload = resource_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):
batch_files = files_to_upload[
i * batch_size : (i + 1) * batch_size
]
with _staged_upload_batch(
dir_to_upload, batch_files
) as upload_dir:
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)
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=xnat_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. It is
# logged whether or not the re-send worked; the line that
# says it DID is repair_completed_message(), after the
# verdict below.
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)
continue
# 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:
# Same gate as the success line, same pass as the repair:
# the only point where a repair is known to have worked.
completed = repair_completed_message(session.name, repaired_on_xnat)
if completed is not None:
logger.info(
completed,
extra={
"event": "xnat_repair_completed",
"session": session.name,
"resources": len(repaired_on_xnat),
},
)
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 repair_completed_message(
session_name: str, repaired_paths: ty.Sequence[str]
) -> ty.Optional[str]:
"""The line saying a repair in this pass completed, or None if none ran.
"Repaired N incomplete resource(s)" is logged when a re-send is ATTEMPTED,
before the verdict, so it appears even when the re-send failed. This one is
logged only after a clean verdict: every re-sent resource uploaded without
error and, with checksum checking on (the default), was re-listed on XNAT
and found complete. With --dont-check-checksums it means only that no
upload call failed. Its wording deliberately matches neither
"Repaired <N> incomplete resource" nor "Successfully uploaded all files in",
which downstream alert rules key on.
Parameters
----------
session_name : str
the session being reported on
repaired_paths : Sequence[str]
resources that were short on XNAT and re-sent in this pass
Returns
-------
str or None
the log line, or None if nothing was repaired
"""
if not repaired_paths:
return None
return (
f"Completed repair of {len(repaired_paths)} incomplete resource(s) on XNAT "
f"in '{session_name}': {sorted(repaired_paths)}"
)
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)