[Core][Flytekit] support async methods in python

[kumare3]

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

[wild-endeavor]

@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

[github-actions[bot]]

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! 🙏

[github-actions[bot]]

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! 🙏

[novahow]

@kumare3 Hi, I have a few questions on this issue.

1. in the example below, is train_ensemble desired to be a workflow or a task?
2. 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()})

[kumare3]

I think this example is wrong and should be @eager - but cc @wild-endeavor as he added that section

[novahow]

@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

[kumare3]

What @task can be async? I guess it is running in sync mode

[novahow]

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?