feat: support imagePullSecrets in TrainJob pod spec overrides (#2806)

Signed-off-by: Tsubasa Nagasawa <[email protected]>
Co-authored-by: Antonin Stefanutti <[email protected]>

api/openapi-spec/swagger.json

@@ -13919,6 +13919,22 @@
             ],
             "x-kubernetes-list-type": "map"
           },
+          "imagePullSecrets": {
+            "description": "ImagePullSecrets overrides the image pull secrets for the Pods in the target job templates.",
+            "type": "array",
+            "items": {
+              "default": {},
+              "allOf": [
+                {
+                  "$ref": "#/components/schemas/io.k8s.api.core.v1.LocalObjectReference"
+                }
+              ]
+            },
+            "x-kubernetes-list-map-keys": [
+              "name"
+            ],
+            "x-kubernetes-list-type": "map"
+          },
           "initContainers": {
             "description": "Overrides for the init container in the target job templates.",
             "type": "array",

api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_pod_spec_override.py

@@ -20,6 +20,7 @@ import json
 from pydantic import BaseModel, ConfigDict, Field, StrictStr
 from typing import Any, ClassVar, Dict, List, Optional
 from kubeflow_trainer_api.models.io_k8s_api_core_v1_affinity import IoK8sApiCoreV1Affinity
+from kubeflow_trainer_api.models.io_k8s_api_core_v1_local_object_reference import IoK8sApiCoreV1LocalObjectReference
 from kubeflow_trainer_api.models.io_k8s_api_core_v1_pod_scheduling_gate import IoK8sApiCoreV1PodSchedulingGate
 from kubeflow_trainer_api.models.io_k8s_api_core_v1_toleration import IoK8sApiCoreV1Toleration
 from kubeflow_trainer_api.models.io_k8s_api_core_v1_volume import IoK8sApiCoreV1Volume
@@ -34,6 +35,7 @@ class TrainerV1alpha1PodSpecOverride(BaseModel):
     """ # noqa: E501
     affinity: Optional[IoK8sApiCoreV1Affinity] = Field(default=None, description="Override for the Pod's affinity.")
     containers: Optional[List[TrainerV1alpha1ContainerOverride]] = Field(default=None, description="Overrides for the containers in the target job templates.")
+    image_pull_secrets: Optional[List[IoK8sApiCoreV1LocalObjectReference]] = Field(default=None, description="ImagePullSecrets overrides the image pull secrets for the Pods in the target job templates.", alias="imagePullSecrets")
     init_containers: Optional[List[TrainerV1alpha1ContainerOverride]] = Field(default=None, description="Overrides for the init container in the target job templates.", alias="initContainers")
     node_selector: Optional[Dict[str, StrictStr]] = Field(default=None, description="Override for the node selector to place Pod on the specific node.", alias="nodeSelector")
     scheduling_gates: Optional[List[IoK8sApiCoreV1PodSchedulingGate]] = Field(default=None, description="SchedulingGates overrides the scheduling gates of the Pods in the target job templates. More info: https://kubernetes.io/docs/concepts/scheduling-eviction/pod-scheduling-readiness/", alias="schedulingGates")
@@ -41,7 +43,7 @@ class TrainerV1alpha1PodSpecOverride(BaseModel):
     target_jobs: List[TrainerV1alpha1PodSpecOverrideTargetJob] = Field(description="TrainJobs is the training job replicas in the training runtime template to apply the overrides.", alias="targetJobs")
     tolerations: Optional[List[IoK8sApiCoreV1Toleration]] = Field(default=None, description="Override for the Pod's tolerations.")
     volumes: Optional[List[IoK8sApiCoreV1Volume]] = Field(default=None, description="Overrides for the Pod volume configurations.")
-    __properties: ClassVar[List[str]] = ["affinity", "containers", "initContainers", "nodeSelector", "schedulingGates", "serviceAccountName", "targetJobs", "tolerations", "volumes"]
+    __properties: ClassVar[List[str]] = ["affinity", "containers", "imagePullSecrets", "initContainers", "nodeSelector", "schedulingGates", "serviceAccountName", "targetJobs", "tolerations", "volumes"]
 
     model_config = ConfigDict(
         populate_by_name=True,
@@ -92,6 +94,13 @@ class TrainerV1alpha1PodSpecOverride(BaseModel):
                 if _item_containers:
                     _items.append(_item_containers.to_dict())
             _dict['containers'] = _items
+        # override the default output from pydantic by calling `to_dict()` of each item in image_pull_secrets (list)
+        _items = []
+        if self.image_pull_secrets:
+            for _item_image_pull_secrets in self.image_pull_secrets:
+                if _item_image_pull_secrets:
+                    _items.append(_item_image_pull_secrets.to_dict())
+            _dict['imagePullSecrets'] = _items
         # override the default output from pydantic by calling `to_dict()` of each item in init_containers (list)
         _items = []
         if self.init_containers:
@@ -141,6 +150,7 @@ class TrainerV1alpha1PodSpecOverride(BaseModel):
         _obj = cls.model_validate({
             "affinity": IoK8sApiCoreV1Affinity.from_dict(obj["affinity"]) if obj.get("affinity") is not None else None,
             "containers": [TrainerV1alpha1ContainerOverride.from_dict(_item) for _item in obj["containers"]] if obj.get("containers") is not None else None,
+            "imagePullSecrets": [IoK8sApiCoreV1LocalObjectReference.from_dict(_item) for _item in obj["imagePullSecrets"]] if obj.get("imagePullSecrets") is not None else None,
             "initContainers": [TrainerV1alpha1ContainerOverride.from_dict(_item) for _item in obj["initContainers"]] if obj.get("initContainers") is not None else None,
             "nodeSelector": obj.get("nodeSelector"),
             "schedulingGates": [IoK8sApiCoreV1PodSchedulingGate.from_dict(_item) for _item in obj["schedulingGates"]] if obj.get("schedulingGates") is not None else None,

charts/kubeflow-trainer/crds/trainer.kubeflow.org_trainjobs.yaml

@@ -1631,6 +1631,29 @@ spec:
                       x-kubernetes-list-map-keys:
                       - name
                       x-kubernetes-list-type: map
+                    imagePullSecrets:
+                      description: ImagePullSecrets overrides the image pull secrets
+                        for the Pods in the target job templates.
+                      items:
+                        description: |-
+                          LocalObjectReference contains enough information to let you locate the
+                          referenced object inside the same namespace.
+                        properties:
+                          name:
+                            default: ""
+                            description: |-
+                              Name of the referent.
+                              This field is effectively required, but due to backwards compatibility is
+                              allowed to be empty. Instances of this type with an empty value here are
+                              almost certainly wrong.
+                              More info: https://kubernetes.io/docs/concepts/overview/working-with-objects/names/#names
+                            type: string
+                        type: object
+                        x-kubernetes-map-type: atomic
+                      type: array
+                      x-kubernetes-list-map-keys:
+                      - name
+                      x-kubernetes-list-type: map
                     initContainers:
                       description: Overrides for the init container in the target
                         job templates.

docs/proposals/2170-kubeflow-trainer-v2/README.md

@@ -796,6 +796,12 @@ type PodSpecOverride struct {
 
 	// Overrides for the containers in the desired job templates.
 	Containers []ContainerOverride `json:"containers,omitempty"`
+
+	// SchedulingGates overrides the scheduling gates of the Pods in the target job templates.
+	SchedulingGates []corev1.PodSchedulingGate `json:"schedulingGates,omitempty"`
+
+	// ImagePullSecrets overrides the image pull secrets for the Pods in the target job templates.
+	ImagePullSecrets []corev1.LocalObjectReference `json:"imagePullSecrets,omitempty"`
 }
 
 // ContainerOverride represents parameters that can be overridden using PodSpecOverrides.

manifests/base/crds/trainer.kubeflow.org_trainjobs.yaml

@@ -1631,6 +1631,29 @@ spec:
                       x-kubernetes-list-map-keys:
                       - name
                       x-kubernetes-list-type: map
+                    imagePullSecrets:
+                      description: ImagePullSecrets overrides the image pull secrets
+                        for the Pods in the target job templates.
+                      items:
+                        description: |-
+                          LocalObjectReference contains enough information to let you locate the
+                          referenced object inside the same namespace.
+                        properties:
+                          name:
+                            default: ""
+                            description: |-
+                              Name of the referent.
+                              This field is effectively required, but due to backwards compatibility is
+                              allowed to be empty. Instances of this type with an empty value here are
+                              almost certainly wrong.
+                              More info: https://kubernetes.io/docs/concepts/overview/working-with-objects/names/#names
+                            type: string
+                        type: object
+                        x-kubernetes-map-type: atomic
+                      type: array
+                      x-kubernetes-list-map-keys:
+                      - name
+                      x-kubernetes-list-type: map
                     initContainers:
                       description: Overrides for the init container in the target
                         job templates.

pkg/apis/trainer/v1alpha1/trainjob_types.go

@@ -271,6 +271,11 @@ type PodSpecOverride struct {
 	// +listType=map
 	// +listMapKey=name
 	SchedulingGates []corev1.PodSchedulingGate `json:"schedulingGates,omitempty"`
+
+	// ImagePullSecrets overrides the image pull secrets for the Pods in the target job templates.
+	// +listType=map
+	// +listMapKey=name
+	ImagePullSecrets []corev1.LocalObjectReference `json:"imagePullSecrets,omitempty"`
 }
 
 type PodSpecOverrideTargetJob struct {

pkg/apis/trainer/v1alpha1/zz_generated.deepcopy.go

@@ -444,6 +444,11 @@ func (in *PodSpecOverride) DeepCopyInto(out *PodSpecOverride) {
 		*out = make([]v1.PodSchedulingGate, len(*in))
 		copy(*out, *in)
 	}
+	if in.ImagePullSecrets != nil {
+		in, out := &in.ImagePullSecrets, &out.ImagePullSecrets
+		*out = make([]v1.LocalObjectReference, len(*in))
+		copy(*out, *in)
+	}
 	return
 }
 

pkg/apis/trainer/v1alpha1/zz_generated.openapi.go

@@ -1110,12 +1110,34 @@ func schema_pkg_apis_trainer_v1alpha1_PodSpecOverride(ref common.ReferenceCallba
 							},
 						},
 					},
+					"imagePullSecrets": {
+						VendorExtensible: spec.VendorExtensible{
+							Extensions: spec.Extensions{
+								"x-kubernetes-list-map-keys": []interface{}{
+									"name",
+								},
+								"x-kubernetes-list-type": "map",
+							},
+						},
+						SchemaProps: spec.SchemaProps{
+							Description: "ImagePullSecrets overrides the image pull secrets for the Pods in the target job templates.",
+							Type:        []string{"array"},
+							Items: &spec.SchemaOrArray{
+								Schema: &spec.Schema{
+									SchemaProps: spec.SchemaProps{
+										Default: map[string]interface{}{},
+										Ref:     ref("k8s.io/api/core/v1.LocalObjectReference"),
+									},
+								},
+							},
+						},
+					},
 				},
 				Required: []string{"targetJobs"},
 			},
 		},
 		Dependencies: []string{
-			"github.com/kubeflow/trainer/v2/pkg/apis/trainer/v1alpha1.ContainerOverride", "github.com/kubeflow/trainer/v2/pkg/apis/trainer/v1alpha1.PodSpecOverrideTargetJob", "k8s.io/api/core/v1.Affinity", "k8s.io/api/core/v1.PodSchedulingGate", "k8s.io/api/core/v1.Toleration", "k8s.io/api/core/v1.Volume"},
+			"github.com/kubeflow/trainer/v2/pkg/apis/trainer/v1alpha1.ContainerOverride", "github.com/kubeflow/trainer/v2/pkg/apis/trainer/v1alpha1.PodSpecOverrideTargetJob", "k8s.io/api/core/v1.Affinity", "k8s.io/api/core/v1.LocalObjectReference", "k8s.io/api/core/v1.PodSchedulingGate", "k8s.io/api/core/v1.Toleration", "k8s.io/api/core/v1.Volume"},
 	}
 }
 

pkg/client/applyconfiguration/trainer/v1alpha1/podspecoverride.go

@@ -33,6 +33,7 @@ type PodSpecOverrideApplyConfiguration struct {
 	InitContainers     []ContainerOverrideApplyConfiguration        `json:"initContainers,omitempty"`
 	Containers         []ContainerOverrideApplyConfiguration        `json:"containers,omitempty"`
 	SchedulingGates    []v1.PodSchedulingGate                       `json:"schedulingGates,omitempty"`
+	ImagePullSecrets   []v1.LocalObjectReference                    `json:"imagePullSecrets,omitempty"`
 }
 
 // PodSpecOverrideApplyConfiguration constructs a declarative configuration of the PodSpecOverride type for use with
@@ -145,3 +146,13 @@ func (b *PodSpecOverrideApplyConfiguration) WithSchedulingGates(values ...v1.Pod
 	}
 	return b
 }
+
+// WithImagePullSecrets adds the given value to the ImagePullSecrets field in the declarative configuration
+// and returns the receiver, so that objects can be build by chaining "With" function invocations.
+// If called multiple times, values provided by each call will be appended to the ImagePullSecrets field.
+func (b *PodSpecOverrideApplyConfiguration) WithImagePullSecrets(values ...v1.LocalObjectReference) *PodSpecOverrideApplyConfiguration {
+	for i := range values {
+		b.ImagePullSecrets = append(b.ImagePullSecrets, values[i])
+	}
+	return b
+}

pkg/runtime/core/trainingruntime_test.go

@@ -857,6 +857,54 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 					Obj(),
 			},
 		},
+		"succeeded to build JobSet with imagePullSecrets overrides from the TrainJob's PodSpecOverrides.": {
+			trainingRuntime: testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").RuntimeSpec(
+				testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").Spec).
+					WithMLPolicy(
+						testingutil.MakeMLPolicyWrapper().
+							WithNumNodes(100).
+							Obj(),
+					).
+					InitContainer(constants.DatasetInitializer, "override-init-container", "test:runtime").
+					InitContainer(constants.Node, "override-init-container", "test:runtime").
+					Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
+					Container(constants.Node, "override-container", "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
+					Obj(),
+			).Obj(),
+			trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
+				UID("uid").
+				RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "test-runtime").
+				Trainer(
+					testingutil.MakeTrainJobTrainerWrapper().
+						Container("test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
+						Obj(),
+				).
+				PodSpecOverrides([]trainer.PodSpecOverride{
+					{
+						TargetJobs: []trainer.PodSpecOverrideTargetJob{{Name: constants.Node}},
+						ImagePullSecrets: []corev1.LocalObjectReference{
+							{Name: "registry-credential"},
+						},
+					},
+				}).
+				Obj(),
+			wantObjs: []runtime.Object{
+				testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
+					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
+					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node).
+					Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer).
+					Completions(1, constants.DatasetInitializer, constants.ModelInitializer).
+					NumNodes(100).
+					InitContainer(constants.DatasetInitializer, "override-init-container", "test:runtime").
+					InitContainer(constants.Node, "override-init-container", "test:runtime").
+					Container(constants.Node, constants.Node, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
+					Container(constants.Node, "override-container", "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
+					ImagePullSecrets(constants.Node,
+						corev1.LocalObjectReference{Name: "registry-credential"},
+					).
+					Obj(),
+			},
+		},
 		"succeeded to build JobSet with dataset and model initializer from the TrainJob.": {
 			trainingRuntime: testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").RuntimeSpec(
 				testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").Spec).

pkg/util/testing/wrapper.go

@@ -314,6 +314,18 @@ func (j *JobSetWrapper) SchedulingGates(rJobName string, schedulingGates ...core
 	return j
 }
 
+func (j *JobSetWrapper) ImagePullSecrets(rJobName string, imagePullSecrets ...corev1.LocalObjectReference) *JobSetWrapper {
+	for i, rJob := range j.Spec.ReplicatedJobs {
+		if rJob.Name == rJobName {
+			j.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.ImagePullSecrets = append(
+				j.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.ImagePullSecrets,
+				imagePullSecrets...,
+			)
+		}
+	}
+	return j
+}
+
 func (j *JobSetWrapper) Tolerations(rJobName string, tolerations ...corev1.Toleration) *JobSetWrapper {
 	for i, rJob := range j.Spec.ReplicatedJobs {
 		if rJob.Name == rJobName {