atdata 0.2.3b1__py3-none-any.whl → 0.3.0b1__py3-none-any.whl

atdata/local/_repo_legacy.py

@@ -0,0 +1,218 @@
+"""Deprecated Repo class for legacy S3 repository operations."""
+
+from atdata import Dataset
+
+from atdata.local._entry import LocalDatasetEntry
+from atdata.local._s3 import _s3_env, _s3_from_credentials, _create_s3_write_callbacks
+
+from pathlib import Path
+from uuid import uuid4
+from tempfile import TemporaryDirectory
+from typing import Any, BinaryIO, TypeVar, cast
+
+from redis import Redis
+import msgpack
+import webdataset as wds
+import warnings
+
+from atdata._protocols import Packable
+
+T = TypeVar("T", bound=Packable)
+
+
+class Repo:
+    """Repository for storing and managing atdata datasets.
+
+    .. deprecated::
+        Use :class:`Index` with :class:`S3DataStore` instead::
+
+            store = S3DataStore(credentials, bucket="my-bucket")
+            index = Index(redis=redis, data_store=store)
+            entry = index.insert_dataset(ds, name="my-dataset")
+
+    Provides storage of datasets in S3-compatible object storage with Redis-based
+    indexing. Datasets are stored as WebDataset tar files with optional metadata.
+
+    Attributes:
+        s3_credentials: S3 credentials dictionary or None.
+        bucket_fs: S3FileSystem instance or None.
+        hive_path: Path within S3 bucket for storing datasets.
+        hive_bucket: Name of the S3 bucket.
+        index: Index instance for tracking datasets.
+    """
+
+    ##
+
+    def __init__(
+        self,
+        s3_credentials: str | Path | dict[str, Any] | None = None,
+        hive_path: str | Path | None = None,
+        redis: Redis | None = None,
+    ) -> None:
+        """Initialize a repository.
+
+        .. deprecated::
+            Use Index with S3DataStore instead.
+
+        Args:
+            s3_credentials: Path to .env file with S3 credentials, or dict with
+                AWS_ENDPOINT, AWS_ACCESS_KEY_ID, and AWS_SECRET_ACCESS_KEY.
+                If None, S3 functionality will be disabled.
+            hive_path: Path within the S3 bucket to store datasets.
+                Required if s3_credentials is provided.
+            redis: Redis connection for indexing. If None, creates a new connection.
+
+        Raises:
+            ValueError: If hive_path is not provided when s3_credentials is set.
+        """
+        warnings.warn(
+            "Repo is deprecated. Use Index with S3DataStore instead:\n"
+            "  store = S3DataStore(credentials, bucket='my-bucket')\n"
+            "  index = Index(redis=redis, data_store=store)\n"
+            "  entry = index.insert_dataset(ds, name='my-dataset')",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+        if s3_credentials is None:
+            self.s3_credentials = None
+        elif isinstance(s3_credentials, dict):
+            self.s3_credentials = s3_credentials
+        else:
+            self.s3_credentials = _s3_env(s3_credentials)
+
+        if self.s3_credentials is None:
+            self.bucket_fs = None
+        else:
+            self.bucket_fs = _s3_from_credentials(self.s3_credentials)
+
+        if self.bucket_fs is not None:
+            if hive_path is None:
+                raise ValueError("Must specify hive path within bucket")
+            self.hive_path = Path(hive_path)
+            self.hive_bucket = self.hive_path.parts[0]
+        else:
+            self.hive_path = None
+            self.hive_bucket = None
+
+        #
+
+        from atdata.local._index import Index
+
+        self.index = Index(redis=redis)
+
+    ##
+
+    def insert(
+        self,
+        ds: Dataset[T],
+        *,
+        name: str,
+        cache_local: bool = False,
+        schema_ref: str | None = None,
+        **kwargs,
+    ) -> tuple[LocalDatasetEntry, Dataset[T]]:
+        """Insert a dataset into the repository.
+
+        Writes the dataset to S3 as WebDataset tar files, stores metadata,
+        and creates an index entry in Redis.
+
+        Args:
+            ds: The dataset to insert.
+            name: Human-readable name for the dataset.
+            cache_local: If True, write to local temporary storage first, then
+                copy to S3. This can be faster for some workloads.
+            schema_ref: Optional schema reference. If None, generates from sample type.
+            **kwargs: Additional arguments passed to wds.ShardWriter.
+
+        Returns:
+            A tuple of (index_entry, new_dataset) where:
+                - index_entry: LocalDatasetEntry for the stored dataset
+                - new_dataset: Dataset object pointing to the stored copy
+
+        Raises:
+            ValueError: If S3 credentials or hive_path are not configured.
+            RuntimeError: If no shards were written.
+        """
+        if self.s3_credentials is None:
+            raise ValueError(
+                "S3 credentials required for insert(). Initialize Repo with s3_credentials."
+            )
+        if self.hive_bucket is None or self.hive_path is None:
+            raise ValueError(
+                "hive_path required for insert(). Initialize Repo with hive_path."
+            )
+
+        new_uuid = str(uuid4())
+
+        hive_fs = _s3_from_credentials(self.s3_credentials)
+
+        # Write metadata
+        metadata_path = (
+            self.hive_path / "metadata" / f"atdata-metadata--{new_uuid}.msgpack"
+        )
+        # Note: S3 doesn't need directories created beforehand - s3fs handles this
+
+        if ds.metadata is not None:
+            # Use s3:// prefix to ensure s3fs treats this as an S3 path
+            with cast(
+                BinaryIO, hive_fs.open(f"s3://{metadata_path.as_posix()}", "wb")
+            ) as f:
+                meta_packed = msgpack.packb(ds.metadata)
+                f.write(cast(bytes, meta_packed))
+
+        # Write data
+        shard_pattern = (self.hive_path / f"atdata--{new_uuid}--%06d.tar").as_posix()
+
+        written_shards: list[str] = []
+        with TemporaryDirectory() as temp_dir:
+            writer_opener, writer_post = _create_s3_write_callbacks(
+                credentials=self.s3_credentials,
+                temp_dir=temp_dir,
+                written_shards=written_shards,
+                fs=hive_fs,
+                cache_local=cache_local,
+                add_s3_prefix=False,
+            )
+
+            with wds.writer.ShardWriter(
+                shard_pattern,
+                opener=writer_opener,
+                post=writer_post,
+                **kwargs,
+            ) as sink:
+                for sample in ds.ordered(batch_size=None):
+                    sink.write(sample.as_wds)
+
+        # Make a new Dataset object for the written dataset copy
+        if len(written_shards) == 0:
+            raise RuntimeError(
+                "Cannot form new dataset entry -- did not write any shards"
+            )
+
+        elif len(written_shards) < 2:
+            new_dataset_url = (
+                self.hive_path / (Path(written_shards[0]).name)
+            ).as_posix()
+
+        else:
+            shard_s3_format = (
+                (self.hive_path / f"atdata--{new_uuid}").as_posix()
+            ) + "--{shard_id}.tar"
+            shard_id_braced = "{" + f"{0:06d}..{len(written_shards) - 1:06d}" + "}"
+            new_dataset_url = shard_s3_format.format(shard_id=shard_id_braced)
+
+        new_dataset = Dataset[ds.sample_type](
+            url=new_dataset_url,
+            metadata_url=metadata_path.as_posix(),
+        )
+
+        # Add to index (use ds._metadata to avoid network requests)
+        new_entry = self.index.add_entry(
+            new_dataset,
+            name=name,
+            schema_ref=schema_ref,
+            metadata=ds._metadata,
+        )
+
+        return new_entry, new_dataset

atdata/local/_s3.py

@@ -0,0 +1,349 @@
+"""S3-compatible data store and helper functions."""
+
+from atdata import Dataset
+
+from pathlib import Path
+from uuid import uuid4
+from tempfile import TemporaryDirectory
+from dotenv import dotenv_values
+from typing import Any, BinaryIO, cast
+
+from s3fs import S3FileSystem
+import webdataset as wds
+
+
+def _s3_env(credentials_path: str | Path) -> dict[str, Any]:
+    """Load S3 credentials from .env file.
+
+    Args:
+        credentials_path: Path to .env file containing AWS_ENDPOINT,
+            AWS_ACCESS_KEY_ID, and AWS_SECRET_ACCESS_KEY.
+
+    Returns:
+        Dict with the three required credential keys.
+
+    Raises:
+        ValueError: If any required key is missing from the .env file.
+    """
+    credentials_path = Path(credentials_path)
+    env_values = dotenv_values(credentials_path)
+
+    required_keys = ("AWS_ENDPOINT", "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY")
+    missing = [k for k in required_keys if k not in env_values]
+    if missing:
+        raise ValueError(
+            f"Missing required keys in {credentials_path}: {', '.join(missing)}"
+        )
+
+    return {k: env_values[k] for k in required_keys}
+
+
+def _s3_from_credentials(creds: str | Path | dict) -> S3FileSystem:
+    """Create S3FileSystem from credentials dict or .env file path."""
+    if not isinstance(creds, dict):
+        creds = _s3_env(creds)
+
+    # Build kwargs, making endpoint_url optional
+    kwargs = {
+        "key": creds["AWS_ACCESS_KEY_ID"],
+        "secret": creds["AWS_SECRET_ACCESS_KEY"],
+    }
+    if "AWS_ENDPOINT" in creds:
+        kwargs["endpoint_url"] = creds["AWS_ENDPOINT"]
+
+    return S3FileSystem(**kwargs)
+
+
+def _create_s3_write_callbacks(
+    credentials: dict[str, Any],
+    temp_dir: str,
+    written_shards: list[str],
+    fs: S3FileSystem | None,
+    cache_local: bool,
+    add_s3_prefix: bool = False,
+) -> tuple:
+    """Create opener and post callbacks for ShardWriter with S3 upload.
+
+    Args:
+        credentials: S3 credentials dict.
+        temp_dir: Temporary directory for local caching.
+        written_shards: List to append written shard paths to.
+        fs: S3FileSystem for direct writes (used when cache_local=False).
+        cache_local: If True, write locally then copy to S3.
+        add_s3_prefix: If True, prepend 's3://' to shard paths.
+
+    Returns:
+        Tuple of (writer_opener, writer_post) callbacks.
+    """
+    if cache_local:
+        import boto3
+
+        s3_client_kwargs = {
+            "aws_access_key_id": credentials["AWS_ACCESS_KEY_ID"],
+            "aws_secret_access_key": credentials["AWS_SECRET_ACCESS_KEY"],
+        }
+        if "AWS_ENDPOINT" in credentials:
+            s3_client_kwargs["endpoint_url"] = credentials["AWS_ENDPOINT"]
+        s3_client = boto3.client("s3", **s3_client_kwargs)
+
+        def _writer_opener(p: str):
+            local_path = Path(temp_dir) / p
+            local_path.parent.mkdir(parents=True, exist_ok=True)
+            return open(local_path, "wb")
+
+        def _writer_post(p: str):
+            local_path = Path(temp_dir) / p
+            path_parts = Path(p).parts
+            bucket = path_parts[0]
+            key = str(Path(*path_parts[1:]))
+
+            with open(local_path, "rb") as f_in:
+                s3_client.put_object(Bucket=bucket, Key=key, Body=f_in.read())
+
+            local_path.unlink()
+            if add_s3_prefix:
+                written_shards.append(f"s3://{p}")
+            else:
+                written_shards.append(p)
+
+        return _writer_opener, _writer_post
+    else:
+        if fs is None:
+            raise ValueError("S3FileSystem required when cache_local=False")
+
+        def _direct_opener(s: str):
+            return cast(BinaryIO, fs.open(f"s3://{s}", "wb"))
+
+        def _direct_post(s: str):
+            if add_s3_prefix:
+                written_shards.append(f"s3://{s}")
+            else:
+                written_shards.append(s)
+
+        return _direct_opener, _direct_post
+
+
+class S3DataStore:
+    """S3-compatible data store implementing AbstractDataStore protocol.
+
+    Handles writing dataset shards to S3-compatible object storage and
+    resolving URLs for reading.
+
+    Attributes:
+        credentials: S3 credentials dictionary.
+        bucket: Target bucket name.
+        _fs: S3FileSystem instance.
+    """
+
+    def __init__(
+        self,
+        credentials: str | Path | dict[str, Any],
+        *,
+        bucket: str,
+    ) -> None:
+        """Initialize an S3 data store.
+
+        Args:
+            credentials: Path to .env file or dict with AWS_ACCESS_KEY_ID,
+                AWS_SECRET_ACCESS_KEY, and optionally AWS_ENDPOINT.
+            bucket: Name of the S3 bucket for storage.
+        """
+        if isinstance(credentials, dict):
+            self.credentials = credentials
+        else:
+            self.credentials = _s3_env(credentials)
+
+        self.bucket = bucket
+        self._fs = _s3_from_credentials(self.credentials)
+
+    def write_shards(
+        self,
+        ds: Dataset,
+        *,
+        prefix: str,
+        cache_local: bool = False,
+        manifest: bool = False,
+        schema_version: str = "1.0.0",
+        source_job_id: str | None = None,
+        parent_shards: list[str] | None = None,
+        pipeline_version: str | None = None,
+        **kwargs,
+    ) -> list[str]:
+        """Write dataset shards to S3.
+
+        Args:
+            ds: The Dataset to write.
+            prefix: Path prefix within bucket (e.g., 'datasets/mnist/v1').
+            cache_local: If True, write locally first then copy to S3.
+            manifest: If True, generate per-shard manifest files alongside
+                each tar shard (``.manifest.json`` + ``.manifest.parquet``).
+            schema_version: Schema version for manifest headers.
+            source_job_id: Optional provenance job identifier for manifests.
+            parent_shards: Optional list of input shard identifiers for provenance.
+            pipeline_version: Optional pipeline version string for provenance.
+            **kwargs: Additional args passed to wds.ShardWriter (e.g., maxcount).
+
+        Returns:
+            List of S3 URLs for the written shards.
+
+        Raises:
+            RuntimeError: If no shards were written.
+        """
+        new_uuid = str(uuid4())
+        shard_pattern = f"{self.bucket}/{prefix}/data--{new_uuid}--%06d.tar"
+
+        written_shards: list[str] = []
+
+        # Manifest tracking state shared with the post callback
+        manifest_builders: list = []
+        current_builder: list = [None]  # mutable ref for closure
+        shard_counter: list[int] = [0]
+
+        if manifest:
+            from atdata.manifest import ManifestBuilder, ManifestWriter
+
+            def _make_builder(shard_idx: int) -> ManifestBuilder:
+                shard_id = f"{self.bucket}/{prefix}/data--{new_uuid}--{shard_idx:06d}"
+                return ManifestBuilder(
+                    sample_type=ds.sample_type,
+                    shard_id=shard_id,
+                    schema_version=schema_version,
+                    source_job_id=source_job_id,
+                    parent_shards=parent_shards,
+                    pipeline_version=pipeline_version,
+                )
+
+            current_builder[0] = _make_builder(0)
+
+        with TemporaryDirectory() as temp_dir:
+            writer_opener, writer_post_orig = _create_s3_write_callbacks(
+                credentials=self.credentials,
+                temp_dir=temp_dir,
+                written_shards=written_shards,
+                fs=self._fs,
+                cache_local=cache_local,
+                add_s3_prefix=True,
+            )
+
+            if manifest:
+
+                def writer_post(p: str):
+                    # Finalize the current manifest builder when a shard completes
+                    builder = current_builder[0]
+                    if builder is not None:
+                        manifest_builders.append(builder)
+                    # Advance to the next shard's builder
+                    shard_counter[0] += 1
+                    current_builder[0] = _make_builder(shard_counter[0])
+                    # Call original post callback
+                    writer_post_orig(p)
+            else:
+                writer_post = writer_post_orig
+
+            offset = 0
+            with wds.writer.ShardWriter(
+                shard_pattern,
+                opener=writer_opener,
+                post=writer_post,
+                **kwargs,
+            ) as sink:
+                for sample in ds.ordered(batch_size=None):
+                    wds_dict = sample.as_wds
+                    sink.write(wds_dict)
+
+                    if manifest and current_builder[0] is not None:
+                        packed_size = len(wds_dict.get("msgpack", b""))
+                        current_builder[0].add_sample(
+                            key=wds_dict["__key__"],
+                            offset=offset,
+                            size=packed_size,
+                            sample=sample,
+                        )
+                        # Approximate tar entry: 512-byte header + data rounded to 512
+                        offset += 512 + packed_size + (512 - packed_size % 512) % 512
+
+            # Finalize the last shard's builder (post isn't called for the last shard
+            # until ShardWriter closes, but we handle it here for safety)
+            if manifest and current_builder[0] is not None:
+                builder = current_builder[0]
+                if builder._rows:  # Only if samples were added
+                    manifest_builders.append(builder)
+
+            # Write all manifest files
+            if manifest:
+                for builder in manifest_builders:
+                    built = builder.build()
+                    writer = ManifestWriter(Path(temp_dir) / Path(built.shard_id))
+                    json_path, parquet_path = writer.write(built)
+
+                    # Upload manifest files to S3 alongside shards
+                    shard_id = built.shard_id
+                    json_key = f"{shard_id}.manifest.json"
+                    parquet_key = f"{shard_id}.manifest.parquet"
+
+                    if cache_local:
+                        import boto3
+
+                        s3_kwargs = {
+                            "aws_access_key_id": self.credentials["AWS_ACCESS_KEY_ID"],
+                            "aws_secret_access_key": self.credentials[
+                                "AWS_SECRET_ACCESS_KEY"
+                            ],
+                        }
+                        if "AWS_ENDPOINT" in self.credentials:
+                            s3_kwargs["endpoint_url"] = self.credentials["AWS_ENDPOINT"]
+                        s3_client = boto3.client("s3", **s3_kwargs)
+
+                        bucket_name = Path(shard_id).parts[0]
+                        json_s3_key = str(Path(*Path(json_key).parts[1:]))
+                        parquet_s3_key = str(Path(*Path(parquet_key).parts[1:]))
+
+                        with open(json_path, "rb") as f:
+                            s3_client.put_object(
+                                Bucket=bucket_name, Key=json_s3_key, Body=f.read()
+                            )
+                        with open(parquet_path, "rb") as f:
+                            s3_client.put_object(
+                                Bucket=bucket_name, Key=parquet_s3_key, Body=f.read()
+                            )
+                    else:
+                        self._fs.put(str(json_path), f"s3://{json_key}")
+                        self._fs.put(str(parquet_path), f"s3://{parquet_key}")
+
+        if len(written_shards) == 0:
+            raise RuntimeError("No shards written")
+
+        return written_shards
+
+    def read_url(self, url: str) -> str:
+        """Resolve an S3 URL for reading/streaming.
+
+        For S3-compatible stores with custom endpoints (like Cloudflare R2,
+        MinIO, etc.), converts s3:// URLs to HTTPS URLs that WebDataset can
+        stream directly.
+
+        For standard AWS S3 (no custom endpoint), URLs are returned unchanged
+        since WebDataset's built-in s3fs integration handles them.
+
+        Args:
+            url: S3 URL to resolve (e.g., 's3://bucket/path/file.tar').
+
+        Returns:
+            HTTPS URL if custom endpoint is configured, otherwise unchanged.
+            Example: 's3://bucket/path' -> 'https://endpoint.com/bucket/path'
+        """
+        endpoint = self.credentials.get("AWS_ENDPOINT")
+        if endpoint and url.startswith("s3://"):
+            # s3://bucket/path -> https://endpoint/bucket/path
+            path = url[5:]  # Remove 's3://' prefix
+            endpoint = endpoint.rstrip("/")
+            return f"{endpoint}/{path}"
+        return url
+
+    def supports_streaming(self) -> bool:
+        """S3 supports streaming reads.
+
+        Returns:
+            True.
+        """
+        return True