[flytepropeller] Watch agent metadata service dynamically

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

@@ -3,6 +3,8 @@ package core
 import (
 	"context"
 	"fmt"
+	"k8s.io/utils/strings/slices"
+	"sync"
 )
 
 //go:generate mockery -all -case=underscore
@@ -55,7 +57,27 @@ type Plugin interface {
 	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

@@ -97,17 +97,11 @@ func getFinalContext(ctx context.Context, operation string, agent *Deployment) (
 	return context.WithTimeout(ctx, timeout)
 }
 
-func updateAgentRegistry(ctx context.Context, cs *ClientSet) {
-	agentRegistry := make(Registry)
+func getUpdatedAgentRegistry(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)
 	}
@@ -144,27 +138,36 @@ func updateAgentRegistry(ctx context.Context, cs *ClientSet) {
 			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
 	}
@@ -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"}
@@ -291,12 +290,17 @@ func newMockAsyncAgentPlugin() webapi.PluginEntry {
 						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}
@@ -331,6 +335,7 @@ func newMockSyncAgentPlugin() webapi.PluginEntry {
 						defaultAgentEndpoint: syncAgentClient,
 					},
 				},
+				registry: agentRegistry,
 			}, nil
 		},
 	}

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

@@ -24,29 +24,16 @@ import (
 	"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,30 @@ type ResourceMetaWrapper struct {
 	TaskCategory      admin.TaskCategory
 }
 
-func (p Plugin) GetConfig() webapi.PluginConfig {
+func (p *Plugin) getRegistryTaskTypes() []core.TaskType {
+	p.mu.RLock()
+	defer p.mu.RUnlock()
+	return maps.Keys(p.registry)
+}
+
+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 +112,7 @@ func (p Plugin) Create(ctx context.Context, taskCtx webapi.TaskExecutionContextR
 	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 +148,7 @@ func (p Plugin) Create(ctx context.Context, taskCtx webapi.TaskExecutionContextR
 	}, nil, nil
 }
 
-func (p Plugin) ExecuteTaskSync(
+func (p *Plugin) ExecuteTaskSync(
 	ctx context.Context,
 	client service.SyncAgentServiceClient,
 	header *admin.CreateRequestHeader,
@@ -209,9 +208,9 @@ func (p Plugin) ExecuteTaskSync(
 	}, 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 {
@@ -239,12 +238,12 @@ func (p Plugin) Get(ctx context.Context, taskCtx webapi.GetContext) (latest weba
 	}, 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 {
@@ -262,7 +261,7 @@ func (p Plugin) Delete(ctx context.Context, taskCtx webapi.DeleteContext) error
 	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}
 
@@ -314,7 +313,7 @@ func (p Plugin) Status(ctx context.Context, taskCtx webapi.StatusContext) (phase
 	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)
@@ -327,7 +326,7 @@ func (p Plugin) getSyncAgentClient(ctx context.Context, agent *Deployment) (serv
 	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)
@@ -340,13 +339,25 @@ func (p Plugin) getAsyncAgentClient(ctx context.Context, agent *Deployment) (ser
 	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 := getUpdatedAgentRegistry(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 {
@@ -369,14 +380,6 @@ func writeOutput(ctx context.Context, taskCtx webapi.StatusContext, outputs *fly
 	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()
 
@@ -391,13 +394,12 @@ func buildTaskExecutionMetadata(taskExecutionMetadata core.TaskExecutionMetadata
 	}
 }
 
-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 := getUpdatedAgentRegistry(ctx, clientSet)
+	supportedTaskTypes := maps.Keys(agentRegistry)
 
 	return webapi.PluginEntry{
 		ID:                 "agent-service",
@@ -407,15 +409,16 @@ func newAgentPlugin() webapi.PluginEntry {
 				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))
 }