feat: Using ormb-storage-initializer (#69)

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

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
+}

pkg/controllers/utils_test.go

@@ -1,13 +1,15 @@
 package controllers
 
 import (
+	"os"
 	"reflect"
 	"testing"
 
 	batchv1 "k8s.io/api/batch/v1"
 	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)
+			}
+		})
+	}
+}

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
 

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