From 16648cb5d5f034528e95422dca5a8c2da99107a7 Mon Sep 17 00:00:00 2001 From: Yuki Iwai Date: Sun, 3 Mar 2024 03:39:22 +0900 Subject: [PATCH] UT: Replace MXNet example with PyTorch example Signed-off-by: Yuki Iwai --- pkg/apis/manager/health/health.pb.go | 2 + pkg/apis/manager/v1beta1/api.pb.go | 12 +++-- .../experiment/experiment_controller_test.go | 18 ++++--- .../experiment/manifest/generator_test.go | 50 +++++++++++-------- .../trial/trial_controller_test.go | 8 +-- .../experiment/validator/validator_test.go | 18 ++++--- 6 files changed, 66 insertions(+), 42 deletions(-) diff --git a/pkg/apis/manager/health/health.pb.go b/pkg/apis/manager/health/health.pb.go index 1f598c1ea49..80f57aa9b55 100644 --- a/pkg/apis/manager/health/health.pb.go +++ b/pkg/apis/manager/health/health.pb.go @@ -5,9 +5,11 @@ Package grpc_health_v1 is a generated protocol buffer package. It is generated from these files: + health.proto It has these top-level messages: + HealthCheckRequest HealthCheckResponse */ diff --git a/pkg/apis/manager/v1beta1/api.pb.go b/pkg/apis/manager/v1beta1/api.pb.go index b73f354974f..f2c20663e20 100644 --- a/pkg/apis/manager/v1beta1/api.pb.go +++ b/pkg/apis/manager/v1beta1/api.pb.go @@ -5,9 +5,11 @@ Package api_v1_beta1 is a generated protocol buffer package. It is generated from these files: + api.proto It has these top-level messages: + Experiment ExperimentSpec ParameterSpec @@ -1138,10 +1140,12 @@ func (m *ValidateAlgorithmSettingsRequest) GetExperiment() *Experiment { type ValidateAlgorithmSettingsReply struct { } -func (m *ValidateAlgorithmSettingsReply) Reset() { *m = ValidateAlgorithmSettingsReply{} } -func (m *ValidateAlgorithmSettingsReply) String() string { return proto.CompactTextString(m) } -func (*ValidateAlgorithmSettingsReply) ProtoMessage() {} -func (*ValidateAlgorithmSettingsReply) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{29} } +func (m *ValidateAlgorithmSettingsReply) Reset() { *m = ValidateAlgorithmSettingsReply{} } +func (m *ValidateAlgorithmSettingsReply) String() string { return proto.CompactTextString(m) } +func (*ValidateAlgorithmSettingsReply) ProtoMessage() {} +func (*ValidateAlgorithmSettingsReply) Descriptor() ([]byte, []int) { + return fileDescriptor0, []int{29} +} type GetEarlyStoppingRulesRequest struct { Experiment *Experiment `protobuf:"bytes,1,opt,name=experiment" json:"experiment,omitempty"` diff --git a/pkg/controller.v1beta1/experiment/experiment_controller_test.go b/pkg/controller.v1beta1/experiment/experiment_controller_test.go index b43f45d7f5c..38f59fddccc 100644 --- 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 index fa3c3f6ff09..3adeb017f74 100644 --- 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 index 6e04813aaa0..dec9bee4851 100644 --- 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 index cef44ca737d..f17b886abe6 100644 --- 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", }, } }