Support object stores: CutSet.from_shar(in_dir='s3://...')

lhotse/ais/__init__.py

@@ -1 +1,2 @@
 from .batch_loader import AISBatchLoader
+from .utils import list_aistore_objects

lhotse/ais/utils.py

@@ -0,0 +1,48 @@
+from typing import Any, Iterable, List, Tuple
+
+from lhotse.serialization import get_aistore_client
+from lhotse.utils import Pathlike, split_object_store_url, object_store_uris
+
+
+def list_aistore_objects(url: Pathlike) -> List[str]:
+    scheme, bucket, prefix = split_object_store_url(url)
+    listing = _list_bucket_objects(scheme, bucket, prefix)
+    keys = _iter_aistore_object_keys(listing)
+    return object_store_uris(scheme, bucket, prefix, keys)
+
+
+def _list_bucket_objects(provider: str, bucket: str, prefix: str) -> Any:
+    client, _ = get_aistore_client()
+    bucket_handle = client.bucket(bucket, provider)
+    listing_prefix = f"{prefix}/" if prefix else ""
+
+    for method_name in ("list_all_objects", "list_objects"):
+        method = getattr(bucket_handle, method_name, None)
+        if method is None:
+            continue
+        for kwargs in (
+            {"prefix": listing_prefix},
+            {"prefix_filter": listing_prefix},
+        ):
+            try:
+                return method(**kwargs)
+            except TypeError:
+                continue
+
+    raise RuntimeError(
+        "The installed AIStore SDK does not expose a supported object listing API for Shar directory scanning."
+    )
+
+
+def _iter_aistore_object_keys(listing: Any) -> Iterable[str]:
+    entries = getattr(listing, "entries", listing)
+    for entry in entries:
+        if isinstance(entry, str):
+            yield entry
+            continue
+        if isinstance(entry, dict):
+            key = entry.get("name")
+        else:
+            key = getattr(entry, "name", None)
+        if isinstance(key, str):
+            yield key

lhotse/shar/readers/lazy.py

@@ -4,14 +4,18 @@
     Callable,
     Dict,
     Generator,
+    Iterable,
     List,
     Literal,
     Optional,
     Sequence,
+    Set,
     Tuple,
     Union,
 )
+from urllib.parse import urlparse
 
+from lhotse.ais.utils import list_aistore_objects
 from lhotse.cut import Cut
 from lhotse.dataset.dataloading import resolve_seed
 from lhotse.lazy import (
@@ -21,9 +25,16 @@
     LazyManifestIterator,
     count_newlines_fast,
 )
-from lhotse.serialization import extension_contains
+from lhotse.serialization import AIStoreIOBackend, extension_contains
 from lhotse.shar.readers.tar import TarIterator
-from lhotse.utils import Pathlike, exactly_one_not_null, ifnone
+from lhotse.utils import Pathlike, exactly_one_not_null, ifnone, list_s3_objects
+
+
+def fail_msg(field: str, expected: int, streams: Dict[str, List[Pathlike]]) -> str:
+    for i, cut in enumerate(streams["cuts"]):
+        print(cut)
+        print(streams[field][i])
+    return f"Expected {expected} shards available for field '{field}' but found {len(streams[field])}"
 
 
 class LazySharIterator(Dillable):
@@ -104,7 +115,9 @@ class LazySharIterator(Dillable):
         and values are lists of shards (either paths or shell commands).
         The field "cuts" pointing to CutSet shards always has to be present.
     :param in_dir: path to a directory created with ``SharWriter`` with
-        all the shards in a single place. Can be used instead of ``fields``.
+        all the shards in a single place, or a supported object store URI
+        prefix such as ``s3://bucket/shar`` or ``ais://bucket/shar``.
+        Can be used instead of ``fields``.
     :param split_for_dataloading: bool, by default ``False`` which does nothing.
         Setting it to ``True`` is intended for PyTorch training with multiple
         dataloader workers and possibly multiple DDP nodes.
@@ -177,9 +190,9 @@ def __init__(
 
         self.num_shards = len(self.streams["cuts"])
         for field in self.fields:
-            assert (
-                len(self.streams[field]) == self.num_shards
-            ), f"Expected {self.num_shards} shards available for field '{field}' but found {len(self.streams[field])}: {self.streams[field]}"
+            assert len(self.streams[field]) == self.num_shards, fail_msg(
+                field, self.num_shards, self.streams
+            )
 
         self.shards = [
             {field: self.streams[field][shard_idx] for field in self.streams}
@@ -197,24 +210,24 @@ def _init_from_inputs(self, fields: Optional[Dict[str, Sequence[str]]] = None):
         self.streams = fields
 
     def _init_from_dir(self, in_dir: Pathlike):
-        self.in_dir = Path(in_dir)
-
-        all_paths = list(self.in_dir.glob("*"))
-        self.fields = set(p.stem.split(".")[0] for p in all_paths)
-        assert "cuts" in self.fields
-        self.fields.remove("cuts")
-
-        self.streams = {
-            "cuts": sorted(
-                p
-                for p in all_paths
-                if p.name.split(".")[0] == "cuts" and extension_contains(".jsonl", p)
+        scheme = urlparse(str(in_dir)).scheme.lower()
+        if scheme in _OBJECT_STORE_SCHEMES:
+            self.in_dir = str(in_dir)
+            self.fields, self.streams = _init_from_object_store_dir(self.in_dir)
+            return
+
+        if scheme:
+            raise ValueError(
+                f"Unsupported object store URI scheme '{scheme}' for Shar directory scanning: {in_dir}. "
+                "Please provide explicit 'fields' or use one of: ais://, s3://, s3a://, s3n://."
             )
-        }
-        for field in self.fields:
-            self.streams[field] = sorted(
-                p for p in all_paths if p.name.split(".")[0] == field
+
+        self.in_dir = Path(in_dir)
+        if not self.in_dir.is_dir():
+            raise ValueError(
+                f"Expected 'in_dir' to be an existing directory or a supported object store URI, got: {in_dir}"
             )
+        self.fields, self.streams = _init_from_local_dir(self.in_dir)
 
     def _maybe_split_for_dataloading(self, shards: List) -> List:
         from .utils import split_by_node, split_by_worker
@@ -309,6 +322,56 @@ def __add__(self, other) -> "LazyIteratorChain":
         return LazyIteratorChain(self, other)
 
 
+_OBJECT_STORE_SCHEMES = frozenset(("ais", "s3", "s3a", "s3n"))
+
+
+def _init_from_local_dir(in_dir: Path) -> Tuple[Set[str], Dict[str, List[Path]]]:
+    return _build_shar_streams(list(in_dir.glob("*")), source=in_dir)
+
+
+def _init_from_object_store_dir(
+    in_dir: Pathlike,
+) -> Tuple[Set[str], Dict[str, List[str]]]:
+    scheme = urlparse(str(in_dir)).scheme.lower()
+    if scheme not in _OBJECT_STORE_SCHEMES:
+        raise ValueError(
+            f"Unsupported object store URI scheme '{scheme}' for Shar directory scanning: {in_dir}. "
+            "Please provide explicit 'fields' or use one of: ais://, s3://, s3a://, s3n://."
+        )
+    if AIStoreIOBackend().is_available():
+        all_paths = list_aistore_objects(in_dir)
+    else:
+        all_paths = list_s3_objects(in_dir)
+    return _build_shar_streams(all_paths, source=in_dir)
+
+
+def _build_shar_streams(
+    all_paths: Iterable[Pathlike], source: Pathlike
+) -> Tuple[Set[str], Dict[str, List[Pathlike]]]:
+    all_paths = list(all_paths)
+    fields = set(Path(p).stem.split(".")[0] for p in all_paths)
+    if not any(field.startswith("cuts") for field in fields):
+        raise ValueError(f"Could not find any Shar 'cuts' shards under: {source}")
+
+    fields = {field for field in fields if not field.startswith("cuts")}
+    streams = {
+        "cuts": sorted(
+            p
+            for p in all_paths
+            if Path(p).name.split(".")[0].startswith("cuts")
+            and extension_contains(".jsonl", p)
+        )
+    }
+    if not streams["cuts"]:
+        raise ValueError(f"Could not find any Shar JSONL manifests under: {source}")
+
+    for field in fields:
+        streams[field] = sorted(
+            p for p in all_paths if Path(p).name.split(".")[0].startswith(field)
+        )
+    return fields, streams
+
+
 def _jsonl_tar_adaptor(
     jsonl_iter: LazyJsonlIterator, field: str
 ) -> Generator[Tuple[Optional[dict], Path], None, None]:

lhotse/utils.py

@@ -138,6 +138,51 @@ def is_valid_url(value: str) -> bool:
         return False
 
 
+def split_object_store_url(url: Pathlike) -> Tuple[str, str, str]:
+    parsed = urlparse(str(url).rstrip("/"))
+    if not (parsed.scheme and parsed.netloc):
+        raise ValueError(f"Not a valid object store URI: {url}")
+    return (
+        parsed.scheme.lower(),
+        parsed.netloc,
+        parsed.path.lstrip("/").rstrip("/"),
+    )
+
+
+def list_s3_objects(url: Pathlike) -> List[str]:
+    try:
+        import boto3
+    except ImportError as e:
+        raise ImportError(
+            "Please run 'pip install boto3' or 'pip install smart_open[s3]' to scan Shar shards in S3."
+        ) from e
+
+    scheme, bucket, prefix = split_object_store_url(url)
+    listing_prefix = f"{prefix}/" if prefix else ""
+
+    paginator = boto3.client("s3").get_paginator("list_objects_v2")
+    keys = (
+        item["Key"]
+        for page in paginator.paginate(
+            Bucket=bucket, Prefix=listing_prefix,
+        )
+        for item in page.get("Contents", [])
+        if isinstance(item.get("Key"), str)
+    )
+    return object_store_uris(scheme, bucket, prefix, keys)
+
+
+def object_store_uris(
+    scheme: str, bucket: str, prefix: str, keys: Iterable[str]
+) -> List[str]:
+    listing_prefix = f"{prefix}/" if prefix else ""
+    return sorted(
+        f"{scheme}://{bucket}/{key}"
+        for key in keys
+        if key.startswith(listing_prefix) and key != listing_prefix
+    )
+
+
 def fix_random_seed(random_seed: int):
     """
     Set the same random seed for the libraries and modules that Lhotse interacts with.

test/shar/test_read_lazy.py

@@ -1,7 +1,10 @@
 import os
 import shutil
+import sys
+import types
 from functools import partial
 from pathlib import Path
+from unittest.mock import MagicMock, patch
 
 import numpy as np
 import pytest
@@ -402,3 +405,94 @@ def test_shar_slice_length(shar_dir: Path):
     # Works via CutSet too
     sliced_cuts2 = list(CutSet.from_shar(in_dir=shar_dir, slice_length=3, seed=6))
     assert sliced_cuts == sliced_cuts2
+
+
+def test_shar_lazy_reader_from_s3_dir_initializes_streams():
+    paginator = MagicMock()
+    paginator.paginate.return_value = [
+        {
+            "Contents": [
+                {"Key": "my-shar/cuts.000000.jsonl.gz"},
+                {"Key": "my-shar/recording.000000.tar"},
+                {"Key": "my-shar/nested/ignored.tar"},
+                {"Key": "my-shar/"},
+            ]
+        }
+    ]
+    s3_client = MagicMock()
+    s3_client.get_paginator.return_value = paginator
+
+    fake_boto3 = types.ModuleType("boto3")
+    fake_boto3.client = MagicMock(return_value=s3_client)
+
+    with patch.dict(sys.modules, {"boto3": fake_boto3}):
+        cuts_iter = LazySharIterator(in_dir="s3://test-bucket/my-shar")
+
+    assert cuts_iter.fields == {"recording"}
+    assert cuts_iter.streams == {
+        "cuts": ["s3://test-bucket/my-shar/cuts.000000.jsonl.gz"],
+        "recording": ["s3://test-bucket/my-shar/recording.000000.tar"],
+    }
+    fake_boto3.client.assert_called_once_with("s3")
+    s3_client.get_paginator.assert_called_once_with("list_objects_v2")
+    paginator.paginate.assert_called_once_with(
+        Bucket="test-bucket", Prefix="my-shar/", Delimiter="/"
+    )
+
+
+def test_shar_lazy_reader_from_aistore_dir_initializes_streams():
+    class FakeBucket:
+        def __init__(self):
+            self.calls = []
+
+        def list_objects(self, **kwargs):
+            self.calls.append(kwargs)
+            return types.SimpleNamespace(
+                entries=[
+                    types.SimpleNamespace(name="my-shar/cuts.000000.jsonl.gz"),
+                    types.SimpleNamespace(name="my-shar/recording.000000.tar"),
+                    types.SimpleNamespace(name="my-shar/nested/ignored.tar"),
+                ]
+            )
+
+    bucket = FakeBucket()
+    client = MagicMock()
+    client.bucket.return_value = bucket
+
+    fake_aistore = types.ModuleType("aistore")
+    fake_aistore.__spec__ = types.SimpleNamespace(
+        name="aistore", origin=None, submodule_search_locations=[]
+    )
+    fake_sdk = types.ModuleType("aistore.sdk")
+    fake_utils = types.ModuleType("aistore.sdk.utils")
+    fake_utils.parse_url = lambda url: ("ais", "test-bucket", "my-shar")
+    fake_sdk.utils = fake_utils
+    fake_aistore.sdk = fake_sdk
+
+    with patch.dict(os.environ, {"AIS_ENDPOINT": "http://localhost:8080"}):
+        with patch.dict(
+            sys.modules,
+            {
+                "aistore": fake_aistore,
+                "aistore.sdk": fake_sdk,
+                "aistore.sdk.utils": fake_utils,
+            },
+        ):
+            with patch(
+                "lhotse.ais.utils.get_aistore_client",
+                return_value=(client, None),
+            ):
+                cuts_iter = LazySharIterator(in_dir="ais://test-bucket/my-shar")
+
+    assert cuts_iter.fields == {"recording"}
+    assert cuts_iter.streams == {
+        "cuts": ["ais://test-bucket/my-shar/cuts.000000.jsonl.gz"],
+        "recording": ["ais://test-bucket/my-shar/recording.000000.tar"],
+    }
+    client.bucket.assert_called_once_with("test-bucket", "ais")
+    assert bucket.calls == [{"prefix": "my-shar/"}]
+
+
+def test_shar_lazy_reader_from_unsupported_object_store_dir():
+    with pytest.raises(ValueError, match="Unsupported object store URI scheme 'gs'"):
+        LazySharIterator(in_dir="gs://test-bucket/my-shar")