Async persistance #488

jernejfrank · 2024-12-27T16:41:09Z

Addressing #484 to have an async persistence interface.

Changes

Adds base classes for saving/loading/running persistor in application via async
Adds async initializer/persister to ApplicationBuilder and lets us build async Application via .abuild()
Adds async validation to Application to warn about async hook being ignored in sync runs and raises error if you try to do sync run with an async Application.
Adds create_async_app in parallelism to create application via AsyncApplicationBuilder (in case of legacy use of sync initializer/persister reverts to old sync Application)
TBD: database specific implementations (only AsyncDevNullPersister for testing)

To think about: (1) fire-and-forget or (2) blocking/transactional options --
Async persistors are naturally (1) and sync persistors are (2). If we want to have both options for both cases, we need to:

effectively block the async save to turn (1)->(2) and
create an event loop / another tread and make sync save a coroutine executing there to go from (2) -> (1).
My initial idea was to have that option as an attribute of the persistor class and handle it within PersisterHook/PersisterHookAsync, but still checking if there is a more natural place for this.

How I tested this

Unit and E2E

Notes

This part is very much for discussion if we want to have another AsyncApplicationBuilder or put this into the existing ApplicationBuilder. My arguments for having two:

It is cleaner to separate async and sync (and maybe less error-prone to use the wrong one?).
There are 3 methods we need to define async: with_state_persister, _load_from_persister, and build.
Related to the above, if with_state_persister is async it gets a bit hairy how to do method chaining -- what I did is to overwrite the original method to just store the state_persister (similar to what initialize_from does) and then pushed all the logic into an async helper function __with_async_state_persister that gets awaited in build to manually chain coroutines.
Having another builder class made it clearer also in the Application class to raise an error when run is used instead of arun.

Having said all of that, I also can push those methods into the existing builder with an abuild() to follow the pattern in the Application class where both sync and async are side-by-side.

Checklist

PR has an informative and human-readable title (this will be pulled into the release notes)
Changes are limited to a single goal (no scope creep)
Code passed the pre-commit check & code is left cleaner/nicer than when first encountered.
Any change in functionality is tested
[] New functions are documented (with a description, list of inputs, and expected output)
Placeholder code is flagged / future TODOs are captured in comments
Project documentation has been updated if adding/changing functionality.

Important

Add asynchronous persistence support with AsyncApplicationBuilder and related async classes and tests.

Behavior:
- Introduces AsyncApplicationBuilder to support asynchronous state persistence.
- Adds create_async_app in parallelism.py for async application creation.
- Implements AsyncDevNullPersister for testing async persistence.
Persistence:
- Adds AsyncBaseStateLoader and AsyncBaseStateSaver in persistence.py for async state operations.
- Introduces PersisterHookAsync for async lifecycle hooks.
Testing:
- Adds tests for async persistence in test_application.py, including test_async_save_and_load_from_persister_end_to_end.
Misc:
- Updates imports in __init__.py and application.py to include async classes.

^{This description was created by}^{for 9e233c8. It will automatically update as commits are pushed.}

ellipsis-dev

👍 Looks good to me! Reviewed everything up to 9e233c8 in 1 minute and 20 seconds

More details

Looked at 803 lines of code in 5 files
Skipped 0 files when reviewing.
Skipped posting 4 drafted comments based on config settings.

1. burr/core/persistence.py:55

Draft comment:
Ensure that the is_async method is overridden in subclasses of BaseStateLoader and BaseStateSaver to accurately reflect their async capabilities.
Reason this comment was not posted:
Comment did not seem useful.

2. burr/core/persistence.py:129

Draft comment:
Ensure that the is_async method is overridden in subclasses of BaseStateLoader and BaseStateSaver to accurately reflect their async capabilities.
Reason this comment was not posted:
Marked as duplicate.

3. burr/core/persistence.py:174

Draft comment:
Ensure that the is_async method is overridden in subclasses of BaseStateLoader and BaseStateSaver to accurately reflect their async capabilities.
Reason this comment was not posted:
Marked as duplicate.

4. burr/core/persistence.py:86

Draft comment:
Ensure that the is_async method is overridden in subclasses of BaseStateLoader and BaseStateSaver to accurately reflect their async capabilities.
Reason this comment was not posted:
Marked as duplicate.

Workflow ID: wflow_QbxcMoFBccy3mX8X

You can customize Ellipsis with 👍 / 👎 feedback, review rules, user-specific overrides, quiet mode, and more.

elijahbenizzy

Lookking good -- some broad comments. I'm falling on the side of making it a single builder to force sharing of code , IMO it makes it a bit simpler + duplicates less logic. That said, i see both ways of doing it.

About this:

To think about: (1) fire-and-forget or (2) blocking/transactional options --
Async persistors are naturally (1) and sync persistors are (2). If we want to have both options for both cases, we need to:

effectively block the async save to turn (1)->(2) and
create an event loop / another tread and make sync save a coroutine executing there to go from (2) -> (1).
My initial idea was to have that option as an attribute of the persistor class and handle it within PersisterHook/PersisterHookAsync, but still checking if there is a more natural place for this.

To clarify -- async persisters are not naturally fire-and-forget -- they'll await completion of the task, which blocks the event loop. asyncio.ensure_future (I think that's the right one) is a way to make it fire/forget. You're right about the thread -- we'll need to tie into some sort of execution. This is also why we can probably push this up in the stack -- E.G. we don't use the result of these for the next execution, so we can put it in a queue.

We would have to worry about how to ensure order, especially in persisters (there's probably an async queue pattern that works -- should be pretty easy, but definitely more complicated). This also gets to some interesting db/distributed systems problems -- I'm not convinced there aren't other ugly pitfalls here.

So I think we should keep it transactional as it is now...

elijahbenizzy · 2024-12-27T19:32:05Z

burr/core/application.py

@@ -1174,6 +1184,22 @@ def iterate(
        :return: Each iteration returns the result of running `step`. This generator also returns a tuple of
            [action, result, current state]
        """
+        # This is a gentle warning for existing users
+        if self._adapter_set.async_hooks:


This check is good to raise -- only in .iterate though (unless I'm missing something?). Maybe take it and refactor out to a specific checking function, then put in the non-iterate functions as well? stream_result, step, run? There's some nuance about things that get called in both, so worth looking out for that.

elijahbenizzy · 2024-12-27T19:35:32Z

burr/core/application.py

+            )
+
+        # Seems fair to raise this if everything is async but the app execution
+        if self._adapter_set.async_hooks and isinstance(self._builder, AsyncApplicationBuilder):


Oh I see -- this makes sense. E.G. why would you ever call ApplicationBuilder for async stuff.

I'm thinking that we want to centralize validation:

If any action is non-async -- we should maybe error out

The other async stuff should call this validation

IMO we should error out when the application is built with .abuild(), has async hooks, and we run it with sync methods, e.g. .run() since this is very easy to miss that the async hooks just don't get called.

Vice versa, if we have an async application with .arun() and sync hooks, I wouldn't error it out since we cannot guarantee that we every adapter has async support / user might be ok for some things to be blocking. Maybe a warning here is more appropriate?

Yes, I agree -- this is backwards compatible. The second case is odd, actually, as you've pointed out. To make it more concrete, not every synchronous hook is incompatible in an async mode. Take, for example, the hook that prints "doing X step" before every step -- this is just going to be a sync hook, not an async hook, but it should work in an async setting. No need to build an async version of something like that, when it never has an await.

I think we shold maintain the design where .arun() allows sync + async hooks, .run() only allows sync hooks (and should break otherwise, althugh that's an edge case...). Also, think this simplifies?

elijahbenizzy · 2024-12-27T19:36:23Z

burr/core/application.py

+        self.state_persister = persister  # track for later
+        return self
+
+    async def __with_async_state_persister(self):


I wonder if we just want to push this into build for both sync/async?

elijahbenizzy · 2024-12-27T19:39:07Z

burr/core/application.py

-                    results=result, final_state=state
+                out = (
+                    action,
+                    StreamingResultContainer.pass_through(results=result, final_state=state),


Formatting -- is there a reason it changed?

No, might have been my local settings..

elijahbenizzy · 2024-12-27T19:40:08Z

burr/core/application.py

@@ -2249,7 +2279,7 @@ def with_tracker(

    def initialize_from(


So I like with_state_persister logic moving to build -- it mirrors this. Might be worth doing for sync side as well. Can always leave as a TODO. But it should simplify logic?

elijahbenizzy · 2024-12-27T19:44:28Z

burr/core/application.py

+                pass
+
+    @telemetry.capture_function_usage
+    async def build(self) -> Application[StateType]:


So this has duplicated logic from the synchronous one. I wonder if this is a reason to combine.

_with_state_persister, _load_from_persister are helper functions for each

_build_common() or something has shared logic. Everything after 2654

build() has synchronous logic -- calls _build_common()

abuild() has async logic -- also calls _build_common

elijahbenizzy · 2024-12-27T19:47:00Z

burr/core/application.py

+        self.state_persister = persister  # track for later
+        return self
+
+    async def __with_async_state_persister(self):


nit -- naming conventions:

no underscore -- public method

one underscore -- private method

two underscores -- name mangling, private for subclasses

two underscores before/after -- system-level concerns

This should have a single underscore

burr/core/persistence.py

elijahbenizzy

Looks really good -- as noted in a DM I think we have to align on behavior. E.G. when using .build() versus .abuild(), sync versus async persister, amethod versus method -- which configurations break/log warning/work? Cartesian product/table is a clean way to think of it.

elijahbenizzy · 2024-12-29T06:14:15Z

burr/core/application.py

+            self.state = State()
+
+        if self.state_persister:
+            await self._set_async_state_persister()  # this is used to save the state during application run


I'm thinking that .abuild() should allow a sync persister but log a warning. But maybe not? Seems fair to force it here...

Scratch that, yes, that case should not work

jernejfrank · 2024-12-29T09:56:06Z

Looks really good -- as noted in a DM I think we have to align on behavior. E.G. when using .build() versus .abuild(), sync versus async persister, amethod versus method -- which configurations break/log warning/work? Cartesian product/table is a clean way to think of it.

This is a current working proposition of which cases should allow what. I think we can replace persister with hooks and put it in the builder docs as well to have clarity on how to use .abuild().

elijahbenizzy · 2024-12-30T03:02:30Z

Looks really good -- as noted in a DM I think we have to align on behavior. E.G. when using .build() versus .abuild(), sync versus async persister, amethod versus method -- which configurations break/log warning/work? Cartesian product/table is a clean way to think of it.

This is a current working proposition of which cases should allow what. I think we can replace persister with hooks and put it in the builder docs as well to have clarity on how to use .abuild().

OK, yeah, I think this is the right way to do it! Nice work enumerating it, you have me convinced. The only case (and I'm not sure how this works in the one above) is when you have a synchronous hook in an async .abuild() method. Not a persister (so not fully covered by the above), but I think it should work. That said, we should have appropriate documentation about it.

elijahbenizzy

Looking good! Some small points:

I think we should document this clearly. Ideally is a sync versus async page in concepts, but this could just be part of the docstring
A bit of a strange case in parallelism, where we call the new function by default. Might be worth adding in a parameter to bypass validation. Or just breaking -- it's an edge case.

Otherwise, let's ship soon! Nice work!

elijahbenizzy · 2024-12-30T03:03:28Z

burr/core/application.py

+            )
+
+        # Seems fair to raise this if everything is async but the app execution is sync
+        if self._adapter_set.async_hooks and self._builder.is_async:


Confused -- this is erroring out if the builder is async and the adapter has async hooks? .is_async?

Yes! So I am calling this in the sync methods and it covers cases 3 and 7 in the above table. Thinking back, maybe it's better to be strict from the get go and only check _builder.is_async and error out.

Previously, if there were no async hooks present (i.e. no async persisters, etc.) it doesn't really matter how the app was built or is run -- since only sync hooks in the app. But might as well just say if you build it async, use it async.

elijahbenizzy · 2024-12-30T03:04:00Z

burr/core/application.py

@@ -2484,3 +2589,51 @@ def build(self) -> Application[StateType]:
            state_persister=self.state_persister,
            state_initializer=self.state_initializer,
        )
+
+    @telemetry.capture_function_usage


Nice this function is extremely clean

elijahbenizzy · 2024-12-30T03:04:15Z

burr/core/application.py

+            self.state = State()
+
+        if self.state_persister:
+            await self._set_async_state_persister()  # this is used to save the state during application run


Scratch that, yes, that case should not work

elijahbenizzy · 2024-12-30T03:08:27Z

burr/core/application.py

+    async def abuild(self) -> Application[StateType]:
+        """Builds the application.
+
+        This function is a bit messy as we iron out the exact logic and rigor we want around things.


No longer messy, fix this. And make it + build() clear about the parameters (good place to have this table?).

elijahbenizzy · 2024-12-30T03:14:14Z

burr/core/persistence.py

+        status: Literal["completed", "failed"],
+        **kwargs,
+    ):
+        # print("I saved something.")


elijahbenizzy · 2024-12-30T03:15:50Z

burr/core/parallelism.py

+        else:
+            builder = builder.with_entrypoint(self.graph.entrypoint).with_state(self.state)
+
+        return await builder.abuild()


OK, complexity, this has the off-chance of making something backwards incompatible. That said, I think that's OK TBH. In this case it's an odd one:

Someone is using the async parallelism piece

Someone is not using the async persister (which they won't be, it isn't there)

It breaks on reloading

I think very few people are doing this, but it might be worth having a parameter in .abuild() to bypass validation, we could do that here...

We can add the parameter, but not sure if it is that helpful. If they are using sync persisters they will have to go through the sync .build() method anyway since .abuild() expects async persisters.

This is what I capture now parallelism .arun(): we check if we have sync persisters and then use .build() to have that backward compatibility, otherwise everything async and ok.

jernejfrank · 2024-12-30T07:49:59Z

OK, yeah, I think this is the right way to do it! Nice work enumerating it, you have me convinced. The only case (and I'm not sure how this works in the one above) is when you have a synchronous hook in an async .abuild() method. Not a persister (so not fully covered by the above), but I think it should work. That said, we should have appropriate documentation about it.

This works as before. You run both sync and async hooks --> sync ones are blocking. Also, if you use sync persister with .build() it will still let you run async (this is backwards compatibility). But if you use the new .abuild() method, it complains in there about using sync persisters and later when you have the application it complains that you are trying to run it sync.

Will document!

elijahbenizzy

Nice work! Really big improvements. Left some small comments but it's there.

elijahbenizzy · 2024-12-31T05:54:20Z

burr/core/application.py

+        # this application is meant to be run in async mode.
+        if self._builder and self._builder.is_async:
+            raise ValueError(
+                "The application was build with async hooks "


Make this say "it was built with abuild" -- think that's clearer/more actionable

elijahbenizzy · 2024-12-31T05:57:25Z

docs/concepts/sync-vs-async.rst

+Sync vs Async Applications
+===========================
+
+TL;DR


Link to the arun, aiterate, etc... methods. Good to tie it back into the actions.

Adding async support for persistance and refactoring builder: - classes for building async persistance adapters / hooks - builder extended to include async initializer/persister, async build - builder refactored - application added async validation, warning/error when async hooks not invoked - automatic built for app in parallelism made backward compatible

Prototyping and testing async persisters: - Add AsyncDevNull and AsyncInMemory persisters for tests - Added support for async sqlite persister - Test Async persister interface, async builder, async application

Docs explaining allowed cases and deprecation warnings.

ellipsis-dev bot reviewed Dec 27, 2024

View reviewed changes

elijahbenizzy reviewed Dec 27, 2024

View reviewed changes

elijahbenizzy reviewed Dec 29, 2024

View reviewed changes

elijahbenizzy reviewed Dec 30, 2024

View reviewed changes

jernejfrank force-pushed the async_persistance branch 2 times, most recently from 5008a12 to b7ba29f Compare December 30, 2024 13:13

elijahbenizzy approved these changes Dec 31, 2024

View reviewed changes

jernejfrank added 3 commits December 31, 2024 19:03

Add local async persisters and tests

5e3e18f

Prototyping and testing async persisters: - Add AsyncDevNull and AsyncInMemory persisters for tests - Added support for async sqlite persister - Test Async persister interface, async builder, async application

Add docs for async persister interface

ea23c98

Docs explaining allowed cases and deprecation warnings.

jernejfrank force-pushed the async_persistance branch from b7ba29f to ea23c98 Compare December 31, 2024 11:04

elijahbenizzy merged commit c36877d into DAGWorks-Inc:main Jan 4, 2025
10 of 11 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Async persistance #488

Async persistance #488

jernejfrank commented Dec 27, 2024 •

edited

Loading

ellipsis-dev bot left a comment

elijahbenizzy left a comment

elijahbenizzy Dec 27, 2024

elijahbenizzy Dec 27, 2024

jernejfrank Dec 28, 2024

elijahbenizzy Dec 29, 2024

elijahbenizzy Dec 27, 2024

elijahbenizzy Dec 27, 2024

jernejfrank Dec 28, 2024

elijahbenizzy Dec 27, 2024

elijahbenizzy Dec 27, 2024

elijahbenizzy Dec 27, 2024

elijahbenizzy left a comment

elijahbenizzy Dec 29, 2024

elijahbenizzy Dec 30, 2024

jernejfrank commented Dec 29, 2024

elijahbenizzy commented Dec 30, 2024

elijahbenizzy left a comment

elijahbenizzy Dec 30, 2024

jernejfrank Dec 30, 2024

elijahbenizzy Dec 30, 2024

elijahbenizzy Dec 30, 2024

elijahbenizzy Dec 30, 2024

elijahbenizzy Dec 30, 2024

elijahbenizzy Dec 30, 2024

jernejfrank Dec 30, 2024

jernejfrank commented Dec 30, 2024 •

edited

Loading

elijahbenizzy left a comment

elijahbenizzy Dec 31, 2024

elijahbenizzy Dec 31, 2024

Async persistance #488

Async persistance #488

Conversation

jernejfrank commented Dec 27, 2024 • edited Loading

Changes

How I tested this

Notes

Checklist

ellipsis-dev bot left a comment

Choose a reason for hiding this comment

elijahbenizzy left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

elijahbenizzy left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jernejfrank commented Dec 29, 2024

elijahbenizzy commented Dec 30, 2024

elijahbenizzy left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jernejfrank commented Dec 30, 2024 • edited Loading

elijahbenizzy left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jernejfrank commented Dec 27, 2024 •

edited

Loading

jernejfrank commented Dec 30, 2024 •

edited

Loading