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.
|
|
||
| class ArtifactStore(abc.ABC): | ||
| """Base class for blob/object storage backends. | ||
|
|
There was a problem hiding this comment.
Blocking: we need to handle the async Python case. Async Python is a first-class execution model in Burr, but this API is synchronous and the S3 implementation performs blocking network I/O through boto3. Using this store from an async action would therefore block the event loop for the duration of an artifact upload or download. Please add an async artifact-store contract and a genuinely async S3 implementation, and define how that store is configured and exposed to async applications. The alternative would be to restrict artifact storage to synchronous Burr explicitly, which would leave a significant part of Burr unsupported.
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