add hostname to input (#4939)

Signed-off-by: novahow <[email protected]>
flyteorg · Mar 4, 2024 · cf21877 · cf21877
1 parent a946070
commit cf21877
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 5 deletions.
diff --git a/flyteplugins/go/tasks/logs/logging_utils.go b/flyteplugins/go/tasks/logs/logging_utils.go
@@ -52,6 +52,7 @@ func GetLogsForContainerInPod(ctx context.Context, logPlugin tasklog.Plugin, tas
 			TaskExecutionID:      taskExecID,
 			ExtraTemplateVars:    extraLogTemplateVars,
 			TaskTemplate:         taskTemplate,
+			HostName:             pod.Spec.Hostname,
 		},
 	)
 

diff --git a/flyteplugins/go/tasks/logs/logging_utils_test.go b/flyteplugins/go/tasks/logs/logging_utils_test.go
@@ -210,6 +210,40 @@ func TestGetLogsForContainerInPod_All(t *testing.T) {
 	assert.Len(t, logs, 2)
 }
 
+func TestGetLogsForContainerInPod_HostName(t *testing.T) {
+	logPlugin, err := InitializeLogPlugins(&LogConfig{
+		IsKubernetesEnabled: true,
+		KubernetesURL:       "k8s.com",
+		IsCloudwatchEnabled: true,
+		CloudwatchRegion:    "us-east-1",
+		CloudwatchLogGroup:  "/kubernetes/flyte-production",
+	})
+	assert.NoError(t, err)
+
+	pod := &v1.Pod{
+		Spec: v1.PodSpec{
+			Containers: []v1.Container{
+				{
+					Name: "ContainerName",
+				},
+			},
+			Hostname: "my-hostname",
+		},
+		Status: v1.PodStatus{
+			ContainerStatuses: []v1.ContainerStatus{
+				{
+					ContainerID: "ContainerID",
+				},
+			},
+		},
+	}
+	pod.Name = podName
+
+	logs, err := GetLogsForContainerInPod(context.TODO(), logPlugin, dummyTaskExecID(), pod, 0, " Suffix", nil, nil)
+	assert.Nil(t, err)
+	assert.Len(t, logs, 2)
+}
+
 func TestGetLogsForContainerInPod_Stackdriver(t *testing.T) {
 	logPlugin, err := InitializeLogPlugins(&LogConfig{
 		IsStackDriverEnabled:       true,
@@ -268,7 +302,7 @@ func TestGetLogsForContainerInPod_LegacyTemplate(t *testing.T) {
 				MessageFormat: core.TaskLog_JSON,
 				Name:          "Stackdriver Logs my-Suffix",
 			},
-		})
+		}, "")
 	})
 
 	t.Run("StackDriver", func(t *testing.T) {
@@ -281,11 +315,11 @@ func TestGetLogsForContainerInPod_LegacyTemplate(t *testing.T) {
 				MessageFormat: core.TaskLog_JSON,
 				Name:          "Stackdriver Logs my-Suffix",
 			},
-		})
+		}, "")
 	})
 }
 
-func assertTestSucceeded(tb testing.TB, config *LogConfig, taskTemplate *core.TaskTemplate, expectedTaskLogs []*core.TaskLog) {
+func assertTestSucceeded(tb testing.TB, config *LogConfig, taskTemplate *core.TaskTemplate, expectedTaskLogs []*core.TaskLog, hostname string) {
 	logPlugin, err := InitializeLogPlugins(config)
 	assert.NoError(tb, err)
 
@@ -300,6 +334,7 @@ func assertTestSucceeded(tb testing.TB, config *LogConfig, taskTemplate *core.Ta
 					Name: "ContainerName",
 				},
 			},
+			Hostname: hostname,
 		},
 		Status: v1.PodStatus{
 			ContainerStatuses: []v1.ContainerStatus{
@@ -347,7 +382,27 @@ func TestGetLogsForContainerInPod_Templates(t *testing.T) {
 			MessageFormat: core.TaskLog_JSON,
 			Name:          "Internal my-Suffix",
 		},
-	})
+	}, "")
+}
+
+func TestGetLogsForContainerInPodTemplates_Hostname(t *testing.T) {
+	assertTestSucceeded(t, &LogConfig{
+		Templates: []tasklog.TemplateLogPlugin{
+			{
+				DisplayName: "StackDriver",
+				TemplateURIs: []string{
+					"{{ .hostname }}/{{ .namespace }}/{{ .podName }}/{{ .containerName }}/{{ .containerId }}",
+				},
+				MessageFormat: core.TaskLog_JSON,
+			},
+		},
+	}, nil, []*core.TaskLog{
+		{
+			Uri:           "my-hostname/my-namespace/my-pod/ContainerName/ContainerID",
+			MessageFormat: core.TaskLog_JSON,
+			Name:          "StackDriver my-Suffix",
+		},
+	}, "my-hostname")
 }
 
 func TestGetLogsForContainerInPod_Flyteinteractive(t *testing.T) {
@@ -507,7 +562,7 @@ func TestGetLogsForContainerInPod_Flyteinteractive(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			assertTestSucceeded(t, tt.config, tt.template, tt.expectedTaskLogs)
+			assertTestSucceeded(t, tt.config, tt.template, tt.expectedTaskLogs, "")
 		})
 	}
 }