wip Async/tasks

[wild-endeavor]

To support the growing interest to support Python async style programming, and to bring the eager mode out of experimental, this pr revamps the call patters in eager to support native Python async semantics. This PR also introduces the concept of async tasks. Note that while documentation refers to eager entities as "workflows", using the @eager decorator produces a task, not a workflow, so they'll continue to show up under the Tasks tab of the Flyte UI.

This is a backwards incompatible change for @eager.

Usage Patterns

Below are some examples of how you might call the new eager workflow.

Simple Example

This is the simplest case. An eager workflow invokes a simple task. The container running the eager function will reach out to the control plane, and kick off a single-task execution.

@task
def add_one(x: int) -> int:  
    return x + 1  
  
  
@eager
async def simple_eager_workflow(x: int) -> int:  
    # This is the normal way of calling tasks. Call normal tasks in an effectively async way by hanging and waiting for  
    # the result.    out = add_one(x=x)  
  
    return out

This is akin to an async function in Python calling a synchronous function. It will automatically block until the results are available.

Controlling Execution

This example shows how you might do work at the same time as a task kicked off by eager is running. Again this follows the same semantics as Python. In Python the executing function also needs to relinquish control of the CPU by calling an await so that other items on the loop can progress.

@eager
async def simple_await() -> int:
    t1 = asyncio.create_task(add_one(x=10))

	# This allows the loop to run in the background, actually kicking
	# off the execution
    await asyncio.sleep(0)
    # <can do more CPU intensive things here while>
  
    # don't forget the comma if just awaiting one
    i1, = await asyncio.gather(t1)  # can have more of course
    return i1

Nested Example

Eager workflows can also be nested. If an eager workflow encounters another eager workflow, it will be launched against Admin as a single task execution, just like any other task.

@task
def add_one(x: int) -> int:
    return x + 1

@task
def double(x: int) -> int:  
    return x * 2  

@dynamic
def level_3_dt(x: int) -> int:
    out = add_one(x=x)
    return out

@eager
async def level_2(x: int) -> int:
    out = add_one(x=x)
    level_3_res = level_3_dt(x=out)
    final_res = double(x=level_3_res)
    return final_res

@eager
async def level_1() -> typing.Tuple[int, int]:
    i1 = add_one(x=10)
    t2 = asyncio.create_task(level_2(x=1))

    i2 = await t2
    return i1, i2

Changes

High level changes

• Addition of a new AsyncPythonFunctionTask and EagerAsyncPythonFunctionTask.
• Changed the main flyte_entity_call_handler to be async.
• Updated main call handler to now allow calling of other flyte entities inside an eager task.
• Introduced two new Execution modes, one for local eager and one for eager backend run.
• Add as_python_native to the LiteralsResolver so it can produce proper un-packable outputs as the result of calling flyte entities.
• Adding two classes for the watching of executions launched by an eager task.

Remaining items

Remaining for this PR

The bar here is to include things that might break other parts of flytekit, or general tidiness for this PR. Eager may or may not be in a completely working state after these are in.

• Correctly raise eager exception
• Proper naming and labeling of executions so that executions are easily findable, including nested eager.
• Hook up flyte Deck again, just using the same template from before.
• Proper signal handling.
• A bunch of unit tests.
• Do a grep for todo:async and do or create follow-up issues for it.
• Update comments about how the call pattern works

UX related

• how to local run properly
• remote/auth story

Cleanup

• Import eager into top level flytekit
• Review of naming of files and objects.
• Documentation in code comments for eager.
• Add error to nested calls, worker queue must be empty.

Other items

These probably make more sense to do after this is merged. Eager should be considered not completely usable until these are in.

• Introduce async workflows that compile the same way as normal workflows, but are callable within other async workflows, and can call eager tasks directly.
• Documentation updates outside of code comments

Investigate

Things that came up in the course of this PR.

• Debugger would hang in the case of an error rather than quit. (search for fut.result(None) in code).
• Not related to eager but the naming of entities changed drastically if i was two folders up from the code rather than one.

How was this patch tested?

Setup process

Screenshots

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Related PRs

Docs link

[codecov]

Codecov Report

Attention: Patch coverage is 31.38075% with 164 lines in your changes missing coverage. Please review.

Project coverage is 51.28%. Comparing base (d4171c3) to head (b45cc87).
Report is 5 commits behind head on master.

Files with missing lines	Patch %	Lines
flytekit/core/worker_queue.py	33.70%	58 Missing and 1 partial ⚠️
flytekit/core/python_function_task.py	33.92%	36 Missing and 1 partial ⚠️
scratch.py	0.00%	26 Missing ⚠️
flytekit/core/promise.py	19.35%	24 Missing and 1 partial ⚠️
flytekit/core/task.py	41.66%	5 Missing and 2 partials ⚠️
flytekit/core/context_manager.py	60.00%	6 Missing ⚠️
flytekit/experimental/eager_function.py	71.42%	1 Missing and 1 partial ⚠️
flytekit/tools/translator.py	33.33%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2927      +/-   ##
==========================================
+ Coverage   47.00%   51.28%   +4.27%     
==========================================
  Files         199      201       +2     
  Lines       20840    21017     +177     
  Branches     2681     2698      +17     
==========================================
+ Hits         9796    10778     +982     
+ Misses      10558     9631     -927     
- Partials      486      608     +122     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.