#1840 [Core][Flytekit] support async methods in python Issue created by 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? ☑︎ Yes Have you read the Code of Conduct? ☑︎ Yes flyteorg/flyte