Skip to content

Commit

Permalink
Add test for sa, fix linting
Browse files Browse the repository at this point in the history
  • Loading branch information
shydefoo committed Feb 21, 2023
1 parent 8eea1f8 commit ff92a6c
Show file tree
Hide file tree
Showing 3 changed files with 154 additions and 29 deletions.
40 changes: 20 additions & 20 deletions api/cmd/api/setup.go
Original file line number Diff line number Diff line change
Expand Up @@ -133,16 +133,16 @@ func initImageBuilder(cfg *config.Config) (webserviceBuilder imagebuilder.ImageB
}

webServiceConfig := imagebuilder.Config{
BaseImages: cfg.ImageBuilderConfig.BaseImages,
BuildNamespace: cfg.ImageBuilderConfig.BuildNamespace,
DockerRegistry: cfg.ImageBuilderConfig.DockerRegistry,
ContextSubPath: cfg.ImageBuilderConfig.ContextSubPath,
BuildTimeoutDuration: timeout,
KanikoImage: cfg.ImageBuilderConfig.KanikoImage,
KanikoServiceAccount: cfg.ImageBuilderConfig.KanikoServiceAccount,
Tolerations: cfg.ImageBuilderConfig.Tolerations,
NodeSelectors: cfg.ImageBuilderConfig.NodeSelectors,
MaximumRetry: cfg.ImageBuilderConfig.MaximumRetry,
BaseImages: cfg.ImageBuilderConfig.BaseImages,
BuildNamespace: cfg.ImageBuilderConfig.BuildNamespace,
DockerRegistry: cfg.ImageBuilderConfig.DockerRegistry,
ContextSubPath: cfg.ImageBuilderConfig.ContextSubPath,
BuildTimeoutDuration: timeout,
KanikoImage: cfg.ImageBuilderConfig.KanikoImage,
KanikoServiceAccount: cfg.ImageBuilderConfig.KanikoServiceAccount,
Tolerations: cfg.ImageBuilderConfig.Tolerations,
NodeSelectors: cfg.ImageBuilderConfig.NodeSelectors,
MaximumRetry: cfg.ImageBuilderConfig.MaximumRetry,

ClusterName: cfg.ImageBuilderConfig.ClusterName,
GcpProject: cfg.ImageBuilderConfig.GcpProject,
Expand All @@ -152,16 +152,16 @@ func initImageBuilder(cfg *config.Config) (webserviceBuilder imagebuilder.ImageB
webserviceBuilder = imagebuilder.NewModelServiceImageBuilder(kubeClient, webServiceConfig)

predJobConfig := imagebuilder.Config{
BaseImages: cfg.ImageBuilderConfig.PredictionJobBaseImages,
BuildNamespace: cfg.ImageBuilderConfig.BuildNamespace,
DockerRegistry: cfg.ImageBuilderConfig.DockerRegistry,
ContextSubPath: cfg.ImageBuilderConfig.PredictionJobContextSubPath,
BuildTimeoutDuration: timeout,
KanikoImage: cfg.ImageBuilderConfig.KanikoImage,
KanikoServiceAccount: cfg.ImageBuilderConfig.KanikoServiceAccount,
Tolerations: cfg.ImageBuilderConfig.Tolerations,
NodeSelectors: cfg.ImageBuilderConfig.NodeSelectors,
MaximumRetry: cfg.ImageBuilderConfig.MaximumRetry,
BaseImages: cfg.ImageBuilderConfig.PredictionJobBaseImages,
BuildNamespace: cfg.ImageBuilderConfig.BuildNamespace,
DockerRegistry: cfg.ImageBuilderConfig.DockerRegistry,
ContextSubPath: cfg.ImageBuilderConfig.PredictionJobContextSubPath,
BuildTimeoutDuration: timeout,
KanikoImage: cfg.ImageBuilderConfig.KanikoImage,
KanikoServiceAccount: cfg.ImageBuilderConfig.KanikoServiceAccount,
Tolerations: cfg.ImageBuilderConfig.Tolerations,
NodeSelectors: cfg.ImageBuilderConfig.NodeSelectors,
MaximumRetry: cfg.ImageBuilderConfig.MaximumRetry,

ClusterName: cfg.ImageBuilderConfig.ClusterName,
GcpProject: cfg.ImageBuilderConfig.GcpProject,
Expand Down
2 changes: 1 addition & 1 deletion api/pkg/imagebuilder/imagebuilder.go
Original file line number Diff line number Diff line change
Expand Up @@ -375,7 +375,7 @@ func (c *imageBuilder) createKanikoJobSpec(project mlp.Project, model *models.Mo
var volumeMount []v1.VolumeMount
var envVar []v1.EnvVar

// Use kaniko service account is not set, use kaniko secret
// If kaniko service account is not set, use kaniko secret
if c.config.KanikoServiceAccount == "" {
kanikoArgs = append(kanikoArgs,
fmt.Sprintf("--build-arg=%s=%s", gacEnvKey, saFilePath))
Expand Down
141 changes: 133 additions & 8 deletions api/pkg/imagebuilder/imagebuilder_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,51 @@ var (
},
MaximumRetry: jobBackOffLimit,
}
configWithSa = Config{
BuildNamespace: buildNamespace,
ContextSubPath: "python/pyfunc-server",
BaseImages: cfg.BaseImageConfigs{
"3.7.*": cfg.BaseImageConfig{
ImageName: "gojek/base-image:1",
BuildContextURI: buildContextURL,
DockerfilePath: "./Dockerfile",
},
"3.8.*": cfg.BaseImageConfig{
ImageName: "gojek/base-image:2",
BuildContextURI: buildContextURL,
DockerfilePath: "./Dockerfile",
},
"3.9.*": cfg.BaseImageConfig{
ImageName: "gojek/base-image:3",
BuildContextURI: buildContextURL,
DockerfilePath: "./Dockerfile",
},
"3.10.*": cfg.BaseImageConfig{
ImageName: "gojek/base-image:4",
BuildContextURI: buildContextURL,
DockerfilePath: "./Dockerfile",
},
},
DockerRegistry: dockerRegistry,
BuildTimeoutDuration: timeout,
ClusterName: "my-cluster",
GcpProject: "test-project",
Environment: "dev",
KanikoImage: "gcr.io/kaniko-project/executor:v1.1.0",
Tolerations: []v1.Toleration{
{
Key: "image-build-job",
Value: "true",
Operator: v1.TolerationOpEqual,
Effect: v1.TaintEffectNoSchedule,
},
},
NodeSelectors: map[string]string{
"cloud.google.com/gke-nodepool": "image-building-job-node-pool",
},
MaximumRetry: jobBackOffLimit,
KanikoServiceAccount: "kaniko-sa",
}
)

func TestBuildImage(t *testing.T) {
Expand Down Expand Up @@ -189,11 +234,11 @@ func TestBuildImage(t *testing.T) {
fmt.Sprintf("--context=%s", config.BaseImages[modelVersion.PythonVersion].BuildContextURI),
fmt.Sprintf("--build-arg=MODEL_URL=%s/model", modelVersion.ArtifactURI),
fmt.Sprintf("--build-arg=BASE_IMAGE=%s", config.BaseImages[modelVersion.PythonVersion].ImageName),
fmt.Sprintf("--build-arg=GOOGLE_APPLICATION_CREDENTIALS=%s", "/secret/kaniko-secret.json"),
fmt.Sprintf("--destination=%s", fmt.Sprintf("%s/%s-%s:%s", config.DockerRegistry, project.Name, model.Name, modelVersion.ID)),
"--cache=true",
"--single-snapshot",
fmt.Sprintf("--context-sub-path=%s", config.ContextSubPath),
fmt.Sprintf("--build-arg=GOOGLE_APPLICATION_CREDENTIALS=%s", "/secret/kaniko-secret.json"),
},
VolumeMounts: []v1.VolumeMount{
{
Expand Down Expand Up @@ -243,6 +288,86 @@ func TestBuildImage(t *testing.T) {
wantImageRef: fmt.Sprintf("%s/%s-%s:%s", config.DockerRegistry, project.Name, model.Name, modelVersion.ID),
config: config,
},
{
name: "success: no existing job, use K8s Service account",
args: args{
project: project,
model: model,
version: modelVersion,
},
existingJob: nil,
wantCreateJob: &batchv1.Job{
ObjectMeta: metav1.ObjectMeta{
Name: fmt.Sprintf("%s-%s-%s", project.Name, model.Name, modelVersion.ID),
Namespace: config.BuildNamespace,
Labels: map[string]string{
"gojek.com/app": model.Name,
"gojek.com/orchestrator": "merlin",
"gojek.com/stream": project.Stream,
"gojek.com/team": project.Team,
"gojek.com/environment": config.Environment,
"gojek.com/component": "image-builder",
},
},
Spec: batchv1.JobSpec{
Completions: &jobCompletions,
BackoffLimit: &jobBackOffLimit,
TTLSecondsAfterFinished: &jobTTLSecondAfterComplete,
ActiveDeadlineSeconds: &timeoutInSecond,
Template: v1.PodTemplateSpec{
ObjectMeta: metav1.ObjectMeta{
Labels: map[string]string{
"gojek.com/app": model.Name,
"gojek.com/orchestrator": "merlin",
"gojek.com/stream": project.Stream,
"gojek.com/team": project.Team,
"gojek.com/environment": config.Environment,
"gojek.com/component": "image-builder",
},
},
Spec: v1.PodSpec{
RestartPolicy: v1.RestartPolicyNever,
Containers: []v1.Container{
{
Name: containerName,
Image: "gcr.io/kaniko-project/executor:v1.1.0",
Args: []string{
fmt.Sprintf("--dockerfile=%s", config.BaseImages[modelVersion.PythonVersion].DockerfilePath),
fmt.Sprintf("--context=%s", config.BaseImages[modelVersion.PythonVersion].BuildContextURI),
fmt.Sprintf("--build-arg=MODEL_URL=%s/model", modelVersion.ArtifactURI),
fmt.Sprintf("--build-arg=BASE_IMAGE=%s", config.BaseImages[modelVersion.PythonVersion].ImageName),
fmt.Sprintf("--destination=%s", fmt.Sprintf("%s/%s-%s:%s", config.DockerRegistry, project.Name, model.Name, modelVersion.ID)),
"--cache=true",
"--single-snapshot",
fmt.Sprintf("--context-sub-path=%s", config.ContextSubPath),
},
Resources: v1.ResourceRequirements{
Requests: defaultResourceRequests,
},
TerminationMessagePolicy: v1.TerminationMessageFallbackToLogsOnError,
},
},
Tolerations: []v1.Toleration{
{
Key: "image-build-job",
Operator: v1.TolerationOpEqual,
Value: "true",
Effect: v1.TaintEffectNoSchedule,
},
},
NodeSelector: map[string]string{
"cloud.google.com/gke-nodepool": "image-building-job-node-pool",
},
ServiceAccountName: "kaniko-sa",
},
},
},
Status: batchv1.JobStatus{},
},
wantDeleteJobName: "",
wantImageRef: fmt.Sprintf("%s/%s-%s:%s", config.DockerRegistry, project.Name, model.Name, modelVersion.ID),
config: configWithSa,
},
{
name: "success: no existing job, tolerations is not set",
args: args{
Expand Down Expand Up @@ -291,11 +416,11 @@ func TestBuildImage(t *testing.T) {
fmt.Sprintf("--context=%s", config.BaseImages[modelVersion.PythonVersion].BuildContextURI),
fmt.Sprintf("--build-arg=MODEL_URL=%s/model", modelVersion.ArtifactURI),
fmt.Sprintf("--build-arg=BASE_IMAGE=%s", config.BaseImages[modelVersion.PythonVersion].ImageName),
fmt.Sprintf("--build-arg=GOOGLE_APPLICATION_CREDENTIALS=%s", "/secret/kaniko-secret.json"),
fmt.Sprintf("--destination=%s", fmt.Sprintf("%s/%s-%s:%s", config.DockerRegistry, project.Name, model.Name, modelVersion.ID)),
"--cache=true",
"--single-snapshot",
fmt.Sprintf("--context-sub-path=%s", config.ContextSubPath),
fmt.Sprintf("--build-arg=GOOGLE_APPLICATION_CREDENTIALS=%s", "/secret/kaniko-secret.json"),
},
VolumeMounts: []v1.VolumeMount{
{
Expand Down Expand Up @@ -420,11 +545,11 @@ func TestBuildImage(t *testing.T) {
fmt.Sprintf("--context=%s", config.BaseImages[modelVersion.PythonVersion].BuildContextURI),
fmt.Sprintf("--build-arg=MODEL_URL=%s/model", modelVersion.ArtifactURI),
fmt.Sprintf("--build-arg=BASE_IMAGE=%s", config.BaseImages[modelVersion.PythonVersion].ImageName),
fmt.Sprintf("--build-arg=GOOGLE_APPLICATION_CREDENTIALS=%s", "/secret/kaniko-secret.json"),
fmt.Sprintf("--destination=%s", fmt.Sprintf("%s/%s-%s:%s", config.DockerRegistry, project.Name, model.Name, modelVersion.ID)),
"--cache=true",
"--single-snapshot",
fmt.Sprintf("--context-sub-path=%s", config.ContextSubPath),
fmt.Sprintf("--build-arg=GOOGLE_APPLICATION_CREDENTIALS=%s", "/secret/kaniko-secret.json"),
},
VolumeMounts: []v1.VolumeMount{
{
Expand Down Expand Up @@ -559,10 +684,10 @@ func TestBuildImage(t *testing.T) {
fmt.Sprintf("--context=%s", config.BaseImages[modelVersion.PythonVersion].BuildContextURI),
fmt.Sprintf("--build-arg=MODEL_URL=%s/model", modelVersion.ArtifactURI),
fmt.Sprintf("--build-arg=BASE_IMAGE=%s", config.BaseImages[modelVersion.PythonVersion].ImageName),
fmt.Sprintf("--build-arg=GOOGLE_APPLICATION_CREDENTIALS=%s", "/secret/kaniko-secret.json"),
fmt.Sprintf("--destination=%s", fmt.Sprintf("%s/%s-%s:%s", config.DockerRegistry, project.Name, model.Name, modelVersion.ID)),
"--cache=true",
"--single-snapshot",
fmt.Sprintf("--build-arg=GOOGLE_APPLICATION_CREDENTIALS=%s", "/secret/kaniko-secret.json"),
},
VolumeMounts: []v1.VolumeMount{
{
Expand Down Expand Up @@ -692,11 +817,11 @@ func TestBuildImage(t *testing.T) {
fmt.Sprintf("--context=%s", config.BaseImages[modelVersion.PythonVersion].BuildContextURI),
fmt.Sprintf("--build-arg=MODEL_URL=%s/model", modelVersion.ArtifactURI),
fmt.Sprintf("--build-arg=BASE_IMAGE=%s", config.BaseImages[modelVersion.PythonVersion].ImageName),
fmt.Sprintf("--build-arg=GOOGLE_APPLICATION_CREDENTIALS=%s", "/secret/kaniko-secret.json"),
fmt.Sprintf("--destination=%s", fmt.Sprintf("%s/%s-%s:%s", config.DockerRegistry, project.Name, model.Name, modelVersion.ID)),
"--cache=true",
"--single-snapshot",
fmt.Sprintf("--context-sub-path=%s", config.ContextSubPath),
fmt.Sprintf("--build-arg=GOOGLE_APPLICATION_CREDENTIALS=%s", "/secret/kaniko-secret.json"),
},
VolumeMounts: []v1.VolumeMount{
{
Expand Down Expand Up @@ -793,11 +918,11 @@ func TestBuildImage(t *testing.T) {
fmt.Sprintf("--context=%s", config.BaseImages[modelVersion.PythonVersion].BuildContextURI),
fmt.Sprintf("--build-arg=MODEL_URL=%s/model", modelVersion.ArtifactURI),
fmt.Sprintf("--build-arg=BASE_IMAGE=%s", config.BaseImages[modelVersion.PythonVersion].ImageName),
fmt.Sprintf("--build-arg=GOOGLE_APPLICATION_CREDENTIALS=%s", "/secret/kaniko-secret.json"),
fmt.Sprintf("--destination=%s", fmt.Sprintf("%s/%s-%s:%s", config.DockerRegistry, project.Name, model.Name, modelVersion.ID)),
"--cache=true",
"--single-snapshot",
fmt.Sprintf("--context-sub-path=%s", config.ContextSubPath),
fmt.Sprintf("--build-arg=GOOGLE_APPLICATION_CREDENTIALS=%s", "/secret/kaniko-secret.json"),
},
VolumeMounts: []v1.VolumeMount{
{
Expand Down Expand Up @@ -897,11 +1022,11 @@ func TestBuildImage(t *testing.T) {
fmt.Sprintf("--context=%s", config.BaseImages[modelVersion.PythonVersion].BuildContextURI),
fmt.Sprintf("--build-arg=MODEL_URL=%s/model", modelVersion.ArtifactURI),
fmt.Sprintf("--build-arg=BASE_IMAGE=%s", config.BaseImages[modelVersion.PythonVersion].ImageName),
fmt.Sprintf("--build-arg=GOOGLE_APPLICATION_CREDENTIALS=%s", "/secret/kaniko-secret.json"),
fmt.Sprintf("--destination=%s", fmt.Sprintf("%s/%s-%s:%s", config.DockerRegistry, project.Name, model.Name, modelVersion.ID)),
"--cache=true",
"--single-snapshot",
fmt.Sprintf("--context-sub-path=%s", config.ContextSubPath),
fmt.Sprintf("--build-arg=GOOGLE_APPLICATION_CREDENTIALS=%s", "/secret/kaniko-secret.json"),
},
VolumeMounts: []v1.VolumeMount{
{
Expand Down Expand Up @@ -989,11 +1114,11 @@ func TestBuildImage(t *testing.T) {
fmt.Sprintf("--context=%s", config.BaseImages[modelVersion.PythonVersion].BuildContextURI),
fmt.Sprintf("--build-arg=MODEL_URL=%s/model", modelVersion.ArtifactURI),
fmt.Sprintf("--build-arg=BASE_IMAGE=%s", config.BaseImages[modelVersion.PythonVersion].ImageName),
fmt.Sprintf("--build-arg=GOOGLE_APPLICATION_CREDENTIALS=%s", "/secret/kaniko-secret.json"),
fmt.Sprintf("--destination=%s", fmt.Sprintf("%s/%s-%s:%s", config.DockerRegistry, project.Name, model.Name, modelVersion.ID)),
"--cache=true",
"--single-snapshot",
fmt.Sprintf("--context-sub-path=%s", config.ContextSubPath),
fmt.Sprintf("--build-arg=GOOGLE_APPLICATION_CREDENTIALS=%s", "/secret/kaniko-secret.json"),
},
VolumeMounts: []v1.VolumeMount{
{
Expand Down

0 comments on commit ff92a6c

Please sign in to comment.