Added VC_TASK_INDEX and added env to initContainers.

Signed-off-by: Klaus Ma <klaus1982.cn@gmail.com>
volcano-sh · Dec 16, 2019 · ef05963 · ef05963
1 parent 33fa0fd
commit ef05963
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 7 deletions.
diff --git a/pkg/controllers/job/plugins/env/const.go b/pkg/controllers/job/plugins/env/const.go
@@ -22,4 +22,7 @@ const (
 
 	// TaskVkIndex  used as key in container env
 	TaskVkIndex = "VK_TASK_INDEX"
+
+	// TaskIndex is used as key in container env
+	TaskIndex = "VC_TASK_INDEX"
 )
diff --git a/pkg/controllers/job/plugins/env/env.go b/pkg/controllers/job/plugins/env/env.go
@@ -43,13 +43,18 @@ func (ep *envPlugin) Name() string {
 }
 
 func (ep *envPlugin) OnPodCreate(pod *v1.Pod, job *batch.Job) error {
-	// add VK_TASK_INDEX env to each container
-	for i, c := range pod.Spec.Containers {
-		vcIndex := v1.EnvVar{
-			Name:  TaskVkIndex,
-			Value: jobhelpers.GetTaskIndex(pod),
-		}
-		pod.Spec.Containers[i].Env = append(c.Env, vcIndex)
+	index := jobhelpers.GetTaskIndex(pod)
+
+	// add VK_TASK_INDEX and VC_TASK_INDEX env to each container
+	for i := range pod.Spec.Containers {
+		pod.Spec.Containers[i].Env = append(pod.Spec.Containers[i].Env, v1.EnvVar{Name: TaskVkIndex, Value: index})
+		pod.Spec.Containers[i].Env = append(pod.Spec.Containers[i].Env, v1.EnvVar{Name: TaskIndex, Value: index})
+	}
+
+	// add VK_TASK_INDEX and VC_TASK_INDEX env to each init container
+	for i := range pod.Spec.InitContainers {
+		pod.Spec.InitContainers[i].Env = append(pod.Spec.InitContainers[i].Env, v1.EnvVar{Name: TaskVkIndex, Value: index})
+		pod.Spec.InitContainers[i].Env = append(pod.Spec.InitContainers[i].Env, v1.EnvVar{Name: TaskIndex, Value: index})
 	}
 
 	return nil