Eager local entrypoint and support for offloaded types

[cosmicBboy]

TL;DR

This PR extends the @eager decorator to support running the workflow locally but executing tasks/subworkflows remotely. Also adds support and tests for preserving offloaded types.

Type

• Bug Fix
• Feature
• Plugin

Are all requirements met?

• Code completed
• Smoke tested
• Unit tests added
• Code documentation added
• Any pending items have an associated Issue

Complete description

@eager(local_entrypoint=True) will treat the eager workflow as a locally executed task that invokes the configured remote object for tasks and subworkflows called inside of it. It also adds support and tests for offloaded types: structured datasets, files, and directories, such that they are not immediately downloaded and can be downloaded/iterated over in the eager workflow function body.

This PR also fixes some formatting issues with the eager workflow flyte deck.

Tracking Issue

NA

Follow-up issue

NA

[codecov]

Codecov Report

Patch coverage: 16.99% and project coverage change: +1.76% 🎉

Comparison is base (89446a1) 18.35% compared to head (ca43b90) 20.12%.
Report is 3 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1833      +/-   ##
==========================================
+ Coverage   18.35%   20.12%   +1.76%     
==========================================
  Files         335      336       +1     
  Lines       32081    32281     +200     
  Branches     3161     5840    +2679     
==========================================
+ Hits         5889     6495     +606     
+ Misses      26107    25620     -487     
- Partials       85      166      +81     

Files Changed	Coverage Δ
...flytekit/unit/experimental/test_eager_workflows.py	0.00% <0.00%> (ø)
flytekit/experimental/eager_function.py	22.94% <6.89%> (+22.94%)	⬆️
...ytekit/integration/experimental/eager_workflows.py	50.47% <43.47%> (ø)
...t/integration/experimental/test_eager_workflows.py	51.28% <50.00%> (+1.28%)	⬆️

... and 37 files with indirect coverage changes

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

[pingsutw]

Maybe we just check if the remote is None?

	local execution	remote execution
AsyncEntity has remote obj	Use flytekit remote to run the task	Use flytekit remote to run the task
AsyncEntity has not remote obj	Run the task function locally	Run the task in the driver (eager) pod

[cosmicBboy]

I think we'd want to run the AsyncEntity locally even if remote is None right? e.g.

# wf.py
@eager(remote=remote)
def eager_wf(): ...

pyflyte run wf.py eager_wf  # runs locally
pyflyte run wf.py --remote eager_wf  # everything, including eager_wf runs remotely

should run locally, even if the user has configured a remote object.

It's only when they say local_entrypoint=True that the behavior changes such that it runs tasks/subworkflows remotely on the configured remote.

# wf.py
@eager(remote=remote, local_entrypoint=True)
def eager_wf(): ...

# eager_wf runs locally, but since local_entrypoint=True,
# tasks and subworkflows run remotely
pyflyte run wf.py eager_wf

[pingsutw]

@eager(remote=remote, local_entrypoint=False)
def wf()
   ...

qq: In the local execution, what happens in this case?, will we also use flytekit remote to execute the task?

[cosmicBboy]

Case 1:

Just specifying a remote shouldn't change any of the local behavior: it'll just run locally as async functions and ignore the `remote object.

@eager(remote=remote, local_entrypoint=False)
# pyflyte run: runs the eager workflow locally, runs the tasks/subworkflows locally
# pyflyte run --remote: runs the eager workflow remotely, runs the tasks/subworkflows remotely

Case 2:

local_entrypoint=True just forces the eager workflow's tasks/subworkflows to run remotely (while running the eager workflow itself locally) as a debugging/development tool.

@eager(remote=remote, local_entrypoint=True)  # runs eager workflow
# pyflyte run: runs the eager workflow locally, but forces tasks and subworkflows to run remotely
# pyflyte run --remote: runs the 

[cosmicBboy]

also happy to change the name local_entrypoint to something clearer. The use case is: "I want to run the eager workflow locally, but run the tasks/subworkflows remotely to that I can debug it, e.g. get access to intermediary results"

[pingsutw]

yes, if we want to add a new arg, it will be better to change local_entrypoint to something clearer.

[cosmicBboy]

any suggestions?

[cosmicBboy]

@eager(remote=remote)
# pyflyte run: runs the eager workflow locally, but forces tasks and subworkflows to run remotely. If fail to read > remote config, Fallback to local execution,
# pyflyte run --remote: runs the eager workflow remotely, runs the tasks/subworkflows remotely

The problem with this is that I have to keep changing my code if I want to run completely locally vs. completely remotely. Think of remote=remote as just configuration, similar to the configuration that you give to @task. I'd like to run my eager workflow locally and remotely by just changing the pyflyte run call (adding --remote flag for remote execution).

The local_entrypoint=True flag is exclusively for the hybrid case of running the eager workflow locally but the tasks/subworkflows remotely.

If it makes the use case clearer we can just introduce a wrapper decorator to isolate this use case:

# this will always just run locally and you can't run this on a remote Flyte cluster
@local_entrypoint(remote=remote)
async def wf():
   ...

(btw I just took the naming from modal, open to suggestions for other names)

[pingsutw]

Got it. Rather than a decorator, I prefer @eager(remote=remote, local_entrypoint=True).
I'm not good at naming. Maybe we merge it first and then rename it once someone finds a better name.

[cosmicBboy]

Cool, will make sure to document this well on flytesnacks