feat(core): add artifact/blob store for large state values

burr/core/application.py

@@ -60,6 +60,7 @@
     StreamingAction,
     StreamingResultContainer,
 )
+from burr.core.artifacts import ArtifactStore
 from burr.core.graph import Graph, GraphBuilder
 from burr.core.persistence import (
     AsyncBaseStateLoader,
@@ -642,6 +643,7 @@ 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
+    object_store: Optional[ArtifactStore] = None
 
     @staticmethod
     def get() -> Optional["ApplicationContext"]:
@@ -849,6 +851,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[ArtifactStore] = None,
     ):
         """Instantiates an Application. This is an internal API -- use the builder!
 
@@ -900,6 +903,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 +947,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 +2229,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 +2550,39 @@ def with_state_persister(
         self.state_persister = persister  # tracks for later; validates in build / abuild
         return self
 
+    def with_object_store(self, object_store: ArtifactStore) -> "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()
+            )
+
+        :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 +2824,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]:

burr/core/artifacts.py

@@ -0,0 +1,242 @@
+# 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)
+
+
+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 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 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"),
+    )

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
 from burr.core.graph import GraphBuilder
 from burr.core.persistence import BaseStateLoader, BaseStateSaver
 from burr.lifecycle import LifecycleAdapter
@@ -102,6 +103,7 @@ class SubGraphTask:
     tracker: Optional[TrackingClient] = None
     state_persister: Optional[BaseStateSaver] = None
     state_initializer: Optional[BaseStateLoader] = None
+    object_store: Optional[ArtifactStore] = None
 
     def _create_app_builder(self, parent_context: ApplicationIdentifiers) -> ApplicationBuilder:
         builder = (
@@ -121,6 +123,9 @@ 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
 
+        if 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)
@@ -513,6 +518,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]:

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.