[WIP] Support task execution of WebAssembly binaries

[cameronraysmith]

This is currently intended to facilitate discussion of some experiments designed to scope out what might be required to support task execution of WebAssembly binaries. If it makes some progress we'll complete the copy of the RFC template below.

RFC template draft

Support for task execution on WebAssembly binaries

Authors:

• @Nickname
• @Nickname

1 Executive Summary

A short paragraph or bullet list that quickly explains what you're trying to do.

2 Motivation

What motivates this proposal, and why is it important?

Here, we aim to get comfortable articulating the value of our actions.

3 Proposed Implementation

This is the core of your proposal, and its purpose is to help you think through the problem because writing is thinking.

Consider:

• using diagrams to help illustrate your ideas.
• including code examples if you're proposing an interface or system contract.
• linking to project briefs or wireframes that are relevant.

4 Metrics & Dashboards

What are the main metrics we should be measuring? For example, when interacting with an external system, it might be the external system latency. When adding a new table, how fast would it fill up?

5 Drawbacks

Are there any reasons why we should not do this? Here we aim to evaluate risk and check ourselves.

6 Alternatives

What are other ways of achieving the same outcome?

7 Potential Impact and Dependencies

Here, we aim to be mindful of our environment and generate empathy towards others who may be impacted by our decisions.

• What other systems or teams are affected by this proposal?
• How could this be exploited by malicious attackers?

8 Unresolved questions

What parts of the proposal are still being defined or not covered by this proposal?

9 Conclusion

Here, we briefly outline why this is the right decision to make at this time and move forward!

10 RFC Process Guide, remove this section when done

By writing an RFC, you're giving insight to your team on the direction you're taking. There may not be a right or better decision in many cases, but we will likely learn from it. By authoring, you're making a decision on where you want us to go and are looking for feedback on this direction from your team members, but ultimately the decision is yours.

This document is a:

• thinking exercise, prototype with words.
• historical record, its value may decrease over time.
• way to broadcast information.
• mechanism to build trust.
• tool to empower.
• communication channel.

This document is not:

• a request for permission.
• the most up to date representation of any process or system

Checklist:

• Copy template
• Draft RFC (think of it as a wireframe)
• Share as WIP with folks you trust to gut-check
• Send pull request when comfortable
• Label accordingly
• Assign reviewers
• Merge PR

Recommendations

• Tag RFC title with [WIP] if you're still ironing out details.
• Tag RFC title with [Newbie] if you're trying out something experimental or you're not entirely convinced of what you're proposing.
• Tag RFC title with [RR] if you'd like to schedule a review request to discuss the RFC.
• If there are areas that you're not convinced on, tag people who you consider may know about this and ask for their input.
• If you have doubts, ask on #feature-discussions for help moving something forward.

[kumare3]

I actually think this is a future that I want to work towards, but IMO the eco system is not mature enough yet. So lets try out an experiement. I will try to provide 2 phases of exploration here, 1. An experiment that proves that it can work and helps us design the API experience 2. What it would need for the final version to work

Phase 1: Setup an experiment.

Caveats / Caution:

• This is not production grade
• Will not work for long running tasks

Beneifts

• Simple to try and get started with
• Isolated without modifying the core

Objective:

One should be able to execute WASM binaries as tasks

Motivating example

wasm_task_1 = WASMTask(
     name=...,
     binary_uri="s3://....",
     inputs=kwtypes(x=int),
     outputs=(int, str),
     metadata=...,
)

wasm_task_2 = ...

@workflow
def my_wasm_wf(x: int) -> int:
     u, v = wasm_task_1(x=x)
    return wasm_task_2(u=u, v=v)

To execute

pyflyte run ... (local and remote)

Overview

We can use something like wasmer-python, as the runtime to execute the wasm binaries. The reason to chose this runtime is so that the runtime is in python and we can easily test things.

Write a New WASM runtime Flyte Agent

1. Using the SyncAgent API.Do method implement the WASM runner that can take the given uri and run it inline

2. Care has to be taken to pass the inputs into the wasm binary. Not sure in the setup how to do it. We can limit the support to only primitive types or even just int and str

Execute locally

[cameronraysmith]

Thanks for this proposal for a minimal experiment! Definitely agree this would be the right place to start.

It probably wouldn't matter too much which language hosted the source for the wasm binary.
However, to create a first example transparent to most flytekit users, it might be nice to use

• componentize-py and/or
• the wasmer fork of the Nuitka compiler > py2wasm | blog post

to write and compile the example wasm binary. We could try to run these on wasmtime-py and wasmer-python as you suggest above.

[cameronraysmith]

This is mostly just collecting a few pointers to relevant source, but let me know if you see any misconceptions as I've not worked with Flyte Agents before.

The SyncAgent interface can be used to implement the class WASMAgent(SyncAgentBase) to be registered AgentRegistry.register(WASMAgent()) to execute something like
class WASMTask(SyncAgentExecutorMixin, ...)
tasks where the

Do method implement the WASM runner that can take the given uri and run it inline

do method of the WASMAgent SyncAgent
https://github.com/flyteorg/flytekit/blob/v1.13.0/flytekit/extend/backend/base_agent.py#L122-L128
that takes at least a TaskTemplate instance
https://github.com/flyteorg/flytekit/blob/v1.13.0/flytekit/models/task.py#L338-L519
and returns a Resource dataclass instance
https://github.com/flyteorg/flytekit/blob/v1.13.0/flytekit/extend/backend/base_agent.py#L77-L92
will call or execute the relevant wasm component or module as the case may be depending on the source and runtime capabilities.

Maybe the sagemaker interface's BotoAgent/BotoTask and ChatGPTAgent/ChatGPTTask are good exampes.

[cameronraysmith]

Care has to be taken to pass the inputs into the wasm binary. Not sure in the setup how to do it. We can limit the support to only primitive types or even just int and str

I think this might be more straightforward and appropriate for the level of abstraction we'd like to achieve if we do indeed use a wasm component with a WIT interface. Maybe it can be done with a wasm module alone as well, but I suspect most users will come to prefer the type system supported by the component model as opposed to the lower-level module types unless there's a specific reason to go there.

[kumare3]

Phase 2: Real runtime

This is all hypothetical.

Steps

1. First convert flytekit core to be rust based (this is already in progress). The entrypoint has to be rust based and all method invocations / task invocations happen using ffi
2. Generate WASM binding for flyteIDL.core.literals / literaltypes
3. Everything before this point should happen in rust (cc @thomasjpfan and @pingsutw and @wild-endeavor ). We should also alter the interface to take input literals / inputs and return (outputs, error)
4. Now the same entrypoint can be used to invoke wasm binaries

What else is needed
At Union we now have a fast async execution engine that goes beyond K8s to make it possible to run super low latency workloads. We might be open to thinking about a future of how that will work with this work to allow running long running tasks