WIP - snakemake v8+ remote storage support

phylogenetic/defaults/config.yaml

@@ -1,14 +1,19 @@
 # Sequences must be FASTA and metadata must be TSV
-# Both files must be zstd compressed
+
+# TODO XXX - remove commented lines below - only for testing the snakemake storage stuff
 inputs:
   - name: ncbi
     metadata: "s3://nextstrain-data/files/workflows/zika/metadata.tsv.zst"
     sequences: "s3://nextstrain-data/files/workflows/zika/sequences.fasta.zst"
+    # metadata: "https://data.nextstrain.org/files/workflows/zika/metadata.tsv.zst"
+    # sequences: "https://data.nextstrain.org/files/workflows/zika/sequences.fasta.zst"
 
 additional_inputs:
   - 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.fasta
 
 # Config files
 exclude: "exclude.txt"
@@ -41,3 +46,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,15 @@ rule filter:
             --min-length {params.min_length:q}
         """
 
+rule upload_filter:
+    """TESTING ONLY TODO XXX REMOVE"""
+    input: "results/filtered.fasta"
+    output: path_or_url("s3://nextstrain-scratch/testing-zika-filtered.fasta")
+    shell:
+        r"""
+        cp {input[0]} {output[0]}
+        """
+
 rule align:
     """
     Aligning sequences to {input.reference}

phylogenetic/rules/remote_files.smk

@@ -0,0 +1,82 @@
+"""
+Helper functions to set-up storage plugins for remote inputs/outputs.
+See the docstring of `path_or_url` for usage instructions.
+
+<https://snakemake.readthedocs.io/en/stable/snakefiles/storage.html>
+"""
+
+from urllib.parse import urlparse
+
+PUBLIC_BUCKETS = set(['nextstrain-data']) # TODO XXX
+
+_storage_registry = {} # keeps track of registered storage plugins to enable reuse
+
+def _storage_s3(*, bucket, keep_local, force_signed):
+    """
+    Returns a Snakemake storage plugin for S3 endpoints
+    """
+    retries=2 # num S3 retries attempted
+
+    # If the bucket is public then we make an unsigned request (i.e. no AWS credentials
+    # need to be set). Note that this won't work for uploads.
+    if not force_signed and bucket in PUBLIC_BUCKETS:
+        if provider:=_storage_registry.get('s3_unsigned', None):
+            return provider
+
+        from botocore import UNSIGNED
+        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']
+
+    # Default: resource fetched via a signed request, which will require AWS credentials
+    if provider:=_storage_registry.get('s3_signed', None):
+        return provider
+
+    # 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 _storage_http(*, keep_local):
+    if provider:=_storage_registry.get('http', None):
+        return provider
+
+    storage:
+        provider="http",
+        allow_redirects=True,
+        supports_head=True,
+        keep_local=keep_local,
+
+    _storage_registry['http'] = storage.http
+    return _storage_registry['http']
+
+def path_or_url(uri, *, keep_local=True, force_signed=False):
+    """
+    Returns the URI wrapped by an applicable storage plugin.
+    Local filepaths will be returned unchanged.
+
+    TODO XXX - document usage more thoroughly
+    """
+    info = urlparse(uri)
+
+    if info.scheme=='': # local
+        return uri      # no storage wrapper
+
+    if info.scheme=='s3':
+        return _storage_s3(bucket=info.netloc, keep_local=keep_local, force_signed=force_signed)(uri)
+
+    if info.scheme in ['http', 'https']:
+        return _storage_http(keep_local=keep_local)(uri)
+
+    # TODO XXX - Google? We allowed this in ncov <https://github.com/nextstrain/ncov/blob/41cf6470d3140963ff3e02c29241f80ae8ed9c33/workflow/snakemake_rules/remote_files.smk#L62>
+
+    raise Exception(f"Input address {uri!r} (scheme={info.scheme!r}) is from a non-supported remote")