KEP-2170: Implement Initializer builders in the JobSet plugin (#2316)

* KEP-2170: Implement Initializer builder in the JobSet plugin Signed-off-by: Andrey Velichkevich <[email protected]> * Update the SDK models Signed-off-by: Andrey Velichkevich <[email protected]> * Remove Info from Initializer builder Signed-off-by: Andrey Velichkevich <[email protected]> * Update manifests Signed-off-by: Andrey Velichkevich <[email protected]> * Update pkg/constants/constants.go Co-authored-by: Yuki Iwai <[email protected]> Signed-off-by: Andrey Velichkevich <[email protected]> * Use var for envs Signed-off-by: Andrey Velichkevich <[email protected]> * Remove check manifests from GitHub actions Signed-off-by: Andrey Velichkevich <[email protected]> * Move consts to JobSet plugin Signed-off-by: Andrey Velichkevich <[email protected]> --------- Signed-off-by: Andrey Velichkevich <[email protected]> Co-authored-by: Yuki Iwai <[email protected]>
kubeflow · Nov 1, 2024 · 2e1e125 · 2e1e125
1 parent 7c5ea70
commit 2e1e125
Show file tree

Hide file tree

Showing 24 changed files with 595 additions and 143 deletions.
diff --git a/.github/workflows/test-go.yaml b/.github/workflows/test-go.yaml
@@ -29,14 +29,10 @@ jobs:
         run: |
           go mod tidy && pushd hack/swagger && go mod tidy && popd && git add go.* &&
           git diff --cached --exit-code || (echo 'Please run "go mod tidy" to sync Go modules' && exit 1);
-      - name: Check manifests
+      - name: Check auto-generated assets
         run: |
-          make manifests && git add manifests &&
-          git diff --cached --exit-code || (echo 'Please run "make manifests" to generate manifests' && exit 1);
-      - name: Check auto-generated codes
-        run: |
-          make generate && git add pkg sdk &&
-          git diff --cached --exit-code || (echo 'Please run "make generate" to generate Go codes' && exit 1);
+          make generate && git add pkg sdk manifests &&
+          git diff --cached --exit-code || (echo 'Please run "make generate" to generate assets' && exit 1);
       - name: Verify gofmt
         run: |
           make fmt && git add pkg cmd &&

diff --git a/Makefile b/Makefile
@@ -47,7 +47,7 @@ manifests: controller-gen ## Generate WebhookConfiguration, ClusterRole and Cust
 		output:rbac:artifacts:config=manifests/v2/base/rbac \
 		output:webhook:artifacts:config=manifests/v2/base/webhook
 
-generate: controller-gen ## Generate apidoc, sdk and code containing DeepCopy, DeepCopyInto, and DeepCopyObject method implementations.
+generate: controller-gen manifests ## Generate apidoc, sdk and code containing DeepCopy, DeepCopyInto, and DeepCopyObject method implementations.
 	$(CONTROLLER_GEN) object:headerFile="hack/boilerplate/boilerplate.go.txt" paths="./pkg/apis/..."
 	hack/update-codegen.sh
 	hack/python-sdk/gen-sdk.sh

diff --git a/api.v2/openapi-spec/swagger.json b/api.v2/openapi-spec/swagger.json
@@ -137,8 +137,8 @@
           }
         },
         "secretRef": {
-          "description": "Reference to the TrainJob's secrets to download dataset.",
-          "$ref": "#/definitions/v1.SecretReference"
+          "description": "Reference to the secret with credentials to download dataset. Secret must be created in the TrainJob's namespace.",
+          "$ref": "#/definitions/v1.LocalObjectReference"
         },
         "storageUri": {
           "description": "Storage uri for the dataset provider.",
@@ -159,8 +159,8 @@
           }
         },
         "secretRef": {
-          "description": "Reference to the TrainJob's secrets to download model.",
-          "$ref": "#/definitions/v1.SecretReference"
+          "description": "Reference to the secret with credentials to download model. Secret must be created in the TrainJob's namespace.",
+          "$ref": "#/definitions/v1.LocalObjectReference"
         },
         "storageUri": {
           "description": "Storage uri for the model provider.",
@@ -315,8 +315,8 @@
           }
         },
         "secretRef": {
-          "description": "Reference to the TrainJob's secrets to export model.",
-          "$ref": "#/definitions/v1.SecretReference"
+          "description": "Reference to the secret with credentials to export model. Secret must be created in the TrainJob's namespace.",
+          "$ref": "#/definitions/v1.LocalObjectReference"
         },
         "storageUri": {
           "description": "Storage uri for the model exporter.",

diff --git a/docs/proposals/2170-kubeflow-training-v2/README.md b/docs/proposals/2170-kubeflow-training-v2/README.md
@@ -597,7 +597,7 @@ type DatasetConfig struct {
 	Env []corev1.EnvVar `json:"env,omitempty"`
 
 	// Reference to the TrainJob's secrets to download dataset.
-	SecretRef *corev1.SecretReference `json:"secretRef,omitempty"`
+	SecretRef *corev1.LocalObjectReference `json:"secretRef,omitempty"`
 }
 ```
 
@@ -665,7 +665,7 @@ type InputModel struct {
 	Env []corev1.EnvVar `json:"env,omitempty"`
 
 	// Reference to the TrainJob's secrets to download model.
-	SecretRef *corev1.SecretReference `json:"secretRef,omitempty"`
+	SecretRef *corev1.LocalObjectReference `json:"secretRef,omitempty"`
 }
 
 type OutputModel struct {
@@ -677,7 +677,7 @@ type OutputModel struct {
 	Env []corev1.EnvVar `json:"env,omitempty"`
 
 	// Reference to the TrainJob's secrets to export model.
-	SecretRef *corev1.SecretReference `json:"secretRef,omitempty"`
+	SecretRef *corev1.LocalObjectReference `json:"secretRef,omitempty"`
 }
 ```
 

diff --git a/manifests/v2/base/crds/kubeflow.org_trainjobs.yaml b/manifests/v2/base/crds/kubeflow.org_trainjobs.yaml
@@ -173,15 +173,15 @@ spec:
                       type: object
                     type: array
                   secretRef:
-                    description: Reference to the TrainJob's secrets to download dataset.
+                    description: |-
+                      Reference to the secret with credentials to download dataset.
+                      Secret must be created in the TrainJob's namespace.
                     properties:
                       name:
-                        description: name is unique within a namespace to reference
-                          a secret resource.
-                        type: string
-                      namespace:
-                        description: namespace defines the space within which the
-                          secret name must be unique.
+                        description: |-
+                          Name of the referent.
+                          More info: https://kubernetes.io/docs/concepts/overview/working-with-objects/names/#names
+                          TODO: Add other useful fields. apiVersion, kind, uid?
                         type: string
                     type: object
                     x-kubernetes-map-type: atomic
@@ -339,16 +339,15 @@ spec:
                           type: object
                         type: array
                       secretRef:
-                        description: Reference to the TrainJob's secrets to download
-                          model.
+                        description: |-
+                          Reference to the secret with credentials to download model.
+                          Secret must be created in the TrainJob's namespace.
                         properties:
                           name:
-                            description: name is unique within a namespace to reference
-                              a secret resource.
-                            type: string
-                          namespace:
-                            description: namespace defines the space within which
-                              the secret name must be unique.
+                            description: |-
+                              Name of the referent.
+                              More info: https://kubernetes.io/docs/concepts/overview/working-with-objects/names/#names
+                              TODO: Add other useful fields. apiVersion, kind, uid?
                             type: string
                         type: object
                         x-kubernetes-map-type: atomic
@@ -479,16 +478,15 @@ spec:
                           type: object
                         type: array
                       secretRef:
-                        description: Reference to the TrainJob's secrets to export
-                          model.
+                        description: |-
+                          Reference to the secret with credentials to export model.
+                          Secret must be created in the TrainJob's namespace.
                         properties:
                           name:
-                            description: name is unique within a namespace to reference
-                              a secret resource.
-                            type: string
-                          namespace:
-                            description: namespace defines the space within which
-                              the secret name must be unique.
+                            description: |-
+                              Name of the referent.
+                              More info: https://kubernetes.io/docs/concepts/overview/working-with-objects/names/#names
+                              TODO: Add other useful fields. apiVersion, kind, uid?
                             type: string
                         type: object
                         x-kubernetes-map-type: atomic

diff --git a/pkg/apis/kubeflow.org/v2alpha1/openapi_generated.go b/pkg/apis/kubeflow.org/v2alpha1/openapi_generated.go
diff --git a/pkg/apis/kubeflow.org/v2alpha1/trainjob_types.go b/pkg/apis/kubeflow.org/v2alpha1/trainjob_types.go
@@ -167,8 +167,9 @@ type DatasetConfig struct {
 	// These values will be merged with the TrainingRuntime's dataset initializer environments.
 	Env []corev1.EnvVar `json:"env,omitempty"`
 
-	// Reference to the TrainJob's secrets to download dataset.
-	SecretRef *corev1.SecretReference `json:"secretRef,omitempty"`
+	// Reference to the secret with credentials to download dataset.
+	// Secret must be created in the TrainJob's namespace.
+	SecretRef *corev1.LocalObjectReference `json:"secretRef,omitempty"`
 }
 
 // ModelConfig represents the desired model configuration.
@@ -193,8 +194,9 @@ type InputModel struct {
 	// These values will be merged with the TrainingRuntime's model initializer environments.
 	Env []corev1.EnvVar `json:"env,omitempty"`
 
-	// Reference to the TrainJob's secrets to download model.
-	SecretRef *corev1.SecretReference `json:"secretRef,omitempty"`
+	// Reference to the secret with credentials to download model.
+	// Secret must be created in the TrainJob's namespace.
+	SecretRef *corev1.LocalObjectReference `json:"secretRef,omitempty"`
 }
 
 // OutputModel represents the desired trained model configuration.
@@ -206,8 +208,9 @@ type OutputModel struct {
 	// These values will be merged with the TrainingRuntime's model exporter environments.
 	Env []corev1.EnvVar `json:"env,omitempty"`
 
-	// Reference to the TrainJob's secrets to export model.
-	SecretRef *corev1.SecretReference `json:"secretRef,omitempty"`
+	// Reference to the secret with credentials to export model.
+	// Secret must be created in the TrainJob's namespace.
+	SecretRef *corev1.LocalObjectReference `json:"secretRef,omitempty"`
 }
 
 // PodSpecOverride represents the custom overrides that will be applied for the TrainJob's resources.

diff --git a/pkg/apis/kubeflow.org/v2alpha1/zz_generated.deepcopy.go b/pkg/apis/kubeflow.org/v2alpha1/zz_generated.deepcopy.go
diff --git a/pkg/controller.v1/pytorch/elastic.go b/pkg/controller.v1/pytorch/elastic.go
@@ -67,7 +67,7 @@ func GetElasticEnvVarGenerator() EnvVarGenerator {
 
 func (e ElasticEnvVarGenerator) Generate(
 	job *kubeflowv1.PyTorchJob) ([]corev1.EnvVar, error) {
-	envVars := []corev1.EnvVar{}
+	var envVars []corev1.EnvVar
 
 	elasticPolicy := job.Spec.ElasticPolicy
 	if elasticPolicy == nil {

diff --git a/pkg/controller.v1/pytorch/master.go b/pkg/controller.v1/pytorch/master.go
@@ -32,7 +32,7 @@ func GetMasterEnvVarGenerator() EnvVarGenerator {
 
 func (e MasterEnvVarGenerator) Generate(
 	job *kubeflowv1.PyTorchJob) ([]corev1.EnvVar, error) {
-	envVars := []corev1.EnvVar{}
+	var envVars []corev1.EnvVar
 	if job.Spec.PyTorchReplicaSpecs[kubeflowv1.PyTorchJobReplicaTypeMaster] != nil {
 		masterPort, err := getPortFromPyTorchJob(job, kubeflowv1.PyTorchJobReplicaTypeMaster)
 		if err != nil {

diff --git a/pkg/runtime.v2/core/clustertrainingruntime_test.go b/pkg/runtime.v2/core/clustertrainingruntime_test.go
@@ -47,9 +47,9 @@ func TestClusterTrainingRuntimeNewObjects(t *testing.T) {
 		"succeeded to build PodGroup and JobSet with NumNodes from the Runtime and container from the Trainer.": {
 			clusterTrainingRuntime: testingutil.MakeClusterTrainingRuntimeWrapper("test-runtime").RuntimeSpec(
 				testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeClusterTrainingRuntimeWrapper("test-runtime").Spec).
+					InitContainerDatasetModelInitializer("test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
 					NumNodes(100).
 					ContainerTrainer("test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
-					ContainerDatasetModelInitializer("test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
 					PodGroupPolicyCoschedulingSchedulingTimeout(120).
 					Obj(),
 			).Obj(),
@@ -59,15 +59,15 @@ func TestClusterTrainingRuntimeNewObjects(t *testing.T) {
 				RuntimeRef(kubeflowv2.SchemeGroupVersion.WithKind(kubeflowv2.ClusterTrainingRuntimeKind), "test-runtime").
 				Trainer(
 					testingutil.MakeTrainJobTrainerWrapper().
-						ContainerTrainer("test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
+						Container("test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
 						Obj(),
 				).
 				Obj(),
 			wantObjs: []client.Object{
 				testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
+					InitContainerDatasetModelInitializer("test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
 					NumNodes(100).
 					ContainerTrainer("test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
-					ContainerDatasetModelInitializer("test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
 					Suspend(true).
 					PodLabel(schedulerpluginsv1alpha1.PodGroupLabel, "test-job").
 					ControllerReference(kubeflowv2.SchemeGroupVersion.WithKind(kubeflowv2.TrainJobKind), "test-job", "uid").