WIP - snakemake v8+ remote storage support

phylogenetic/defaults/config.yaml

@@ -1,14 +1,38 @@
 # Sequences must be FASTA and metadata must be TSV
-# Both files must be zstd compressed
+
+# TODO XXX - fixup this mess!
 inputs:
+  # - name: ncbi
+  #   metadata: "s3://nextstrain-scratch/zika-pr-89/metadata.tsv.zst"
+  #   sequences: "s3://nextstrain-scratch/zika-pr-89/sequences.fasta.zst"
+
   - name: ncbi
     metadata: "s3://nextstrain-data/files/workflows/zika/metadata.tsv.zst"
     sequences: "s3://nextstrain-data/files/workflows/zika/sequences.fasta.zst"
 
-additional_inputs:
+  # - name: ncbi
+  #   metadata: "https://data.nextstrain.org/files/workflows/zika/metadata.tsv.zst"
+  #   sequences: "https://data.nextstrain.org/workflows/zika/sequences.fasta.zst"
+
+  # - name: usvi
+  #   metadata: s3://nextstrain-data/files/workflows/zika/metadata_usvi.tsv.zst
+  #   sequences: s3://nextstrain-data/files/workflows/zika/sequences_usvi.fasta.zst
+
+  # - name: usvi
+  #   metadata: https://data.nextstrain.org/files/workflows/zika/metadata_usvi.tsv.zst
+  #   sequences: https://data.nextstrain.org/files/workflows/zika/sequences_usvi.fasta.zst
+
+  # - name: usvi
+  #   metadata: metadata_usvi.tsv
+  #   sequences: sequences_usvi.fasta.zst
+
   - name: usvi
-    metadata: "data/metadata_usvi.tsv"
-    sequences: "data/sequences_usvi.fasta"
+    metadata: https://raw.githubusercontent.com/nextstrain/zika/refs/heads/main/phylogenetic/data/metadata_usvi.tsv
+    sequences: https://raw.githubusercontent.com/nextstrain/zika/refs/heads/main/phylogenetic/data/sequences_usvi.fasta
+
+  # - name: usvi
+  #   metadata: gs://nextstrain-data/files/workflows/zika/metadata_usvi.tsv.zst
+  #   sequences: gs://nextstrain-data/files/workflows/zika/sequences_usvi.fasta.zst
 
 # Config files
 exclude: "exclude.txt"
@@ -41,3 +65,4 @@ traits:
     - region
     - country
   sampling_bias_correction: 3
+# 

phylogenetic/rules/merge_inputs.smk

@@ -22,55 +22,7 @@ This part of the workflow usually includes the following steps:
 
 """
 
-
-# ------------- helper functions to collect, merge & download input files ------------------- #
-NEXTSTRAIN_PUBLIC_BUCKET = "s3://nextstrain-data/"
-
-def _parse_config_input(input):
-    """
-    Parses information from an individual config-defined input, i.e. an element within `config.inputs` or `config.additional_inputs`
-    and returns information snakemake rules can use to obtain the underlying data.
-
-    The structure of `input` is a dictionary with keys:
-    - name:string (required)
-    - metadata:string (optional) - a s3 URI or a local file path
-    - sequences:string (optional) - a s3 URI or a local file path
-
-    Returns a dictionary with optional keys:
-    - metadata:string - the relative path to the metadata file. If the original data was remote then this represents
-      the output of a rule which downloads the file
-    - metadata_location:string - the URI for the remote file if applicable else `None`
-    - sequences:string - the relative path to the sequences FASTA. If the original data was remote then this represents
-      the output of a rule which downloads the file
-    - sequences_location:string - the URI for the remote file if applicable else `None`
-
-    Raises InvalidConfigError
-    """
-    name = input['name']
-
-    info = {'metadata': None, 'metadata_location': None, 'sequences': None, 'sequences_location': None}
-
-    def _source(uri, *,  s3, local):
-        if uri.startswith('s3://'):
-            return s3
-        elif uri.lower().startswith(('http://','https://')):
-            raise InvalidConfigError("Workflow cannot yet handle HTTP[S] inputs")
-        # USVI files are expected to be part of the workflow source,
-        # and are _not_ expected to be provided via the analysis directory
-        elif uri.startswith('data/metadata_usvi.tsv') or uri.startswith('data/sequences_usvi.fasta'):
-            return workflow.source_path("../" + uri)
-        return local
-
-    if location:=input.get('metadata', False):
-        info['metadata'] = _source(location,  s3=f"data/{name}/metadata.tsv", local=location)
-        info['metadata_location'] = _source(location,  s3=location, local=None)
-
-    if location:=input.get('sequences', False):
-        info['sequences'] = _source(location,  s3=f"data/{name}/sequences.fasta", local=location)
-        info['sequences_location'] = _source(location,  s3=location, local=None)
-
-    return info
-
+include: "remote_files.smk"
 
 def _gather_inputs():
     all_inputs = [*config['inputs'], *config.get('additional_inputs', [])]
@@ -85,7 +37,11 @@ def _gather_inputs():
     if not all(['name' in i and ('sequences' in i or 'metadata' in i) for i in all_inputs]):
         raise InvalidConfigError("Each input (config.inputs and config.additional_inputs) must have a 'name' and 'metadata' and/or 'sequences'")
 
-    return {i['name']: _parse_config_input(i) for i in all_inputs}
+    available_keys = set(['name', 'metadata', 'sequences'])
+    if any([len(set(el.keys())-available_keys)>0 for el in all_inputs]):
+        raise InvalidConfigError(f"Each input (config.inputs and config.additional_inputs) can only include keys of {', '.join(available_keys)}")
+
+    return {el['name']: {k:(v if k=='name' else path_or_url(v)) for k,v in el.items()} for el in all_inputs}
 
 input_sources = _gather_inputs()
 
@@ -97,44 +53,6 @@ def input_sequences(wildcards):
     inputs = [info['sequences'] for info in input_sources.values() if info.get('sequences', None)]
     return inputs[0] if len(inputs)==1 else "results/sequences_merged.fasta"
 
-rule download_s3_sequences:
-    output:
-        sequences = "data/{input_name}/sequences.fasta",
-    log:
-        "logs/{input_name}/download_s3_sequences.txt",
-    benchmark:
-        "benchmarks/{input_name}/download_s3_sequences.txt"
-    params:
-        address = lambda w: input_sources[w.input_name]['sequences_location'],
-        no_sign_request=lambda w: "--no-sign-request" \
-            if input_sources[w.input_name]['sequences_location'].startswith(NEXTSTRAIN_PUBLIC_BUCKET) \
-            else "",
-    shell:
-        r"""
-        exec &> >(tee {log:q})
-
-        aws s3 cp {params.no_sign_request:q} {params.address:q} - | zstd -d > {output.sequences}
-        """
-
-rule download_s3_metadata:
-    output:
-        metadata = "data/{input_name}/metadata.tsv",
-    log:
-        "logs/{input_name}/download_s3_metadata.txt",
-    benchmark:
-        "benchmarks/{input_name}/download_s3_metadata.txt"
-    params:
-        address = lambda w: input_sources[w.input_name]['metadata_location'],
-        no_sign_request=lambda w: "--no-sign-request" \
-            if input_sources[w.input_name]['metadata_location'].startswith(NEXTSTRAIN_PUBLIC_BUCKET) \
-            else "",
-    shell:
-        r"""
-        exec &> >(tee {log:q})
-
-        aws s3 cp {params.no_sign_request:q} {params.address:q} - | zstd -d > {output.metadata}
-        """
-
 rule merge_metadata:
     """
     This rule should only be invoked if there are multiple defined metadata inputs
@@ -180,5 +98,3 @@ rule merge_sequences:
 
         seqkit rmdup {input:q} > {output.sequences:q}
         """
-
-# -------------------------------------------------------------------------------------------- #

phylogenetic/rules/prepare_sequences.smk

@@ -61,6 +61,42 @@ rule filter:
             --min-length {params.min_length:q}
         """
 
+
+# rule upload_filter:
+#     """TESTING ONLY TODO XXX REMOVE"""
+#     input:
+#         sequences = input_sequences,
+#         metadata = input_metadata,
+#         exclude = resolve_config_path(config["exclude"]),
+#     output:
+#         sequences = path_or_url("s3://nextstrain-scratch/zika-pr-89/filtered.fasta")
+#     params:
+#         group_by = as_list(config["filter"]["group_by"]),
+#         sequences_per_group = config["filter"]["sequences_per_group"],
+#         min_date = config["filter"]["min_date"],
+#         min_length = config["filter"]["min_length"],
+#         strain_id = config.get("strain_id_field", "strain"),
+#     log:
+#         "logs/filter.txt",
+#     benchmark:
+#         "benchmarks/filter.txt"
+#     shell:
+#         r"""
+#         exec &> >(tee {log:q})
+
+#         augur filter \
+#             --sequences {input.sequences:q} \
+#             --metadata {input.metadata:q} \
+#             --metadata-id-columns {params.strain_id:q} \
+#             --exclude {input.exclude:q} \
+#             --output-sequences {output.sequences:q} \
+#             --group-by {params.group_by:q} \
+#             --sequences-per-group {params.sequences_per_group:q} \
+#             --min-date {params.min_date:q} \
+#             --min-length {params.min_length:q}
+#         """
+
+
 rule align:
     """
     Aligning sequences to {input.reference}

phylogenetic/rules/remote_files.smk

@@ -0,0 +1,144 @@
+"""
+Helper functions to set-up storage plugins for remote inputs/outputs.
+See the docstring of `path_or_url` for usage instructions.
+
+The errors raised by storage plugins are often confusing. For instance,
+a HTTP 404 error will result in a `MissingInputException` with little
+hint as to the underlying issue. S3 credentials errors are similarly
+confusing and we attempt to check these ourselves to improve UX here.
+"""
+
+from urllib.parse import urlparse
+
+# Keep a list of known public buckets, which we'll allow uncredentialled (unsigned) access to
+# We could make this config-definable in the future
+PUBLIC_BUCKETS = set(['nextstrain-data'])
+
+# Keep track of registered storage plugins to enable reuse
+_storage_registry = {}
+
+class RemoteFilesMissingCredentials(Exception):
+    pass
+
+def _storage_s3(*, bucket, keep_local, retries) -> snakemake.storage.StorageProviderProxy:
+    """
+    Registers and returns an instance of snakemake-storage-plugin-s3
+    """
+    # If the bucket is public then we may use an unsigned request which has the nice UX
+    # of not needing credentials to be present. If we've made other signed requests _or_
+    # credentials are present then we just sign everything. This has implications for upload:
+    # if you attempt to upload to a public bucket without credentials then we allow that here
+    # and you'll get a subsequent `AccessDenied` error when the upload is attempted.
+    if bucket in PUBLIC_BUCKETS and \
+        "s3_signed" not in _storage_registry and \
+        ("s3_unsigned" in _storage_registry or not _aws_credentials_present()):
+
+        if provider:=_storage_registry.get('s3_unsigned', None):
+            return provider
+
+        from botocore import UNSIGNED # dependency of snakemake-storage-plugin-s3
+        storage s3_unsigned:
+            provider="s3",
+            signature_version=UNSIGNED,
+            retries=retries,
+            keep_local=keep_local,
+
+        _storage_registry['s3_unsigned'] = storage.s3_unsigned
+        return _storage_registry['s3_unsigned']
+
+    # Resource fetched/uploaded via a signed request, which will require AWS credentials
+    if provider:=_storage_registry.get('s3_signed', None):
+        return provider
+
+    # Enforce the presence of credentials to paper over <https://github.com/snakemake/snakemake/issues/3663>
+    if not _aws_credentials_present():
+        raise RemoteFilesMissingCredentials()
+
+    # the tag appears in the local file path, so reference 'signed' to give a hint about credential errors
+    storage s3_signed:
+        provider="s3",
+        retries=retries,
+        keep_local=keep_local,
+
+    _storage_registry['s3_signed'] = storage.s3_signed
+    return _storage_registry['s3_signed']
+
+def _aws_credentials_present() -> bool:
+    import boto3 # dependency of snakemake-storage-plugin-s3
+    session = boto3.Session()
+    creds = session.get_credentials()
+    return creds is not None
+
+def _storage_http(*, keep_local, retries) -> snakemake.storage.StorageProviderProxy:
+    """
+    Registers and returns an instance of snakemake-storage-plugin-http
+    """
+    if provider:=_storage_registry.get('http', None):
+        return provider
+
+    storage:
+        provider="http",
+        allow_redirects=True,
+        supports_head=True,
+        keep_local=keep_local,
+        retries=retries,
+
+    _storage_registry['http'] = storage.http
+    return _storage_registry['http']
+
+
+def path_or_url(uri, *, keep_local=True, retries=2) -> str:
+    """
+    Intended for use in Snakemake inputs / outputs to transparently use remote
+    resources. Returns the URI wrapped by an applicable storage plugin. Local
+    filepaths will be returned unchanged.
+
+    For example, the following rule will download inputs from HTTPs and upload
+    the output to S3:
+
+        rule filter:
+            input:
+                sequences = path_or_url("https://data.nextstrain.org/files/workflows/zika/sequences.fasta.zst"),
+                metadata = path_or_url("https://data.nextstrain.org/files/workflows/zika/metadata.tsv.zst"),
+            output:
+                sequences = path_or_url("s3://nextstrain-scratch/filtered.fasta")
+            shell:
+                r'''
+                augur filter \
+                    --sequences {input.sequences:q} \
+                    --metadata {input.metadata:q} \
+                    --metadata-id-columns accession \
+                    --output-sequences {output.sequences:q}
+                '''
+
+    If `keep_local` is True (the default) then downloaded/uploaded files will
+    remain in `.snakemake/storage/`. The presence of a previously downloaded
+    file (via `keep_local=True`) does not guarantee that the file will not be
+    re-downloaded if the storage plugin decides the local file is out of date.
+
+    See <https://snakemake.readthedocs.io/en/stable/snakefiles/storage.html> for
+    more information on Snakemake storage plugins. Note: various snakemake
+    plugins will be required depending on the URIs provided.
+    """
+    info = urlparse(uri)
+
+    if info.scheme=='': # local
+        return uri      # no storage wrapper
+
+    if info.scheme=='s3':
+        try:
+            return _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)
+    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.")
+
+    if info.scheme in ['gs', 'gcs']:
+        raise Exception(f"Google Storage is not yet implemented for nextstrain workflows (attempting to access {uri!r}).\n"
+            "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")