Re: [PR] `feature/issue-67200`: Adding `AssetState` Task SDK mechanism [airflow]
jroachgolf84 commented on PR #67248:
URL: https://github.com/apache/airflow/pull/67248#issuecomment-4547765634
@amoghrajesh - got it, that makes sense to me. If that's the case, then we
shouldn't need this PR, correct? We could use the `from
airflow.sdk.execution_time.context import AssetStateAccessor` as the keys for
the return value of `self.asset_states`? i.e.
```python
self.asset_states = {
"asset_a": AssetStateAccessor(...),
"asset_b": AssetStateAccessor(...)
}
```
Thoughts?
cc: @vincbeck
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] `feature/issue-67200`: Adding `AssetState` Task SDK mechanism [airflow]
vincbeck commented on PR #67248:
URL: https://github.com/apache/airflow/pull/67248#issuecomment-4544853225
> Hey @jroachgolf84, I think the plumbing fix here will be correct.
>
> `AssetState(name=self.asset_name)` works but it is inconsistent with how
tasks access asset state as Vincent also mentioned, and nothing prevents a
trigger from doing `AssetState(name="some_other_asset") `even if it's not
associated with that asset.
>
> Looking at `BaseTrigger`, the triggerer already injects
`self.task_instance` before calling `run()`. We could use the same pattern here
and the triggerer could populate `self.asset_states` on `BaseEventTrigger`
before `run()`, keyed by asset name, using the assets associated with the TI.
The trigger author can then do:
>
> ```python
> async def run(self):
> watermark = self.asset_states["orders"].get("watermark")
> self.asset_states["orders"].set("watermark", new_watermark)
> yield TriggerEvent(...)
> ```
>
> This will also provide some benefits like:
>
> * Make the triggers scoped automatically (only the trigger's associated
assets are present)
> * Should also work cleanly for multi-asset triggers(just use the asset
name while accessing)
> * The framework handles scoping, not the author
>
> The triggerer already knows which assets are associated with the TI at the
point it sets `trigger.task_instance = ti`, so populating `asset_states` there
is probably straightforward.
I really like that!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] `feature/issue-67200`: Adding `AssetState` Task SDK mechanism [airflow]
jroachgolf84 commented on PR #67248: URL: https://github.com/apache/airflow/pull/67248#issuecomment-4539117423 @cmarteepants, @vikramkoka -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] `feature/issue-67200`: Adding `AssetState` Task SDK mechanism [airflow]
jroachgolf84 commented on PR #67248:
URL: https://github.com/apache/airflow/pull/67248#issuecomment-4539111717
> I just looked at #67376 and I meant exactly this same mechanism but for
asset. Cannot we plumb through automatically the state of an asset to its
related triggers? That would avoid forcing Dag author to manually get the asset
state with `AssetState(name=self.asset_name)`
@vincbeck - I think that makes sense. The only caveat is that a
`BaseEventTrigger` can have multiple Assets "plumbed" through to it (as we
discussed here: https://github.com/apache/airflow/pull/66595). That would make
the pattern for access the state for a specific Asset a bit more difficult. I
guess the syntax would be something more like this:
```python
# Contains the asset states
self.asset_states = ...
asset_a_state = self.asset_states.get("asset_a")
asset_b_state = self.asset_states.get("asset_b")
```
Thoughts? I'm kinda caught in the middle on this one.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] `feature/issue-67200`: Adding `AssetState` Task SDK mechanism [airflow]
vincbeck commented on PR #67248: URL: https://github.com/apache/airflow/pull/67248#issuecomment-4535711253 I just looked at https://github.com/apache/airflow/pull/67376 and I meant exactly this same mechanism but for asset. Cannot we plumb through automatically the state of an asset to its related triggers? That would avoid forcing Dag author to manually get the asset state with `AssetState(name=self.asset_name)` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] `feature/issue-67200`: Adding `AssetState` Task SDK mechanism [airflow]
vincbeck commented on PR #67248: URL: https://github.com/apache/airflow/pull/67248#issuecomment-4535623264 > > From [AIP-103](https://cwiki.apache.org/confluence/display/AIRFLOW/AIP-103%3A+Task+State+Management) I understood this new task state management could be used across multiple different use cases such as [intra-task progress checkpointing ](https://cwiki.apache.org/confluence/pages/viewpage.action?pageId=406623137#AIP103:TaskStateManagement-Pattern3:Intra-TaskProgressCheckpointing). If that's the case, I am not sure about the name `AssetState`, it feels like something specific to asset although it should be agnostic. What do you think? > > AIP-103 addresses both Task and Asset state. Here are some of the PR's that have added Asset State. > > * [AIP-103: Add Core API endpoints for task state and asset state #67041](https://github.com/apache/airflow/pull/67041) > * [AIP-103: Add Execution API endpoints for task and asset states #66073](https://github.com/apache/airflow/pull/66073) > * [AIP-103: Refactor asset state routes to by-name/by-uri and add AssetUriRef support #66336](https://github.com/apache/airflow/pull/66336) > > cc: @amoghrajesh The APIs look good to me, I am only questioning the way to access the state. In your example you do `asset_state = AssetState(name="generic_asset")`, but in reality the task instance/asset state is scoped to its task instance/asset. It feels weird to me to be able to specify the asset name. If I do `asset_state = AssetState(name="another_asset")`, would I get the state from this other asset even though my code s not scoped to `another_asset`? Would it be protected by a JWT? In my mental model I was expecting a getter and setter like `set_state`, `get_state` that would use the runtime context (task instance ID/asset name) to retrieve the state -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] `feature/issue-67200`: Adding `AssetState` Task SDK mechanism [airflow]
jroachgolf84 commented on PR #67248: URL: https://github.com/apache/airflow/pull/67248#issuecomment-4527046132 @amoghrajesh - when you get a chance, can you look at this? I'm going to work on getting these checks green. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] `feature/issue-67200`: Adding `AssetState` Task SDK mechanism [airflow]
jroachgolf84 commented on PR #67248: URL: https://github.com/apache/airflow/pull/67248#issuecomment-4525365064 > From [AIP-103](https://cwiki.apache.org/confluence/display/AIRFLOW/AIP-103%3A+Task+State+Management) I understood this new task state management could be used across multiple different use cases such as [intra-task progress checkpointing ](https://cwiki.apache.org/confluence/pages/viewpage.action?pageId=406623137#AIP103:TaskStateManagement-Pattern3:Intra-TaskProgressCheckpointing). If that's the case, I am not sure about the name `AssetState`, it feels like something specific to asset although it should be agnostic. What do you think? AIP-103 addresses both Task and Asset state. Here are some of the PR's that have added Asset State. - https://github.com/apache/airflow/pull/67041 - https://github.com/apache/airflow/pull/66073 - https://github.com/apache/airflow/pull/66336 cc: @amoghrajesh -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] `feature/issue-67200`: Adding `AssetState` Task SDK mechanism [airflow]
vincbeck commented on PR #67248: URL: https://github.com/apache/airflow/pull/67248#issuecomment-4509281993 From [AIP-103](https://cwiki.apache.org/confluence/display/AIRFLOW/AIP-103%3A+Task+State+Management) I understood this new task state management could be used across multiple different use cases such as [intra-task progress checkpointing ](https://cwiki.apache.org/confluence/pages/viewpage.action?pageId=406623137#AIP103:TaskStateManagement-Pattern3:Intra-TaskProgressCheckpointing). If that's the case, I am not sure about the name `AssetState`, it feels like something specific to asset although it should be agnostic. What do you think? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
[PR] `feature/issue-67200`: Adding `AssetState` Task SDK mechanism [airflow]
jroachgolf84 opened a new pull request, #67248:
URL: https://github.com/apache/airflow/pull/67248
## Description
This PR adds the `AssetState` mechanism to use the foundations put in place
in AIP-103 within a Trigger or Task.
## Testing
These changes were unit-tested, as well as tested E2E. See below for more
information.
### Unit Tests
Existing unit tests were updated and new tests were added to validate the
changes that were made in this branch. *Note, there are additional unit-tests
to be added.*
```bash
# Updating existing unit testes
breeze testing core-tests airflow-core/tests/unit/jobs/test_triggerer_job.py
# New unit tests
breeze testing task-sdk-tests
task-sdk/tests/task_sdk/definitions/test_asset_state.py
```
### E2E Testing
This `GenericEventTrigger` was used for E2E testing. Note that `AssetState`
is used in the `run` method of the Trigger. This code properly stores and
retrieves the generated number, and logs the output accordingly.
```python
from airflow.sdk import AssetState
from airflow.triggers.base import BaseEventTrigger, TriggerEvent
from collections.abc import AsyncIterator
from typing import Any
import asyncio
import logging
import random
class GenericEventTrigger(BaseEventTrigger):
def __init__(
self,
random_number,
waiter_delay,
asset_name,
**kwargs
):
super().__init__(**kwargs)
self.random_number = random_number
self.waiter_delay = waiter_delay
self.asset_name = asset_name
def serialize(self) -> tuple[str, dict[str, Any]]:
"""Serialize the Trigger, including the func, params, and
waiter_delay."""
return (
self.__class__.__module__ + "." + self.__class__.__qualname__,
{
"random_number": self.random_number,
"waiter_delay": self.waiter_delay,
"asset_name": self.asset_name,
},
)
async def run(self) -> AsyncIterator[TriggerEvent]:
"""Logic that fires a TriggerEvent."""
# Here's where the AssetState is actually being used
asset_state = AssetState(name=self.asset_name)
logging.info(f"* asset_state: {asset_state}")
while True:
result = random.randint(0, 5)
logging.info(f"result: {result}")
asset_state.set("result", str(result))
_result = asset_state.get("result")
logging.info(f"_result: {_result}")
if result == self.random_number:
logging.info("yield'ing TriggerEvent")
yield TriggerEvent({"status": "success", "result": result})
break
logging.info(f"Sleeping for {self.waiter_delay} seconds")
await asyncio.sleep(self.waiter_delay)
```
## TODO
The following items still need to be completed:
- [ ] Validate that `AssetState` works when called within a Task.
- [ ] Ensure complete test coverage of the changes that were made.
* closes: #67200
# Was generative AI tooling used to co-author this PR?
No, generative AI was not used to generate this PR.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
