UT: Replace MXNet example with PyTorch example

Signed-off-by: Yuki Iwai <[email protected]>
kubeflow · Mar 2, 2024 · 16648cb · 16648cb
1 parent 8df3c5c
commit 16648cb
Show file tree

Hide file tree

Showing 6 changed files with 66 additions and 42 deletions.
diff --git a/pkg/apis/manager/health/health.pb.go b/pkg/apis/manager/health/health.pb.go
diff --git a/pkg/apis/manager/v1beta1/api.pb.go b/pkg/apis/manager/v1beta1/api.pb.go
diff --git a/pkg/controller.v1beta1/experiment/experiment_controller_test.go b/pkg/controller.v1beta1/experiment/experiment_controller_test.go
@@ -477,12 +477,14 @@ func newFakeInstance() *experimentsv1beta1.Experiment {
 					Containers: []corev1.Container{
 						{
 							Name:  primaryContainer,
-							Image: "docker.io/kubeflowkatib/mxnet-mnist",
+							Image: "docker.io/kubeflowkatib/pytorch-mnist-cpu",
 							Command: []string{
 								"python3",
-								"/opt/mxnet-mnist/mnist.py",
+								"/opt/pytorch-mnist/mnist.py",
+								"--epochs=1",
+								"--batch-size=16",
 								"--lr=${trialParameters.learningRate}",
-								"--num-layers=${trialParameters.numberLayers}",
+								"--momentum=${trialParameters.momentum}",
 							},
 						},
 					},
@@ -611,12 +613,14 @@ func newFakeBatchJob() *batchv1.Job {
 					Containers: []corev1.Container{
 						{
 							Name:  primaryContainer,
-							Image: "docker.io/kubeflowkatib/mxnet-mnist",
+							Image: "docker.io/kubeflowkatib/pytorch-mnist-cpu",
 							Command: []string{
 								"python3",
-								"/opt/mxnet-mnist/mnist.py",
-								"--lr=0.01",
-								"--num-layers=5",
+								"/opt/pytorch-mnist/mnist.py",
+								"--epochs=1",
+								"--batch-size=16",
+								"--lr=${trialParameters.learningRate}",
+								"--momentum=${trialParameters.momentum}",
 							},
 						},
 					},

diff --git a/pkg/controller.v1beta1/experiment/manifest/generator_test.go b/pkg/controller.v1beta1/experiment/manifest/generator_test.go
@@ -61,12 +61,14 @@ func TestGetRunSpecWithHP(t *testing.T) {
 					Containers: []v1.Container{
 						{
 							Name:  "training-container",
-							Image: "docker.io/kubeflowkatib/mxnet-mnist",
+							Image: "docker.io/kubeflowkatib/pytorch-mnist-cpu",
 							Command: []string{
 								"python3",
-								"/opt/mxnet-mnist/mnist.py",
+								"/opt/pytorch-mnist/mnist.py",
+								"--epochs=1",
+								"--batch-size=16",
 								"--lr=0.05",
-								"--num-layers=5",
+								"--momentum=0.9",
 							},
 							Env: []v1.EnvVar{
 								{Name: consts.TrialTemplateMetaKeyOfName, Value: "trial-name"},
@@ -176,12 +178,14 @@ spec:
     spec:
       containers:
         - name: training-container
-          image: docker.io/kubeflowkatib/mxnet-mnist
+          image: docker.io/kubeflowkatib/pytorch-mnist-cpu
           command:
             - "python3"
-            - "/opt/mxnet-mnist/mnist.py"
+            - "/opt/pytorch-mnist/mnist.py"
+            - "--epochs=1"
+            - "--batch-size=16"
             - "--lr=${trialParameters.learningRate}"
-            - "--num-layers=${trialParameters.numberLayers}"`
+            - "--momentum=${trialParameters.momentum}"`
 
 	invalidTrialSpec := `apiVersion: batch/v1
 kind: Job
@@ -190,12 +194,14 @@ spec:
     spec:
       containers:
         - name: training-container
-          image: docker.io/kubeflowkatib/mxnet-mnist
+          image: docker.io/kubeflowkatib/pytorch-mnist-cpu
           command:
             - python3
-            - /opt/mxnet-mnist/mnist.py
+            - /opt/pytorch-mnist/mnist.py
+            - --epochs=1
+            - --batch-size=16
             - --lr=${trialParameters.learningRate}
-            - --num-layers=${trialParameters.numberLayers}
+            - --momentum=${trialParameters.momentum}
             - --invalidParameter={'num_layers': 2, 'input_sizes': [32, 32, 3]}`
 
 	validGetConfigMap1 := c.EXPECT().GetConfigMap(gomock.Any(), gomock.Any()).Return(
@@ -228,12 +234,14 @@ spec:
     spec:
       containers:
         - name: training-container
-          image: docker.io/kubeflowkatib/mxnet-mnist
+          image: docker.io/kubeflowkatib/pytorch-mnist-cpu
           command:
             - "python3"
-            - "/opt/mxnet-mnist/mnist.py"
+            - "/opt/pytorch-mnist/mnist.py"
+            - "--epochs=1"
+            - "--batch-size=16"
             - "--lr=0.05"
-            - "--num-layers=5"`
+            - "--momentum=0.9"`
 
 	expectedRunSpec, err := util.ConvertStringToUnstructured(expectedStr)
 	if err != nil {
@@ -347,12 +355,14 @@ func newFakeInstance() *experimentsv1beta1.Experiment {
 					Containers: []v1.Container{
 						{
 							Name:  "training-container",
-							Image: "docker.io/kubeflowkatib/mxnet-mnist",
+							Image: "docker.io/kubeflowkatib/pytorch-mnist-cpu",
 							Command: []string{
 								"python3",
-								"/opt/mxnet-mnist/mnist.py",
+								"/opt/pytorch-mnist/mnist.py",
+								"--epochs=1",
+								"--batch-size=16",
 								"--lr=${trialParameters.learningRate}",
-								"--num-layers=${trialParameters.numberLayers}",
+								"--momentum=${trialParameters.momentum}",
 							},
 							Env: []v1.EnvVar{
 								{Name: consts.TrialTemplateMetaKeyOfName, Value: "${trialParameters.trialName}"},
@@ -381,9 +391,9 @@ func newFakeInstance() *experimentsv1beta1.Experiment {
 						Reference:   "lr",
 					},
 					{
-						Name:        "numberLayers",
-						Description: "Number of layers",
-						Reference:   "num-layers",
+						Name:        "momentum",
+						Description: "Momentum for the training model",
+						Reference:   "momentum",
 					},
 					{
 						Name:        "trialName",
@@ -418,8 +428,8 @@ func newFakeParameterAssignment() []commonapiv1beta1.ParameterAssignment {
 			Value: "0.05",
 		},
 		{
-			Name:  "num-layers",
-			Value: "5",
+			Name:  "momentum",
+			Value: "0.9",
 		},
 	}
 }
diff --git a/pkg/controller.v1beta1/trial/trial_controller_test.go b/pkg/controller.v1beta1/trial/trial_controller_test.go
@@ -388,12 +388,14 @@ func newFakeTrialBatchJob() *trialsv1beta1.Trial {
 					Containers: []corev1.Container{
 						{
 							Name:  primaryContainer,
-							Image: "docker.io/kubeflowkatib/mxnet-mnist",
+							Image: "docker.io/kubeflowkatib/pytorch-mnist-cpu",
 							Command: []string{
 								"python3",
-								"/opt/mxnet-mnist/mnist.py",
+								"/opt/pytorch-mnist/mnist.py",
+								"--epochs=1",
+								"--batch-size=16",
 								"--lr=0.01",
-								"--num-layers=5",
+								"--momentum=0.9",
 							},
 						},
 					},

diff --git a/pkg/webhook/v1beta1/experiment/validator/validator_test.go b/pkg/webhook/v1beta1/experiment/validator/validator_test.go
@@ -1273,10 +1273,10 @@ func newFakeInstance() *experimentsv1beta1.Experiment {
 					},
 				},
 				{
-					Name:          "num-layers",
+					Name:          "momentum",
 					ParameterType: experimentsv1beta1.ParameterTypeCategorical,
 					FeasibleSpace: experimentsv1beta1.FeasibleSpace{
-						List: []string{"1", "2", "3"},
+						List: []string{"0.95", "0.85", "0.75"},
 					},
 				},
 			},
@@ -1298,12 +1298,14 @@ func newFakeBatchJob() *batchv1.Job {
 					Containers: []v1.Container{
 						{
 							Name:  "training-container",
-							Image: "docker.io/kubeflowkatib/mxnet-mnist",
+							Image: "docker.io/kubeflowkatib/pytorch-mnist-cpu",
 							Command: []string{
 								"python3",
-								"/opt/mxnet-mnist/mnist.py",
+								"--epochs=1",
+								"--batch-size=16",
+								"/opt/pytorch-mnist/mnist.py",
 								"--lr=${trialParameters.learningRate}",
-								"--num-layers=${trialParameters.numberLayers}",
+								"--momentum=${trialParameters.momentum}",
 							},
 						},
 					},
@@ -1321,9 +1323,9 @@ func newFakeTrialParamters() []experimentsv1beta1.TrialParameterSpec {
 			Reference:   "lr",
 		},
 		{
-			Name:        "numberLayers",
-			Description: "Number of layers",
-			Reference:   "num-layers",
+			Name:        "momentum",
+			Description: "Momentum for the training model",
+			Reference:   "momentum",
 		},
 	}
 }