[flytepropeller] Watch agent metadata service dynamically

flyteplugins/go/tasks/pluginmachinery/core/plugin.go

@@ -3,6 +3,9 @@
 import (
 	"context"
 	"fmt"
+	"sync"
+
+	"k8s.io/utils/strings/slices"
 )
 
 //go:generate mockery -all -case=underscore
@@ -55,7 +58,27 @@
 	Finalize(ctx context.Context, tCtx TaskExecutionContext) error
 }
 
-// Loads and validates a plugin.
+type AgentService struct {
+	mu                 sync.RWMutex
+	supportedTaskTypes []TaskType
+	CorePlugin         Plugin
+}
+
+// ContainTaskType check if agent supports this task type.
+func (p *AgentService) ContainTaskType(taskType TaskType) bool {
+	p.mu.RLock()
+	defer p.mu.RUnlock()
+	return slices.Contains(p.supportedTaskTypes, taskType)
+}
+
+// SetSupportedTaskType set supportTaskType in the agent service.
+func (p *AgentService) SetSupportedTaskType(taskTypes []TaskType) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.supportedTaskTypes = taskTypes
+}
+
+// LoadPlugin Loads and validates a plugin.
 func LoadPlugin(ctx context.Context, iCtx SetupContext, entry PluginEntry) (Plugin, error) {
 	plugin, err := entry.LoadPlugin(ctx, iCtx)
 	if err != nil {

flyteplugins/go/tasks/plugins/webapi/agent/client.go

@@ -90,17 +90,11 @@
 	return context.WithTimeout(ctx, timeout)
 }
 
-func updateAgentRegistry(ctx context.Context, cs *ClientSet) {
-	agentRegistry := make(Registry)
+func getAgentRegistry(ctx context.Context, cs *ClientSet) Registry {
+	newAgentRegistry := make(Registry)
 	cfg := GetConfig()
 	var agentDeployments []*Deployment
 
-	// Ensure that the old configuration is backward compatible
-	for taskType, agentDeploymentID := range cfg.AgentForTaskTypes {
-		agent := Agent{AgentDeployment: cfg.AgentDeployments[agentDeploymentID], IsSync: false}
-		agentRegistry[taskType] = map[int32]*Agent{defaultTaskTypeVersion: &agent}
-	}
-
 	if len(cfg.DefaultAgent.Endpoint) != 0 {
 		agentDeployments = append(agentDeployments, &cfg.DefaultAgent)
 	}
@@ -137,27 +131,36 @@
 			deprecatedSupportedTaskTypes := agent.SupportedTaskTypes
 			for _, supportedTaskType := range deprecatedSupportedTaskTypes {
 				agent := &Agent{AgentDeployment: agentDeployment, IsSync: agent.IsSync}
-				agentRegistry[supportedTaskType] = map[int32]*Agent{defaultTaskTypeVersion: agent}
+				newAgentRegistry[supportedTaskType] = map[int32]*Agent{defaultTaskTypeVersion: agent}
 			}
 
 			supportedTaskCategories := agent.SupportedTaskCategories
 			for _, supportedCategory := range supportedTaskCategories {
 				agent := &Agent{AgentDeployment: agentDeployment, IsSync: agent.IsSync}
-				agentRegistry[supportedCategory.GetName()] = map[int32]*Agent{supportedCategory.GetVersion(): agent}
+				newAgentRegistry[supportedCategory.GetName()] = map[int32]*Agent{supportedCategory.GetVersion(): agent}
 			}
 		}
-		// If the agent doesn't implement the metadata service, we construct the registry based on the configuration
-		for taskType, agentDeploymentID := range cfg.AgentForTaskTypes {
-			if agentDeployment, ok := cfg.AgentDeployments[agentDeploymentID]; ok {
-				if _, ok := agentRegistry[taskType]; !ok {
-					agent := &Agent{AgentDeployment: agentDeployment, IsSync: false}
-					agentRegistry[taskType] = map[int32]*Agent{defaultTaskTypeVersion: agent}
-				}
+	}
+
+	// If the agent doesn't implement the metadata service, we construct the registry based on the configuration
+	for taskType, agentDeploymentID := range cfg.AgentForTaskTypes {
+		if agentDeployment, ok := cfg.AgentDeployments[agentDeploymentID]; ok {
+			if _, ok := newAgentRegistry[taskType]; !ok {
+				agent := &Agent{AgentDeployment: agentDeployment, IsSync: false}
+				newAgentRegistry[taskType] = map[int32]*Agent{defaultTaskTypeVersion: agent}
 			}
 		}
 	}
-	logger.Debugf(ctx, "AgentDeployment service supports task types: %v", maps.Keys(agentRegistry))
-	setAgentRegistry(agentRegistry)
+
+	// Ensure that the old configuration is backward compatible
+	for _, taskType := range cfg.SupportedTaskTypes {
+		if _, ok := newAgentRegistry[taskType]; !ok {
+			agent := &Agent{AgentDeployment: &cfg.DefaultAgent, IsSync: false}
+			newAgentRegistry[taskType] = map[int32]*Agent{defaultTaskTypeVersion: agent}
+		}
+	}
+
+	return newAgentRegistry
 }
 
 func getAgentClientSets(ctx context.Context) *ClientSet {

flyteplugins/go/tasks/plugins/webapi/agent/integration_test.go

@@ -35,10 +35,6 @@ import (
 )
 
 func TestEndToEnd(t *testing.T) {
-	agentRegistry = Registry{
-		"openai": {defaultTaskTypeVersion: {AgentDeployment: &Deployment{Endpoint: defaultAgentEndpoint}, IsSync: true}},
-		"spark":  {defaultTaskTypeVersion: {AgentDeployment: &Deployment{Endpoint: defaultAgentEndpoint}, IsSync: false}},
-	}
 	iter := func(ctx context.Context, tCtx pluginCore.TaskExecutionContext) error {
 		return nil
 	}
@@ -117,7 +113,7 @@ func TestEndToEnd(t *testing.T) {
 	t.Run("failed to create a job", func(t *testing.T) {
 		agentPlugin := newMockAsyncAgentPlugin()
 		agentPlugin.PluginLoader = func(ctx context.Context, iCtx webapi.PluginSetupContext) (webapi.AsyncPlugin, error) {
-			return Plugin{
+			return &Plugin{
 				metricScope: iCtx.MetricsScope(),
 				cfg:         GetConfig(),
 				cs: &ClientSet{
@@ -259,6 +255,9 @@ func getTaskContext(t *testing.T) *pluginCoreMocks.TaskExecutionContext {
 
 func newMockAsyncAgentPlugin() webapi.PluginEntry {
 	asyncAgentClient := new(agentMocks.AsyncAgentServiceClient)
+	agentRegistry := Registry{
+		"spark": {defaultTaskTypeVersion: {AgentDeployment: &Deployment{Endpoint: defaultAgentEndpoint}, IsSync: false}},
+	}
 
 	mockCreateRequestMatcher := mock.MatchedBy(func(request *admin.CreateTaskRequest) bool {
 		expectedArgs := []string{"pyflyte-fast-execute", "--output-prefix", "/tmp/123"}
@@ -283,20 +282,25 @@ func newMockAsyncAgentPlugin() webapi.PluginEntry {
 		ID:                 "agent-service",
 		SupportedTaskTypes: []core.TaskType{"bigquery_query_job_task", "spark"},
 		PluginLoader: func(ctx context.Context, iCtx webapi.PluginSetupContext) (webapi.AsyncPlugin, error) {
-			return Plugin{
+			return &Plugin{
 				metricScope: iCtx.MetricsScope(),
 				cfg:         &cfg,
 				cs: &ClientSet{
 					asyncAgentClients: map[string]service.AsyncAgentServiceClient{
 						defaultAgentEndpoint: asyncAgentClient,
 					},
 				},
+				registry: agentRegistry,
 			}, nil
 		},
 	}
 }
 
 func newMockSyncAgentPlugin() webapi.PluginEntry {
+	agentRegistry := Registry{
+		"openai": {defaultTaskTypeVersion: {AgentDeployment: &Deployment{Endpoint: defaultAgentEndpoint}, IsSync: true}},
+	}
+
 	syncAgentClient := new(agentMocks.SyncAgentServiceClient)
 	output, _ := coreutils.MakeLiteralMap(map[string]interface{}{"x": 1})
 	resource := &admin.Resource{Phase: flyteIdlCore.TaskExecution_SUCCEEDED, Outputs: output}
@@ -323,14 +327,15 @@ func newMockSyncAgentPlugin() webapi.PluginEntry {
 		ID:                 "agent-service",
 		SupportedTaskTypes: []core.TaskType{"openai"},
 		PluginLoader: func(ctx context.Context, iCtx webapi.PluginSetupContext) (webapi.AsyncPlugin, error) {
-			return Plugin{
+			return &Plugin{
 				metricScope: iCtx.MetricsScope(),
 				cfg:         &cfg,
 				cs: &ClientSet{
 					syncAgentClients: map[string]service.SyncAgentServiceClient{
 						defaultAgentEndpoint: syncAgentClient,
 					},
 				},
+				registry: agentRegistry,
 			}, nil
 		},
 	}

flyteplugins/go/tasks/plugins/webapi/agent/plugin.go

@@ -24,29 +24,16 @@
 	"github.com/flyteorg/flyte/flytestdlib/promutils"
 )
 
-type Registry map[string]map[int32]*Agent // map[taskTypeName][taskTypeVersion] => Agent
-
-var (
-	agentRegistry Registry
-	mu            sync.RWMutex
-)
+const ID = "agent-service"
 
-func getAgentRegistry() Registry {
-	mu.Lock()
-	defer mu.Unlock()
-	return agentRegistry
-}
-
-func setAgentRegistry(r Registry) {
-	mu.Lock()
-	defer mu.Unlock()
-	agentRegistry = r
-}
+type Registry map[string]map[int32]*Agent // map[taskTypeName][taskTypeVersion] => Agent
 
 type Plugin struct {
 	metricScope promutils.Scope
 	cfg         *Config
 	cs          *ClientSet
+	registry    Registry
+	mu          sync.RWMutex
 }
 
 type ResourceWrapper struct {
@@ -69,18 +56,24 @@
 	TaskCategory      admin.TaskCategory
 }
 
-func (p Plugin) GetConfig() webapi.PluginConfig {
+func (p *Plugin) setRegistry(r Registry) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.registry = r
+}
+
+func (p *Plugin) GetConfig() webapi.PluginConfig {
 	return GetConfig().WebAPI
 }
 
-func (p Plugin) ResourceRequirements(_ context.Context, _ webapi.TaskExecutionContextReader) (
+func (p *Plugin) ResourceRequirements(_ context.Context, _ webapi.TaskExecutionContextReader) (
 	namespace core.ResourceNamespace, constraints core.ResourceConstraintsSpec, err error) {
 
 	// Resource requirements are assumed to be the same.
 	return "default", p.cfg.ResourceConstraints, nil
 }
 
-func (p Plugin) Create(ctx context.Context, taskCtx webapi.TaskExecutionContextReader) (webapi.ResourceMeta,
+func (p *Plugin) Create(ctx context.Context, taskCtx webapi.TaskExecutionContextReader) (webapi.ResourceMeta,
 	webapi.Resource, error) {
 	taskTemplate, err := taskCtx.TaskReader().Read(ctx)
 	if err != nil {
@@ -113,7 +106,7 @@
 	outputPrefix := taskCtx.OutputWriter().GetOutputPrefixPath().String()
 
 	taskCategory := admin.TaskCategory{Name: taskTemplate.Type, Version: taskTemplate.TaskTypeVersion}
-	agent, isSync := getFinalAgent(&taskCategory, p.cfg)
+	agent, isSync := p.getFinalAgent(&taskCategory, p.cfg)
 
 	taskExecutionMetadata := buildTaskExecutionMetadata(taskCtx.TaskExecutionMetadata())
 
@@ -149,7 +142,7 @@
 	}, nil, nil
 }
 
-func (p Plugin) ExecuteTaskSync(
+func (p *Plugin) ExecuteTaskSync(
 	ctx context.Context,
 	client service.SyncAgentServiceClient,
 	header *admin.CreateRequestHeader,
@@ -206,9 +199,9 @@
 	}, err
 }
 
-func (p Plugin) Get(ctx context.Context, taskCtx webapi.GetContext) (latest webapi.Resource, err error) {
+func (p *Plugin) Get(ctx context.Context, taskCtx webapi.GetContext) (latest webapi.Resource, err error) {
 	metadata := taskCtx.ResourceMeta().(ResourceMetaWrapper)
-	agent, _ := getFinalAgent(&metadata.TaskCategory, p.cfg)
+	agent, _ := p.getFinalAgent(&metadata.TaskCategory, p.cfg)
 
 	client, err := p.getAsyncAgentClient(ctx, agent)
 	if err != nil {
@@ -236,12 +229,12 @@
 	}, nil
 }
 
-func (p Plugin) Delete(ctx context.Context, taskCtx webapi.DeleteContext) error {
+func (p *Plugin) Delete(ctx context.Context, taskCtx webapi.DeleteContext) error {
 	if taskCtx.ResourceMeta() == nil {
 		return nil
 	}
 	metadata := taskCtx.ResourceMeta().(ResourceMetaWrapper)
-	agent, _ := getFinalAgent(&metadata.TaskCategory, p.cfg)
+	agent, _ := p.getFinalAgent(&metadata.TaskCategory, p.cfg)
 
 	client, err := p.getAsyncAgentClient(ctx, agent)
 	if err != nil {
@@ -259,7 +252,7 @@
 	return err
 }
 
-func (p Plugin) Status(ctx context.Context, taskCtx webapi.StatusContext) (phase core.PhaseInfo, err error) {
+func (p *Plugin) Status(ctx context.Context, taskCtx webapi.StatusContext) (phase core.PhaseInfo, err error) {
 	resource := taskCtx.Resource().(ResourceWrapper)
 	taskInfo := &core.TaskInfo{Logs: resource.LogLinks}
 
@@ -311,7 +304,7 @@
 	return core.PhaseInfoUndefined, pluginErrors.Errorf(core.SystemErrorCode, "unknown execution state [%v].", resource.State)
 }
 
-func (p Plugin) getSyncAgentClient(ctx context.Context, agent *Deployment) (service.SyncAgentServiceClient, error) {
+func (p *Plugin) getSyncAgentClient(ctx context.Context, agent *Deployment) (service.SyncAgentServiceClient, error) {
 	client, ok := p.cs.syncAgentClients[agent.Endpoint]
 	if !ok {
 		conn, err := getGrpcConnection(ctx, agent)
@@ -324,7 +317,7 @@
 	return client, nil
 }
 
-func (p Plugin) getAsyncAgentClient(ctx context.Context, agent *Deployment) (service.AsyncAgentServiceClient, error) {
+func (p *Plugin) getAsyncAgentClient(ctx context.Context, agent *Deployment) (service.AsyncAgentServiceClient, error) {
 	client, ok := p.cs.asyncAgentClients[agent.Endpoint]
 	if !ok {
 		conn, err := getGrpcConnection(ctx, agent)
@@ -337,13 +330,25 @@
 	return client, nil
 }
 
-func (p Plugin) watchAgents(ctx context.Context) {
+func (p *Plugin) watchAgents(ctx context.Context, agentService *core.AgentService) {
 	go wait.Until(func() {
 		clientSet := getAgentClientSets(ctx)
-		updateAgentRegistry(ctx, clientSet)
+		agentRegistry := getAgentRegistry(ctx, clientSet)
+		p.setRegistry(agentRegistry)
+		agentService.SetSupportedTaskType(maps.Keys(agentRegistry))
 	}, p.cfg.PollInterval.Duration, ctx.Done())
 }
 
+func (p *Plugin) getFinalAgent(taskCategory *admin.TaskCategory, cfg *Config) (*Deployment, bool) {
+	p.mu.RLock()
+	defer p.mu.RUnlock()
+
+	if agent, exists := p.registry[taskCategory.Name][taskCategory.Version]; exists {
+		return agent.AgentDeployment, agent.IsSync
+	}
+	return &cfg.DefaultAgent, false
+}
+
 func writeOutput(ctx context.Context, taskCtx webapi.StatusContext, outputs *flyteIdl.LiteralMap) error {
 	taskTemplate, err := taskCtx.TaskReader().Read(ctx)
 	if err != nil {
@@ -366,14 +371,6 @@
 	return taskCtx.OutputWriter().Put(ctx, opReader)
 }
 
-func getFinalAgent(taskCategory *admin.TaskCategory, cfg *Config) (*Deployment, bool) {
-	r := getAgentRegistry()
-	if agent, exists := r[taskCategory.Name][taskCategory.Version]; exists {
-		return agent.AgentDeployment, agent.IsSync
-	}
-	return &cfg.DefaultAgent, false
-}
-
 func buildTaskExecutionMetadata(taskExecutionMetadata core.TaskExecutionMetadata) admin.TaskExecutionMetadata {
 	taskExecutionID := taskExecutionMetadata.GetTaskExecutionID().GetID()
 
@@ -388,13 +385,12 @@
 	}
 }
 
-func newAgentPlugin() webapi.PluginEntry {
+func newAgentPlugin(agentService *core.AgentService) webapi.PluginEntry {
 	ctx := context.Background()
 	cfg := GetConfig()
-
 	clientSet := getAgentClientSets(ctx)
-	updateAgentRegistry(ctx, clientSet)
-	supportedTaskTypes := append(maps.Keys(getAgentRegistry()), cfg.SupportedTaskTypes...)
+	agentRegistry := getAgentRegistry(ctx, clientSet)
+	supportedTaskTypes := maps.Keys(agentRegistry)
 
 	return webapi.PluginEntry{
 		ID:                 "agent-service",
@@ -404,15 +400,16 @@
 				metricScope: iCtx.MetricsScope(),
 				cfg:         cfg,
 				cs:          clientSet,
+				registry:    agentRegistry,
 			}
-			plugin.watchAgents(ctx)
+			plugin.watchAgents(ctx, agentService)
 			return plugin, nil
 		},
 	}
 }
 
-func RegisterAgentPlugin() {
+func RegisterAgentPlugin(agentService *core.AgentService) {
 	gob.Register(ResourceMetaWrapper{})
 	gob.Register(ResourceWrapper{})
-	pluginmachinery.PluginRegistry().RegisterRemotePlugin(newAgentPlugin())
+	pluginmachinery.PluginRegistry().RegisterRemotePlugin(newAgentPlugin(agentService))
 }