Replace 'aws s3 ls' shell-out in dataset.py with boto3

Author: goanpeca

.gitignore

@@ -163,4 +163,6 @@ cython_debug/
 *.wav
 wandb/*
 *.out
-test_*
+test_*
+# macOS
+.DS_Store

README.md

@@ -171,6 +171,20 @@ The following properties are defined in the top level of the model configuration
 ## Dataset config
 `stable-audio-tools` currently supports two kinds of data sources: local directories of audio files, and WebDataset datasets stored in Amazon S3. More information can be found in [the dataset config documentation](docs/datasets.md)
 
+## S3-compatible storage (Backblaze B2)
+The S3 dataset loader uses `boto3`, which ships in the `train` extra. If you installed without that extra, add it with `pip install boto3` (or `pip install "stable-audio-tools[train]"`).
+
+The loader honors the `AWS_ENDPOINT_URL` environment variable, so you can point it at any S3-compatible host without changing the dataset config.
+
+Example for [Backblaze B2](https://www.backblaze.com/cloud-storage):
+```bash
+export AWS_ENDPOINT_URL=https://s3.us-west-004.backblazeb2.com
+export AWS_ACCESS_KEY_ID=<B2 application key ID>
+export AWS_SECRET_ACCESS_KEY=<B2 application key>
+```
+
+When `AWS_ENDPOINT_URL` is unset, the loader uses default AWS S3, so existing setups are unaffected.
+
 # Todo
 - [ ] Add troubleshooting section
 - [ ] Add contribution guidelines 

pyproject.toml

@@ -37,6 +37,7 @@ dependencies = [
 [project.optional-dependencies]
 train = [
     "auraloss==0.4.0",
+    "boto3",
     "descript-audio-codec==1.0.0",
     "encodec==0.1.1",
     "inf-cl",

stable_audio_tools/data/dataset.py

@@ -1,13 +1,13 @@
 import importlib
+from importlib.metadata import PackageNotFoundError, version
+from functools import lru_cache
 import numpy as np
 import io
 import json
 import os
 import dill
 import posixpath
 import random
-import re
-import subprocess
 import time
 import torch
 import torchaudio
@@ -483,51 +483,108 @@ def __getitem__(self, idx):
 
 # S3 code and WDS preprocessing code based on implementation by Scott Hawley originally in https://github.com/zqevans/audio-diffusion/blob/main/dataset/dataset.py
 
+try:
+    _PACKAGE_VERSION = version("stable-audio-tools")
+except PackageNotFoundError:  # source/editable checkout without dist metadata
+    _PACKAGE_VERSION = "dev"
+
+_USER_AGENT = f"stable-audio-tools/{_PACKAGE_VERSION}"
+
+
+def _build_user_agent_extra(user_agent_extra=None):
+    """``stable-audio-tools/<version>`` with any caller- or env-provided
+    (``STABLE_AUDIO_TOOLS_USER_AGENT_EXTRA``) value appended, not replacing it."""
+    extra = user_agent_extra or os.environ.get("STABLE_AUDIO_TOOLS_USER_AGENT_EXTRA")
+    return f"{_USER_AGENT} {extra}" if extra else _USER_AGENT
+
+
+@lru_cache(maxsize=32)
+def _build_s3_client(profile, endpoint_url, user_agent_extra):
+    import boto3  # local import so boto3 is only required when S3 is used
+    from botocore.config import Config
+
+    session = boto3.Session(profile_name=profile) if profile else boto3.Session()
+    return session.client(
+        "s3",
+        endpoint_url=endpoint_url,
+        config=Config(user_agent_extra=user_agent_extra),
+    )
+
+
+def _get_s3_client(profile=None, user_agent_extra=None):
+    """
+    Build (and reuse) a boto3 S3 client. Honors AWS_ENDPOINT_URL when set so the
+    same code path works against any S3-compatible endpoint (AWS S3 by default;
+    set AWS_ENDPOINT_URL to a Backblaze B2 endpoint to point it at B2). When the
+    env var is unset, behavior matches the default AWS client.
+
+    Clients are cached per (profile, endpoint, user-agent) so listing and
+    presigning share one client instead of building a new one on each call.
+    """
+    endpoint_url = os.environ.get("AWS_ENDPOINT_URL") or None
+    return _build_s3_client(
+        profile, endpoint_url, _build_user_agent_extra(user_agent_extra)
+    )
+
+
+def _parse_s3_url(url):
+    "Split an ``s3://bucket/key`` URL into (bucket, key). Raises ValueError otherwise."
+    if not url.startswith("s3://"):
+        raise ValueError(f"expected an s3:// URL, got: {url!r}")
+    bucket, _, key = url[len("s3://"):].partition("/")
+    return bucket, key
+
+
 def get_s3_contents(dataset_path, s3_url_prefix=None, filter='', recursive=True, debug=False, profile=None):
     """
-    Returns a list of full S3 paths to files in a given S3 bucket and directory path.
+    Returns a list of S3 keys (relative to ``dataset_path``) for objects in a
+    given S3 bucket and directory path. Uses boto3 directly so it works
+    against any S3-compatible endpoint when ``AWS_ENDPOINT_URL`` is set.
     """
     # Ensure dataset_path ends with a trailing slash
     if dataset_path != '' and not dataset_path.endswith('/'):
         dataset_path += '/'
-    # Use posixpath to construct the S3 URL path
+    # Use posixpath to construct the S3 URL path (e.g. "s3://bucket/prefix/")
     bucket_path = posixpath.join(s3_url_prefix or '', dataset_path)
-    # Construct the `aws s3 ls` command
-    cmd = ['aws', 's3', 'ls', bucket_path]
 
-    if profile is not None:
-        cmd.extend(['--profile', profile])
+    bucket, prefix = _parse_s3_url(bucket_path)
+
+    s3 = _get_s3_client(profile=profile)
+    paginator = s3.get_paginator("list_objects_v2")
+    list_kwargs = {"Bucket": bucket, "Prefix": prefix}
+    if not recursive:
+        list_kwargs["Delimiter"] = "/"
+
+    keys = []
+    for page in paginator.paginate(**list_kwargs):
+        for obj in page.get("Contents", []) or []:
+            key = obj.get("Key", "")
+            if not key or key.endswith("/"):
+                continue
+            keys.append(key)
 
-    if recursive:
-        # Add the --recursive flag if requested
-        cmd.append('--recursive')
-    
-    # Run the `aws s3 ls` command and capture the output
-    run_ls = subprocess.run(cmd, capture_output=True, check=True)
-    # Split the output into lines and strip whitespace from each line
-    contents = run_ls.stdout.decode('utf-8').split('\n')
-    contents = [x.strip() for x in contents if x]
-    # Remove the timestamp from lines that begin with a timestamp
-    contents = [re.sub(r'^\S+\s+\S+\s+\d+\s+', '', x)
-                if re.match(r'^\S+\s+\S+\s+\d+\s+', x) else x for x in contents]
-    # Construct a full S3 path for each file in the contents list
-    contents = [posixpath.join(s3_url_prefix or '', x)
-                for x in contents if not x.endswith('/')]
     # Apply the filter, if specified
     if filter:
-        contents = [x for x in contents if filter in x]
-    # Remove redundant directory names in the S3 URL
-    if recursive:
-        # Get the main directory name from the S3 URL
-        main_dir = "/".join(bucket_path.split('/')[3:])
-        # Remove the redundant directory names from each file path
-        contents = [x.replace(f'{main_dir}', '').replace(
-            '//', '/') for x in contents]
-    # Print debugging information, if requested
+        keys = [k for k in keys if filter in k]
+
+    # Match the legacy `aws s3 ls` output shape: paths relative to dataset_path.
+    # The legacy CLI emitted basenames in non-recursive mode and full keys
+    # (which it then stripped) in recursive mode; both paths ended up
+    # relative to dataset_path. boto3 always returns full keys, so strip
+    # the prefix unconditionally.
+    if prefix:
+        keys = [k[len(prefix):] if k.startswith(prefix) else k for k in keys]
+        keys = [k.lstrip('/') for k in keys]
+
     if debug:
-        print("contents = \n", contents)
-    # Return the list of S3 paths to files
-    return contents
+        print("contents = \n", keys)
+
+    return keys
+
+
+# 7 days (SigV4 max) so shard URLs outlast long training runs. Override per
+# call or via STABLE_AUDIO_TOOLS_S3_PRESIGN_EXPIRY.
+_DEFAULT_PRESIGN_EXPIRY_SECONDS = 7 * 24 * 3600
 
 
 def get_all_s3_urls(
@@ -540,8 +597,21 @@ def get_all_s3_urls(
     # print debugging info -- note: info displayed likely to change at dev's whims
     debug=False,
     profiles={},        # dictionary of profiles for each item in names, e.g. {'dataset1': 'profile1', 'dataset2': 'profile2'}
+    presign_expiry_seconds=None,  # presigned-URL lifetime; None -> env var or default
 ):
     "get urls of shards (tar files) for multiple datasets in one s3 bucket"
+    if presign_expiry_seconds is None:
+        raw = os.environ.get("STABLE_AUDIO_TOOLS_S3_PRESIGN_EXPIRY")
+        if raw is None:
+            presign_expiry_seconds = _DEFAULT_PRESIGN_EXPIRY_SECONDS
+        else:
+            try:
+                presign_expiry_seconds = int(raw)
+            except ValueError:
+                raise ValueError(
+                    "STABLE_AUDIO_TOOLS_S3_PRESIGN_EXPIRY must be an integer number "
+                    f"of seconds, got: {raw!r}"
+                )
     urls = []
     for name in names:
         # If s3_url_prefix is not specified, assume the full S3 path is included in each element of the names list
@@ -560,22 +630,30 @@ def get_all_s3_urls(
             profile = profiles.get(name, None)
             tar_list = get_s3_contents(
                 subset_str, s3_url_prefix=None, recursive=recursive, filter=filter_str, debug=debug, profile=profile)
+            # Reuse the cached S3 client (shared with get_s3_contents) for presigning.
+            s3_client = _get_s3_client(profile=profile)
             for tar in tar_list:
-                # Escape spaces and parentheses in the tar filename for use in the shell command
-                tar = tar.replace(" ", "\ ").replace(
-                    "(", "\(").replace(")", "\)")
-                # Construct the S3 path to the current tar file
-                s3_path = posixpath.join(name, subset, tar) + " -"
-                # Construct the AWS CLI command to download the current tar file
+                # Construct the full s3:// URL for the current tar file.
                 if s3_url_prefix is None:
-                    request_str = f"pipe:aws s3 --cli-connect-timeout 0 cp {s3_path}"
+                    full_s3_url = posixpath.join(name, subset, tar)
                 else:
-                    request_str = f"pipe:aws s3 --cli-connect-timeout 0 cp {posixpath.join(s3_url_prefix, s3_path)}"
-                if profiles.get(name):
-                    request_str += f" --profile {profiles.get(name)}"
+                    full_s3_url = posixpath.join(s3_url_prefix, name, subset, tar)
+
+                bucket, key = _parse_s3_url(full_s3_url)
+
+                # Presigned GET URL works against AWS and any S3-compatible
+                # endpoint when AWS_ENDPOINT_URL is set. Expiry is configurable
+                # so long training runs do not outlive their shard URLs.
+                presigned = s3_client.generate_presigned_url(
+                    "get_object",
+                    Params={"Bucket": bucket, "Key": key},
+                    ExpiresIn=presign_expiry_seconds,
+                )
+                request_str = f'pipe:curl -fsSL "{presigned}"'
                 if debug:
-                    print("request_str = ", request_str)
-                # Add the constructed URL to the list of URLs
+                    # Strip the signed query string so signatures are not logged.
+                    redacted = presigned.split("?", 1)[0]
+                    print(f'request_str = pipe:curl -fsSL "{redacted}?<redacted>"')
                 urls.append(request_str)
     return urls