refactor: Adjust CoPilot init and Pod status logic (#6523)

Signed-off-by: Kevin Su <[email protected]>

Author: pingsutw

flyteplugins/go/tasks/pluginmachinery/flytek8s/copilot.go

@@ -271,7 +271,8 @@ func AddCoPilotToPod(ctx context.Context, cfg config.FlyteCoPilotConfig, coPilot
 			if err != nil {
 				return primaryInitContainerName, err
 			}
-			coPilotPod.InitContainers = append(coPilotPod.InitContainers, sidecar)
+			// Let the sidecar container start before the downloader; it will ensure the signal watcher is started before the main container finishes.
+			coPilotPod.InitContainers = append([]v1.Container{sidecar}, coPilotPod.InitContainers...)
 
 			coPilotPod.TerminationGracePeriodSeconds = (*int64)(&cfg.Timeout.Duration)
 		}

flyteplugins/go/tasks/pluginmachinery/flytek8s/copilot_test.go

@@ -228,7 +228,8 @@ func assertContainerHasVolumeMounts(t *testing.T, cfg config.FlyteCoPilotConfig,
 }
 
 func assertPodHasCoPilot(t *testing.T, cfg config.FlyteCoPilotConfig, pilot *core.DataLoadingConfig, iFace *core.TypedInterface, pod *v1.PodSpec) {
-	for _, c := range pod.Containers {
+	containers := append(pod.Containers, pod.InitContainers...)
+	for _, c := range containers {
 		if c.Name == "test" {
 			cntr := c
 			assertContainerHasVolumeMounts(t, cfg, pilot, iFace, &cntr)
@@ -510,6 +511,8 @@ func TestAddCoPilotToPod(t *testing.T) {
 		primaryInitContainerName, err := AddCoPilotToPod(ctx, cfg, &pod, iface, taskMetadata, inputPaths, opath, pilot)
 		assert.NoError(t, err)
 		assert.Equal(t, "test-downloader", primaryInitContainerName)
+		assert.Equal(t, pod.InitContainers[0].Name, cfg.NamePrefix+flyteSidecarContainerName)
+		assert.Equal(t, pod.InitContainers[1].Name, cfg.NamePrefix+flyteDownloaderContainerName)
 		assertPodHasCoPilot(t, cfg, pilot, iface, &pod)
 	})
 

flyteplugins/go/tasks/plugins/k8s/pod/plugin.go

@@ -208,7 +208,7 @@ func DemystifyPodStatus(ctx context.Context, pod *v1.Pod, info pluginsCore.TaskI
 		// DO NOTHING
 	default:
 		if !primaryContainerExists {
-			// if all of the containers in the Pod are complete, as an optimization, we can declare the task as
+			// if all of the containers and sidecars in the Pod are complete, as an optimization, we can declare the task as
 			// succeeded rather than waiting for the Pod to be marked completed.
 			allSuccessfullyTerminated := len(pod.Status.ContainerStatuses) > 0
 			for _, s := range pod.Status.ContainerStatuses {
@@ -217,6 +217,14 @@ func DemystifyPodStatus(ctx context.Context, pod *v1.Pod, info pluginsCore.TaskI
 				}
 			}
 
+			// Init container will become sidecar if the restart policy is set to Always
+			// https://kubernetes.io/docs/concepts/workloads/pods/sidecar-containers/#sidecar-containers-and-pod-lifecycle
+			for _, s := range pod.Status.InitContainerStatuses {
+				if s.State.Waiting != nil || s.State.Running != nil || (s.State.Terminated != nil && s.State.Terminated.ExitCode != 0) {
+					allSuccessfullyTerminated = false
+				}
+			}
+
 			if allSuccessfullyTerminated {
 				return flytek8s.DemystifySuccess(pod.Status, info)
 			}