flyte
flyte copied to clipboard
[Core][Flytekit] support async methods in python
Motivating Example 1
Motivation: Why do you think this is important?
It is possible that users may define their tasks as async
@task
async def foo():
pass
Since the invocation of this call is hijacked by flytekit (pyflyte), in the default call pattern - asyncio.run(foo)
is not invoked.
Goal: What should the final outcome look like, ideally?
pyflyte execute - eventually the execute
method should invoke the underlying task function using asyncio.run, if it is declared to be async
. This can be detected easily using
inspect.iscoroutinefunction(object)
Care should be taken to allow all task extensions to use this basic construct. Also, it is possible that users potentially do no wait for their own downstream tasks. This is not really a problem with flytekit, but Flytekit should do its best to help the users identify the problem. This can be achieved by listing all pending tasks
asyncio.Task.all_tasks().
Describe alternatives you've considered
It is possible for users to invoke their async methods today
async def async_foo():
pass
@task
def foo():
asyncio.run(async_foo)
But this is less desirable.
Motivating Example 2
I would love an async/await
API where there's just tasks which are async
, and they return a special Promise type which can be awaited to get the real value. As an example, I just wanted to write something like
@task
def ensemble(trained_models_and_weights: Dict[ModelID, float]) -> ModelID:
...
@workflow
def train_ensemble(models_and_weights: Dict[str, Tuple[float, ModelConfig]) -> ModelID:
model_ids = [response["model_id"] for response in train_model(cfg) for _, cfg in models_and_weights.values()]
return ensemble({model_id: weight for model_id, (weight, _) in zip(model_ids, models_and_weights.values()})
but because of the way promises are resolved I need to do something like
@task
def ensemble(trained_models_and_weights: Dict[ModelID, float]) -> ModelID:
...
@dynamic
def _resolve_model_ids_and_ensemble(model_ids: List[ModelID], weights: List[float]) -> ModelID:
return ensemble({model_id: weight for model_id, weight in zip(model_ids, weights})
@task
def model_id(response: TrainResponse) -> ModelID:
return response["model_id"]
@workflow
def train_ensemble(models_and_weights: Dict[str, Tuple[float, ModelConfig]) -> ModelID:
model_ids = [model_id(response) for response in train_model(cfg) for _, cfg in models_and_weights.values()]
return _resolve_model_ids_and_ensemble(model_ids, [weight for _, weight in models_and_weights.values()])
with async/await
I could instead write something like
@task
async def ensemble(trained_models_and_weights: Dict[ModelID, float]) -> ModelID:
...
@task
async def train_ensemble(models_and_weights: Dict[str, Tuple[float, ModelConfig]) -> ModelID:
responses = await gather([train_model(cfg) for _, cfg in models_and_weights.values()])
return await ensemble({response["model_id"]: weight for response, (weight, _) in zip(responses, models_and_weights.values()})
Misc
Are you sure this issue hasn't been raised already?
- [X] Yes
Have you read the Code of Conduct?
- [X] Yes
@bethebunny this was the original issue for async support. I added your example to the main description.
Also related is this ticket https://github.com/flyteorg/flyte/issues/2483
Hello 👋, This issue has been inactive for over 9 months. To help maintain a clean and focused backlog, we'll be marking this issue as stale and will close the issue if we detect no activity in the next 7 days. Thank you for your contribution and understanding! 🙏
Hello 👋, This issue has been inactive for over 9 months and hasn't received any updates since it was marked as stale. We'll be closing this issue for now, but if you believe this issue is still relevant, please feel free to reopen it. Thank you for your contribution and understanding! 🙏
@kumare3 Hi, I have a few questions on this issue.
- in the example below, is
train_ensemble
desired to be a workflow or a task? - I think maybe this example can be handled with eager workflows, are eager workflows less desired?
thanks
@task
async def ensemble(trained_models_and_weights: Dict[ModelID, float]) -> ModelID:
...
@task
async def train_ensemble(models_and_weights: Dict[str, Tuple[float, ModelConfig]) -> ModelID:
responses = await gather([train_model(cfg) for _, cfg in models_and_weights.values()])
return await ensemble({response["model_id"]: weight for response, (weight, _) in zip(responses, models_and_weights.values()})
I think this example is wrong and should be @eager - but cc @wild-endeavor as he added that section
@kumare3 Thanks, I ran the following with current flytekit and seems that the eager version works.
@task(enable_deck=True)
async def async_add_one(x: int) -> int:
await asyncio.sleep(7.5)
return x + 1
@eager(
remote=FlyteRemote(
config=Config.auto(),
default_project="flytesnacks",
default_domain="development",
)
)
async def simple_async_workflow(models: list[int] = [1,2,3]) -> list[int]:
coros = [async_add_one(x=x) for x in models]
res = await asyncio.gather(*coros)
return res
What @task can be async? I guess it is running in sync mode
Sorry, can you elaborate more on "running in sync mode"? Do you mean that coroutines are executed before awaited or the coroutines didn't execute in parallel?