feat(core): add artifact/blob store for large state values - #931
vignesh-manel wants to merge 2 commits into
Conversation
skrawcz
left a comment
There was a problem hiding this comment.
Thanks for putting this together. The abstraction is promising, the documentation is thorough, the S3 tests are appropriately isolated, and propagation into parallel subgraphs looks consistent. I found several functional correctness issues that need to be addressed before merge:
1. ArtifactRef deserialization fails in a fresh process
The serde hooks are registered by module-import side effects in burr/core/artifacts.py:181-202. A process loading persisted state does not necessarily import burr.core.artifacts first: application.py imports ArtifactStore only under TYPE_CHECKING, and the normal State.deserialize() path does not eagerly import this module.
As a result, the main persistence/replay use case fails in a fresh process with:
ValueError: No deserializer registered for key: 'ArtifactRef'
The current test passes because tests/core/test_artifacts.py imports ArtifactRef at module scope before testing deserialization, which performs the registration and masks the real behavior. Please make the deserializer available on the normal state-deserialization path and add a fresh-process/subprocess test that loads a serialized ArtifactRef without first importing burr.core.artifacts. This is also required for the documentation's “out of the box, no extra setup” guarantee to be true.
2. Reusing an explicit key returns a reference that does not describe stored content
ArtifactStore.put_artifact() computes a digest for the new bytes and returns that digest in the ref (burr/core/artifacts.py:111-117), but both implementations skip the write whenever the explicit key already exists:
LocalFileSystemArtifactStore.put():burr/core/artifacts.py:161-168S3ArtifactStore.put():burr/integrations/artifacts/s3.py:75-79
For example:
first = store.put_artifact(b"first", key="mutable")
second = store.put_artifact(b"second", key="mutable")
store.get_artifact(second) # Digest mismatchThe second call returns a ref whose digest is for b"second", while the key still contains b"first". A successful write API must not return an immediately invalid reference. Please define and enforce collision semantics consistently: overwrite atomically, verify that existing content matches before treating the operation as idempotent, or reject a conflicting explicit key.
3. Local path containment can be bypassed through symlinks
LocalFileSystemArtifactStore._path_for_key() uses abspath and lexical commonpath checks (burr/core/artifacts.py:151-159). It does not resolve symlinks. If a path within root_dir is a symlink to a directory outside the root, a key beneath that symlink passes validation and writes outside the configured store.
I reproduced this with a root/link -> /outside symlink and store.put(..., "link/outside.txt"). Please enforce containment against resolved paths and account for creation of a final path that does not exist yet. Add tests covering symlinked parent directories for put, get, and exists.
4. Local writes are not atomic and can permanently poison a content-addressed key
LocalFileSystemArtifactStore.put() writes directly to the final destination (burr/core/artifacts.py:166-168). Concurrent readers can observe a partial file; a process crash can leave a truncated file. Future calls then see that the destination exists and skip writing it, so the key remains corrupt and every verified read fails indefinitely.
I confirmed that the final destination becomes visible at size zero while a large write is in progress. Please write to a temporary file in the destination directory, flush/close it, and publish it atomically with clearly defined collision behavior. This should have concurrency/interrupted-write coverage.
5. Adding object_store breaks direct construction of public ApplicationContext
burr/core/application.py:645 adds object_store as a required dataclass field. ApplicationContext is exported from burr.core and is public API, so existing code constructing it directly now fails with:
TypeError: ApplicationContext.__init__() missing 1 required positional argument: 'object_store'
Please preserve compatibility by making this a trailing field with a None default (and adjust trailing defaults/order as necessary), with a regression test using the previous constructor shape.
This is a functional correctness review only. A separate code-quality and architecture/API-design review will follow, so resolving these findings should not yet be treated as final approval of the overall design.
| } | ||
|
|
||
|
|
||
| @serde.deserializer.register(_ARTIFACT_REF_SERDE_KEY) |
There was a problem hiding this comment.
This deserializer is registered only as an import side effect of burr.core.artifacts. A fresh process loading persisted state through State.deserialize() does not necessarily import this module, so it fails with No deserializer registered for key: ArtifactRef. The current test masks this by importing ArtifactRef at module scope first. Please register this on the normal deserialization path and add a fresh-process/subprocess regression test.
| """ | ||
| digest = hashlib.sha256(data).hexdigest() | ||
| resolved_key = key if key is not None else digest | ||
| self.put(data, resolved_key) |
There was a problem hiding this comment.
With an explicit key, this returns a digest for the new data even when a backend skips the write because that key already exists. For example, storing b"first" and then b"second" under "mutable" returns a second ref that immediately fails get_artifact() verification because the store still contains the first bytes. Please define consistent collision semantics: atomically overwrite, verify identical existing content, or reject conflicting keys.
| os.makedirs(root_dir, exist_ok=True) | ||
|
|
||
| def _path_for_key(self, key: str) -> str: | ||
| root = os.path.abspath(self.root_dir) |
There was a problem hiding this comment.
abspath plus lexical commonpath does not enforce containment through symlinks. I reproduced an escape by creating root/link as a symlink to an outside directory and writing key link/outside.txt; this check accepted it and wrote outside root_dir. Please validate resolved parent paths safely, including when the destination does not yet exist, and test symlinked parents for put, get, and exists.
| if os.path.exists(path): | ||
| # content-addressed keys make writes idempotent -- skip re-writing existing data. | ||
| return | ||
| os.makedirs(os.path.dirname(path), exist_ok=True) |
There was a problem hiding this comment.
Writing directly to the final destination publishes a partial file to concurrent readers and can leave a truncated artifact after interruption. Because the earlier existence check then skips all later writes, that content-addressed key can remain permanently poisoned. I confirmed the final path becomes visible at size zero during a large write. Please write to a temporary file in the destination directory and publish atomically, with concurrency/interruption coverage.
| parallel_executor_factory: Callable[[], Executor] | ||
| state_initializer: Optional[BaseStateLoader] | ||
| state_persister: Optional[BaseStateSaver] | ||
| object_store: Optional["ArtifactStore"] |
There was a problem hiding this comment.
ApplicationContext is exported public API, but this new field has no default. Existing callers constructing the previous context shape now fail with TypeError: ... missing ... object_store. Please preserve backwards compatibility by making this a trailing optional field defaulting to None, with a regression test for the previous constructor shape.
Adds an optional object/blob storage abstraction for large values (files,
images, dataframes, etc...) that shouldn't be embedded directly in
State. Today, State is fully re-serialized to the persister/trackerafter every action, so putting a large binary blob in state means it
gets re-written on every subsequent step even when it never changes.
This implements the "push to the persistence layer" approach discussed
in #467 -- a store for the bytes, and a small content-addressed
reference kept in state instead.
Changes
burr.core.artifactswithArtifactStore(abstractput/get/exists, withput_artifact/get_artifactprovided for free tohandle digest computation/verification and content-addressed keys),
ArtifactRef(frozen, content-addressed reference), andLocalFileSystemArtifactStore(stdlib-only, local disk backend)burr.integrations.artifacts.s3.S3ArtifactStore, gated behinda new optional
s3extra (pip install "apache-burr[s3]")ArtifactRefwithburr.core.serdeso it serializes /deserializes automatically as part of normal
Statehandling --deserializing a ref never eagerly re-fetches the underlying bytes
ApplicationBuilder.with_object_store(...), threaded throughApplication/ApplicationContextso actions can access theconfigured store via
__context.object_storewithout constructingor importing their own
writes to the store on its own, actions call
put_artifact/get_artifactexplicitlydocs/concepts/artifact-storage.rst(conceptual overviewdocs/reference/artifacts.rst/docs/reference/integrations/s3.rst(API reference)How I tested this
python -m pytest tests/core/test_artifacts.py tests/integrations/artifacts/test_s3.py tests/core/test_application.py tests/core/test_parallelism.py -qbotocore.stub.Stubber-- no live AWS account or network access requiredpre-commit run --all-filesNotes
with_object_storeis purely a way to configure and share a store(e.g. local disk in dev, S3 in prod) across actions -- it's not a
replacement for
with_state_persister, which handles the statemachine's own bookkeeping.
put/get/exists;digest computation/verification is provided by the base class.
Checklist