from __future__ import annotations
import io
import threading
from abc import ABC, abstractmethod
from collections.abc import Mapping, Sequence
from concurrent.futures import ThreadPoolExecutor
import boto3
import numpy as np
from botocore.config import Config
from botocore.exceptions import ClientError
DEFAULT_MAX_ARTIFACTS_PER_BATCH = 100
[docs]
class ArtifactReadError(RuntimeError):
"""Raised when an artifact cannot be read from its configured storage location.
Args:
artifact_id: Identifier of the artifact that could not be read.
bucket: Effective bucket used for the read.
"""
def __init__(self, artifact_id: str, bucket: str | None):
self.artifact_id = artifact_id
self.bucket = bucket
super().__init__(f"could not read artifact '{artifact_id}' from bucket '{bucket}'")
[docs]
class ArtifactStorage(ABC):
"""Abstract base for artifact key-value stores.
Implementations must support put/get/remove by string artifact_id and
optionally list all stored IDs. All mutating methods accept an optional
``bucket`` override; when omitted they use the backend's default bucket.
"""
def __init__(self, max_artifacts_per_batch: int = DEFAULT_MAX_ARTIFACTS_PER_BATCH):
if max_artifacts_per_batch < 1:
raise ValueError("max_artifacts_per_batch must be at least 1")
self._max_artifacts_per_batch = max_artifacts_per_batch
@property
def default_bucket(self) -> str | None:
"""Default bucket name for this storage backend, or None if not applicable."""
return None
[docs]
@abstractmethod
def put(self, artifact_id: str, data: np.ndarray, bucket: str | None = None) -> str | None:
"""Store a numpy array under the given artifact_id.
Args:
artifact_id: Unique key (UUID string).
data: Array to store.
bucket: Bucket override; uses backend default when None.
Returns:
The effective bucket name used for storage, or None for backends
without a bucket concept (e.g. in-memory).
"""
[docs]
@abstractmethod
def get(self, artifact_id: str, bucket: str | None = None) -> np.ndarray:
"""Retrieve a numpy array by artifact_id.
Args:
artifact_id: Key to look up.
bucket: Bucket override; uses backend default when None.
Returns:
The stored numpy array.
"""
[docs]
@abstractmethod
def remove(self, artifact_id: str, bucket: str | None = None) -> None:
"""Delete the artifact with the given artifact_id.
Args:
artifact_id: Key to delete.
bucket: Bucket override; uses backend default when None.
"""
[docs]
@abstractmethod
def list(self, bucket: str | None = None) -> list[str]:
"""Return all artifact_ids currently stored.
Args:
bucket: Bucket override; uses backend default when None.
Returns:
List of artifact_id strings.
"""
[docs]
def get_many(self, artifact_ids: list[str], bucket: str | None = None) -> list[np.ndarray]:
"""Retrieve multiple arrays by artifact_id, in the same order as the input.
Args:
artifact_ids: Keys to look up.
bucket: Bucket override forwarded to each :meth:`get` call.
Returns:
List of numpy arrays in the same order as artifact_ids.
"""
return [self.get(aid, bucket=bucket) for aid in artifact_ids]
[docs]
def get_artifact(self, artifact_id: str, bucket: str | None = None) -> np.ndarray:
"""Retrieve one decoded artifact by its ID and optional bucket.
Args:
artifact_id: Identifier of the artifact to retrieve.
bucket: Bucket override; uses :attr:`default_bucket` when None.
Returns:
The decoded numpy array.
Raises:
ArtifactReadError: If the artifact cannot be read.
"""
return self.get(artifact_id, bucket=bucket)
[docs]
def get_artifacts(self, locations: Sequence[tuple[str, str | None]]) -> Mapping[str, np.ndarray]:
"""Retrieve a bounded set of artifacts addressed by ID and bucket.
The returned mapping is keyed only by artifact ID. An ID may therefore
occur at one effective bucket only within a request. Identical locations
are read once and represented once in the result.
Args:
locations: Artifact ID and optional bucket pairs.
Returns:
Mapping of artifact IDs to decoded numpy arrays.
Raises:
ValueError: If the request exceeds the batch limit or one artifact
ID is requested from multiple effective buckets.
ArtifactReadError: If an artifact cannot be read.
"""
if len(locations) > self._max_artifacts_per_batch:
raise ValueError(f"requested {len(locations)} artifacts; maximum is {self._max_artifacts_per_batch}")
buckets_by_id: dict[str, str | None] = {}
ids_by_bucket: dict[str | None, list[str]] = {}
ordered_ids: list[str] = []
for artifact_id, bucket in locations:
effective_bucket = bucket if bucket is not None else self.default_bucket
if artifact_id in buckets_by_id:
if buckets_by_id[artifact_id] != effective_bucket:
raise ValueError(f"artifact '{artifact_id}' was requested from multiple buckets")
continue
buckets_by_id[artifact_id] = effective_bucket
ids_by_bucket.setdefault(effective_bucket, []).append(artifact_id)
ordered_ids.append(artifact_id)
arrays_by_id: dict[str, np.ndarray] = {}
for bucket, artifact_ids in ids_by_bucket.items():
arrays_by_id.update(zip(artifact_ids, self.get_many(artifact_ids, bucket=bucket)))
return {artifact_id: arrays_by_id[artifact_id] for artifact_id in ordered_ids}
[docs]
def ensure_bucket(self, bucket: str | None) -> None:
"""Ensure a bucket is available for writes when the backend requires it.
Args:
bucket: Bucket to ensure, or None for backends without buckets.
"""
[docs]
def put_many(self, artifact_ids: list[str], arrays: list[np.ndarray], bucket: str | None = None) -> None:
"""Store multiple arrays under their respective artifact_ids.
Args:
artifact_ids: Unique keys.
arrays: Arrays to store, paired with artifact_ids by position.
bucket: Bucket override forwarded to each :meth:`put` call.
"""
for aid, arr in zip(artifact_ids, arrays):
self.put(aid, arr, bucket=bucket)
[docs]
def remove_many(self, artifact_ids: list[str], bucket: str | None = None) -> None:
"""Delete multiple artifacts by artifact_id.
Args:
artifact_ids: Keys to delete.
bucket: Bucket override forwarded to each :meth:`remove` call.
"""
for aid in artifact_ids:
self.remove(aid, bucket=bucket)
[docs]
class InMemoryArtifactStorage(ArtifactStorage):
"""Dict-backed in-memory ArtifactStorage — useful for testing and dry runs.
The ``bucket`` parameter is accepted on all methods for interface compatibility
but is ignored; all artifacts share a single in-memory dict regardless of bucket.
"""
def __init__(self, max_artifacts_per_batch: int = DEFAULT_MAX_ARTIFACTS_PER_BATCH):
super().__init__(max_artifacts_per_batch=max_artifacts_per_batch)
self._store: dict[str, np.ndarray] = {}
[docs]
def put(self, artifact_id: str, data: np.ndarray, bucket: str | None = None) -> str | None:
"""Store array in memory.
Args:
artifact_id: Unique key.
data: Array to store.
bucket: Ignored.
Returns:
None (no bucket concept for in-memory storage).
"""
self._store[artifact_id] = data
return None
[docs]
def get(self, artifact_id: str, bucket: str | None = None) -> np.ndarray:
"""Retrieve array from memory.
Args:
artifact_id: Key to look up.
bucket: Ignored.
Returns:
The stored numpy array.
"""
try:
return self._store[artifact_id]
except KeyError as exc:
effective_bucket = bucket if bucket is not None else self.default_bucket
raise ArtifactReadError(artifact_id, effective_bucket) from exc
[docs]
def remove(self, artifact_id: str, bucket: str | None = None) -> None:
"""Remove array from memory.
Args:
artifact_id: Key to delete.
bucket: Ignored.
"""
del self._store[artifact_id]
[docs]
def list(self, bucket: str | None = None) -> list[str]:
"""Return all stored keys.
Args:
bucket: Ignored.
Returns:
List of artifact_id strings.
"""
return list(self._store)
[docs]
class S3ArtifactStorage(ArtifactStorage):
"""ArtifactStorage backed by any S3-compatible object store.
Numpy arrays are serialised via np.save/np.load into BytesIO and stored
as S3 objects keyed by artifact_id (UUID string). The bucket is created
on first use if it does not already exist.
Defaults point to the RustFS instance defined in docker-compose.yml
(endpoint ``http://localhost:9000``, credentials ``rustfs``/``rustfsdev``).
To use AWS S3, pass ``endpoint_url=None`` and supply real IAM credentials::
S3ArtifactStorage(
endpoint_url=None,
bucket="my-bucket",
access_key="MY_ACCESS_KEY",
secret_key="MY_SECRET_KEY",
)
Args:
endpoint_url: S3 API URL for self-hosted backends, e.g.
``"http://localhost:9000"`` for local RustFS. Pass ``None`` to
use the default AWS S3 endpoint.
bucket: Default bucket name.
access_key: Access key id (``RUSTFS_ACCESS_KEY`` for RustFS,
AWS access key id for AWS S3).
secret_key: Secret access key (``RUSTFS_SECRET_KEY`` for RustFS,
AWS secret access key for AWS S3).
create_buckets: Create missing buckets on construction and before writes.
Set False for read-only clients.
"""
def __init__(
self,
endpoint_url: str | None = "http://localhost:9000",
bucket: str = "artifacts",
access_key: str = "heros",
secret_key: str = "heros",
max_workers: int = 10,
max_artifacts_per_batch: int = DEFAULT_MAX_ARTIFACTS_PER_BATCH,
create_buckets: bool = True,
):
super().__init__(max_artifacts_per_batch=max_artifacts_per_batch)
self._bucket = bucket
self._max_workers = max_workers
self._create_buckets = create_buckets
self._known_buckets: set[str] = set()
self._bucket_lock = threading.Lock()
self._client = boto3.client(
"s3",
endpoint_url=endpoint_url,
aws_access_key_id=access_key,
aws_secret_access_key=secret_key,
config=Config(signature_version="s3v4", max_pool_connections=max_workers),
)
if create_buckets:
self.ensure_bucket(bucket)
@property
def default_bucket(self) -> str:
"""Default bucket name configured at construction time."""
return self._bucket
[docs]
def ensure_bucket(self, bucket: str | None) -> None:
"""Ensure an S3 bucket exists, tolerating concurrent creation.
Args:
bucket: Bucket to ensure. None uses :attr:`default_bucket`.
Raises:
ClientError: If the bucket cannot be inspected or created.
"""
effective_bucket = bucket if bucket is not None else self._bucket
if not self._create_buckets:
raise PermissionError("bucket creation is disabled for this storage instance")
if effective_bucket in self._known_buckets:
return
with self._bucket_lock:
if effective_bucket in self._known_buckets:
return
try:
self._client.head_bucket(Bucket=effective_bucket)
except ClientError:
try:
self._client.create_bucket(Bucket=effective_bucket)
except ClientError as create_error:
try:
self._client.head_bucket(Bucket=effective_bucket)
except ClientError:
raise create_error
self._known_buckets.add(effective_bucket)
[docs]
def put(self, artifact_id: str, data: np.ndarray, bucket: str | None = None) -> str:
"""Serialise array to npy bytes and upload as an S3 object.
Args:
artifact_id: S3 object key.
data: Array to store.
bucket: Bucket override; uses :attr:`default_bucket` when None.
Returns:
The effective bucket name used for storage.
"""
effective = bucket or self._bucket
self.ensure_bucket(effective)
buf = io.BytesIO()
np.save(buf, data)
self._client.put_object(Bucket=effective, Key=artifact_id, Body=buf.getvalue())
return effective
[docs]
def get(self, artifact_id: str, bucket: str | None = None) -> np.ndarray:
"""Download and deserialise an S3 object to a numpy array.
Args:
artifact_id: S3 object key.
bucket: Bucket override; uses :attr:`default_bucket` when None.
Returns:
The stored numpy array.
"""
effective_bucket = bucket if bucket is not None else self._bucket
try:
response = self._client.get_object(Bucket=effective_bucket, Key=artifact_id)
return np.load(io.BytesIO(response["Body"].read()))
except ClientError as exc:
raise ArtifactReadError(artifact_id, effective_bucket) from exc
[docs]
def remove(self, artifact_id: str, bucket: str | None = None) -> None:
"""Delete an S3 object.
Args:
artifact_id: S3 object key to delete.
bucket: Bucket override; uses :attr:`default_bucket` when None.
"""
self._client.delete_object(Bucket=bucket or self._bucket, Key=artifact_id)
[docs]
def list(self, bucket: str | None = None) -> list[str]:
"""Return all object keys in the bucket via paginated listing.
Args:
bucket: Bucket override; uses :attr:`default_bucket` when None.
Returns:
List of artifact_id strings.
"""
paginator = self._client.get_paginator("list_objects_v2")
keys = []
for page in paginator.paginate(Bucket=bucket or self._bucket):
keys.extend(obj["Key"] for obj in page.get("Contents", []))
return keys
[docs]
def get_many(self, artifact_ids: list[str], bucket: str | None = None) -> list[np.ndarray]:
"""Download and deserialise multiple S3 objects in parallel.
Args:
artifact_ids: S3 object keys to fetch.
bucket: Bucket override forwarded to each :meth:`get` call.
Returns:
List of numpy arrays in the same order as artifact_ids.
"""
effective = bucket or self._bucket
with ThreadPoolExecutor(max_workers=self._max_workers) as ex:
return list(ex.map(lambda aid: self.get(aid, bucket=effective), artifact_ids))
[docs]
def put_many(self, artifact_ids: list[str], arrays: list[np.ndarray], bucket: str | None = None) -> None:
"""Serialise and upload multiple arrays to S3 in parallel.
Args:
artifact_ids: S3 object keys.
arrays: Arrays to store, paired with artifact_ids by position.
bucket: Bucket override forwarded to each :meth:`put` call.
"""
effective = bucket or self._bucket
with ThreadPoolExecutor(max_workers=self._max_workers) as ex:
list(ex.map(lambda p: self.put(p[0], p[1], bucket=effective), zip(artifact_ids, arrays)))
[docs]
def remove_many(self, artifact_ids: list[str], bucket: str | None = None) -> None:
"""Delete multiple S3 objects using the native batch delete API.
Sends at most 1000 keys per request as required by the S3 API.
Args:
artifact_ids: S3 object keys to delete.
bucket: Bucket override; uses :attr:`default_bucket` when None.
"""
effective = bucket or self._bucket
chunk_size = 1000
for i in range(0, len(artifact_ids), chunk_size):
chunk = artifact_ids[i : i + chunk_size]
self._client.delete_objects(
Bucket=effective,
Delete={"Objects": [{"Key": k} for k in chunk]},
)