From 61e75da759bbbdaa0cfd3ccbff0cd495fcfc5910 Mon Sep 17 00:00:00 2001 From: Future-Outlier Date: Tue, 3 Sep 2024 13:11:26 +0800 Subject: [PATCH] [flyteagent] Add Logging for Agent Supported Task Types (#5718) * Add Logging for Agent-Supported Task Types Signed-off-by: Future-Outlier * use make(map[string]struct{}) Signed-off-by: Future-Outlier --------- Signed-off-by: Future-Outlier --- flyteplugins/go/tasks/plugins/webapi/agent/client.go | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/flyteplugins/go/tasks/plugins/webapi/agent/client.go b/flyteplugins/go/tasks/plugins/webapi/agent/client.go index 35e6662107..148113fb38 100644 --- a/flyteplugins/go/tasks/plugins/webapi/agent/client.go +++ b/flyteplugins/go/tasks/plugins/webapi/agent/client.go @@ -3,6 +3,7 @@ package agent import ( "context" "crypto/x509" + "strings" "golang.org/x/exp/maps" "google.golang.org/grpc" @@ -127,19 +128,26 @@ func getAgentRegistry(ctx context.Context, cs *ClientSet) Registry { continue } + agentSupportedTaskCategories := make(map[string]struct{}) for _, agent := range res.GetAgents() { deprecatedSupportedTaskTypes := agent.SupportedTaskTypes for _, supportedTaskType := range deprecatedSupportedTaskTypes { agent := &Agent{AgentDeployment: agentDeployment, IsSync: agent.IsSync} newAgentRegistry[supportedTaskType] = map[int32]*Agent{defaultTaskTypeVersion: agent} + agentSupportedTaskCategories[supportedTaskType] = struct{}{} } supportedTaskCategories := agent.SupportedTaskCategories for _, supportedCategory := range supportedTaskCategories { agent := &Agent{AgentDeployment: agentDeployment, IsSync: agent.IsSync} - newAgentRegistry[supportedCategory.GetName()] = map[int32]*Agent{supportedCategory.GetVersion(): agent} + supportedCategoryName := supportedCategory.GetName() + newAgentRegistry[supportedCategoryName] = map[int32]*Agent{supportedCategory.GetVersion(): agent} + agentSupportedTaskCategories[supportedCategoryName] = struct{}{} } + } + logger.Infof(ctx, "AgentDeployment [%v] supports the following task types: [%v]", agentDeployment.Endpoint, + strings.Join(maps.Keys(agentSupportedTaskCategories), ", ")) } // If the agent doesn't implement the metadata service, we construct the registry based on the configuration @@ -160,6 +168,7 @@ func getAgentRegistry(ctx context.Context, cs *ClientSet) Registry { } } + logger.Infof(ctx, "AgentDeployments support the following task types: [%v]", strings.Join(maps.Keys(newAgentRegistry), ", ")) return newAgentRegistry }