From 6ca8fece3f53de020d7901fb1b164ff9761888fc Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario Date: Fri, 26 Jan 2024 17:03:43 -0800 Subject: [PATCH] Evaluate templatizable taskConfig values at runtime Signed-off-by: Eduardo Apolinario --- .../tasks/pluginmachinery/tasklog/template.go | 21 ++++----- .../pluginmachinery/tasklog/template_test.go | 45 +++++++------------ 2 files changed, 26 insertions(+), 40 deletions(-) diff --git a/flyteplugins/go/tasks/pluginmachinery/tasklog/template.go b/flyteplugins/go/tasks/pluginmachinery/tasklog/template.go index aa386ef6a9..0de13d432c 100644 --- a/flyteplugins/go/tasks/pluginmachinery/tasklog/template.go +++ b/flyteplugins/go/tasks/pluginmachinery/tasklog/template.go @@ -41,7 +41,6 @@ type templateRegexes struct { ExecutionProject *regexp.Regexp ExecutionDomain *regexp.Regexp GeneratedName *regexp.Regexp - Port *regexp.Regexp } func initDefaultRegexes() templateRegexes { @@ -67,7 +66,6 @@ func initDefaultRegexes() templateRegexes { MustCreateRegex("executionProject"), MustCreateRegex("executionDomain"), MustCreateRegex("generatedName"), - MustCreateDynamicLogRegex("port"), } } @@ -82,7 +80,7 @@ func replaceAll(template string, vars []TemplateVar) string { return template } -func (input Input) templateVarsForScheme() []TemplateVar { +func (input Input) templateVars() []TemplateVar { vars := []TemplateVar{ TemplateVar{defaultRegexes.LogName, input.LogName}, } @@ -92,14 +90,6 @@ func (input Input) templateVarsForScheme() []TemplateVar { vars = append(vars, input.ExtraTemplateVars...) } - port := input.TaskTemplate.GetConfig()["port"] - if port != "" { - vars = append( - vars, - TemplateVar{defaultRegexes.Port, port}, - ) - } - // Container IDs are prefixed with docker://, cri-o://, etc. which is stripped by fluentd before pushing to a log // stream. Therefore, we must also strip the prefix. containerID := input.ContainerID @@ -193,6 +183,13 @@ func (input Input) templateVarsForScheme() []TemplateVar { }, ) + // Add values from task template config as dynamic regexes (i.e. templates prefixed by .taskConfig) + for key, value := range input.TaskTemplate.GetConfig() { + if value != "" { + vars = append(vars, TemplateVar{MustCreateDynamicLogRegex(key), value}) + } + } + return vars } @@ -212,7 +209,7 @@ func getDynamicLogLinkTypes(taskTemplate *core.TaskTemplate) []string { } func (p TemplateLogPlugin) GetTaskLogs(input Input) (Output, error) { - templateVars := input.templateVarsForScheme() + templateVars := input.templateVars() taskLogs := make([]*core.TaskLog, 0, len(p.TemplateURIs)) for _, templateURI := range p.TemplateURIs { taskLogs = append(taskLogs, &core.TaskLog{ diff --git a/flyteplugins/go/tasks/pluginmachinery/tasklog/template_test.go b/flyteplugins/go/tasks/pluginmachinery/tasklog/template_test.go index e7001003dc..a94af5e23f 100644 --- a/flyteplugins/go/tasks/pluginmachinery/tasklog/template_test.go +++ b/flyteplugins/go/tasks/pluginmachinery/tasklog/template_test.go @@ -43,19 +43,23 @@ func dummyTaskExecID() pluginCore.TaskExecutionID { return tID } -func Test_Input_templateVarsForScheme(t *testing.T) { +func Test_Input_templateVars(t *testing.T) { testRegexes := struct { - Foo *regexp.Regexp - Bar *regexp.Regexp - Baz *regexp.Regexp - Ham *regexp.Regexp - Spam *regexp.Regexp + Foo *regexp.Regexp + Bar *regexp.Regexp + Baz *regexp.Regexp + Ham *regexp.Regexp + Spam *regexp.Regexp + LinkType *regexp.Regexp + Port *regexp.Regexp }{ MustCreateRegex("foo"), MustCreateRegex("bar"), MustCreateRegex("baz"), MustCreateRegex("ham"), MustCreateRegex("spam"), + MustCreateDynamicLogRegex("link_type"), + MustCreateDynamicLogRegex("port"), } podBase := Input{ HostName: "my-host", @@ -93,7 +97,8 @@ func Test_Input_templateVarsForScheme(t *testing.T) { PodUnixFinishTime: 12345, TaskTemplate: &core.TaskTemplate{ Config: map[string]string{ - "port": "1234", + "link_type": "vscode", + "port": "1234", }, }, } @@ -188,30 +193,13 @@ func Test_Input_templateVarsForScheme(t *testing.T) { }, nil, }, - // { - // "task execution with unused extra vars", - // TemplateSchemeTaskExecution, - // taskExecutionBase, - // &TemplateVarsByScheme{ - // Pod: TemplateVars{ - // {testRegexes.Bar, "bar"}, - // {testRegexes.Baz, "baz"}, - // }, - // }, - // nil, - // nil, - // TemplateVars{ - // {testRegexes.Bar, "bar"}, - // {testRegexes.Baz, "baz"}, - // }, - // }, { "flyin happy path", flyinBase, nil, nil, []TemplateVar{ - {defaultRegexes.Port, "1234"}, + {testRegexes.Port, "1234"}, }, nil, }, @@ -221,7 +209,6 @@ func Test_Input_templateVarsForScheme(t *testing.T) { nil, []TemplateVar{ {defaultRegexes.LogName, "main_logs"}, - {defaultRegexes.Port, "1234"}, {defaultRegexes.PodName, "my-pod"}, {defaultRegexes.PodUID, "my-pod-uid"}, {defaultRegexes.Namespace, "my-namespace"}, @@ -232,6 +219,8 @@ func Test_Input_templateVarsForScheme(t *testing.T) { {defaultRegexes.PodRFC3339FinishTime, "1970-01-01T04:25:45+01:00"}, {defaultRegexes.PodUnixStartTime, "123"}, {defaultRegexes.PodUnixFinishTime, "12345"}, + {testRegexes.LinkType, "vscode"}, + {testRegexes.Port, "1234"}, }, nil, nil, @@ -243,7 +232,7 @@ func Test_Input_templateVarsForScheme(t *testing.T) { nil, nil, []TemplateVar{ - {defaultRegexes.Port, "1234"}, + {testRegexes.Port, "1234"}, }, }, } @@ -251,7 +240,7 @@ func Test_Input_templateVarsForScheme(t *testing.T) { t.Run(tt.name, func(t *testing.T) { base := tt.baseVars base.ExtraTemplateVars = tt.extraVars - got := base.templateVarsForScheme() + got := base.templateVars() if tt.exact != nil { assert.Equal(t, got, tt.exact) }