Add podUID as a template input in tasklog plugin (flyteorg#304)

Signed-off-by: Ankit Goyal <[email protected]> Signed-off-by: Ankit Goyal <[email protected]>
eapolinario · Jan 4, 2023 · 9eeb694 · 9eeb694
1 parent 8efeda7
commit 9eeb694
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 4 deletions.
diff --git a/go/tasks/logs/logging_utils.go b/go/tasks/logs/logging_utils.go
@@ -41,6 +41,7 @@ func GetLogsForContainerInPod(ctx context.Context, logPlugin tasklog.Plugin, pod
 	logs, err := logPlugin.GetTaskLogs(
 		tasklog.Input{
 			PodName:           pod.Name,
+			PodUID:            string(pod.GetUID()),
 			Namespace:         pod.Namespace,
 			ContainerName:     pod.Spec.Containers[index].Name,
 			ContainerID:       pod.Status.ContainerStatuses[index].ContainerID,

diff --git a/go/tasks/pluginmachinery/tasklog/plugin.go b/go/tasks/pluginmachinery/tasklog/plugin.go
@@ -13,6 +13,7 @@ type Input struct {
 	LogName           string `json:"logName"`
 	PodUnixStartTime  int64  `json:"podUnixStartTime"`
 	PodUnixFinishTime int64  `json:"podUnixFinishTime"`
+	PodUID            string `json:"podUID"`
 }
 
 // Output contains all task logs a plugin generates for a given Input.

diff --git a/go/tasks/pluginmachinery/tasklog/template.go b/go/tasks/pluginmachinery/tasklog/template.go
@@ -30,6 +30,7 @@ type regexValPair struct {
 
 type templateRegexes struct {
 	PodName           *regexp.Regexp
+	PodUID            *regexp.Regexp
 	Namespace         *regexp.Regexp
 	ContainerName     *regexp.Regexp
 	ContainerID       *regexp.Regexp
@@ -42,6 +43,7 @@ type templateRegexes struct {
 func mustInitTemplateRegexes() templateRegexes {
 	return templateRegexes{
 		PodName:           mustCreateRegex("podName"),
+		PodUID:            mustCreateRegex("podUID"),
 		Namespace:         mustCreateRegex("namespace"),
 		ContainerName:     mustCreateRegex("containerName"),
 		ContainerID:       mustCreateRegex("containerID"),
@@ -66,11 +68,12 @@ func replaceAll(template string, values []regexValPair) string {
 	return template
 }
 
-func (s TemplateLogPlugin) GetTaskLog(podName, namespace, containerName, containerID, logName string, podUnixStartTime, podUnixFinishTime int64) (core.TaskLog, error) {
+func (s TemplateLogPlugin) GetTaskLog(podName, podUID, namespace, containerName, containerID, logName string, podUnixStartTime, podUnixFinishTime int64) (core.TaskLog, error) {
 	o, err := s.GetTaskLogs(Input{
 		LogName:           logName,
 		Namespace:         namespace,
 		PodName:           podName,
+		PodUID:            podUID,
 		ContainerName:     containerName,
 		ContainerID:       containerID,
 		PodUnixStartTime:  podUnixStartTime,
@@ -103,6 +106,10 @@ func (s TemplateLogPlugin) GetTaskLogs(input Input) (Output, error) {
 						regex: regexes.PodName,
 						val:   input.PodName,
 					},
+					{
+						regex: regexes.PodUID,
+						val:   input.PodUID,
+					},
 					{
 						regex: regexes.Namespace,
 						val:   input.Namespace,

diff --git a/go/tasks/pluginmachinery/tasklog/template_test.go b/go/tasks/pluginmachinery/tasklog/template_test.go
@@ -11,9 +11,10 @@ import (
 )
 
 func TestTemplateLog(t *testing.T) {
-	p := NewTemplateLogPlugin([]string{"https://console.aws.amazon.com/cloudwatch/home?region=us-east-1#logEventViewer:group=/flyte-production/kubernetes;stream=var.log.containers.{{.podName}}_{{.namespace}}_{{.containerName}}-{{.containerId}}.log"}, core.TaskLog_JSON)
+	p := NewTemplateLogPlugin([]string{"https://console.aws.amazon.com/cloudwatch/home?region=us-east-1#logEventViewer:group=/flyte-production/kubernetes;stream=var.log.containers.{{.podName}}_{{.podUID}}_{{.namespace}}_{{.containerName}}-{{.containerId}}.log"}, core.TaskLog_JSON)
 	tl, err := p.GetTaskLog(
 		"f-uuid-driver",
+		"pod-uid",
 		"flyteexamples-production",
 		"spark-kubernetes-driver",
 		"cri-o://abc",
@@ -24,7 +25,7 @@ func TestTemplateLog(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, tl.GetName(), "main_logs")
 	assert.Equal(t, tl.GetMessageFormat(), core.TaskLog_JSON)
-	assert.Equal(t, "https://console.aws.amazon.com/cloudwatch/home?region=us-east-1#logEventViewer:group=/flyte-production/kubernetes;stream=var.log.containers.f-uuid-driver_flyteexamples-production_spark-kubernetes-driver-abc.log", tl.Uri)
+	assert.Equal(t, "https://console.aws.amazon.com/cloudwatch/home?region=us-east-1#logEventViewer:group=/flyte-production/kubernetes;stream=var.log.containers.f-uuid-driver_pod-uid_flyteexamples-production_spark-kubernetes-driver-abc.log", tl.Uri)
 }
 
 // Latest Run: Benchmark_mustInitTemplateRegexes-16    	   45960	     26914 ns/op
@@ -41,6 +42,7 @@ func Test_templateLogPlugin_Regression(t *testing.T) {
 	}
 	type args struct {
 		podName           string
+		podUID            string
 		namespace         string
 		containerName     string
 		containerID       string
@@ -63,6 +65,7 @@ func Test_templateLogPlugin_Regression(t *testing.T) {
 			},
 			args{
 				podName:           "f-uuid-driver",
+				podUID:            "pod-uid",
 				namespace:         "flyteexamples-production",
 				containerName:     "spark-kubernetes-driver",
 				containerID:       "cri-o://abc",
@@ -85,6 +88,7 @@ func Test_templateLogPlugin_Regression(t *testing.T) {
 			},
 			args{
 				podName:           "podName",
+				podUID:            "pod-uid",
 				namespace:         "flyteexamples-production",
 				containerName:     "spark-kubernetes-driver",
 				containerID:       "cri-o://abc",
@@ -107,6 +111,7 @@ func Test_templateLogPlugin_Regression(t *testing.T) {
 			},
 			args{
 				podName:           "flyteexamples-development-task-name",
+				podUID:            "pod-uid",
 				namespace:         "flyteexamples-development",
 				containerName:     "ignore",
 				containerID:       "ignore",
@@ -129,7 +134,7 @@ func Test_templateLogPlugin_Regression(t *testing.T) {
 				messageFormat: tt.fields.messageFormat,
 			}
 
-			got, err := s.GetTaskLog(tt.args.podName, tt.args.namespace, tt.args.containerName, tt.args.containerID, tt.args.logName, tt.args.podUnixStartTime, tt.args.podUnixFinishTime)
+			got, err := s.GetTaskLog(tt.args.podName, tt.args.podUID, tt.args.namespace, tt.args.containerName, tt.args.containerID, tt.args.logName, tt.args.podUnixStartTime, tt.args.podUnixFinishTime)
 			if (err != nil) != tt.wantErr {
 				t.Errorf("GetTaskLog() error = %v, wantErr %v", err, tt.wantErr)
 				return