[Core feature] Rust Flytekit Remote Client

[pingsutw]

Motivation: Why do you think this is important?

The flytekit remote client relies on gRPC for communication with the control plane (FlyteAdmin). However, our experience has revealed numerous issues within Python gRPC, leading to test failures. Additionally, users frequently encounter challenges when working with specific versions of gRPC.

In order to eliminate the gRPC Python dependency from flytekit, our goal is to develop a Rust crate to serve as a replacement for the flytekit remote client. To achieve this, we plan to utilize Pyo3 to generate a Python binding, allowing flytekit to seamlessly integrate and utilize the new implementation.

Goal: What should the final outcome look like, ideally?

flyteorg/flytekit#2377 is a POC of a rust remote client that allows you to get/create tasks without Python grpc

Add a flag (enable_rust) to the FlyteRemote class. When this flag is true, FlyteRemote will use the rust implementation.

remote_rs = FlyteRemote(Config.auto(), default_project=PROJECT, default_domain=DOMAIN, enable_rust=True)
flyte_task = remote_rs.register_task(
    entity=my_test_task,
    serialization_settings=SerializationSettings(
        image_config=ImageConfig.auto(img_name="flyte-cr.io/image-name:tag")
    ),
    version=VERSION_ID,
)
flyte_task = remote_rs.fetch_task(name=TASK_NAME, version=VERSION_ID)

Follow up Issues

• Add a test for the rust remote client
• Add a GA workflow to run rust remote client without installing Python grpc
• Add a GA workflow to make sure the output of the Rust remote and Python remote are the same
• Pass credentials from python to rust
• OAuth flow in the rust implementation

Describe alternatives you've considered

Use existing flytekit remote client

Propose: Link/Inline OR Additional context

• Fix: Fix GCP Identity Aware Proxy integration errors occurring with grpcio>1.56.0 flytekit#2034
• Skip version 1.45.0 of grpcio and grpcio-status flytekit#913
• bug(python): Regression in grpcio>=1.59.0 related to grpc thread pool hanging in grpc_google_default_credentials_create grpc/grpc#36265

Are you sure this issue hasn't been raised already?

• Yes

Have you read the Code of Conduct?

• Yes

[runllm]

To address the issues with Python gRPC in the Flytekit remote client, a Rust-based replacement is proposed. This approach aims to eliminate Python gRPC dependencies by using Pyo3 for Python bindings, ensuring seamless integration with Flytekit. A proof of concept is available in this pull request, demonstrating task operations without Python gRPC. The proposed solution includes an enable_rust flag in FlyteRemote to switch to the Rust implementation. Future tasks involve testing the Rust client, ensuring output consistency between Rust and Python implementations, and integrating OAuth and credential passing.

remote_rs = FlyteRemote(Config.auto(), default_project="flytesnacks", default_domain="development", enable_rust=True)
flyte_task = remote_rs.register_task(
    entity=my_test_task,
    serialization_settings=SerializationSettings(
        image_config=ImageConfig.auto(img_name="flyte-cr.io/image-name:tag")
    ),
    version="v1"
)
flyte_task = remote_rs.fetch_task(name="task_name", version="v1")

Sources:

• Discuss Flyte
• Flytekit Pull Request #2377
• gRPC Issue #36265

---

This is an AI-generated response and your feedback is appreciated! Please leave a 👍 if this is helpful and 👎 if it is not.

Mention @RunLLM to ask follow-up questions. See runllm.com for more!

[austin362667]

@pingsutw , could you please assist me in adding one more follow-up issue: 'Add secure channel if flyteadmin has SSL enabled'? Thank you!