cdm_data_loaders.pipelines.ncbi_ftp_download ¶

def download_batch(  # noqa: PLR0913
    manifest_path: Path,
    output_dir: Path,
    threads: int = 4,
    ftp_host: str = FTP_HOST,
    limit: int | None = None,
    *,
    on_assembly_downloaded: Callable[[PurePosixPath], None] | None = None,
    progress_desc: str = "Downloading from NCBI FTP",
) -> dict[str, Any]:
    """Download all assemblies listed in the manifest.

    :param manifest_path: path to the transfer manifest file
    :type manifest_path: Path
    :param output_dir: base output directory
    :type output_dir: Path
    :param threads: number of parallel download threads
    :type threads: int
    :param ftp_host: FTP hostname
    :type ftp_host: str
    :param limit: optional limit for testing
    :type limit: int | None
    :param on_assembly_downloaded: optional callback invoked with each assembly's FTP path after download succeeds
    :type on_assembly_downloaded: Callable[[PurePosixPath], None] | None
    :param progress_desc: description shown on the tqdm progress bar
    :type progress_desc: str
    :return: report dict with overall stats
    """
    with manifest_path.open() as f:
        assembly_paths = [PurePosixPath(line.strip()) for line in f if line.strip() and not line.startswith("#")]

    if limit:
        assembly_paths = assembly_paths[:limit]

    logger.info("Starting download of %d assemblies with %d threads", len(assembly_paths), threads)

    pool = ThreadLocalFTP(ftp_host)
    lock = threading.Lock()
    success_count = 0
    failed: list[dict[str, str]] = []
    all_stats: list[dict[str, Any]] = []

    def _download_one(path: PurePosixPath) -> tuple[str, Exception | None]:
        nonlocal success_count

        @retry(
            retry=retry_if_exception_type((error_temp, BrokenPipeError, ConnectionResetError, EOFError)),
            stop=stop_after_attempt(5),
            wait=wait_exponential(multiplier=1, min=5, max=60),
            reraise=True,
            before_sleep=before_sleep_log(logger, logging.WARNING),
        )
        def _attempt() -> dict[str, Any]:
            return download_assembly_to_local(path, output_dir, ftp_host=ftp_host, ftp=pool.get())

        try:
            stats = _attempt()
        except Exception as e:  # noqa: BLE001
            return str(path), e

        if on_assembly_downloaded is not None:
            on_assembly_downloaded(path)

        with lock:
            success_count += 1
            all_stats.append(stats)
        return str(path), None

    try:
        with (
            tqdm.tqdm(total=len(assembly_paths), unit="assembly", desc=progress_desc, smoothing=0.01) as pbar,
            ThreadPoolExecutor(max_workers=threads) as executor,
        ):
            futures = {executor.submit(_download_one, p): p for p in assembly_paths}
            for future in as_completed(futures):
                path, error = future.result()
                if error:
                    logger.error("FAILED: %s: %s", path, error)
                    with lock:
                        failed.append({"path": path, "error": str(error)})
                pbar.update(1)
    finally:
        pool.close_all()

    report: dict[str, Any] = {
        "timestamp": datetime.now(UTC).isoformat(),
        "total_attempted": len(assembly_paths),
        "succeeded": success_count,
        "failed": len(failed),
        "failures": failed,
        "assembly_stats": all_stats,
    }

    report_path = output_dir / "download_report.json"
    report_path.parent.mkdir(parents=True, exist_ok=True)
    with report_path.open("w") as f:
        json.dump(report, f, indent=2)
    logger.info("Download report written to: %s", report_path)

    logger.info(
        "SUMMARY: %d attempted, %d succeeded, %d failed",
        len(assembly_paths),
        success_count,
        len(failed),
    )

    return report

def run_download(config: DownloadSettings) -> None:
    """Main CTS entry point for Phase 2 download.

    :param config: validated download settings
    :type config: DownloadSettings
    """
    report = download_batch(
        manifest_path=config.manifest,
        output_dir=config.output_dir,
        threads=config.threads,
        ftp_host=config.ftp_host,
        limit=config.limit,
    )
    if report["failed"] > 0:
        msg = f"Download completed with {report['failed']} failures"
        raise RuntimeError(msg)

def download_and_stage(  # noqa: PLR0913
    *,
    bucket: PurePosixPath,
    staging_key_prefix: PurePosixPath,
    manifest_s3_key: PurePosixPath | None = None,
    manifest_local_path: Path | None = None,
    threads: int = 4,
    ftp_host: str = FTP_HOST,
    limit: int | None = None,
    dry_run: bool = False,
) -> dict[str, Any]:
    """Download assemblies from NCBI FTP and stage them to S3 (Phase 2).

    Exactly one of *manifest_s3_key* or *manifest_local_path* must be given.

    Downloads and uploads are pipelined per assembly: each worker downloads one
    assembly, immediately uploads its files to S3, then deletes the local copies
    before picking up the next assembly.  At most *threads* assembly directories
    exist on disk simultaneously, preventing disk exhaustion on large batches.

    :param bucket: destination S3 bucket name
    :type bucket: PurePosixPath
    :param staging_key_prefix: key prefix inside the bucket (e.g. ``"staging/run1/"``)
    :type staging_key_prefix: PurePosixPath
    :param manifest_s3_key: S3 object key of the transfer manifest within *bucket*
    :type manifest_s3_key: PurePosixPath | None
    :param manifest_local_path: local path to the transfer manifest file
    :type manifest_local_path: Path | None
    :param threads: number of parallel download-and-upload workers
    :type threads: int
    :param ftp_host: NCBI FTP hostname
    :type ftp_host: str
    :param limit: optional limit for testing
    :type limit: int | None
    :param dry_run: when ``True``, download but skip all S3 uploads
    :type dry_run: bool
    :return: download report extended with ``staged_objects``, ``staging_key_prefix``, ``dry_run``
    """
    if manifest_s3_key is not None and manifest_local_path is not None:
        msg = "Provide exactly one of manifest_s3_key or manifest_local_path, not both"
        raise ValueError(msg)
    if manifest_s3_key is None and manifest_local_path is None:
        msg = "One of manifest_s3_key or manifest_local_path must be provided"
        raise ValueError(msg)

    with tempfile.TemporaryDirectory() as _tmpdir:
        tmp = Path(_tmpdir)
        manifest_dest = tmp / "transfer_manifest.txt"

        if manifest_s3_key is not None:
            s3 = get_s3_client()
            response = s3.get_object(Bucket=str(bucket), Key=str(manifest_s3_key))
            manifest_dest.write_bytes(response["Body"].read())
            logger.info("Manifest read from S3: s3://%s/%s", bucket, manifest_s3_key)
        elif manifest_local_path:
            manifest_dest.write_bytes(manifest_local_path.read_bytes())
            logger.info("Manifest read from local path: %s", manifest_local_path)

        staged_objects = 0
        stage_lock = threading.Lock()

        def _stage_assembly(path: PurePosixPath) -> None:
            nonlocal staged_objects
            _db, assembly_dir_name, _accession = parse_assembly_path(path)
            assembly_local_dir = tmp / build_accession_path(assembly_dir_name)
            count = _upload_assembly_dir(assembly_local_dir, tmp, bucket, staging_key_prefix)
            with stage_lock:
                staged_objects += count

        report = download_batch(
            manifest_path=manifest_dest,
            output_dir=tmp,
            threads=threads,
            ftp_host=ftp_host,
            limit=limit,
            on_assembly_downloaded=None if dry_run else _stage_assembly,
            progress_desc="Downloading & staging",
        )

        if not dry_run:
            report_path = tmp / "download_report.json"
            if upload_file(report_path, str(bucket / staging_key_prefix), show_progress=False):
                staged_objects += 1
            else:
                logger.warning("Failed to upload download report to s3://%s/%s", bucket, staging_key_prefix)
            logger.info("Staged %d objects to s3://%s/%s", staged_objects, bucket, staging_key_prefix)

    return {
        **report,
        "staged_objects": staged_objects,
        "staging_key_prefix": staging_key_prefix,
        "dry_run": dry_run,
    }