Agent ClientSet

[Future-Outlier]

Tracking issue

Fixes #3936

Why are the changes needed?

For better readability and scalability in package agent in the future.

What changes were proposed in this pull request?

• rename getClientFunc to getAgentClientFunc
• use a struct ClientFuncSet to store getAgentClientFunc and getAgentMetadataClientFunc
• pass ClientFuncSet to all related functions and objects
• add related tests

How was this patch tested?

unit test, integration test, and single binary mode.

Setup process

0. start the agent server
1. start the single binary mode
2. run an agent workflow to test it in remote environment

pyflyte serve agent

make compile
flyte start --config ../flyte-single-binary-local-dev.yaml

pyflyte run --remote sensor_example.py wf

flyte-single-binary-local-dev.yaml

tasks:
  task-plugins:
    enabled-plugins:
      - container
      - sidecar
      - K8S-ARRAY
      - agent-service

    default-for-task-types:
      sensor: agent-service
      container: container
      container_array: K8S-ARRAY

plugins:
  agent-service:
    supportedTaskTypes:
      - sensor
    defaultAgent:
      endpoint: "dns:///localhost:8000"
      insecure: true
      timeouts:
        GetTask: 100s
      defaultTimeout: 100s

sensor_example.py

from flytekit.sensor.file_sensor import FileSensor
from flytekit import ImageSpec, task, workflow

sensor = FileSensor(name="test_sensor")

@task()
def t1():
    return

@workflow()
def wf():
    # sensor(path="s3://my-s3-bucket") >> t1()
    sensor(path="/tmp") >> t1()

if __name__ == "__main__":
    wf()

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: 67 lines in your changes are missing coverage. Please review.

Comparison is base (c379b71) 59.02% compared to head (34a4049) 58.99%.

Files	Patch %	Lines
...yteplugins/go/tasks/plugins/webapi/agent/client.go	45.54%	51 Missing and 4 partials ⚠️
...yteplugins/go/tasks/plugins/webapi/agent/plugin.go	25.00%	12 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #4718      +/-   ##
==========================================
- Coverage   59.02%   58.99%   -0.04%     
==========================================
  Files         643      644       +1     
  Lines       55153    55148       -5     
==========================================
- Hits        32555    32535      -20     
- Misses      20018    20039      +21     
+ Partials     2580     2574       -6     

Flag	Coverage Δ
unittests	58.99% <42.73%> (-0.04%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

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

[Future-Outlier]

generate mock file by mockery

cd flyte/flyteidl/gen/pb-go/flyteidl/service

mockery --name=AsyncAgentServiceClient --output=/mnt/c/code/dev/flyte/flyteplugins/go/tasks/plugins/webapi/agent/mocks/ --outpkg=mocks

[Future-Outlier]

Only integration test needs to be fixed, other test files are correct!

[Future-Outlier]

• Improve Error Message
  When you set up the agent-service, and run an agent task, while don't have any connected agent, you will have an error.

  agent-service:
    supportedTaskTypes:
      - sensor
      - spark
      - default_task
      - custom_task
      - chatgpt
      - sensor
      - airflow
    # By default, all the request will be sent to the default agent.
    # defaultAgent:
    #   endpoint: "dns:///localhost:8000"
    #   insecure: true
    #   timeouts:
    #     GetTask: 100s
    #   defaultTimeout: 100s

Before Error Message

After Error Message

[honnix]

Could you please help me understand what happens if one of the agent endpoints is not available when flytepropeller is booting up? Would it just crash?

[honnix]

If so, it is a very strong assumption that I'm afraid we cannot take.

[honnix]

Even the plugin loading is guarded so a single failing plugin would not impact propeller booting, it is still undesired that one failed agent endpoint (out of for example 50 healthy ones) would fail the whole plugin loading. I think some kind of late binding would be nicer.

[Future-Outlier]

We have to wait for every endpoint to start; otherwise, FlytePropeller will keep crashing.

[Future-Outlier]

Really nice suggestions, I will discuss with @pingsutw and reply the result here.
I think late binding or some kind of endpoint or service detect mechanism will be really helpful.
For example, in k8s, maybe readiness probe can help us realize lazy binding?
I am not 100% familiar with how agent is deployed in flyte cluster, but yes, I will try to help.

Thank you for always providing great advices.

[Future-Outlier]

Or maybe, only need 1 agent is connected is enough

[honnix]

I think this plugin used to work in a late binding way, meaning the connection was not established until the very first RPC. In this way the propeller does not depend on agents in the wild when booting up.

For example, in k8s, maybe readiness probe can help us realize lazy binding?

I think we cannot make any assumption how agents are deployed, and only interface they expose to the plugin is gRPC endpoints, so we can only make sense out of those.

[Future-Outlier]

Will update the solution to your case with Kevin this week, thank you!

[Future-Outlier]

@honnix, We will add a watcher to solve your case!
Thank you for the patience, we will mention you to see the new update, thank you very much.

[honnix]

@Future-Outlier That is fantastic. Thank you!