From fa534cd4f28d372fa5aa17e73674b628690925f2 Mon Sep 17 00:00:00 2001 From: Codex Date: Mon, 20 Jul 2026 21:43:52 -0400 Subject: [PATCH] Handle reference genome download failures --- .../controllers/cli_controller.py | 3 +- .../scripts/reference_genome_manager.py | 104 ++++++++++-- .../scripts/test_reference_genome_manager.py | 156 ++++++++++++++++++ 3 files changed, 252 insertions(+), 11 deletions(-) create mode 100644 tests/scripts/test_reference_genome_manager.py diff --git a/SigProfilerMatrixGenerator/controllers/cli_controller.py b/SigProfilerMatrixGenerator/controllers/cli_controller.py index 0c43333..2f10485 100644 --- a/SigProfilerMatrixGenerator/controllers/cli_controller.py +++ b/SigProfilerMatrixGenerator/controllers/cli_controller.py @@ -60,7 +60,8 @@ def parse_arguments_install(args: List[str]) -> argparse.Namespace: "--local_genome", help=""" Install an offline reference genome downloaded from the Alexandrov Lab's FTP server. - Provide the absolute path to the locally-stored genome file. + Provide the absolute path to the directory containing the locally-stored + .tar.gz file. For downloads, visit AlexandrovLab's ftp server: ftp://alexandrovlab-ftp.ucsd.edu/pub/tools/SigProfilerMatrixGenerator/ """, diff --git a/SigProfilerMatrixGenerator/scripts/reference_genome_manager.py b/SigProfilerMatrixGenerator/scripts/reference_genome_manager.py index 660f2ac..df6d1c5 100644 --- a/SigProfilerMatrixGenerator/scripts/reference_genome_manager.py +++ b/SigProfilerMatrixGenerator/scripts/reference_genome_manager.py @@ -348,6 +348,10 @@ logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s") +class GenomeDownloadError(RuntimeError): + """Raised when a reference genome archive cannot be downloaded or installed.""" + + class ReferenceGenomeManager: """ A class for downloading and managing reference genomes. @@ -370,38 +374,77 @@ def download_genome(self, genome_name): file_name = f"{genome_name}.tar.gz" local_filepath = self.reference_dir.get_tsb_dir() / file_name + local_filepath.parent.mkdir(parents=True, exist_ok=True) + + downloaded = False + failures = [] # First try to download using FTP for server_key in FTP_SERVERS: server_info = FTP_SERVERS[server_key] try: + self._remove_partial_archive(local_filepath) self._download_via_ftplib( server_info["address"], server_info["path"], file_name, local_filepath, ) - logging.info(f"Downloaded {genome_name} from {server_key} using FTP.") - break # Exit the loop if download is successful + if self._downloaded_archive_exists(local_filepath): + logging.info( + f"Downloaded {genome_name} from {server_key} using FTP." + ) + downloaded = True + break # Exit the loop if download is successful + failures.append( + f"{server_key} FTP: download completed but did not create " + f"{local_filepath}." + ) except ftplib.all_errors as e: + failures.append(f"{server_key} FTP: {e}") logging.info(f"Attempt to download from {server_key} failed: {e}") + if downloaded: + break if shutil.which("curl"): try: + self._remove_partial_archive(local_filepath) self._download_via_curl( server_info["address"], server_info["path"], file_name, local_filepath, ) - logging.info(f"Downloaded {genome_name} using curl.") - break # Exit the loop if download is successful - except (subprocess.CalledProcessError, KeyboardInterrupt) as e: + if self._downloaded_archive_exists(local_filepath): + logging.info(f"Downloaded {genome_name} using curl.") + downloaded = True + break # Exit the loop if download is successful + failures.append( + f"{server_key} curl: download completed but did not create " + f"{local_filepath}." + ) + except (subprocess.CalledProcessError, OSError) as e: + failures.append(f"{server_key} curl: {e}") logging.error(f"Curl download failed with error: {e}") else: + failures.append(f"{server_key} curl: curl is not available.") logging.error("Curl is not available. Unable to download the file.") - return - self._unzip_file(local_filepath) + if not downloaded: + self._remove_partial_archive(local_filepath) + raise GenomeDownloadError( + self._format_download_failure_message(file_name, failures) + ) + + try: + self._unzip_file(local_filepath) + except (tarfile.TarError, OSError) as e: + self._remove_partial_archive(local_filepath) + raise GenomeDownloadError( + f"Downloaded archive {local_filepath} could not be extracted. " + "Please retry the installation; if the problem persists, manually " + "download a fresh archive and install it with --local_genome in the " + "CLI or offline_files_path in the Python API." + ) from e local_filepath.unlink() logging.info(f"{genome_name} has been successfully installed.") @@ -568,16 +611,57 @@ def _download_via_curl(self, ftp_server, ftp_path, filename, local_filepath): Helper function for download_genome. """ try: - url = f"ftp://{ftp_server}/{ftp_path}/{filename}" - command = ["curl", "-o", str(local_filepath), url] + local_filepath.parent.mkdir(parents=True, exist_ok=True) + url = f"ftp://{ftp_server}/{ftp_path.rstrip('/')}/{filename}" + command = [ + "curl", + "--fail", + "--location", + "--retry", + "3", + "--connect-timeout", + "30", + "-o", + str(local_filepath), + url, + ] subprocess.run(command, check=True) except subprocess.CalledProcessError as e: logging.error(f"Curl download failed with error: {e}.") raise + def _downloaded_archive_exists(self, local_filepath): + return local_filepath.exists() and local_filepath.stat().st_size > 0 + + def _remove_partial_archive(self, local_filepath): + try: + local_filepath.unlink() + except FileNotFoundError: + pass + + def _format_download_failure_message(self, file_name, failures): + message = [ + f"Unable to download {file_name} from the configured reference genome mirrors.", + ] + if failures: + message.append("Download attempts failed with:") + message.extend(f" - {failure}" for failure in failures) + message.append( + "Please check FTP/network access or manually download the archive from " + "ftp://alexandrovlab-ftp.ucsd.edu/pub/tools/SigProfilerMatrixGenerator/ " + "and install it with --local_genome in the CLI or offline_files_path in " + "the Python API. Both options expect the directory containing the archive." + ) + return "\n".join(message) + def _unzip_file(self, file_path): with tarfile.open(file_path, "r:gz") as tar: - tar.extractall(path=self.reference_dir.get_tsb_dir()) + extraction_options = {} + if hasattr(tarfile, "data_filter"): + extraction_options["filter"] = "data" + tar.extractall( + path=self.reference_dir.get_tsb_dir(), **extraction_options + ) def _verify_checksum(self, file_path, checksum): """ diff --git a/tests/scripts/test_reference_genome_manager.py b/tests/scripts/test_reference_genome_manager.py new file mode 100644 index 0000000..feda408 --- /dev/null +++ b/tests/scripts/test_reference_genome_manager.py @@ -0,0 +1,156 @@ +import ftplib +import io +import subprocess +import tarfile + +import pytest + +from SigProfilerMatrixGenerator.scripts import reference_genome_manager + + +def write_test_archive(archive_path, genome_name="test_genome"): + archive_path.parent.mkdir(parents=True, exist_ok=True) + contents = b"test chromosome contents" + member = tarfile.TarInfo(f"{genome_name}/1.txt") + member.size = len(contents) + with tarfile.open(archive_path, "w:gz") as archive: + archive.addfile(member, io.BytesIO(contents)) + + +def test_download_genome_installs_archive_from_ftp(monkeypatch, tmp_path): + manager = reference_genome_manager.ReferenceGenomeManager(reference_dir=tmp_path) + archive_path = tmp_path.resolve() / "tsb" / "test_genome.tar.gz" + installed_file = tmp_path.resolve() / "tsb" / "test_genome" / "1.txt" + + monkeypatch.setattr(manager, "is_genome_installed", lambda genome: False) + monkeypatch.setattr( + manager, + "_download_via_ftplib", + lambda *args: write_test_archive(args[-1]), + ) + + def unexpected_curl(*args): + raise AssertionError("curl should not run after a successful FTP download") + + monkeypatch.setattr(manager, "_download_via_curl", unexpected_curl) + + manager.download_genome("test_genome") + + assert installed_file.read_bytes() == b"test chromosome contents" + assert not archive_path.exists() + + +def test_download_genome_falls_back_to_curl_on_same_mirror(monkeypatch, tmp_path): + manager = reference_genome_manager.ReferenceGenomeManager(reference_dir=tmp_path) + calls = [] + + monkeypatch.setattr(manager, "is_genome_installed", lambda genome: False) + monkeypatch.setattr(reference_genome_manager.shutil, "which", lambda name: "curl") + + def fail_ftp(server, path, filename, local_filepath): + calls.append(("ftp", server)) + raise ftplib.error_temp("Connection reset by peer") + + def successful_curl(server, path, filename, local_filepath): + calls.append(("curl", server)) + write_test_archive(local_filepath) + + monkeypatch.setattr(manager, "_download_via_ftplib", fail_ftp) + monkeypatch.setattr(manager, "_download_via_curl", successful_curl) + + manager.download_genome("test_genome") + + assert calls == [ + ("ftp", "alexandrovlab-ftp.ucsd.edu"), + ("curl", "alexandrovlab-ftp.ucsd.edu"), + ] + assert (tmp_path.resolve() / "tsb" / "test_genome" / "1.txt").exists() + + +def test_download_via_curl_uses_retrying_ftp_command(monkeypatch, tmp_path): + manager = reference_genome_manager.ReferenceGenomeManager(reference_dir=tmp_path) + archive_path = tmp_path / "downloads" / "ebv.tar.gz" + observed = {} + + def capture_run(command, check): + observed["command"] = command + observed["check"] = check + + monkeypatch.setattr(reference_genome_manager.subprocess, "run", capture_run) + + manager._download_via_curl( + "alexandrovlab-ftp.ucsd.edu", + "pub/tools/SigProfilerMatrixGenerator/", + "ebv.tar.gz", + archive_path, + ) + + assert observed == { + "command": [ + "curl", + "--fail", + "--location", + "--retry", + "3", + "--connect-timeout", + "30", + "-o", + str(archive_path), + "ftp://alexandrovlab-ftp.ucsd.edu/pub/tools/" + "SigProfilerMatrixGenerator/ebv.tar.gz", + ], + "check": True, + } + + +def test_download_genome_raises_clear_error_when_all_downloads_fail( + monkeypatch, tmp_path +): + manager = reference_genome_manager.ReferenceGenomeManager(reference_dir=tmp_path) + archive_path = tmp_path.resolve() / "tsb" / "mm10.tar.gz" + + monkeypatch.setattr(manager, "is_genome_installed", lambda genome: False) + monkeypatch.setattr(reference_genome_manager.shutil, "which", lambda name: "curl") + + def fail_ftp(*args): + raise ftplib.error_temp("Connection reset by peer") + + def fail_curl(*args): + local_filepath = args[-1] + local_filepath.write_bytes(b"partial archive") + raise subprocess.CalledProcessError(56, ["curl"]) + + monkeypatch.setattr(manager, "_download_via_ftplib", fail_ftp) + monkeypatch.setattr(manager, "_download_via_curl", fail_curl) + + with pytest.raises(reference_genome_manager.GenomeDownloadError) as error: + manager.download_genome("mm10") + + message = str(error.value) + assert "Unable to download mm10.tar.gz" in message + assert "alexandrovlab FTP" in message + assert "sanger curl" in message + assert "--local_genome" in message + assert "offline_files_path" in message + assert "directory containing the archive" in message + assert not archive_path.exists() + + +def test_download_genome_removes_bad_archive_when_extract_fails(monkeypatch, tmp_path): + manager = reference_genome_manager.ReferenceGenomeManager(reference_dir=tmp_path) + archive_path = tmp_path.resolve() / "tsb" / "mm10.tar.gz" + + monkeypatch.setattr(manager, "is_genome_installed", lambda genome: False) + + def write_bad_archive(*args): + local_filepath = args[-1] + local_filepath.write_bytes(b"not a tar.gz archive") + + monkeypatch.setattr(manager, "_download_via_ftplib", write_bad_archive) + + with pytest.raises(reference_genome_manager.GenomeDownloadError) as error: + manager.download_genome("mm10") + + assert "could not be extracted" in str(error.value) + assert "offline_files_path" in str(error.value) + assert not archive_path.exists()