diff --git a/burr/core/application.py b/burr/core/application.py index 415e0984e..2806dd00a 100644 --- a/burr/core/application.py +++ b/burr/core/application.py @@ -60,6 +60,7 @@ StreamingAction, StreamingResultContainer, ) +from burr.core.artifacts import ArtifactStore, AsyncArtifactStore from burr.core.graph import Graph, GraphBuilder from burr.core.persistence import ( AsyncBaseStateLoader, @@ -642,6 +643,9 @@ def my_action(state: State, __context: ApplicationContext) -> State: state_initializer: Optional[BaseStateLoader] state_persister: Optional[BaseStateSaver] action_name: Optional[str] # Store just the action name + # Trailing field with a `None` default so directly constructing this public dataclass the + # pre-existing way (without `object_store`) keeps working -- see with_object_store. + object_store: Optional[Union[ArtifactStore, AsyncArtifactStore]] = None @staticmethod def get() -> Optional["ApplicationContext"]: @@ -849,6 +853,7 @@ def __init__( parallel_executor_factory: Optional[Executor] = None, state_persister: Union[BaseStateSaver, LifecycleAdapter, None] = None, state_initializer: Union[BaseStateLoader, LifecycleAdapter, None] = None, + object_store: Optional[Union[ArtifactStore, AsyncArtifactStore]] = None, ): """Instantiates an Application. This is an internal API -- use the builder! @@ -900,6 +905,7 @@ def __init__( self._spawning_parent_pointer = spawning_parent_pointer self._state_initializer = state_initializer self._state_persister = state_persister + self._object_store = object_store self._adapter_set.call_all_lifecycle_hooks_sync( "post_application_create", state=self._state, @@ -943,6 +949,7 @@ def _context_factory(self, action: Action, sequence_id: int) -> ApplicationConte parallel_executor_factory=self._parallel_executor_factory, state_initializer=self._state_initializer, state_persister=self._state_persister, + object_store=self._object_store, action_name=action.name if action else None, # Pass just the action name ) @@ -2224,6 +2231,7 @@ def __init__(self): self.typing_system = None self.parallel_executor_factory = None self.state_persister = None + self.object_store = None self._is_async: bool = False def with_identifiers( @@ -2544,6 +2552,46 @@ def with_state_persister( self.state_persister = persister # tracks for later; validates in build / abuild return self + def with_object_store( + self, object_store: Union[ArtifactStore, AsyncArtifactStore] + ) -> "ApplicationBuilder[StateType]": + """Adds an object/blob store to the application, for storing large values (files, + dataframes, images, etc...) outside of ``State`` -- see :py:mod:`burr.core.artifacts`. + + Unlike :py:meth:`with_state_persister`, this does not add a lifecycle hook -- Burr never + reads from or writes to the store on its own. It is purely made available to actions + through :py:class:`ApplicationContext` (``__context.object_store``) so they don't each + need to construct/import their own store, and so it can be swapped (e.g. local disk in + dev, S3 in prod) in one place. + + .. code-block:: python + + from burr.core import action, State, ApplicationContext + from burr.core.artifacts import ArtifactStore + + @action(reads=[], writes=["pdf_doc"]) + def ingest_pdf(state: State, pdf_bytes: bytes, __context: ApplicationContext) -> State: + ref = __context.object_store.put_artifact(pdf_bytes, media_type="application/pdf") + return state.update(pdf_doc=ref) + + app = ( + ApplicationBuilder() + .with_actions(ingest_pdf, ...) + .with_object_store(LocalFileSystemArtifactStore(root_dir="./blobs")) + .build() + ) + + An :py:class:`AsyncArtifactStore` (e.g. an ``aiobotocore``-backed S3 store) may be passed + too, for use from async actions without blocking the event loop -- build the application + with :py:meth:`abuild` in that case, since :py:meth:`build` rejects an async store (it + can't be awaited from a sync action). + + :param object_store: The artifact store to make available to actions. + :return: The application builder for future chaining. + """ + self.object_store = object_store + return self + def with_spawning_parent( self, app_id: str, sequence_id: int, partition_key: Optional[str] = None ) -> "ApplicationBuilder[StateType]": @@ -2785,6 +2833,7 @@ def _build_common(self) -> Application: parallel_executor_factory=self.parallel_executor_factory, state_persister=self.state_persister, state_initializer=self.state_initializer, + object_store=self.object_store, ) def build(self) -> Application[StateType]: @@ -2815,6 +2864,12 @@ def build(self) -> Application[StateType]: :return: The application object. """ _validate_app_id(self.app_id) + if self.object_store is not None and self.object_store.is_async(): + raise ValueError( + "You are building the sync application, but have used an " + "async object store. Please use a sync object store (ArtifactStore) or " + "use the .abuild() method to build an async application." + ) if self.state is None: self.state = State() @@ -2836,6 +2891,9 @@ async def abuild(self) -> Application[StateType]: Note: When you run an async application you can still use the normal sync functionalities, i.e. sync hooks of other adapters, but they will block the async event loop until finished. + Likewise, a synchronous :py:class:`~burr.core.artifacts.ArtifactStore` passed to + :py:meth:`with_object_store` still works here, but its calls will block the event loop -- + prefer an :py:class:`~burr.core.artifacts.AsyncArtifactStore` for I/O-bound backends. In case you are using state initializers and persisters, the asynchronous application should be used in the following cases: diff --git a/burr/core/artifacts.py b/burr/core/artifacts.py new file mode 100644 index 000000000..78287cea0 --- /dev/null +++ b/burr/core/artifacts.py @@ -0,0 +1,349 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Blob/object storage abstraction for large values held in :py:class:`burr.core.state.State`. + +Rather than embedding large binary payloads (files, dataframes, images, etc...) directly in +state -- which gets fully re-serialized to the persistence layer after every action -- store the +payload in an :py:class:`ArtifactStore` and keep a small, content-addressed :py:class:`ArtifactRef` +in state instead. Deserializing a ref does *not* eagerly fetch the underlying bytes -- callers +must explicitly call :py:meth:`ArtifactRef.read` (or ``store.get_artifact``) to do so. + +See :doc:`../reference/serde` and :doc:`../concepts/state-persistence` for background on how +state serialization interacts with this module. +""" + +import abc +import hashlib +import os +import tempfile +from dataclasses import dataclass +from typing import Optional + +from burr.core import serde + + +@dataclass(frozen=True) +class ArtifactRef: + """A lightweight, content-addressed reference to a blob stored in an :py:class:`ArtifactStore`. + + This is what you should store in :py:class:`burr.core.state.State` -- *not* the raw bytes/object. + It is registered with :py:mod:`burr.core.serde` so it serializes/deserializes automatically as + part of normal state (de)serialization; deserializing only recreates this reference, it never + re-fetches the underlying bytes on its own. + + :param key: The store-specific key/path the blob is stored under. + :param size_bytes: Size of the blob, in bytes. + :param digest: SHA-256 hex digest of the blob's contents, used to verify integrity on read + and to detect duplicate content. + :param media_type: Optional MIME type (e.g. ``"application/pdf"``) for the caller's own use. + """ + + key: str + size_bytes: int + digest: str + media_type: Optional[str] = None + + def read(self, store: "ArtifactStore", verify: bool = True) -> bytes: + """Lazily fetches the underlying bytes for this reference from ``store``. + + :param store: The store this artifact was originally written to. + :param verify: Whether to verify the fetched bytes match ``self.digest``. Defaults to True. + :return: The raw bytes of the artifact. + """ + return store.get_artifact(self, verify=verify) + + async def aread(self, store: "AsyncArtifactStore", verify: bool = True) -> bytes: + """Asynchronous counterpart to :py:meth:`read`, for use with an + :py:class:`AsyncArtifactStore` (e.g. from an async action). + + :param store: The async store this artifact was originally written to. + :param verify: Whether to verify the fetched bytes match ``self.digest``. Defaults to True. + :return: The raw bytes of the artifact. + """ + return await store.get_artifact(self, verify=verify) + + +class ArtifactStore(abc.ABC): + """Base class for blob/object storage backends. + + Implement :py:meth:`put`, :py:meth:`get`, and :py:meth:`exists` for a new backend -- + :py:meth:`put_artifact` and :py:meth:`get_artifact` are provided for free on top of those, + so digest computation/verification is not duplicated across backends. + """ + + @abc.abstractmethod + def put(self, data: bytes, key: str) -> None: + """Stores ``data`` under ``key``. Should be idempotent -- writing the same key twice + with the same content should not error. + + :param data: The raw bytes to store. + :param key: The key/path to store the data under. + """ + + @abc.abstractmethod + def get(self, key: str) -> bytes: + """Retrieves the raw bytes stored under ``key``. + + :param key: The key/path to load. + :raises FileNotFoundError: If no data is stored under ``key``. + """ + + @abc.abstractmethod + def exists(self, key: str) -> bool: + """Returns whether ``key`` is present in the store.""" + + def is_async(self) -> bool: + """Whether this store's I/O methods are async. Always False here -- see + :py:class:`AsyncArtifactStore` for the async contract. Mirrors + :py:meth:`burr.core.persistence.BaseStateSaver.is_async`.""" + return False + + def put_artifact( + self, data: bytes, media_type: Optional[str] = None, key: Optional[str] = None + ) -> ArtifactRef: + """Computes the digest/size for ``data``, writes it to the store, and returns a + fully-formed :py:class:`ArtifactRef`. This is the recommended way to write artifacts -- + it ensures every backend computes/verifies digests the same way. + + With the default content-addressed ``key`` (the hex digest), writing identical content + twice is always a safe no-op -- the key can only match if the bytes match. + + With an explicit ``key``, a write is only skipped if the *existing* content stored under + that key has the same digest as ``data``. If the key already exists with *different* + content, this raises -- it never silently returns a ref that doesn't describe what's + actually stored. + + :param data: The raw bytes to store. + :param media_type: Optional MIME type to record on the returned ref. + :param key: Optional explicit key to store under. If not provided, a content-addressed + key (the hex digest) is used -- so writing identical content twice is a no-op. + :raises ValueError: If ``key`` is explicitly provided and already exists with content + that does not match the digest of ``data``. + :return: An :py:class:`ArtifactRef` describing the stored artifact. + """ + digest = hashlib.sha256(data).hexdigest() + resolved_key = key if key is not None else digest + if key is not None and self.exists(resolved_key): + existing_digest = hashlib.sha256(self.get(resolved_key)).hexdigest() + if existing_digest != digest: + raise ValueError( + f"Key '{resolved_key}' already exists with different content (existing " + f"digest {existing_digest}, new digest {digest}). Use a different explicit " + f"key, or omit `key` to use content-addressed storage." + ) + # Identical content is already stored under this key -- nothing to do. + else: + self.put(data, resolved_key) + return ArtifactRef( + key=resolved_key, size_bytes=len(data), digest=digest, media_type=media_type + ) + + def get_artifact(self, ref: ArtifactRef, verify: bool = True) -> bytes: + """Fetches the bytes described by ``ref``, optionally verifying its digest. + + :param ref: The reference describing what to fetch. + :param verify: Whether to verify the fetched bytes match ``ref.digest``. Defaults to True. + :raises ValueError: If ``verify`` is True and the fetched bytes' digest does not match. + """ + data = self.get(ref.key) + if verify: + actual_digest = hashlib.sha256(data).hexdigest() + if actual_digest != ref.digest: + raise ValueError( + f"Digest mismatch for artifact '{ref.key}': " + f"expected {ref.digest}, got {actual_digest}. The stored content may have " + f"been overwritten or corrupted." + ) + return data + + +class AsyncArtifactStore(abc.ABC): + """Asynchronous counterpart to :py:class:`ArtifactStore`, for backends whose I/O should not + block the event loop (e.g. network calls to a cloud object store). Implement this instead of + :py:class:`ArtifactStore` if your backend's ``put``/``get``/``exists`` are I/O-bound and you + want it to be usable from async actions without blocking other concurrent work -- mirrors the + sync/async split already used for state persisters (see + :py:class:`burr.core.persistence.BaseStateSaver` / + :py:class:`burr.core.persistence.AsyncBaseStateSaver`). + + A synchronous store (e.g. :py:class:`LocalFileSystemArtifactStore`) can still be used from an + async action, but its calls will block the event loop for their duration -- prefer an async + store (e.g. an ``aiobotocore``-backed S3 store) for I/O-bound backends used from async code. + """ + + @abc.abstractmethod + async def put(self, data: bytes, key: str) -> None: + """Stores ``data`` under ``key``. Should be idempotent -- writing the same key twice + with the same content should not error. + + :param data: The raw bytes to store. + :param key: The key/path to store the data under. + """ + + @abc.abstractmethod + async def get(self, key: str) -> bytes: + """Retrieves the raw bytes stored under ``key``. + + :param key: The key/path to load. + :raises FileNotFoundError: If no data is stored under ``key``. + """ + + @abc.abstractmethod + async def exists(self, key: str) -> bool: + """Returns whether ``key`` is present in the store.""" + + def is_async(self) -> bool: + """Whether this store's I/O methods are async. Always True here -- see + :py:class:`ArtifactStore` for the synchronous contract.""" + return True + + async def put_artifact( + self, data: bytes, media_type: Optional[str] = None, key: Optional[str] = None + ) -> ArtifactRef: + """Asynchronous counterpart to :py:meth:`ArtifactStore.put_artifact` -- same digest + computation and explicit-key collision semantics (see there for details), just awaited + instead of blocking. + + :param data: The raw bytes to store. + :param media_type: Optional MIME type to record on the returned ref. + :param key: Optional explicit key to store under. If not provided, a content-addressed + key (the hex digest) is used -- so writing identical content twice is a no-op. + :raises ValueError: If ``key`` is explicitly provided and already exists with content + that does not match the digest of ``data``. + :return: An :py:class:`ArtifactRef` describing the stored artifact. + """ + digest = hashlib.sha256(data).hexdigest() + resolved_key = key if key is not None else digest + if key is not None and await self.exists(resolved_key): + existing_digest = hashlib.sha256(await self.get(resolved_key)).hexdigest() + if existing_digest != digest: + raise ValueError( + f"Key '{resolved_key}' already exists with different content (existing " + f"digest {existing_digest}, new digest {digest}). Use a different explicit " + f"key, or omit `key` to use content-addressed storage." + ) + # Identical content is already stored under this key -- nothing to do. + else: + await self.put(data, resolved_key) + return ArtifactRef( + key=resolved_key, size_bytes=len(data), digest=digest, media_type=media_type + ) + + async def get_artifact(self, ref: ArtifactRef, verify: bool = True) -> bytes: + """Fetches the bytes described by ``ref``, optionally verifying its digest. + + :param ref: The reference describing what to fetch. + :param verify: Whether to verify the fetched bytes match ``ref.digest``. Defaults to True. + :raises ValueError: If ``verify`` is True and the fetched bytes' digest does not match. + """ + data = await self.get(ref.key) + if verify: + actual_digest = hashlib.sha256(data).hexdigest() + if actual_digest != ref.digest: + raise ValueError( + f"Digest mismatch for artifact '{ref.key}': " + f"expected {ref.digest}, got {actual_digest}. The stored content may have " + f"been overwritten or corrupted." + ) + return data + + +class LocalFileSystemArtifactStore(ArtifactStore): + """An :py:class:`ArtifactStore` backed by a directory on the local filesystem. + Has no third-party dependencies. For cloud-backed stores (S3, GCS, ...) see + ``burr.integrations.artifacts``. + """ + + def __init__(self, root_dir: str): + """Constructor. + + :param root_dir: Directory to store artifacts under. Created if it does not exist. + """ + self.root_dir = root_dir + os.makedirs(root_dir, exist_ok=True) + + def _path_for_key(self, key: str) -> str: + # Resolve symlinks (os.path.realpath) rather than just normalizing (os.path.abspath) -- + # otherwise a symlink inside root_dir pointing outside of it would let a key escape the + # store undetected. realpath resolves symlinks in existing path components even when the + # final path segment (e.g. a not-yet-written artifact file) does not exist yet. + root = os.path.realpath(self.root_dir) + path = os.path.realpath(os.path.join(root, key)) + if os.path.commonpath([root, path]) != root: + raise ValueError( + f"Invalid artifact key '{key}': resolves outside of the store's root directory." + ) + return path + + def put(self, data: bytes, key: str) -> None: + path = self._path_for_key(key) + if os.path.exists(path): + # content-addressed keys make writes idempotent -- skip re-writing existing data. + return + directory = os.path.dirname(path) + os.makedirs(directory, exist_ok=True) + # Write to a temp file in the same directory and atomically rename it into place, so + # concurrent readers/crashes never observe a partial/truncated file at `path` -- either + # `path` doesn't exist yet, or it's fully written. + fd, tmp_path = tempfile.mkstemp(dir=directory, prefix=".tmp-artifact-") + try: + with os.fdopen(fd, "wb") as f: + f.write(data) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp_path, path) + except BaseException: + try: + os.remove(tmp_path) + except OSError: + pass + raise + + def get(self, key: str) -> bytes: + path = self._path_for_key(key) + if not os.path.exists(path): + raise FileNotFoundError(f"Artifact '{key}' not found in store at '{self.root_dir}'.") + with open(path, "rb") as f: + return f.read() + + def exists(self, key: str) -> bool: + return os.path.exists(self._path_for_key(key)) + + +_ARTIFACT_REF_SERDE_KEY = "ArtifactRef" + + +@serde.serialize.register(ArtifactRef) +def _serialize_artifact_ref(value: ArtifactRef, **kwargs) -> dict: + return { + serde.KEY: _ARTIFACT_REF_SERDE_KEY, + "key": value.key, + "size_bytes": value.size_bytes, + "digest": value.digest, + "media_type": value.media_type, + } + + +@serde.deserializer.register(_ARTIFACT_REF_SERDE_KEY) +def _deserialize_artifact_ref(value: dict, **kwargs) -> ArtifactRef: + return ArtifactRef( + key=value["key"], + size_bytes=value["size_bytes"], + digest=value["digest"], + media_type=value.get("media_type"), + ) diff --git a/burr/core/parallelism.py b/burr/core/parallelism.py index 857fed333..228dff408 100644 --- a/burr/core/parallelism.py +++ b/burr/core/parallelism.py @@ -40,6 +40,7 @@ from burr.core import Action, ApplicationBuilder, ApplicationContext, Graph, State from burr.core.action import SingleStepAction from burr.core.application import ApplicationIdentifiers +from burr.core.artifacts import ArtifactStore, AsyncArtifactStore from burr.core.graph import GraphBuilder from burr.core.persistence import BaseStateLoader, BaseStateSaver from burr.lifecycle import LifecycleAdapter @@ -102,8 +103,11 @@ class SubGraphTask: tracker: Optional[TrackingClient] = None state_persister: Optional[BaseStateSaver] = None state_initializer: Optional[BaseStateLoader] = None + object_store: Optional[Union[ArtifactStore, AsyncArtifactStore]] = None - def _create_app_builder(self, parent_context: ApplicationIdentifiers) -> ApplicationBuilder: + def _create_app_builder( + self, parent_context: ApplicationIdentifiers, attach_object_store: bool = True + ) -> ApplicationBuilder: builder = ( ApplicationBuilder() .with_graph(self.graph.graph) @@ -121,6 +125,15 @@ def _create_app_builder(self, parent_context: ApplicationIdentifiers) -> Applica if self.tracker is not None: builder = builder.with_tracker(self.tracker) # TODO -- move this into the adapter + # attach_object_store=False is used by arun() -- it attaches the object store directly + # to the built Application afterwards instead, since build()/abuild() choice there is + # driven entirely by state_persister/state_initializer sync-ness and may not agree with + # the object store's sync-ness (e.g. an async object store cascaded alongside a sync + # persister still works fine, since the sub-app is always executed via app.arun() + # regardless of whether build() or abuild() was used to construct it). + if attach_object_store and self.object_store is not None: + builder = builder.with_object_store(self.object_store) + # In this case we want to persist the state for the app if self.state_persister is not None: builder = builder.with_state_persister(self.state_persister) @@ -154,7 +167,14 @@ def run( return state async def arun(self, parent_context: ApplicationContext): - # Here for backwards compatibility, not ideal + # Here for backwards compatibility, not ideal. Note the object store's sync/async-ness + # is intentionally *not* a factor in choosing build() vs abuild() below -- the sub-app is + # always run via app.arun() in this method regardless of which one is used, so either + # store type works either way. It's attached directly to the built app afterwards + # (attach_object_store=False + the explicit assignment below) instead of going through + # with_object_store(), since build()/abuild() would otherwise reject an object store + # whose sync/async-ness doesn't match -- a check that only makes sense for the public + # build()/abuild() -> run()/arun() pairing, not this internal always-async codepath. if (self.state_initializer is not None and not self.state_initializer.is_async()) or ( self.state_persister is not None and not self.state_persister.is_async() ): @@ -162,9 +182,11 @@ async def arun(self, parent_context: ApplicationContext): "You are using sync persisters for an async application which is not optimal. " "Consider switching to an async persister implementation. We will make this an error soon." ) - app = self._create_app_builder(parent_context).build() + app = self._create_app_builder(parent_context, attach_object_store=False).build() else: - app = await self._create_app_builder(parent_context).abuild() + app = await self._create_app_builder(parent_context, attach_object_store=False).abuild() + if self.object_store is not None: + app._object_store = self.object_store action, result, state = await app.arun( halt_after=self.graph.halt_after, inputs={key: value for key, value in self.inputs.items() if not key.startswith("__")}, @@ -513,6 +535,7 @@ def _create_task(key: str, action: Action, substate: State) -> SubGraphTask: tracker=tracker, state_persister=state_persister, state_initializer=state_initializer, + object_store=context.object_store, ) def _tasks() -> Generator[SubGraphTask, None, None]: diff --git a/burr/integrations/artifacts/__init__.py b/burr/integrations/artifacts/__init__.py new file mode 100644 index 000000000..13a83393a --- /dev/null +++ b/burr/integrations/artifacts/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/burr/integrations/artifacts/s3.py b/burr/integrations/artifacts/s3.py new file mode 100644 index 000000000..b77cec620 --- /dev/null +++ b/burr/integrations/artifacts/s3.py @@ -0,0 +1,211 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""An :py:class:`~burr.core.artifacts.ArtifactStore` backed by AWS S3. + +Requires the ``boto3`` package -- install with ``pip install "apache-burr[s3]"``. + +Also provides :py:class:`AsyncS3ArtifactStore`, a non-blocking counterpart backed by +``aiobotocore`` (installed alongside ``boto3`` via the same ``s3`` extra) for use from async +actions -- ``S3ArtifactStore`` performs blocking network I/O and should not be called from async +actions, since it would block the event loop for the duration of each upload/download. +""" + +from typing import Any, Optional + +from burr.core.artifacts import ArtifactStore, AsyncArtifactStore +from burr.integrations import base + +try: + import boto3 + import botocore.exceptions +except ImportError as e: + base.require_plugin(e, "s3") + +try: + import aiobotocore.session +except ImportError as e: + base.require_plugin(e, "s3") + + +class S3ArtifactStore(ArtifactStore): + """Stores artifacts as objects in an S3 bucket, optionally under a key prefix.""" + + @classmethod + def from_config(cls, config: dict) -> "S3ArtifactStore": + """Creates a new instance from a configuration dictionary. See :py:meth:`from_values` + for the accepted keys.""" + return cls.from_values(**config) + + @classmethod + def from_values( + cls, + bucket: str, + prefix: str = "", + client_kwargs: Optional[dict] = None, + ) -> "S3ArtifactStore": + """Creates a new instance, constructing its own boto3 S3 client. + + :param bucket: Name of the S3 bucket to store artifacts in. + :param prefix: Optional key prefix (e.g. ``"my-app/artifacts"``) applied to every key. + :param client_kwargs: Optional kwargs passed to ``boto3.client("s3", ...)``. + """ + client = boto3.client("s3", **(client_kwargs or {})) + return cls(bucket, prefix=prefix, client=client) + + def __init__(self, bucket: str, prefix: str = "", client: Optional[Any] = None): + """Constructor. + + :param bucket: Name of the S3 bucket to store artifacts in. + :param prefix: Optional key prefix (e.g. ``"my-app/artifacts"``) applied to every key. + :param client: An existing boto3 S3 client. If not provided, one is created with + default credentials/region resolution (``boto3.client("s3")``). + """ + self.bucket = bucket + self.prefix = prefix + self.client = client if client is not None else boto3.client("s3") + + def _object_key(self, key: str) -> str: + if self.prefix: + return f"{self.prefix.rstrip('/')}/{key}" + return key + + def put(self, data: bytes, key: str) -> None: + # content-addressed keys make writes idempotent -- skip re-uploading existing data. + if self.exists(key): + return + self.client.put_object(Bucket=self.bucket, Key=self._object_key(key), Body=data) + + def get(self, key: str) -> bytes: + try: + response = self.client.get_object(Bucket=self.bucket, Key=self._object_key(key)) + except botocore.exceptions.ClientError as e: + error_code = e.response.get("Error", {}).get("Code") + if error_code in ("NoSuchKey", "404"): + raise FileNotFoundError( + f"Artifact '{key}' not found in bucket '{self.bucket}' " + f"(prefix='{self.prefix}')." + ) from e + raise + return response["Body"].read() + + def exists(self, key: str) -> bool: + try: + self.client.head_object(Bucket=self.bucket, Key=self._object_key(key)) + return True + except botocore.exceptions.ClientError as e: + error_code = e.response.get("Error", {}).get("Code") + if error_code in ("404", "NoSuchKey"): + return False + raise + + +class AsyncS3ArtifactStore(AsyncArtifactStore): + """Non-blocking counterpart to :py:class:`S3ArtifactStore`, backed by ``aiobotocore`` so it + can safely be used from async actions without blocking the event loop. Prefer this over + :py:class:`S3ArtifactStore` whenever the store is configured on an application built with + :py:meth:`~burr.core.application.ApplicationBuilder.abuild`. + """ + + @classmethod + async def acreate( + cls, + bucket: str, + prefix: str = "", + client_kwargs: Optional[dict] = None, + ) -> "AsyncS3ArtifactStore": + """Async factory that opens its own ``aiobotocore`` session/client. Must be awaited -- + unlike ``boto3.client(...)``, an ``aiobotocore`` client is opened via an async context + manager, so this can't be done in a plain (synchronous) constructor/``from_values``. + + :param bucket: Name of the S3 bucket to store artifacts in. + :param prefix: Optional key prefix (e.g. ``"my-app/artifacts"``) applied to every key. + :param client_kwargs: Optional kwargs passed to ``session.create_client("s3", ...)``. + """ + session = aiobotocore.session.get_session() + client_cm = session.create_client("s3", **(client_kwargs or {})) + client = await client_cm.__aenter__() + return cls(bucket, prefix=prefix, client=client, _client_cm=client_cm) + + def __init__( + self, + bucket: str, + prefix: str = "", + client: Optional[Any] = None, + _client_cm: Optional[Any] = None, + ): + """Constructor. + + :param bucket: Name of the S3 bucket to store artifacts in. + :param prefix: Optional key prefix (e.g. ``"my-app/artifacts"``) applied to every key. + :param client: An already-open ``aiobotocore`` S3 client (e.g. entered via + ``async with session.create_client("s3") as client``). Prefer :py:meth:`acreate` if + you don't already have one -- this is intentionally not optional/self-constructing + since opening an ``aiobotocore`` client requires an ``await``. + """ + if client is None: + raise ValueError( + "AsyncS3ArtifactStore requires an already-open aiobotocore client. Use " + "`await AsyncS3ArtifactStore.acreate(...)` to construct one, or open your own " + "and pass it as `client`." + ) + self.bucket = bucket + self.prefix = prefix + self.client = client + self._client_cm = _client_cm + + async def aclose(self) -> None: + """Closes the underlying client, if this instance opened one itself via + :py:meth:`acreate`. No-op if a pre-opened ``client`` was passed to the constructor + directly -- that client is owned by the caller.""" + if self._client_cm is not None: + await self._client_cm.__aexit__(None, None, None) + + def _object_key(self, key: str) -> str: + if self.prefix: + return f"{self.prefix.rstrip('/')}/{key}" + return key + + async def put(self, data: bytes, key: str) -> None: + # content-addressed keys make writes idempotent -- skip re-uploading existing data. + if await self.exists(key): + return + await self.client.put_object(Bucket=self.bucket, Key=self._object_key(key), Body=data) + + async def get(self, key: str) -> bytes: + try: + response = await self.client.get_object(Bucket=self.bucket, Key=self._object_key(key)) + except botocore.exceptions.ClientError as e: + error_code = e.response.get("Error", {}).get("Code") + if error_code in ("NoSuchKey", "404"): + raise FileNotFoundError( + f"Artifact '{key}' not found in bucket '{self.bucket}' " + f"(prefix='{self.prefix}')." + ) from e + raise + async with response["Body"] as stream: + return await stream.read() + + async def exists(self, key: str) -> bool: + try: + await self.client.head_object(Bucket=self.bucket, Key=self._object_key(key)) + return True + except botocore.exceptions.ClientError as e: + error_code = e.response.get("Error", {}).get("Code") + if error_code in ("404", "NoSuchKey"): + return False + raise diff --git a/docs/concepts/artifact-storage.rst b/docs/concepts/artifact-storage.rst new file mode 100644 index 000000000..dee9000b2 --- /dev/null +++ b/docs/concepts/artifact-storage.rst @@ -0,0 +1,252 @@ +.. + Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + + +================ +Artifact Storage +================ + +.. _artifact-storage: + +.. note:: + + Burr comes with an optional object/blob storage abstraction for large values (files, + images, dataframes, model weights, etc...) that you don't want to embed directly in + :py:class:`State `. This is separate from :doc:`state-persistence` -- + it's for the *values* your actions produce/consume, not the state machine's bookkeeping. + +TL;DR +----- + +:py:class:`State ` is fully re-serialized to your persister/tracker +after every action. If you put a large binary blob (a PDF, an image, a big dataframe) directly +in state, that entire blob gets re-written on every subsequent step, even if it never changes. + +Burr's artifact storage API solves this by giving you a place to write the blob once +(an :py:class:`ArtifactStore `), and a small, +content-addressed handle to keep in state instead (an +:py:class:`ArtifactRef `). The ref serializes/deserializes like +any other state value, but loading it back does **not** eagerly fetch the underlying bytes -- +you explicitly call :py:meth:`ArtifactRef.read ` (or its +async counterpart, :py:meth:`ArtifactRef.aread `, for an +``AsyncArtifactStore``) when you actually need them. + +This is an opt-in, explicit mechanism. Burr never reads from or writes to an artifact store on +its own -- your actions are always the ones calling ``put_artifact``/``get_artifact``. + +Why not just put the bytes in State? +------------------------------------- + +Consider an action that ingests an uploaded PDF and a later action that summarizes it: + +.. code-block:: python + + @action(reads=[], writes=["pdf_bytes"]) + def ingest_pdf(state: State, pdf_bytes: bytes) -> State: + return state.update(pdf_bytes=pdf_bytes) # the whole PDF is now in State + + @action(reads=["pdf_bytes"], writes=["summary"]) + def summarize(state: State) -> State: + summary = call_llm(state["pdf_bytes"]) + return state.update(summary=summary) + +Every time ``State`` is persisted or tracked after this point (including for actions that have +nothing to do with the PDF), the full PDF bytes are re-serialized along with everything else. +This gets expensive quickly, and it also means large binary data ends up embedded in your +persistence layer and shown as raw JSON in the tracking UI. + +Storing an artifact reference instead +-------------------------------------- + +Use an :py:class:`ArtifactStore ` to store the blob once, and +keep only the small :py:class:`ArtifactRef ` -- a key, size, +SHA-256 digest, and optional media type -- in state: + +.. code-block:: python + + from burr.core import action, State, ApplicationContext + from burr.core.artifacts import ArtifactStore + + @action(reads=[], writes=["pdf_doc"]) + def ingest_pdf(state: State, pdf_bytes: bytes, __context: ApplicationContext) -> State: + store: ArtifactStore = __context.object_store + ref = store.put_artifact(pdf_bytes, media_type="application/pdf") + return state.update(pdf_doc=ref) # only the small ref is in State now + + @action(reads=["pdf_doc"], writes=["summary"]) + def summarize(state: State, __context: ApplicationContext) -> State: + store: ArtifactStore = __context.object_store + pdf_bytes = state["pdf_doc"].read(store) # explicit, lazy fetch + summary = call_llm(pdf_bytes) + return state.update(summary=summary) + +``put_artifact`` computes a SHA-256 digest of the data and, unless you pass an explicit ``key``, +uses that digest as the storage key. This makes writes idempotent and content-addressed -- +storing the same bytes twice (even across app runs) is a no-op, so you get de-duplication for +free without any bookkeeping of your own. + +``get_artifact`` (and ``ArtifactRef.read``) re-computes the digest of whatever comes back from +the store and compares it to ``ref.digest`` by default, raising a ``ValueError`` if they don't +match -- this catches corrupted or overwritten data early rather than silently returning bad +bytes. + +If you pass an explicit ``key`` (rather than relying on content-addressing), reusing that key is +only safe if the content is the same: if the key already exists with *different* content, +``put_artifact`` raises a ``ValueError`` instead of silently overwriting or returning a ref that +doesn't match what's actually stored. Stick to content-addressed (default) keys unless you have a +specific reason to name your own -- e.g. a stable "latest" pointer you intend to update +out-of-band. + +:py:class:`LocalFileSystemArtifactStore ` writes +are atomic (via a temp file + rename) and reject keys that would resolve outside ``root_dir`` +(including through symlinks), so concurrent writers/readers never observe partial files and +artifacts can't be written outside the configured store directory. + +Making a store available to actions +------------------------------------ + +Rather than having every action construct or import its own store, configure one on the +:py:class:`ApplicationBuilder ` with +:py:meth:`with_object_store `. It's +then available to any action via :py:class:`ApplicationContext ` +(``__context.object_store``), the same way :py:meth:`with_state_persister ` +and :py:meth:`with_tracker ` expose the +persister/tracker. + +.. code-block:: python + + from burr.core import ApplicationBuilder + from burr.core.artifacts import LocalFileSystemArtifactStore + + app = ( + ApplicationBuilder() + .with_actions(ingest_pdf, summarize, ...) + .with_transitions(...) + .with_state(...) + .with_entrypoint(...) + .with_object_store(LocalFileSystemArtifactStore(root_dir="./blobs")) + .build() + ) + +Unlike ``with_state_persister``, ``with_object_store`` does not register a lifecycle hook -- +Burr never calls into the store itself. It's purely a way to configure and share the store, so +you can swap implementations (local disk in dev, S3 in prod) in one place. + +.. note:: + + ``__context`` is only injected if it appears in your action's signature -- see + :ref:`State Persistence ` for more on ``ApplicationContext``. + +Using an artifact store from async actions +-------------------------------------------- + +:py:class:`ArtifactStore ` is synchronous -- calling it from +an async action blocks the event loop for the duration of each ``put``/``get`` (e.g. a network +round-trip to S3). For async applications, use an +:py:class:`AsyncArtifactStore ` instead (e.g. +:py:class:`AsyncS3ArtifactStore `), and +build the application with :py:meth:`abuild `: + +.. code-block:: python + + from burr.core import action, State, ApplicationContext + from burr.core.artifacts import AsyncArtifactStore + from burr.integrations.artifacts.s3 import AsyncS3ArtifactStore + + @action(reads=[], writes=["pdf_doc"]) + async def ingest_pdf(state: State, pdf_bytes: bytes, __context: ApplicationContext) -> State: + store: AsyncArtifactStore = __context.object_store + ref = await store.put_artifact(pdf_bytes, media_type="application/pdf") + return state.update(pdf_doc=ref) + + object_store = await AsyncS3ArtifactStore.acreate(bucket="my-bucket") + app = await ( + ApplicationBuilder() + .with_actions(ingest_pdf, ...) + .with_transitions(...) + .with_state(...) + .with_entrypoint(...) + .with_object_store(object_store) + .abuild() + ) + +:py:meth:`build ` (the synchronous builder) +rejects an async store with a clear error at build time, since it has no way to await it. An +``ArtifactStore`` (sync) can still be used from ``abuild()``-built applications, but, as with a +sync state persister, it will block the event loop while it runs. + +Supported Backends +------------------- + +.. list-table:: Burr Implemented Artifact Stores + :header-rows: 1 + :widths: auto + + * - Backend + - Class + - Extra dependency + * - Local disk + - :py:class:`LocalFileSystemArtifactStore ` + - none (stdlib only) + * - AWS S3 + - :py:class:`S3ArtifactStore ` + - ``pip install "apache-burr[s3]"`` + * - AWS S3 (async) + - :py:class:`AsyncS3ArtifactStore ` + - ``pip install "apache-burr[s3]"`` (includes ``aiobotocore``) + +See :ref:`the API reference ` for full details, and +:ref:`the S3 integration reference ` for setup instructions. + +Implementing your own backend +------------------------------ + +To back artifacts with a different store (GCS, Azure Blob Storage, a database, ...), +subclass :py:class:`ArtifactStore ` and implement three +methods -- ``put``, ``get``, and ``exists``: + +.. code-block:: python + + from burr.core.artifacts import ArtifactStore + + class MyCustomArtifactStore(ArtifactStore): + def put(self, data: bytes, key: str) -> None: + ... # write data under key -- must be idempotent + + def get(self, key: str) -> bytes: + ... # raise FileNotFoundError if key does not exist + + def exists(self, key: str) -> bool: + ... + +``put_artifact`` and ``get_artifact`` (digest computation/verification, content-addressed keys) +are provided for free by the base class, so you don't need to reimplement them for a new backend. + +If your backend's I/O is itself async (e.g. a cloud SDK with a native async client), subclass +:py:class:`AsyncArtifactStore ` instead and implement +``async def put/get/exists`` -- it provides the same ``put_artifact``/``get_artifact`` helpers, +just awaited. + +Interaction with serialization +-------------------------------- + +``ArtifactRef`` is a plain, frozen dataclass registered with :doc:`serde` out of the box, so it +serializes/deserializes as part of normal :py:class:`State ` (de)serialization +with no extra setup required. Deserializing a ref only recreates the small dataclass -- it never +fetches the underlying bytes on its own, so browsing history, forking, or replaying an +application never triggers unexpected I/O against your artifact store. diff --git a/docs/concepts/index.rst b/docs/concepts/index.rst index 8762d8518..e0ebec28f 100644 --- a/docs/concepts/index.rst +++ b/docs/concepts/index.rst @@ -38,6 +38,7 @@ Overview of the concepts -- read these to get a mental model for how Burr works. tracking s3-tracking-aws state-persistence + artifact-storage serde streaming-actions state-typing diff --git a/docs/reference/artifacts.rst b/docs/reference/artifacts.rst new file mode 100644 index 000000000..2296cb70b --- /dev/null +++ b/docs/reference/artifacts.rst @@ -0,0 +1,76 @@ +.. + Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + + +================= +Artifact Storage +================= + +.. _artifactsref: + +Burr provides a set of tools for storing large values (files, images, dataframes, ...) outside +of :py:class:`State `, keeping only a small reference in state. See +:doc:`../concepts/artifact-storage` for a conceptual overview and usage examples. + +We currently support the following backends: + +.. list-table:: Burr Implemented Artifact Stores + :header-rows: 1 + :widths: auto + + * - Backend + - Class + - Extra dependency + * - Local disk + - :ref:`LocalFileSystemArtifactStore ` + - none (stdlib only) + * - AWS S3 + - :ref:`S3ArtifactStore ` + - ``pip install "apache-burr[s3]"`` + * - AWS S3 (async) + - :ref:`AsyncS3ArtifactStore ` + - ``pip install "apache-burr[s3]"`` + +If you want to implement your own artifact store (to bridge it with a new backend), you should +implement the ``ArtifactStore`` interface -- or ``AsyncArtifactStore`` if your backend's I/O is +itself async (e.g. a cloud SDK with a native async client) and you want to use it from async +actions without blocking the event loop. + +.. autoclass:: burr.core.artifacts.ArtifactStore + :members: + :show-inheritance: + +.. autoclass:: burr.core.artifacts.AsyncArtifactStore + :members: + :show-inheritance: + +.. autoclass:: burr.core.artifacts.ArtifactRef + :members: + +Supported Implementations +========================== + +.. _localartifactstoreref: + +.. autoclass:: burr.core.artifacts.LocalFileSystemArtifactStore + :members: + + .. automethod:: __init__ + +See :doc:`integrations/index` for cloud-backed implementations such as +:ref:`S3ArtifactStore `. diff --git a/docs/reference/index.rst b/docs/reference/index.rst index d734972e1..bc7fc40ec 100644 --- a/docs/reference/index.rst +++ b/docs/reference/index.rst @@ -36,6 +36,7 @@ need functionality that is not publicly exposed, please open an issue and we can state serde persister + artifacts conditions tracking visibility diff --git a/docs/reference/integrations/index.rst b/docs/reference/integrations/index.rst index e5b288363..64d861aef 100644 --- a/docs/reference/integrations/index.rst +++ b/docs/reference/integrations/index.rst @@ -36,3 +36,4 @@ Integrations -- we will be adding more haystack ray bedrock + s3 diff --git a/docs/reference/integrations/s3.rst b/docs/reference/integrations/s3.rst new file mode 100644 index 000000000..afae3c9f9 --- /dev/null +++ b/docs/reference/integrations/s3.rst @@ -0,0 +1,130 @@ +.. + Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + + +.. _s3-artifacts-integration: + +=================== +S3 Artifact Storage +=================== + +An :py:class:`ArtifactStore ` backed by an +`AWS S3 `_ bucket. See :doc:`../../concepts/artifact-storage` for +background on artifact storage in general. + +Install the optional extra (pulls ``boto3`` and ``aiobotocore``): + +.. code-block:: bash + + pip install "apache-burr[s3]" + +------------- +Quick Example +------------- + +.. code-block:: python + + from burr.core import ApplicationBuilder + from burr.integrations.artifacts.s3 import S3ArtifactStore + + object_store = S3ArtifactStore.from_values( + bucket="my-bucket", + prefix="burr-artifacts", # optional + ) + + app = ( + ApplicationBuilder() + .with_actions(...) + .with_transitions(...) + .with_state(...) + .with_entrypoint(...) + .with_object_store(object_store) + .build() + ) + +If you already have a boto3 client (e.g. reusing a connection, custom endpoint for testing with +`moto `_ or a local S3-compatible server), construct the store +directly and pass it in: + +.. code-block:: python + + import boto3 + from burr.integrations.artifacts.s3 import S3ArtifactStore + + client = boto3.client("s3", region_name="us-east-1") + object_store = S3ArtifactStore(bucket="my-bucket", prefix="burr-artifacts", client=client) + +-------------------------- +Async (non-blocking) Usage +-------------------------- + +``S3ArtifactStore`` performs blocking network I/O through ``boto3`` -- calling it from an async +action blocks the event loop for the duration of each upload/download. For async applications +(built with :py:meth:`abuild `), use +:py:class:`AsyncS3ArtifactStore ` instead, +which is backed by `aiobotocore `_: + +.. code-block:: python + + from burr.core import ApplicationBuilder + from burr.integrations.artifacts.s3 import AsyncS3ArtifactStore + + object_store = await AsyncS3ArtifactStore.acreate( + bucket="my-bucket", + prefix="burr-artifacts", # optional + ) + + app = await ( + ApplicationBuilder() + .with_actions(...) + .with_transitions(...) + .with_state(...) + .with_entrypoint(...) + .with_object_store(object_store) + .abuild() + ) + +``acreate`` opens and owns its underlying ``aiobotocore`` client -- call +``await object_store.aclose()`` when you're done with it (e.g. on app shutdown) to release it. +If you already have an open ``aiobotocore`` client, pass it directly via the constructor's +``client`` argument instead. + +--------------- +IAM Permissions +--------------- + +At minimum, the runtime identity used by Burr needs permission to read/write/check objects in +the target bucket (optionally scoped to the configured prefix): + +* ``s3:PutObject`` +* ``s3:GetObject`` +* ``s3:HeadObject`` + +---------------- +API Reference +---------------- + +.. autoclass:: burr.integrations.artifacts.s3.S3ArtifactStore + :members: + + .. automethod:: __init__ + +.. autoclass:: burr.integrations.artifacts.s3.AsyncS3ArtifactStore + :members: + + .. automethod:: __init__ diff --git a/pyproject.toml b/pyproject.toml index 7fba971cc..f2c7eec98 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -83,6 +83,11 @@ redis = [ "redis" ] +s3 = [ + "boto3", + "aiobotocore" +] + release = [ "jinja2", ] @@ -105,12 +110,14 @@ tests = [ "apache-burr[opentelemetry]", "apache-burr[langfuse]", "apache-burr[haystack]", - "apache-burr[ray]" + "apache-burr[ray]", + "apache-burr[s3]" ] documentation = [ "apache-burr[tests]", "apache-burr[bedrock]", + "apache-burr[s3]", "sphinx<9", "sphinx-autobuild", "myst-nb", diff --git a/tests/core/test_application.py b/tests/core/test_application.py index 5a252f778..43af4745e 100644 --- a/tests/core/test_application.py +++ b/tests/core/test_application.py @@ -17,6 +17,7 @@ import asyncio import collections +import concurrent.futures import datetime import logging import typing @@ -4148,6 +4149,157 @@ def load( builder.with_state_persister(persister) +def test_with_object_store_defaults_to_none(): + builder = ApplicationBuilder() + assert builder.object_store is None + + +def test_with_object_store_sets_builder_field(): + from burr.core.artifacts import LocalFileSystemArtifactStore + + store = LocalFileSystemArtifactStore(root_dir="/tmp/burr-test-object-store") + builder = ApplicationBuilder().with_object_store(store) + assert builder.object_store is store + + +def test_with_object_store_exposed_through_application_context(tmp_path): + from burr.core.artifacts import ArtifactRef, LocalFileSystemArtifactStore + + store = LocalFileSystemArtifactStore(root_dir=str(tmp_path)) + seen_context: Dict[str, Optional[ApplicationContext]] = {"context": None} + + @action(reads=[], writes=["doc"]) + def ingest(state: State, __context: ApplicationContext) -> State: + seen_context["context"] = __context + ref = __context.object_store.put_artifact(b"hello world", media_type="text/plain") + return state.update(doc=ref) + + app = ( + ApplicationBuilder() + .with_actions(ingest=ingest, terminal=Result("doc")) + .with_transitions(("ingest", "terminal")) + .with_entrypoint("ingest") + .with_state() + .with_object_store(store) + .build() + ) + + *_, state = app.run(halt_after=["terminal"]) + + assert seen_context["context"] is not None + assert seen_context["context"].object_store is store + + ref = state["doc"] + assert isinstance(ref, ArtifactRef) + assert ref.read(store) == b"hello world" + + +def test_without_object_store_application_context_has_none(): + @action(reads=[], writes=["ran"]) + def check_no_store(state: State, __context: ApplicationContext) -> State: + assert __context.object_store is None + return state.update(ran=True) + + app = ( + ApplicationBuilder() + .with_actions(check_no_store=check_no_store, terminal=Result("ran")) + .with_transitions(("check_no_store", "terminal")) + .with_entrypoint("check_no_store") + .with_state() + .build() + ) + + *_, state = app.run(halt_after=["terminal"]) + assert state["ran"] is True + + +def test_application_context_backwards_compatible_without_object_store_kwarg(): + """ApplicationContext is public, exported API. Constructing it the way code did before + `object_store` was added (i.e. omitting the kwarg entirely) must keep working -- the field + must have a trailing `None` default, not become a newly-required positional/keyword arg.""" + context = ApplicationContext( + app_id="app_id", + partition_key=None, + sequence_id=0, + tracker=None, + parallel_executor_factory=lambda: concurrent.futures.ThreadPoolExecutor(), + state_initializer=None, + state_persister=None, + action_name="some_action", + # note: no object_store kwarg passed at all + ) + assert context.object_store is None + + +def _make_in_memory_async_store(): + """Builds a minimal `AsyncArtifactStore` test double -- in-memory, so tests don't need + network/disk. Imported lazily to match this file's convention of importing + `burr.core.artifacts` within individual tests rather than at module scope.""" + from burr.core.artifacts import AsyncArtifactStore + + class InMemoryAsyncArtifactStore(AsyncArtifactStore): + def __init__(self): + self._data: Dict[str, bytes] = {} + + async def put(self, data: bytes, key: str) -> None: + self._data[key] = data + + async def get(self, key: str) -> bytes: + if key not in self._data: + raise FileNotFoundError(key) + return self._data[key] + + async def exists(self, key: str) -> bool: + return key in self._data + + return InMemoryAsyncArtifactStore() + + +def test_build_rejects_async_object_store(): + """A sync `.build()` can't await an async store's `put`/`get`/`exists` -- using one should + fail fast at build time with a clear error, not silently produce a broken sync app.""" + store = _make_in_memory_async_store() + builder = ( + ApplicationBuilder() + .with_actions(terminal=Result()) + .with_transitions() + .with_entrypoint("terminal") + .with_state() + .with_object_store(store) + ) + with pytest.raises(ValueError, match="async object store"): + builder.build() + + +async def test_abuild_exposes_async_object_store_through_application_context(): + store = _make_in_memory_async_store() + seen_context: Dict[str, Optional[ApplicationContext]] = {"context": None} + + @action(reads=[], writes=["doc"]) + async def ingest(state: State, __context: ApplicationContext) -> State: + seen_context["context"] = __context + ref = await __context.object_store.put_artifact(b"hello async world") + return state.update(doc=ref) + + app = await ( + ApplicationBuilder() + .with_actions(ingest=ingest, terminal=Result("doc")) + .with_transitions(("ingest", "terminal")) + .with_entrypoint("ingest") + .with_state() + .with_object_store(store) + .abuild() + ) + + *_, state = await app.arun(halt_after=["terminal"]) + + assert seen_context["context"] is not None + assert seen_context["context"].object_store is store + + ref = state["doc"] + assert await ref.aread(store) == b"hello async world" + + class ActionWithoutContext(Action): def run(self, other_param, foo): pass diff --git a/tests/core/test_artifacts.py b/tests/core/test_artifacts.py new file mode 100644 index 000000000..919a83222 --- /dev/null +++ b/tests/core/test_artifacts.py @@ -0,0 +1,371 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import hashlib +import subprocess +import sys +import textwrap + +import pytest + +from burr.core import serde, state +from burr.core.artifacts import ArtifactRef, AsyncArtifactStore, LocalFileSystemArtifactStore + + +@pytest.fixture +def store(tmp_path): + return LocalFileSystemArtifactStore(root_dir=str(tmp_path)) + + +def test_put_and_get_roundtrip(store): + data = b"hello world large payload" + store.put(data, key="my-key") + + assert store.exists("my-key") + assert store.get("my-key") == data + + +def test_get_missing_key_raises(store): + with pytest.raises(FileNotFoundError): + store.get("does-not-exist") + + +def test_exists_false_for_missing_key(store): + assert not store.exists("does-not-exist") + + +def test_put_and_get_nested_key_creates_parent_dirs(store): + data = b"nested payload" + store.put(data, key="sub/dir/file.txt") + + assert store.exists("sub/dir/file.txt") + assert store.get("sub/dir/file.txt") == data + + +@pytest.mark.parametrize( + "bad_key", + ["../escape.txt", "../../etc/passwd", "sub/../../escape.txt"], +) +def test_put_rejects_keys_that_escape_root_dir(store, bad_key): + with pytest.raises(ValueError): + store.put(b"malicious", key=bad_key) + + +def test_get_rejects_keys_that_escape_root_dir(store): + with pytest.raises(ValueError): + store.get("../escape.txt") + + +def test_put_rejects_symlinked_escape(store, tmp_path): + outside_dir = tmp_path.parent / f"{tmp_path.name}-outside" + outside_dir.mkdir() + (tmp_path / "link").symlink_to(outside_dir) + + with pytest.raises(ValueError): + store.put(b"malicious", key="link/escape.txt") + assert not (outside_dir / "escape.txt").exists() + + +def test_get_rejects_symlinked_escape(store, tmp_path): + outside_dir = tmp_path.parent / f"{tmp_path.name}-outside" + outside_dir.mkdir() + (outside_dir / "escape.txt").write_bytes(b"outside data") + (tmp_path / "link").symlink_to(outside_dir) + + with pytest.raises(ValueError): + store.get("link/escape.txt") + + +def test_exists_rejects_symlinked_escape(store, tmp_path): + outside_dir = tmp_path.parent / f"{tmp_path.name}-outside" + outside_dir.mkdir() + (tmp_path / "link").symlink_to(outside_dir) + + with pytest.raises(ValueError): + store.exists("link/escape.txt") + + +def test_put_writes_atomically_and_recovers_from_interrupted_write(store, tmp_path, monkeypatch): + def failing_fsync(fd): + raise OSError("simulated crash mid-write") + + monkeypatch.setattr("os.fsync", failing_fsync) + + with pytest.raises(OSError): + store.put(b"x" * 1000, key="flaky-key") + + # the destination must not exist -- the write failed before the atomic rename into place + assert not store.exists("flaky-key") + # no leftover temp files should be left behind in the store directory + leftovers = [p.name for p in tmp_path.iterdir() if p.name.startswith(".tmp-artifact-")] + assert leftovers == [] + + monkeypatch.undo() + + # a subsequent, uninterrupted write for the same key must succeed normally + store.put(b"good data", key="flaky-key") + assert store.get("flaky-key") == b"good data" + + +def test_put_artifact_computes_digest_and_size(store): + data = b"some binary content" + ref = store.put_artifact(data, media_type="application/octet-stream") + + assert isinstance(ref, ArtifactRef) + assert ref.size_bytes == len(data) + assert ref.digest == hashlib.sha256(data).hexdigest() + assert ref.media_type == "application/octet-stream" + # content-addressed by default -- key derived from digest + assert ref.key == ref.digest + + +def test_put_artifact_with_explicit_key(store): + data = b"some binary content" + ref = store.put_artifact(data, key="explicit-key") + + assert ref.key == "explicit-key" + assert store.get("explicit-key") == data + + +def test_put_artifact_explicit_key_identical_content_is_noop(store): + data = b"same content, written twice" + ref1 = store.put_artifact(data, key="stable-key") + ref2 = store.put_artifact(data, key="stable-key") + + assert ref1 == ref2 + assert store.get("stable-key") == data + + +def test_put_artifact_explicit_key_conflicting_content_raises(store): + store.put_artifact(b"first", key="mutable") + + with pytest.raises(ValueError, match="already exists with different content"): + store.put_artifact(b"second", key="mutable") + + # the original content must be left untouched -- a rejected write must not corrupt the key + assert store.get("mutable") == b"first" + + +def test_put_artifact_is_idempotent_for_identical_content(store): + data = b"duplicate content" + ref1 = store.put_artifact(data) + ref2 = store.put_artifact(data) + + # same content -> same content-addressed key, written only once + assert ref1.key == ref2.key + + +def test_get_artifact_roundtrip(store): + data = b"round trip me" + ref = store.put_artifact(data) + + assert store.get_artifact(ref) == data + + +def test_get_artifact_verifies_digest_by_default(store, tmp_path): + data = b"original content" + ref = store.put_artifact(data, key="mutable-key") + + # simulate corruption/tampering of the underlying stored bytes on disk directly -- + # store.put() is a no-op for an already-existing key, so we bypass it here. + (tmp_path / "mutable-key").write_bytes(b"tampered content") + + with pytest.raises(ValueError, match="Digest mismatch"): + store.get_artifact(ref) + + +def test_get_artifact_skips_verification_when_disabled(store, tmp_path): + data = b"original content" + ref = store.put_artifact(data, key="mutable-key") + (tmp_path / "mutable-key").write_bytes(b"tampered content") + + # no error, just returns the (tampered) bytes as-is + assert store.get_artifact(ref, verify=False) == b"tampered content" + + +def test_artifact_ref_read_delegates_to_store(store): + data = b"delegate read" + ref = store.put_artifact(data) + + assert ref.read(store) == data + + +def test_artifact_ref_serde_roundtrip(): + ref = ArtifactRef(key="abc", size_bytes=10, digest="0" * 64, media_type="text/plain") + + serialized = serde.serialize(ref) + assert serialized[serde.KEY] == "ArtifactRef" + assert serialized["key"] == "abc" + + deserialized = serde.deserialize(serialized) + assert deserialized == ref + + +def test_state_with_artifact_ref_serializes_as_reference_only(store): + data = b"a very large file we do not want in the DB" + ref = store.put_artifact(data, media_type="application/pdf") + s = state.State({"document": ref}) + + serialized_state = s.serialize() + # state.serialize() must only contain the lightweight reference, not the raw bytes + assert serialized_state["document"] == { + serde.KEY: "ArtifactRef", + "key": ref.key, + "size_bytes": ref.size_bytes, + "digest": ref.digest, + "media_type": "application/pdf", + } + + +def test_state_deserialize_restores_ref_without_fetching_bytes(store): + data = b"a very large file we do not want eagerly loaded" + ref = store.put_artifact(data) + s = state.State({"document": ref}) + serialized_state = s.serialize() + + restored_state = state.State.deserialize(serialized_state) + + restored_ref = restored_state["document"] + assert isinstance(restored_ref, ArtifactRef) + assert restored_ref == ref + # deserializing state must not read artifact bytes -- only .read()/get_artifact() should + restored_ref.read(store) # exercised separately/explicitly, proving it's a distinct step + + +def test_artifact_ref_deserializes_in_fresh_process_without_explicit_import(): + """Regression test: the ArtifactRef (de)serializer must be registered as a side effect of + importing ``burr.core`` (transitively, e.g. via ``burr.core.state``) -- not only when a + caller happens to import ``burr.core.artifacts`` directly first. Runs in a fresh subprocess + so no other test in this process can have already triggered the registration.""" + script = textwrap.dedent( + """ + import sys + assert "burr.core.artifacts" not in sys.modules, ( + "burr.core.artifacts must not already be imported for this to be a valid regression " + "test" + ) + + from burr.core import serde + from burr.core.state import State + + serialized = { + "document": { + serde.KEY: "ArtifactRef", + "key": "abc123", + "size_bytes": 5, + "digest": "abc123", + "media_type": "text/plain", + } + } + restored = State.deserialize(serialized) + ref = restored["document"] + assert type(ref).__name__ == "ArtifactRef" + assert ref.key == "abc123" + print("OK") + """ + ) + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, result.stderr + assert "OK" in result.stdout + + +class InMemoryAsyncArtifactStore(AsyncArtifactStore): + """Minimal in-memory `AsyncArtifactStore` used to exercise the async base-class contract + (digest computation, explicit-key collision semantics, verification) without needing a real + I/O-bound backend.""" + + def __init__(self): + self._data: dict = {} + + async def put(self, data: bytes, key: str) -> None: + self._data[key] = data + + async def get(self, key: str) -> bytes: + if key not in self._data: + raise FileNotFoundError(key) + return self._data[key] + + async def exists(self, key: str) -> bool: + return key in self._data + + +@pytest.fixture +def async_store(): + return InMemoryAsyncArtifactStore() + + +def test_artifact_store_is_async_false(store): + assert store.is_async() is False + + +def test_async_artifact_store_is_async_true(async_store): + assert async_store.is_async() is True + + +async def test_async_put_artifact_computes_digest_and_size(async_store): + data = b"async payload" + ref = await async_store.put_artifact(data) + + assert ref.digest == hashlib.sha256(data).hexdigest() + assert ref.size_bytes == len(data) + assert ref.key == ref.digest # content-addressed by default + + +async def test_async_put_artifact_explicit_key_identical_content_is_noop(async_store): + data = b"same bytes" + first = await async_store.put_artifact(data, key="mutable") + second = await async_store.put_artifact(data, key="mutable") + + assert first == second + assert await async_store.get("mutable") == data + + +async def test_async_put_artifact_explicit_key_conflicting_content_raises(async_store): + await async_store.put_artifact(b"first", key="mutable") + + with pytest.raises(ValueError, match="already exists with different content"): + await async_store.put_artifact(b"second", key="mutable") + + # the original content must be untouched after the rejected write + assert await async_store.get("mutable") == b"first" + + +async def test_async_get_artifact_roundtrip(async_store): + data = b"async roundtrip" + ref = await async_store.put_artifact(data) + + assert await async_store.get_artifact(ref) == data + + +async def test_async_get_artifact_verifies_digest_by_default(async_store): + ref = await async_store.put_artifact(b"original") + await async_store.put(b"tampered", ref.key) # simulate corruption/overwrite out-of-band + + with pytest.raises(ValueError, match="Digest mismatch"): + await async_store.get_artifact(ref) + + +async def test_artifact_ref_aread_delegates_to_async_store(async_store): + data = b"delegate aread" + ref = await async_store.put_artifact(data) + + assert await ref.aread(async_store) == data diff --git a/tests/core/test_parallelism.py b/tests/core/test_parallelism.py index 25d37cc24..fc54bf99f 100644 --- a/tests/core/test_parallelism.py +++ b/tests/core/test_parallelism.py @@ -19,6 +19,7 @@ import concurrent.futures import dataclasses import datetime +import tempfile from random import random from typing import Any, AsyncGenerator, Callable, Dict, Generator, List, Literal, Optional, Union @@ -34,6 +35,8 @@ action, ) from burr.core.action import Input, Result +from burr.core.application import ApplicationIdentifiers +from burr.core.artifacts import AsyncArtifactStore, LocalFileSystemArtifactStore from burr.core.graph import GraphBuilder from burr.core.parallelism import ( MapActions, @@ -83,6 +86,21 @@ async def simple_single_fn_subgraph_async( ) +# Single action/callable subgraph that reads/writes an artifact through the +# (async) object store exposed on the sub-application's context. +@action(reads=["input_number", "number_to_add"], writes=["output_number"]) +async def simple_single_fn_subgraph_async_object_store( + state: State, __context: ApplicationContext, additional_number: int = 1 +) -> State: + assert __context.object_store is not None + ref = await __context.object_store.put_artifact(b"hello from subtask") + data = await ref.aread(__context.object_store) + assert data == b"hello from subtask" + return state.update( + output_number=state["input_number"] + state["number_to_add"] + additional_number + ) + + class ClassBasedAction(Action): def __init__(self, identifying_number: int, name: str = "class_based_action"): super().__init__() @@ -1217,6 +1235,7 @@ def reads(self) -> list[str]: tracker=tracker, state_persister=persister, state_initializer=persister, + object_store=None, parallel_executor_factory=lambda: concurrent.futures.ThreadPoolExecutor(), action_name=action.name, ), @@ -1227,3 +1246,127 @@ def reads(self) -> list[str]: assert task.state_initializer is not None assert task.tracker is not None assert task.state_persister is task.state_initializer # This ensures they're the same + + +def test_map_actions_and_states_cascades_object_store(): + """Tests that the object_store on the parent ApplicationContext is cascaded to the + SubGraphTask created for each sub-application, so actions running inside a parallel + task can still access it through __context.object_store.""" + + class SimpleMapStates(MapActionsAndStates): + def actions( + self, state: State, context: ApplicationContext, inputs: Dict[str, Any] + ) -> Generator[Union[Action, Callable, RunnableGraph], None, None]: + for graph_ in [ + simple_single_fn_subgraph.bind(identifying_number=1000), + ]: + yield graph_ + + def states( + self, state: State, context: ApplicationContext, inputs: Dict[str, Any] + ) -> Generator[State, None, None]: + yield state.update(input_number=0, number_to_add=0) + + def reduce(self, state: State, states: Generator[State, None, None]) -> State: + new_state = state + for output_state in states: + new_state = new_state.append(output_numbers_in_state=output_state["output_number"]) + return new_state + + @property + def writes(self) -> list[str]: + return ["output_numbers_in_state"] + + @property + def reads(self) -> list[str]: + return ["input_numbers_in_state"] + + action = SimpleMapStates() + object_store = LocalFileSystemArtifactStore(tempfile.mkdtemp()) + + task_generator = action.tasks( + state=State(), + context=ApplicationContext( + app_id="app_id", + partition_key="partition_key", + sequence_id=0, + tracker=None, + state_persister=None, + state_initializer=None, + object_store=object_store, + parallel_executor_factory=lambda: concurrent.futures.ThreadPoolExecutor(), + action_name=action.name, + ), + inputs={}, + ) + (task,) = task_generator # one task + assert task.object_store is object_store + + # And ensure it's actually threaded into the built sub-application's context + builder = task._create_app_builder( + ApplicationIdentifiers(app_id="app_id", partition_key="partition_key", sequence_id=0) + ) + assert builder.object_store is object_store + + +class InMemoryAsyncArtifactStore(AsyncArtifactStore): + """Minimal async object store test double -- used to exercise SubGraphTask.arun() with + an async object store, independent of any real backend.""" + + def __init__(self): + self._data: Dict[str, bytes] = {} + + async def put(self, data: bytes, key: str) -> None: + self._data[key] = data + + async def get(self, key: str) -> bytes: + return self._data[key] + + async def exists(self, key: str) -> bool: + return key in self._data + + +@pytest.mark.asyncio +async def test_sub_graph_task_arun_with_async_object_store_and_sync_persister(): + """Regression test: SubGraphTask.arun() must succeed when the cascaded object_store is + async even though the state_persister/state_initializer are sync (a supported, documented + combination). Previously this crashed because the sync persister forced `.build()` to be + used internally, which unconditionally rejects an async object_store. + """ + object_store = InMemoryAsyncArtifactStore() + persister = DummyPersister() + + task = SubGraphTask( + graph=RunnableGraph( + graph=GraphBuilder() + .with_actions( + subgraph=simple_single_fn_subgraph_async_object_store, + terminal=Result("output_number"), + ) + .with_transitions(("subgraph", "terminal")) + .build(), + entrypoint="subgraph", + halt_after=["terminal"], + ), + inputs={}, + state=State({"input_number": 1, "number_to_add": 2}), + application_id="sub-app-object-store-test", + state_persister=persister, + state_initializer=None, + object_store=object_store, + ) + + parent_context = ApplicationContext( + app_id="parent", + partition_key="partition_key", + sequence_id=0, + tracker=None, + state_persister=None, + state_initializer=None, + object_store=object_store, + parallel_executor_factory=lambda: concurrent.futures.ThreadPoolExecutor(), + action_name="parent_action", + ) + + state = await task.arun(parent_context) + assert state["output_number"] == 1 + 2 + 1 # input + number_to_add + additional_number default diff --git a/tests/integrations/artifacts/__init__.py b/tests/integrations/artifacts/__init__.py new file mode 100644 index 000000000..13a83393a --- /dev/null +++ b/tests/integrations/artifacts/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/tests/integrations/artifacts/test_s3.py b/tests/integrations/artifacts/test_s3.py new file mode 100644 index 000000000..1fd63a64d --- /dev/null +++ b/tests/integrations/artifacts/test_s3.py @@ -0,0 +1,386 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Tests for S3ArtifactStore and AsyncS3ArtifactStore. + +S3ArtifactStore's tests use botocore's built-in ``Stubber`` (a dependency of boto3, already +required by this module) to mock S3 responses. AsyncS3ArtifactStore's tests mock an aiobotocore +client directly with ``AsyncMock`` instead -- the same approach already used elsewhere in this +repo for aiobotocore-backed code (see tests/tracking/test_bip0042_s3_buffering.py). Neither needs +a live AWS account or network access. +""" + +import hashlib +import io +from unittest.mock import AsyncMock, MagicMock + +import boto3 +import botocore.exceptions +import pytest +from botocore.response import StreamingBody +from botocore.stub import Stubber + +from burr.core.artifacts import ArtifactRef +from burr.integrations.artifacts.s3 import AsyncS3ArtifactStore, S3ArtifactStore + + +def _streaming_body(data: bytes) -> StreamingBody: + return StreamingBody(io.BytesIO(data), len(data)) + + +@pytest.fixture +def s3_client(): + # dummy credentials/region -- Stubber intercepts calls before any network access happens. + client = boto3.client( + "s3", + region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + with Stubber(client) as stubber: + yield client, stubber + + +@pytest.fixture +def store(s3_client): + client, _stubber = s3_client + return S3ArtifactStore(bucket="my-bucket", prefix="artifacts", client=client) + + +def test_put_uploads_when_key_does_not_exist(store, s3_client): + _client, stubber = s3_client + data = b"hello from s3" + + stubber.add_client_error("head_object", service_error_code="404") + stubber.add_response( + "put_object", + {}, + expected_params={"Bucket": "my-bucket", "Key": "artifacts/my-key", "Body": data}, + ) + + store.put(data, key="my-key") + + stubber.assert_no_pending_responses() + + +def test_put_skips_upload_when_key_already_exists(store, s3_client): + _client, stubber = s3_client + data = b"hello from s3" + + stubber.add_response( + "head_object", + {}, + expected_params={"Bucket": "my-bucket", "Key": "artifacts/my-key"}, + ) + # no put_object stubbed -- if the store tried to upload, Stubber would raise. + + store.put(data, key="my-key") + + stubber.assert_no_pending_responses() + + +def test_get_returns_object_bytes(store, s3_client): + _client, stubber = s3_client + data = b"the stored artifact bytes" + + stubber.add_response( + "get_object", + {"Body": _streaming_body(data)}, + expected_params={"Bucket": "my-bucket", "Key": "artifacts/my-key"}, + ) + + assert store.get("my-key") == data + stubber.assert_no_pending_responses() + + +def test_get_missing_key_raises_file_not_found(store, s3_client): + _client, stubber = s3_client + stubber.add_client_error("get_object", service_error_code="NoSuchKey", http_status_code=404) + + with pytest.raises(FileNotFoundError): + store.get("missing-key") + + +def test_exists_true(store, s3_client): + _client, stubber = s3_client + stubber.add_response( + "head_object", + {}, + expected_params={"Bucket": "my-bucket", "Key": "artifacts/my-key"}, + ) + + assert store.exists("my-key") + + +def test_exists_false_for_missing_key(store, s3_client): + _client, stubber = s3_client + stubber.add_client_error("head_object", service_error_code="404") + + assert not store.exists("missing-key") + + +def test_no_prefix_uses_bare_key(s3_client): + client, stubber = s3_client + store = S3ArtifactStore(bucket="my-bucket", client=client) + + stubber.add_client_error("head_object", service_error_code="404") + stubber.add_response( + "put_object", + {}, + expected_params={"Bucket": "my-bucket", "Key": "my-key", "Body": b"data"}, + ) + + store.put(b"data", key="my-key") + stubber.assert_no_pending_responses() + + +def test_put_artifact_and_get_artifact_end_to_end(store, s3_client): + """Exercises the shared ArtifactStore.put_artifact/get_artifact helpers against S3.""" + _client, stubber = s3_client + data = b"end to end artifact content" + digest = hashlib.sha256(data).hexdigest() + + # put_artifact() -> content-addressed key == digest + stubber.add_client_error("head_object", service_error_code="404") + stubber.add_response( + "put_object", + {}, + expected_params={"Bucket": "my-bucket", "Key": f"artifacts/{digest}", "Body": data}, + ) + + ref = store.put_artifact(data, media_type="text/plain") + assert isinstance(ref, ArtifactRef) + assert ref.key == digest + assert ref.digest == digest + assert ref.size_bytes == len(data) + + # get_artifact() -> verifies digest matches after fetch + stubber.add_response( + "get_object", + {"Body": _streaming_body(data)}, + expected_params={"Bucket": "my-bucket", "Key": f"artifacts/{digest}"}, + ) + assert store.get_artifact(ref) == data + stubber.assert_no_pending_responses() + + +def test_put_artifact_explicit_key_identical_content_is_noop(store, s3_client): + """Same collision-safety guarantee as the local backend: re-`put_artifact`-ing identical + content under an explicit key that already exists must not re-upload.""" + _client, stubber = s3_client + data = b"same content" + + stubber.add_response( + "head_object", + {}, + expected_params={"Bucket": "my-bucket", "Key": "artifacts/stable-key"}, + ) + stubber.add_response( + "get_object", + {"Body": _streaming_body(data)}, + expected_params={"Bucket": "my-bucket", "Key": "artifacts/stable-key"}, + ) + # no put_object stubbed -- if the store tried to re-upload, Stubber would raise. + + ref = store.put_artifact(data, key="stable-key") + + assert ref.key == "stable-key" + stubber.assert_no_pending_responses() + + +def test_put_artifact_explicit_key_conflicting_content_raises(store, s3_client): + """An explicit key that already holds *different* content must not be silently overwritten + or skipped -- put_artifact must surface a clear conflict instead of returning a ref that + doesn't describe what's actually stored.""" + _client, stubber = s3_client + existing_data = b"first" + new_data = b"second" + + stubber.add_response( + "head_object", + {}, + expected_params={"Bucket": "my-bucket", "Key": "artifacts/mutable"}, + ) + stubber.add_response( + "get_object", + {"Body": _streaming_body(existing_data)}, + expected_params={"Bucket": "my-bucket", "Key": "artifacts/mutable"}, + ) + # no put_object stubbed -- the conflicting write must be rejected before any upload attempt. + + with pytest.raises(ValueError, match="already exists with different content"): + store.put_artifact(new_data, key="mutable") + + stubber.assert_no_pending_responses() + + +# --- AsyncS3ArtifactStore tests below --- + + +def _client_error(code: str, operation: str = "HeadObject") -> botocore.exceptions.ClientError: + return botocore.exceptions.ClientError({"Error": {"Code": code}}, operation) + + +def _async_streaming_body(data: bytes) -> MagicMock: + stream = AsyncMock() + stream.read = AsyncMock(return_value=data) + body_cm = MagicMock() + body_cm.__aenter__ = AsyncMock(return_value=stream) + body_cm.__aexit__ = AsyncMock(return_value=None) + return body_cm + + +@pytest.fixture +def async_client(): + return AsyncMock() + + +@pytest.fixture +def async_store(async_client): + return AsyncS3ArtifactStore(bucket="my-bucket", prefix="artifacts", client=async_client) + + +def test_async_store_requires_already_open_client(): + with pytest.raises(ValueError, match="already-open aiobotocore client"): + AsyncS3ArtifactStore(bucket="my-bucket") + + +async def test_async_put_uploads_when_key_does_not_exist(async_store, async_client): + data = b"hello from async s3" + async_client.head_object = AsyncMock(side_effect=_client_error("404")) + async_client.put_object = AsyncMock(return_value={}) + + await async_store.put(data, key="my-key") + + async_client.put_object.assert_awaited_once_with( + Bucket="my-bucket", Key="artifacts/my-key", Body=data + ) + + +async def test_async_put_skips_upload_when_key_already_exists(async_store, async_client): + async_client.head_object = AsyncMock(return_value={}) + async_client.put_object = AsyncMock() + + await async_store.put(b"hello from async s3", key="my-key") + + async_client.put_object.assert_not_awaited() + + +async def test_async_get_returns_object_bytes(async_store, async_client): + data = b"the stored artifact bytes" + async_client.get_object = AsyncMock(return_value={"Body": _async_streaming_body(data)}) + + assert await async_store.get("my-key") == data + async_client.get_object.assert_awaited_once_with(Bucket="my-bucket", Key="artifacts/my-key") + + +async def test_async_get_missing_key_raises_file_not_found(async_store, async_client): + async_client.get_object = AsyncMock( + side_effect=_client_error("NoSuchKey", operation="GetObject") + ) + + with pytest.raises(FileNotFoundError): + await async_store.get("missing-key") + + +async def test_async_exists_true(async_store, async_client): + async_client.head_object = AsyncMock(return_value={}) + + assert await async_store.exists("my-key") + + +async def test_async_exists_false_for_missing_key(async_store, async_client): + async_client.head_object = AsyncMock(side_effect=_client_error("404")) + + assert not await async_store.exists("missing-key") + + +async def test_async_put_artifact_and_get_artifact_end_to_end(async_store, async_client): + """Exercises the shared AsyncArtifactStore.put_artifact/get_artifact helpers against S3.""" + data = b"end to end async artifact content" + digest = hashlib.sha256(data).hexdigest() + + async_client.head_object = AsyncMock(side_effect=_client_error("404")) + async_client.put_object = AsyncMock(return_value={}) + + ref = await async_store.put_artifact(data, media_type="text/plain") + assert isinstance(ref, ArtifactRef) + assert ref.key == digest + assert ref.digest == digest + assert ref.size_bytes == len(data) + async_client.put_object.assert_awaited_once_with( + Bucket="my-bucket", Key=f"artifacts/{digest}", Body=data + ) + + async_client.get_object = AsyncMock(return_value={"Body": _async_streaming_body(data)}) + assert await async_store.get_artifact(ref) == data + + +async def test_async_put_artifact_explicit_key_identical_content_is_noop(async_store, async_client): + """Same collision-safety guarantee as the sync backend: re-`put_artifact`-ing identical + content under an explicit key that already exists must not re-upload.""" + data = b"same content" + async_client.head_object = AsyncMock(return_value={}) + async_client.get_object = AsyncMock(return_value={"Body": _async_streaming_body(data)}) + async_client.put_object = AsyncMock() + + ref = await async_store.put_artifact(data, key="stable-key") + + assert ref.key == "stable-key" + async_client.put_object.assert_not_awaited() + + +async def test_async_put_artifact_explicit_key_conflicting_content_raises( + async_store, async_client +): + """An explicit key that already holds *different* content must not be silently overwritten + or skipped -- put_artifact must surface a clear conflict instead of returning a ref that + doesn't describe what's actually stored.""" + existing_data = b"first" + new_data = b"second" + async_client.head_object = AsyncMock(return_value={}) + async_client.get_object = AsyncMock(return_value={"Body": _async_streaming_body(existing_data)}) + async_client.put_object = AsyncMock() + + with pytest.raises(ValueError, match="already exists with different content"): + await async_store.put_artifact(new_data, key="mutable") + + async_client.put_object.assert_not_awaited() + + +async def test_async_store_acreate_opens_client_via_session_and_aclose_closes_it(monkeypatch): + """`acreate` must open the aiobotocore client via the async-context-manager protocol (it + can't just call the constructor directly, unlike boto3), and `aclose` must release it.""" + import burr.integrations.artifacts.s3 as s3_module + + fake_client = AsyncMock() + client_cm = MagicMock() + client_cm.__aenter__ = AsyncMock(return_value=fake_client) + client_cm.__aexit__ = AsyncMock(return_value=None) + + fake_session = MagicMock() + fake_session.create_client = MagicMock(return_value=client_cm) + + monkeypatch.setattr(s3_module.aiobotocore.session, "get_session", lambda: fake_session) + + store = await AsyncS3ArtifactStore.acreate(bucket="my-bucket", prefix="artifacts") + + assert store.client is fake_client + fake_session.create_client.assert_called_once_with("s3") + + await store.aclose() + client_cm.__aexit__.assert_awaited_once() diff --git a/tests/tracking/test_bip0042_s3_buffering.py b/tests/tracking/test_bip0042_s3_buffering.py index 4cbd151e2..badbb1b7d 100644 --- a/tests/tracking/test_bip0042_s3_buffering.py +++ b/tests/tracking/test_bip0042_s3_buffering.py @@ -23,6 +23,9 @@ import pytest pytest.importorskip("aiobotocore") +# tortoise-orm is only pulled in by the `tracking-server-s3` extra, not the (now +# aiobotocore-containing) `s3`/`tests` extras, so it needs its own skip-guard here. +pytest.importorskip("tortoise") from burr.tracking.server.backend import EventDrivenBackendMixin from burr.tracking.server.s3.backend import (