From 09e12404df443cadf31a7649a2ed6ff5cca97365 Mon Sep 17 00:00:00 2001 From: Jover Lee Date: Thu, 1 Oct 2026 14:42:29 -0700 Subject: [PATCH] remote_files: Add support for directly using local cache The Snakemake storage usually checks the existence and revalidates remote files before using the local cache, which can lead to errors when running the workflow offline even if the local cache already exists. This change adds support to directly use the local cache if it exists and bypass the remote file check when 1. the config parameter `use_remote_files` is False 2. the environment variable `NEXTSTRAIN_REMOTE_FILES` is False 3. the workflow is running offline Based on discussions in and changes tested in measles repo --- snakemake/remote_files.smk | 93 ++++++++++++++++++++++++++++++++++++-- 1 file changed, 89 insertions(+), 4 deletions(-) diff --git a/snakemake/remote_files.smk b/snakemake/remote_files.smk index 844f80e..6943e75 100644 --- a/snakemake/remote_files.smk +++ b/snakemake/remote_files.smk @@ -7,8 +7,43 @@ The errors raised by storage plugins are often confusing. For instance, a HTTP underlying issue. S3 credentials errors are similarly confusing and we attempt to check these ourselves to improve UX here. """ - +import os +import socket from urllib.parse import urlparse +from snakemake.io import get_flag_value, AnnotatedString +from snakemake.logging import logger + + +def is_online(host="8.8.8.8", port=53, timeout=3) -> bool: + """ + Check if workflow has network connection. + Based on + Host: 8.8.8.8 (google-public-dns-a.google.com) + OpenPort: 53/tcp + Service: domain (DNS/TCP) + """ + try: + socket.setdefaulttimeout(timeout) + socket.socket(socket.AF_INET, socket.SOCK_STREAM).connect((host, port)) + return True + except socket.error as err: + logger.debug(f"Assuming no internet connection with error: {err}") + return False + + +# Determine whether workflow should try to fetch or revalidate remote input +# files based on prioritization order: +# 1. Per workflow config parameter `use_remote_files` +# 2. Environment variable `NEXTSTRAIN_REMOTE_FILES` +# 3. Whether the workflow has network connection +USE_REMOTE_FILES = True +if "use_remote_files" in config: + USE_REMOTE_FILES = bool(config["use_remote_files"]) +elif "NEXTSTRAIN_USE_REMOTE_FILES" in os.environ: + USE_REMOTE_FILES = bool(os.environ["NEXTSTRAIN_USE_REMOTE_FILES"]) +else: + USE_REMOTE_FILES = is_online() + # Keep a list of known public buckets, which we'll allow uncredentialled (unsigned) access to # We could make this config-definable in the future @@ -99,7 +134,7 @@ def _storage_http(*, keep_local, retries) -> snakemake.storage.StorageProviderPr return _storage_registry['http'] -def path_or_url(uri, *, keep_local=True, retries=2) -> str: +def path_or_url(uri, *, keep_local=True, retries=2) -> str | AnnotatedString: """ Intended for use in Snakemake inputs / outputs to transparently use remote resources. Returns the URI wrapped by an applicable storage plugin. Local @@ -142,12 +177,12 @@ def path_or_url(uri, *, keep_local=True, retries=2) -> str: if info.scheme=='s3': try: - return _storage_s3(bucket=info.netloc, keep_local=keep_local, retries=retries)(uri) + return _local_or_storage(_storage_s3(bucket=info.netloc, keep_local=keep_local, retries=retries)(uri)) except RemoteFilesMissingCredentials as e: raise Exception(f"AWS credentials are required to access {uri!r}") from e if info.scheme=='https': - return _storage_http(keep_local=keep_local, retries=retries)(uri) + return _local_or_storage(_storage_http(keep_local=keep_local, retries=retries)(uri)) elif info.scheme=='http': raise Exception(f"HTTP remote file support is not implemented in nextstrain workflows (attempting to access {uri!r}).\n" "Please use an HTTPS address instead.") @@ -157,3 +192,53 @@ def path_or_url(uri, *, keep_local=True, retries=2) -> str: "Please get in touch if you require this functionality and we can add it to our workflows") raise Exception(f"Input address {uri!r} (scheme={info.scheme!r}) is from a non-supported remote") + + +def _local_or_storage(wrapped: AnnotatedString) -> str | AnnotatedString: + """ + Given a storage-wrapped input (an AnnotatedString carrying a `storage_object` + flag), transparently fall back to the local cache copy as determined by + USE_REMOTE_FILES: + * If USE_REMOTE_FILES == True, then return object unchanged and let the + allow normal Snakemake retrieval, revalidation, or + MissingInputException behaviour. + * If USE_REMOTE_FILES == False and cache does not exist, raise an error. + * If USE_REMOTE_FILES == False and cache does exist, then return the + plain local path. Snakemake then treats it as an ordinary local input + (existence via `os.path.exists`, no network), allowing the workflow to + run offline from cache. + """ + so = get_flag_value(wrapped, "storage_object") + if so is None: + raise Exception("_local_or_storage must be called with a storage object AnnotatedString. Provided argument: ", wrapped) + + # Return StorageObject as-is since remote files are expected and allow it + # to fetch or revalidate cached files as needed. + # Any connection errors will be raised downstream + if USE_REMOTE_FILES: + return wrapped + + # The keep_local flag is managed by the workflow author and this cannot be + # changed by the user. So output warning for user and return the StorageObject as-is. + if not so.keep_local: + logger.warning( + f"Unable to use cached copy of {so.query!r} because the workflow " + "author has opted out of keeping a local copy for this input. " + "Please contact the workflow author to enable use of local cached files." + ) + return wrapped + + local = so.local_path() + if not local.exists(): + raise Exception( + f"Unable to find local cached file for {so.query!r}. " + "Please make sure you have run the workflow at least once with " + "internet connection and config.use_remote_files set to True to " + "to fetch and save the remote file locally.") + + logger.warning( + f"Using the cached copy of {so.query} at {local} without revalidation. " + "Set config.use_remote_files to True if you want to force revalidation." + ) + + return str(local)