From bc7c66f42acd2aa4a1994613fe7b08c46bd207d1 Mon Sep 17 00:00:00 2001 From: Jun Chen Date: Sun, 9 Aug 2020 14:24:12 +0800 Subject: [PATCH] feat: Using ormb-storage-initializer (#69) --- .../templates/preset-image-configmap.yaml | 1 + pkg/controllers/utils.go | 62 +++++++++++++++++++ pkg/controllers/utils_test.go | 43 +++++++++++++ scripts/shell/run.sh | 8 --- testutil/testutil.go | 3 +- 5 files changed, 108 insertions(+), 9 deletions(-) diff --git a/manifests/klever-modeljob-operator/templates/preset-image-configmap.yaml b/manifests/klever-modeljob-operator/templates/preset-image-configmap.yaml index 93debe3e..20c7c717 100644 --- a/manifests/klever-modeljob-operator/templates/preset-image-configmap.yaml +++ b/manifests/klever-modeljob-operator/templates/preset-image-configmap.yaml @@ -12,6 +12,7 @@ data: savedmodel-extract: '{{ .Values.ormb.domain }}/release/savedmodel:v0.2' torchscript-extract: '{{ .Values.ormb.domain }}/release/torchscript:v0.2' pmml-extract: '{{ .Values.ormb.domain }}/release/openscoring:v2.0.1-1.0' + ormb-storage-initializer: "{{ .Values.ormb.domain }}/release/clever-ormb-storage-initializer:v0.0.1" kind: ConfigMap metadata: name: modeljob-image-config diff --git a/pkg/controllers/utils.go b/pkg/controllers/utils.go index 610f2521..813d53a2 100644 --- a/pkg/controllers/utils.go +++ b/pkg/controllers/utils.go @@ -39,6 +39,11 @@ func generateJobResource(modeljob *modeljobsv1alpha1.ModelJob) (*batchv1.Job, er return nil, fmt.Errorf("%v", "not support source") } + initContainers, err := generateInitContainers(modeljob) + if err != nil { + return nil, err + } + job := &batchv1.Job{ ObjectMeta: metav1.ObjectMeta{ Namespace: modeljob.Namespace, @@ -47,6 +52,7 @@ func generateJobResource(modeljob *modeljobsv1alpha1.ModelJob) (*batchv1.Job, er Spec: batchv1.JobSpec{ Template: corev1.PodTemplateSpec{ Spec: corev1.PodSpec{ + InitContainers: initContainers, Containers: []corev1.Container{ { Name: "executor", @@ -99,6 +105,20 @@ func generateJobResource(modeljob *modeljobsv1alpha1.ModelJob) (*batchv1.Job, er Value: viper.GetString(common.ORMBPasswordEnvKey), }, }, + VolumeMounts: []corev1.VolumeMount{ + { + Name: "model", + MountPath: modeljobsv1alpha1.SourceModelPath, + }, + }, + }, + }, + Volumes: []corev1.Volume{ + { + Name: "model", + VolumeSource: corev1.VolumeSource{ + EmptyDir: &corev1.EmptyDirVolumeSource{}, + }, }, }, RestartPolicy: corev1.RestartPolicyNever, @@ -109,3 +129,45 @@ func generateJobResource(modeljob *modeljobsv1alpha1.ModelJob) (*batchv1.Job, er return job, nil } + +// generateInitContainers will pull model from harbor and export the model to /models/input path +func generateInitContainers(modeljob *modeljobsv1alpha1.ModelJob) ([]corev1.Container, error) { + ormbDomain := viper.GetString(common.ORMBDomainEnvKey) + ormbUsername := viper.GetString(common.ORMBUsernameEnvkey) + ormbPassword := viper.GetString(common.ORMBPasswordEnvKey) + if ormbDomain == "" || ormbUsername == "" || ormbPassword == "" { + return nil, nil + } + + image, ok := PresetAnalyzeImageConfig.Data["ormb-storage-initializer"] + if !ok { + return nil, fmt.Errorf("failed get ormb-storage-initializer image") + } + + initContainers := []corev1.Container{ + { + Name: "model-initializer", + Image: image, + Args: []string{modeljob.Spec.Model, modeljobsv1alpha1.SourceModelPath}, + WorkingDir: "/models", + Env: []corev1.EnvVar{ + corev1.EnvVar{ + Name: "AWS_ACCESS_KEY_ID", + Value: ormbUsername, + }, + corev1.EnvVar{ + Name: "AWS_SECRET_ACCESS_KEY", + Value: ormbPassword, + }, + }, + VolumeMounts: []corev1.VolumeMount{ + { + Name: "model", + MountPath: modeljobsv1alpha1.SourceModelPath, + }, + }, + }, + } + + return initContainers, nil +} diff --git a/pkg/controllers/utils_test.go b/pkg/controllers/utils_test.go index 54aa0fc2..361ccf5d 100644 --- a/pkg/controllers/utils_test.go +++ b/pkg/controllers/utils_test.go @@ -1,6 +1,7 @@ package controllers import ( + "os" "reflect" "testing" @@ -8,6 +9,7 @@ import ( corev1 "k8s.io/api/core/v1" modeljobsv1alpha1 "github.com/kleveross/klever-model-registry/pkg/apis/modeljob/v1alpha1" + "github.com/kleveross/klever-model-registry/pkg/common" test "github.com/kleveross/klever-model-registry/testutil" ) @@ -262,3 +264,44 @@ func Test_generateJobResource(t *testing.T) { }) } } + +func Test_generateInitContainers(t *testing.T) { + os.Setenv(common.ORMBDomainEnvKey, "demo.goharbo.com") + os.Setenv(common.ORMBUsernameEnvkey, "ormbtest") + os.Setenv(common.ORMBPasswordEnvKey, "ORMBtest12345") + PresetAnalyzeImageConfig = test.InitPresetModelImageConfigMap() + + type args struct { + modeljob *modeljobsv1alpha1.ModelJob + } + tests := []struct { + name string + args args + want []corev1.Container + wantErr bool + }{ + { + name: "generateInitContainers successfully", + args: args{ + modeljob: &modeljobsv1alpha1.ModelJob{ + Spec: modeljobsv1alpha1.ModelJobSpec{ + Model: "demo.goharbor.com/release/testmodel:v1", + }, + }, + }, + wantErr: false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := generateInitContainers(tt.args.modeljob) + if (err != nil) != tt.wantErr { + t.Errorf("generateInitContainers() error = %v, wantErr %v", err, tt.wantErr) + return + } + if got[0].Args[0] != tt.args.modeljob.Spec.Model { + t.Errorf("generateInitContainers() = %v, want %v", got, tt.want) + } + }) + } +} diff --git a/scripts/shell/run.sh b/scripts/shell/run.sh index ab5e2b1c..eb915abe 100755 --- a/scripts/shell/run.sh +++ b/scripts/shell/run.sh @@ -40,14 +40,6 @@ mkdir -p $output_dir/model ormb login --insecure $ORMB_DOMAIN -u $ORMB_USERNAME -p $ORMB_PASSWORD checkOrExit $? $ormb_login_err -# pull source model. -ormb pull $src_tag --plain-http -checkOrExit $? $ormb_pull_model_err - -# export source model. -ormb export -d $input_dir $src_tag -checkOrExit $? $ormb_export_model_err - if [ $dst_tag == "empty" ] then diff --git a/testutil/testutil.go b/testutil/testutil.go index 6a921e52..4269305a 100644 --- a/testutil/testutil.go +++ b/testutil/testutil.go @@ -12,7 +12,8 @@ import ( func InitPresetModelImageConfigMap() *corev1.ConfigMap { presetImageConfigMap := &corev1.ConfigMap{} presetImageConfigMap.Data = map[string]string{ - "savedmodel-extract": "cargo.dev.caicloud.xyz/release/savedmodel:v0.2", + "savedmodel-extract": "demo.goharbor.com/release/savedmodel:v0.2", + "ormb-storage-initializer": "demo.goharbor.com/release/savedmodel:v0.2", } return presetImageConfigMap