Add name to TemplateLogPlugin and use that to identify dynamic log links

Signed-off-by: Eduardo Apolinario <[email protected]>
flyteorg · Jan 26, 2024 · 0a42ca1 · 0a42ca1
1 parent 59dd968
commit 0a42ca1
Show file tree

Hide file tree

Showing 7 changed files with 32 additions and 41 deletions.
diff --git a/flyteplugins/go/tasks/logs/config.go b/flyteplugins/go/tasks/logs/config.go
@@ -28,7 +28,7 @@ type LogConfig struct {
 	StackdriverLogResourceName string              `json:"stackdriver-logresourcename" pflag:",Name of the logresource in stackdriver"`
 	StackDriverTemplateURI     tasklog.TemplateURI `json:"stackdriver-template-uri" pflag:",Template Uri to use when building stackdriver log links"`
 
-	DynamicLogLinks          map[string]tasklog.TemplateURI `json:"dynamic-log-links" pflag:",Map of dynamic log links"`
+	DynamicLogLinks map[string]tasklog.TemplateURI `json:"dynamic-log-links" pflag:",Map of dynamic log links"`
 
 	Templates []tasklog.TemplateLogPlugin `json:"templates" pflag:"-,"`
 }

diff --git a/flyteplugins/go/tasks/logs/logging_utils.go b/flyteplugins/go/tasks/logs/logging_utils.go
@@ -115,10 +115,12 @@ func InitializeLogPlugins(cfg *LogConfig) (tasklog.Plugin, error) {
 		dynamicPlugins = append(
 			dynamicPlugins,
 			tasklog.TemplateLogPlugin{
+				Name:        logLinkType,
 				DisplayName: fmt.Sprintf("%s logs", logLinkType),
 				Scheme:      tasklog.TemplateSchemeDynamic,
-				DynamicTemplateURIs: []tasklog.DynamicTemplateURI{
-					{TemplateURI: dynamicLogLink, Kind: logLinkType}},
+				DynamicTemplateURIs: []tasklog.TemplateURI{
+					dynamicLogLink,
+				},
 				MessageFormat: core.TaskLog_JSON,
 			})
 	}

diff --git a/flyteplugins/go/tasks/logs/logging_utils_test.go b/flyteplugins/go/tasks/logs/logging_utils_test.go
@@ -362,7 +362,7 @@ func TestGetLogsForContainerInPod_Flyin(t *testing.T) {
 			"Flyin enabled but no task template",
 			&LogConfig{
 				DynamicLogLinks: map[string]tasklog.TemplateURI{
-					"flyin": "https://flyin.mydomain.com:{{ .taskConfig.port }}/{{ .namespace }}/{{ .podName }}/{{ .containerName }}/{{ .containerId }}",
+					"vscode": "https://flyin.mydomain.com:{{ .taskConfig.port }}/{{ .namespace }}/{{ .podName }}/{{ .containerName }}/{{ .containerId }}",
 				},
 			},
 			nil,
@@ -372,7 +372,7 @@ func TestGetLogsForContainerInPod_Flyin(t *testing.T) {
 			"Flyin enabled but config not found in task template",
 			&LogConfig{
 				DynamicLogLinks: map[string]tasklog.TemplateURI{
-					"flyin": "https://flyin.mydomain.com:{{ .taskConfig.port }}/{{ .namespace }}/{{ .podName }}/{{ .containerName }}/{{ .containerId }}",
+					"vscode": "https://flyin.mydomain.com:{{ .taskConfig.port }}/{{ .namespace }}/{{ .podName }}/{{ .containerName }}/{{ .containerId }}",
 				},
 			},
 			&core.TaskTemplate{},
@@ -382,7 +382,7 @@ func TestGetLogsForContainerInPod_Flyin(t *testing.T) {
 			"Flyin enabled but no port in task template",
 			&LogConfig{
 				DynamicLogLinks: map[string]tasklog.TemplateURI{
-					"flyin": "https://flyin.mydomain.com:{{ .taskConfig.port }}/{{ .namespace }}/{{ .podName }}/{{ .containerName }}/{{ .containerId }}",
+					"vscode": "https://flyin.mydomain.com:{{ .taskConfig.port }}/{{ .namespace }}/{{ .podName }}/{{ .containerName }}/{{ .containerId }}",
 				},
 			},
 			&core.TaskTemplate{
@@ -394,14 +394,13 @@ func TestGetLogsForContainerInPod_Flyin(t *testing.T) {
 				{
 					Uri:           "https://flyin.mydomain.com:8080/my-namespace/my-pod/ContainerName/ContainerID",
 					MessageFormat: core.TaskLog_JSON,
-					Name:          "flyin logs my-Suffix",
+					Name:          "vscode logs my-Suffix",
 				},
 			},
 		},
 		{
 			"Flyin disabled but config present in TaskTemplate",
-			&LogConfig{
-			},
+			&LogConfig{},
 			&core.TaskTemplate{
 				Config: map[string]string{
 					"link_type": "vscode",
@@ -413,8 +412,8 @@ func TestGetLogsForContainerInPod_Flyin(t *testing.T) {
 		{
 			"Flyin disabled and K8s enabled and flyin config present in TaskTemplate",
 			&LogConfig{
-				IsKubernetesEnabled:      true,
-				KubernetesTemplateURI:    "https://k8s.com/{{ .namespace }}/{{ .podName }}/{{ .containerName }}/{{ .containerId }}",
+				IsKubernetesEnabled:   true,
+				KubernetesTemplateURI: "https://k8s.com/{{ .namespace }}/{{ .podName }}/{{ .containerName }}/{{ .containerId }}",
 			},
 			&core.TaskTemplate{
 				Config: map[string]string{
@@ -433,10 +432,10 @@ func TestGetLogsForContainerInPod_Flyin(t *testing.T) {
 		{
 			"Flyin and K8s enabled",
 			&LogConfig{
-				IsKubernetesEnabled:      true,
-				KubernetesTemplateURI:    "https://k8s.com/{{ .namespace }}/{{ .podName }}/{{ .containerName }}/{{ .containerId }}",
+				IsKubernetesEnabled:   true,
+				KubernetesTemplateURI: "https://k8s.com/{{ .namespace }}/{{ .podName }}/{{ .containerName }}/{{ .containerId }}",
 				DynamicLogLinks: map[string]tasklog.TemplateURI{
-					"flyin": "https://flyin.mydomain.com:{{ .taskConfig.port }}/{{ .namespace }}/{{ .podName }}/{{ .containerName }}/{{ .containerId }}",
+					"vscode": "https://flyin.mydomain.com:{{ .taskConfig.port }}/{{ .namespace }}/{{ .podName }}/{{ .containerName }}/{{ .containerId }}",
 				},
 			},
 			&core.TaskTemplate{
@@ -454,7 +453,7 @@ func TestGetLogsForContainerInPod_Flyin(t *testing.T) {
 				{
 					Uri:           "https://flyin.mydomain.com:65535/my-namespace/my-pod/ContainerName/ContainerID",
 					MessageFormat: core.TaskLog_JSON,
-					Name:          "flyin logs my-Suffix",
+					Name:          "vscode logs my-Suffix",
 				},
 			},
 		},

diff --git a/flyteplugins/go/tasks/pluginmachinery/tasklog/plugin.go b/flyteplugins/go/tasks/pluginmachinery/tasklog/plugin.go
@@ -20,11 +20,6 @@ const (
 // TemplateURI is a URI that accepts templates. See: go/tasks/pluginmachinery/tasklog/template.go for available templates.
 type TemplateURI = string
 
-type DynamicTemplateURI struct {
-	TemplateURI TemplateURI
-	Kind        string // Kind describes the kind of the dynamic template. Currently only "flyin" is supported.
-}
-
 type TemplateVar struct {
 	Regex *regexp.Regexp
 	Value string
@@ -69,9 +64,10 @@ type Plugin interface {
 }
 
 type TemplateLogPlugin struct {
+	Name                string                     `json:"name" pflag:",Name of the plugin."`
 	DisplayName         string                     `json:"displayName" pflag:",Display name for the generated log when displayed in the console."`
 	TemplateURIs        []TemplateURI              `json:"templateUris" pflag:",URI Templates for generating task log links."`
-	DynamicTemplateURIs []DynamicTemplateURI       `json:"dynamicTemplateUris" pflag:",Dynamic URI Templates for generating task log links."`
+	DynamicTemplateURIs []TemplateURI              `json:"dynamictemplateUris" pflag:",URI Templates for generating dynamic task log links."`
 	MessageFormat       core.TaskLog_MessageFormat `json:"messageFormat" pflag:",Log Message Format."`
 	Scheme              TemplateScheme             `json:"scheme" pflag:",Templating scheme to use. Supported values are Pod and TaskExecution."`
 }
diff --git a/flyteplugins/go/tasks/pluginmachinery/tasklog/template.go b/flyteplugins/go/tasks/pluginmachinery/tasklog/template.go
@@ -212,17 +212,7 @@ func getDynamicLogLinkTypes(taskTemplate *core.TaskTemplate) []string {
 	if linkType == "" {
 		return nil
 	}
-	dynamicLogLinkTypes := []string{}
-	for _, linkType := range strings.Split(linkType, ",") {
-		// NB: we are going to grandfather in flyin.
-		if linkType == "vscode" {
-			dynamicLogLinkTypes = append(dynamicLogLinkTypes, "flyin")
-		} else {
-			dynamicLogLinkTypes = append(dynamicLogLinkTypes, linkType)
-		}
-	}
-
-	return dynamicLogLinkTypes
+	return strings.Split(linkType, ",")
 }
 
 func (p TemplateLogPlugin) GetTaskLogs(input Input) (Output, error) {
@@ -238,9 +228,9 @@ func (p TemplateLogPlugin) GetTaskLogs(input Input) (Output, error) {
 
 	for _, dynamicLogLinkType := range getDynamicLogLinkTypes(input.TaskTemplate) {
 		for _, dynamicTemplateURI := range p.DynamicTemplateURIs {
-			if dynamicTemplateURI.Kind == dynamicLogLinkType {
+			if p.Name == dynamicLogLinkType {
 				taskLogs = append(taskLogs, &core.TaskLog{
-					Uri:           replaceAll(dynamicTemplateURI.TemplateURI, templateVars),
+					Uri:           replaceAll(dynamicTemplateURI, templateVars),
 					Name:          p.DisplayName + input.LogName,
 					MessageFormat: p.MessageFormat,
 				})

diff --git a/flyteplugins/go/tasks/pluginmachinery/tasklog/template_test.go b/flyteplugins/go/tasks/pluginmachinery/tasklog/template_test.go
@@ -548,8 +548,9 @@ func TestTemplateLogPlugin(t *testing.T) {
 		{
 			"flyin",
 			TemplateLogPlugin{
+				Name:                "vscode",
 				Scheme:              TemplateSchemeDynamic,
-				DynamicTemplateURIs: []DynamicTemplateURI{{"vscode://flyin:{{ .taskConfig.port }}/{{ .podName }}", "flyin"}},
+				DynamicTemplateURIs: []TemplateURI{"vscode://flyin:{{ .taskConfig.port }}/{{ .podName }}"},
 				MessageFormat:       core.TaskLog_JSON,
 			},
 			args{
@@ -575,9 +576,10 @@ func TestTemplateLogPlugin(t *testing.T) {
 		{
 			"flyin - default port",
 			TemplateLogPlugin{
-				Scheme:        TemplateSchemeDynamic,
-				TemplateURIs:  []TemplateURI{"vscode://flyin:{{ .taskConfig.port }}/{{ .podName }}"},
-				MessageFormat: core.TaskLog_JSON,
+				Name:                "vscode",
+				Scheme:              TemplateSchemeDynamic,
+				DynamicTemplateURIs: []TemplateURI{"vscode://flyin:{{ .taskConfig.port }}/{{ .podName }}"},
+				MessageFormat:       core.TaskLog_JSON,
 			},
 			args{
 				input: Input{
@@ -601,8 +603,9 @@ func TestTemplateLogPlugin(t *testing.T) {
 		{
 			"flyin - no link_type in task template",
 			TemplateLogPlugin{
+				Name:                "vscode",
 				Scheme:              TemplateSchemeDynamic,
-				DynamicTemplateURIs: []DynamicTemplateURI{{"vscode://flyin:{{ .taskConfig.port }}/{{ .podName }}", "flyin"}},
+				DynamicTemplateURIs: []TemplateURI{"vscode://flyin:{{ .taskConfig.port }}/{{ .podName }}"},
 				MessageFormat:       core.TaskLog_JSON,
 				DisplayName:         "Flyin Logs",
 			},

diff --git a/flyteplugins/go/tasks/plugins/k8s/kfoperators/common/common_operator_test.go b/flyteplugins/go/tasks/plugins/k8s/kfoperators/common/common_operator_test.go
@@ -211,8 +211,9 @@ func TestGetLogsTemplateUri(t *testing.T) {
 
 	taskCtx := dummyTaskContext()
 	pytorchJobObjectMeta := meta_v1.ObjectMeta{
-		Name:      "test",
-		Namespace: "pytorch-namespace",
+		Name: "test",
+		Namespace: "pytorch-" +
+			"namespace",
 		CreationTimestamp: meta_v1.Time{
 			Time: time.Date(2022, time.January, 1, 12, 0, 0, 0, time.UTC),
 		},