From a433b334ff919195ab7618feeb9e2b1dbdea13b8 Mon Sep 17 00:00:00 2001 From: Future-Outlier Date: Fri, 17 Jan 2025 16:51:53 +0800 Subject: [PATCH] Add FlyteAgent Default Service Config Using Round Robin Mechanism Signed-off-by: Future-Outlier --- .../go/tasks/plugins/webapi/agent/config.go | 7 ++++--- .../tasks/plugins/webapi/agent/config_test.go | 20 +++++++++++++++++++ 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/flyteplugins/go/tasks/plugins/webapi/agent/config.go b/flyteplugins/go/tasks/plugins/webapi/agent/config.go index f26499f320..3f9c7bb5fe 100644 --- a/flyteplugins/go/tasks/plugins/webapi/agent/config.go +++ b/flyteplugins/go/tasks/plugins/webapi/agent/config.go @@ -40,9 +40,10 @@ var ( }, }, DefaultAgent: Deployment{ - Endpoint: "", - Insecure: true, - DefaultTimeout: config.Duration{Duration: 10 * time.Second}, + Endpoint: "", + Insecure: true, + DefaultTimeout: config.Duration{Duration: 10 * time.Second}, + DefaultServiceConfig: `{"loadBalancingConfig": [{"round_robin":{}}]}`, }, // AsyncPlugin should be registered to at least one task type. // Reference: https://github.com/flyteorg/flyte/blob/master/flyteplugins/go/tasks/pluginmachinery/registry.go#L27 diff --git a/flyteplugins/go/tasks/plugins/webapi/agent/config_test.go b/flyteplugins/go/tasks/plugins/webapi/agent/config_test.go index 09abffb83c..27a17746dc 100644 --- a/flyteplugins/go/tasks/plugins/webapi/agent/config_test.go +++ b/flyteplugins/go/tasks/plugins/webapi/agent/config_test.go @@ -39,3 +39,23 @@ func TestGetAndSetConfig(t *testing.T) { assert.NoError(t, err) assert.Equal(t, &cfg, GetConfig()) } + +func TestDefaultAgentConfig(t *testing.T) { + cfg := defaultConfig + + assert.Equal(t, "", cfg.DefaultAgent.Endpoint) + assert.True(t, cfg.DefaultAgent.Insecure) + assert.Equal(t, 10*time.Second, cfg.DefaultAgent.DefaultTimeout.Duration) + assert.Equal(t, `{"loadBalancingConfig": [{"round_robin":{}}]}`, cfg.DefaultAgent.DefaultServiceConfig) + + assert.Empty(t, cfg.DefaultAgent.Timeouts) + + expectedTaskTypes := []string{"task_type_1", "task_type_2"} + assert.Equal(t, expectedTaskTypes, cfg.SupportedTaskTypes) + + assert.Empty(t, cfg.AgentDeployments) + + assert.Empty(t, cfg.AgentForTaskTypes) + + assert.Equal(t, 10*time.Second, cfg.PollInterval.Duration) +}