Introduce new constants for commonly used artifact and registry types…

… and fix incorrect if-else logic
caraml-dev · Nov 7, 2024 · fe5de0d · fe5de0d
1 parent 35ad11e
commit fe5de0d
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 27 deletions.
diff --git a/api/pkg/imagebuilder/imagebuilder.go b/api/pkg/imagebuilder/imagebuilder.go
@@ -104,6 +104,10 @@ const (
 	// jobDeletionTickDurationMilliseconds is the interval at which the API server checks if a job has been deleted
 	jobDeletionTickDurationMilliseconds = 100
 
+	dockerRegistryPushRegistryType        = "docker"
+	googleCloudRegistryPushRegistryType   = "gcr"
+	googleCloudStorageArtifactServiceType = "gcs"
+
 	gacEnvKey  = "GOOGLE_APPLICATION_CREDENTIALS"
 	saFilePath = "/secret/kaniko-secret.json"
 
@@ -715,7 +719,8 @@ func (c *imageBuilder) createKanikoJobSpec(
 }
 
 func (c *imageBuilder) configureKanikoArgsToAddCredentials(kanikoArgs []string) []string {
-	if c.config.KanikoPushRegistryType == "gcr" || c.artifactService.GetType() == "gcs" {
+	if c.config.KanikoPushRegistryType == googleCloudRegistryPushRegistryType ||
+		c.artifactService.GetType() == googleCloudStorageArtifactServiceType {
 		if c.config.KanikoServiceAccount == "" {
 			kanikoArgs = append(kanikoArgs,
 				fmt.Sprintf("--build-arg=%s=%s", gacEnvKey, saFilePath))
@@ -728,7 +733,8 @@ func (c *imageBuilder) configureVolumesAndVolumeMountsToAddCredentials(
 	volumes []v1.Volume,
 	volumeMounts []v1.VolumeMount,
 ) ([]v1.Volume, []v1.VolumeMount) {
-	if c.config.KanikoPushRegistryType == "gcr" || c.artifactService.GetType() == "gcs" {
+	if c.config.KanikoPushRegistryType == googleCloudRegistryPushRegistryType ||
+		c.artifactService.GetType() == googleCloudStorageArtifactServiceType {
 		if c.config.KanikoServiceAccount == "" {
 			volumes = append(volumes, v1.Volume{
 				Name: kanikoSecretName,
@@ -743,7 +749,8 @@ func (c *imageBuilder) configureVolumesAndVolumeMountsToAddCredentials(
 				MountPath: "/secret",
 			})
 		}
-	} else if c.config.KanikoPushRegistryType == "docker" {
+	}
+	if c.config.KanikoPushRegistryType == dockerRegistryPushRegistryType {
 		volumes = append(volumes, v1.Volume{
 			Name: kanikoSecretName,
 			VolumeSource: v1.VolumeSource{
@@ -761,7 +768,8 @@ func (c *imageBuilder) configureVolumesAndVolumeMountsToAddCredentials(
 }
 
 func (c *imageBuilder) configureEnvVarsToAddCredentials(envVar []v1.EnvVar) []v1.EnvVar {
-	if c.config.KanikoPushRegistryType == "gcr" || c.artifactService.GetType() == "gcs" {
+	if c.config.KanikoPushRegistryType == googleCloudRegistryPushRegistryType ||
+		c.artifactService.GetType() == googleCloudStorageArtifactServiceType {
 		if c.config.KanikoServiceAccount == "" {
 			envVar = append(envVar, v1.EnvVar{
 				Name:  gacEnvKey,

diff --git a/api/pkg/imagebuilder/imagebuilder_test.go b/api/pkg/imagebuilder/imagebuilder_test.go
@@ -136,7 +136,7 @@ var (
 		ClusterName:             "my-cluster",
 		GcpProject:              "test-project",
 		Environment:             testEnvironmentName,
-		KanikoPushRegistryType:  "gcr",
+		KanikoPushRegistryType:  googleCloudRegistryPushRegistryType,
 		KanikoImage:             "gcr.io/kaniko-project/executor:v1.1.0",
 		KanikoAdditionalArgs:    defaultKanikoAdditionalArgs,
 		SupportedPythonVersions: defaultSupportedPythonVersions,
@@ -176,7 +176,7 @@ var (
 		ClusterName:                      "my-cluster",
 		GcpProject:                       "test-project",
 		Environment:                      testEnvironmentName,
-		KanikoPushRegistryType:           "docker",
+		KanikoPushRegistryType:           dockerRegistryPushRegistryType,
 		KanikoImage:                      "gcr.io/kaniko-project/executor:v1.1.0",
 		KanikoAdditionalArgs:             defaultKanikoAdditionalArgs,
 		KanikoDockerCredentialSecretName: "docker-secret",
@@ -393,7 +393,7 @@ func TestBuildImage(t *testing.T) {
 										fmt.Sprintf("--dockerfile=%s", configWithGCRPushRegistry.BaseImage.DockerfilePath),
 										fmt.Sprintf("--context=%s", configWithGCRPushRegistry.BaseImage.BuildContextURI),
 										fmt.Sprintf("--build-arg=BASE_IMAGE=%s", configWithGCRPushRegistry.BaseImage.ImageName),
-										fmt.Sprintf("--build-arg=MLFLOW_ARTIFACT_STORAGE_TYPE=%s", "gcs"),
+										fmt.Sprintf("--build-arg=MLFLOW_ARTIFACT_STORAGE_TYPE=%s", googleCloudStorageArtifactServiceType),
 										fmt.Sprintf("--build-arg=MODEL_DEPENDENCIES_URL=gs%s", modelDependenciesURLSuffix),
 										fmt.Sprintf("--build-arg=MODEL_ARTIFACTS_URL=gs%s/model", testArtifactURISuffix),
 										fmt.Sprintf("--destination=%s", fmt.Sprintf("%s/%s-%s:%s", configWithGCRPushRegistry.DockerRegistry, project.Name, model.Name, modelVersionWithGCSArtifact.ID)),
@@ -449,7 +449,7 @@ func TestBuildImage(t *testing.T) {
 			wantDeleteJobName:        "",
 			wantImageRef:             fmt.Sprintf("%s/%s-%s:%s", configWithGCRPushRegistry.DockerRegistry, project.Name, model.Name, modelVersionWithGCSArtifact.ID),
 			config:                   configWithGCRPushRegistry,
-			artifactServiceType:      "gcs",
+			artifactServiceType:      googleCloudStorageArtifactServiceType,
 			artifactServiceURLScheme: "gs",
 		},
 		{
@@ -618,7 +618,7 @@ func TestBuildImage(t *testing.T) {
 										fmt.Sprintf("--dockerfile=%s", configWithGCRPushRegistry.BaseImage.DockerfilePath),
 										fmt.Sprintf("--context=%s", configWithGCRPushRegistry.BaseImage.BuildContextURI),
 										fmt.Sprintf("--build-arg=BASE_IMAGE=%s", configWithGCRPushRegistry.BaseImage.ImageName),
-										fmt.Sprintf("--build-arg=MLFLOW_ARTIFACT_STORAGE_TYPE=%s", "gcs"),
+										fmt.Sprintf("--build-arg=MLFLOW_ARTIFACT_STORAGE_TYPE=%s", googleCloudStorageArtifactServiceType),
 										fmt.Sprintf("--build-arg=MODEL_DEPENDENCIES_URL=gs%s", modelDependenciesURLSuffix),
 										fmt.Sprintf("--build-arg=MODEL_ARTIFACTS_URL=gs%s/model", testArtifactURISuffix),
 										fmt.Sprintf("--destination=%s", fmt.Sprintf("%s/%s-%s:%s", configWithGCRPushRegistry.DockerRegistry, project.Name, model.Name, modelVersionWithGCSArtifact.ID)),
@@ -652,7 +652,7 @@ func TestBuildImage(t *testing.T) {
 			wantDeleteJobName:        "",
 			wantImageRef:             fmt.Sprintf("%s/%s-%s:%s", configWithGCRPushRegistry.DockerRegistry, project.Name, model.Name, modelVersionWithGCSArtifact.ID),
 			config:                   configWithSa,
-			artifactServiceType:      "gcs",
+			artifactServiceType:      googleCloudStorageArtifactServiceType,
 			artifactServiceURLScheme: "gs",
 		},
 		{
@@ -712,7 +712,7 @@ func TestBuildImage(t *testing.T) {
 										fmt.Sprintf("--dockerfile=%s", configWithGCRPushRegistry.BaseImage.DockerfilePath),
 										fmt.Sprintf("--context=%s", configWithGCRPushRegistry.BaseImage.BuildContextURI),
 										fmt.Sprintf("--build-arg=BASE_IMAGE=%s", configWithGCRPushRegistry.BaseImage.ImageName),
-										fmt.Sprintf("--build-arg=MLFLOW_ARTIFACT_STORAGE_TYPE=%s", "gcs"),
+										fmt.Sprintf("--build-arg=MLFLOW_ARTIFACT_STORAGE_TYPE=%s", googleCloudStorageArtifactServiceType),
 										fmt.Sprintf("--build-arg=MODEL_DEPENDENCIES_URL=gs%s", modelDependenciesURLSuffix),
 										fmt.Sprintf("--build-arg=MODEL_ARTIFACTS_URL=gs%s/model", testArtifactURISuffix),
 										fmt.Sprintf("--destination=%s", fmt.Sprintf("%s/%s-%s:%s", configWithGCRPushRegistry.DockerRegistry, project.Name, model.Name, modelVersionWithGCSArtifact.ID)),
@@ -773,7 +773,7 @@ func TestBuildImage(t *testing.T) {
 				GcpProject:              "test-project",
 				Environment:             testEnvironmentName,
 				KanikoImage:             "gcr.io/kaniko-project/executor:v1.1.0",
-				KanikoPushRegistryType:  "gcr",
+				KanikoPushRegistryType:  googleCloudRegistryPushRegistryType,
 				KanikoAdditionalArgs:    defaultKanikoAdditionalArgs,
 				SupportedPythonVersions: defaultSupportedPythonVersions,
 				DefaultResources:        configWithGCRPushRegistry.DefaultResources,
@@ -782,7 +782,7 @@ func TestBuildImage(t *testing.T) {
 				},
 				MaximumRetry: jobBackOffLimit,
 			},
-			artifactServiceType:      "gcs",
+			artifactServiceType:      googleCloudStorageArtifactServiceType,
 			artifactServiceURLScheme: "gs",
 		},
 		{
@@ -842,7 +842,7 @@ func TestBuildImage(t *testing.T) {
 										fmt.Sprintf("--dockerfile=%s", configWithGCRPushRegistry.BaseImage.DockerfilePath),
 										fmt.Sprintf("--context=%s", configWithGCRPushRegistry.BaseImage.BuildContextURI),
 										fmt.Sprintf("--build-arg=BASE_IMAGE=%s", configWithGCRPushRegistry.BaseImage.ImageName),
-										fmt.Sprintf("--build-arg=MLFLOW_ARTIFACT_STORAGE_TYPE=%s", "gcs"),
+										fmt.Sprintf("--build-arg=MLFLOW_ARTIFACT_STORAGE_TYPE=%s", googleCloudStorageArtifactServiceType),
 										fmt.Sprintf("--build-arg=MODEL_DEPENDENCIES_URL=gs%s", modelDependenciesURLSuffix),
 										fmt.Sprintf("--build-arg=MODEL_ARTIFACTS_URL=gs%s/model", testArtifactURISuffix),
 										fmt.Sprintf("--destination=%s", fmt.Sprintf("%s/%s-%s:%s", configWithGCRPushRegistry.DockerRegistry, project.Name, model.Name, modelVersionWithGCSArtifact.ID)),
@@ -908,7 +908,7 @@ func TestBuildImage(t *testing.T) {
 				GcpProject:              "test-project",
 				Environment:             testEnvironmentName,
 				KanikoImage:             "gcr.io/kaniko-project/executor:v1.1.0",
-				KanikoPushRegistryType:  "gcr",
+				KanikoPushRegistryType:  googleCloudRegistryPushRegistryType,
 				KanikoAdditionalArgs:    defaultKanikoAdditionalArgs,
 				SupportedPythonVersions: defaultSupportedPythonVersions,
 				DefaultResources:        configWithGCRPushRegistry.DefaultResources,
@@ -922,7 +922,7 @@ func TestBuildImage(t *testing.T) {
 				},
 				MaximumRetry: jobBackOffLimit,
 			},
-			artifactServiceType:      "gcs",
+			artifactServiceType:      googleCloudStorageArtifactServiceType,
 			artifactServiceURLScheme: "gs",
 		},
 		{
@@ -982,7 +982,7 @@ func TestBuildImage(t *testing.T) {
 										fmt.Sprintf("--dockerfile=%s", configWithGCRPushRegistry.BaseImage.DockerfilePath),
 										fmt.Sprintf("--context=%s", configWithGCRPushRegistry.BaseImage.BuildContextURI),
 										fmt.Sprintf("--build-arg=BASE_IMAGE=%s", configWithGCRPushRegistry.BaseImage.ImageName),
-										fmt.Sprintf("--build-arg=MLFLOW_ARTIFACT_STORAGE_TYPE=%s", "gcs"),
+										fmt.Sprintf("--build-arg=MLFLOW_ARTIFACT_STORAGE_TYPE=%s", googleCloudStorageArtifactServiceType),
 										fmt.Sprintf("--build-arg=MODEL_DEPENDENCIES_URL=gs%s", modelDependenciesURLSuffix),
 										fmt.Sprintf("--build-arg=MODEL_ARTIFACTS_URL=gs%s/model", testArtifactURISuffix),
 										fmt.Sprintf("--destination=%s", fmt.Sprintf("%s/%s-%s:%s", configWithGCRPushRegistry.DockerRegistry, project.Name, model.Name, modelVersionWithGCSArtifact.ID)),
@@ -1057,7 +1057,7 @@ func TestBuildImage(t *testing.T) {
 				NodeSelectors:           configWithGCRPushRegistry.NodeSelectors,
 				Tolerations:             configWithGCRPushRegistry.Tolerations,
 			},
-			artifactServiceType:      "gcs",
+			artifactServiceType:      googleCloudStorageArtifactServiceType,
 			artifactServiceURLScheme: "gs",
 		},
 		{
@@ -1116,7 +1116,7 @@ func TestBuildImage(t *testing.T) {
 										fmt.Sprintf("--dockerfile=%s", configWithGCRPushRegistry.BaseImage.DockerfilePath),
 										fmt.Sprintf("--context=%s", configWithGCRPushRegistry.BaseImage.BuildContextURI),
 										fmt.Sprintf("--build-arg=BASE_IMAGE=%s", configWithGCRPushRegistry.BaseImage.ImageName),
-										fmt.Sprintf("--build-arg=MLFLOW_ARTIFACT_STORAGE_TYPE=%s", "gcs"),
+										fmt.Sprintf("--build-arg=MLFLOW_ARTIFACT_STORAGE_TYPE=%s", googleCloudStorageArtifactServiceType),
 										fmt.Sprintf("--build-arg=MODEL_DEPENDENCIES_URL=gs%s", modelDependenciesURLSuffix),
 										fmt.Sprintf("--build-arg=MODEL_ARTIFACTS_URL=gs%s/model", testArtifactURISuffix),
 										fmt.Sprintf("--destination=%s", fmt.Sprintf("%s/%s-%s:%s", configWithGCRPushRegistry.DockerRegistry, project.Name, model.Name, modelVersionWithGCSArtifact.ID)),
@@ -1172,7 +1172,7 @@ func TestBuildImage(t *testing.T) {
 			wantCreateJob:            nil,
 			wantImageRef:             fmt.Sprintf("%s/%s-%s:%s", configWithGCRPushRegistry.DockerRegistry, project.Name, model.Name, modelVersionWithGCSArtifact.ID),
 			config:                   configWithGCRPushRegistry,
-			artifactServiceType:      "gcs",
+			artifactServiceType:      googleCloudStorageArtifactServiceType,
 			artifactServiceURLScheme: "gs",
 		},
 		{
@@ -1231,7 +1231,7 @@ func TestBuildImage(t *testing.T) {
 										fmt.Sprintf("--dockerfile=%s", configWithGCRPushRegistry.BaseImage.DockerfilePath),
 										fmt.Sprintf("--context=%s", configWithGCRPushRegistry.BaseImage.BuildContextURI),
 										fmt.Sprintf("--build-arg=BASE_IMAGE=%s", configWithGCRPushRegistry.BaseImage.ImageName),
-										fmt.Sprintf("--build-arg=MLFLOW_ARTIFACT_STORAGE_TYPE=%s", "gcs"),
+										fmt.Sprintf("--build-arg=MLFLOW_ARTIFACT_STORAGE_TYPE=%s", googleCloudStorageArtifactServiceType),
 										fmt.Sprintf("--build-arg=MODEL_DEPENDENCIES_URL=gs%s", modelDependenciesURLSuffix),
 										fmt.Sprintf("--build-arg=MODEL_ARTIFACTS_URL=gs%s/model", testArtifactURISuffix),
 										fmt.Sprintf("--destination=%s", fmt.Sprintf("%s/%s-%s:%s", configWithGCRPushRegistry.DockerRegistry, project.Name, model.Name, modelVersionWithGCSArtifact.ID)),
@@ -1290,7 +1290,7 @@ func TestBuildImage(t *testing.T) {
 			wantDeleteJobName:        "",
 			wantImageRef:             fmt.Sprintf("%s/%s-%s:%s", configWithGCRPushRegistry.DockerRegistry, project.Name, model.Name, modelVersionWithGCSArtifact.ID),
 			config:                   configWithGCRPushRegistry,
-			artifactServiceType:      "gcs",
+			artifactServiceType:      googleCloudStorageArtifactServiceType,
 			artifactServiceURLScheme: "gs",
 		},
 		{
@@ -1349,7 +1349,7 @@ func TestBuildImage(t *testing.T) {
 										fmt.Sprintf("--dockerfile=%s", configWithGCRPushRegistry.BaseImage.DockerfilePath),
 										fmt.Sprintf("--context=%s", configWithGCRPushRegistry.BaseImage.BuildContextURI),
 										fmt.Sprintf("--build-arg=BASE_IMAGE=%s", configWithGCRPushRegistry.BaseImage.ImageName),
-										fmt.Sprintf("--build-arg=MLFLOW_ARTIFACT_STORAGE_TYPE=%s", "gcs"),
+										fmt.Sprintf("--build-arg=MLFLOW_ARTIFACT_STORAGE_TYPE=%s", googleCloudStorageArtifactServiceType),
 										fmt.Sprintf("--build-arg=MODEL_DEPENDENCIES_URL=gs%s", modelDependenciesURLSuffix),
 										fmt.Sprintf("--build-arg=MODEL_ARTIFACTS_URL=gs%s/model", testArtifactURISuffix),
 										fmt.Sprintf("--destination=%s", fmt.Sprintf("%s/%s-%s:%s", configWithGCRPushRegistry.DockerRegistry, project.Name, model.Name, modelVersionWithGCSArtifact.ID)),
@@ -1453,7 +1453,7 @@ func TestBuildImage(t *testing.T) {
 										fmt.Sprintf("--dockerfile=%s", configWithGCRPushRegistry.BaseImage.DockerfilePath),
 										fmt.Sprintf("--context=%s", configWithGCRPushRegistry.BaseImage.BuildContextURI),
 										fmt.Sprintf("--build-arg=BASE_IMAGE=%s", configWithGCRPushRegistry.BaseImage.ImageName),
-										fmt.Sprintf("--build-arg=MLFLOW_ARTIFACT_STORAGE_TYPE=%s", "gcs"),
+										fmt.Sprintf("--build-arg=MLFLOW_ARTIFACT_STORAGE_TYPE=%s", googleCloudStorageArtifactServiceType),
 										fmt.Sprintf("--build-arg=MODEL_DEPENDENCIES_URL=gs%s", modelDependenciesURLSuffix),
 										fmt.Sprintf("--build-arg=MODEL_ARTIFACTS_URL=gs%s/model", testArtifactURISuffix),
 										fmt.Sprintf("--destination=%s", fmt.Sprintf("%s/%s-%s:%s", configWithGCRPushRegistry.DockerRegistry, project.Name, model.Name, modelVersionWithGCSArtifact.ID)),
@@ -1509,7 +1509,7 @@ func TestBuildImage(t *testing.T) {
 			wantDeleteJobName:        fmt.Sprintf("%s-%s-%s", project.Name, model.Name, modelVersionWithGCSArtifact.ID),
 			wantImageRef:             fmt.Sprintf("%s/%s-%s:%s", configWithGCRPushRegistry.DockerRegistry, project.Name, model.Name, modelVersionWithGCSArtifact.ID),
 			config:                   configWithGCRPushRegistry,
-			artifactServiceType:      "gcs",
+			artifactServiceType:      googleCloudStorageArtifactServiceType,
 			artifactServiceURLScheme: "gs",
 		},
 		{
@@ -1573,7 +1573,7 @@ func TestBuildImage(t *testing.T) {
 										fmt.Sprintf("--dockerfile=%s", configWithGCRPushRegistry.BaseImage.DockerfilePath),
 										fmt.Sprintf("--context=%s", configWithGCRPushRegistry.BaseImage.BuildContextURI),
 										fmt.Sprintf("--build-arg=BASE_IMAGE=%s", configWithGCRPushRegistry.BaseImage.ImageName),
-										fmt.Sprintf("--build-arg=MLFLOW_ARTIFACT_STORAGE_TYPE=%s", "gcs"),
+										fmt.Sprintf("--build-arg=MLFLOW_ARTIFACT_STORAGE_TYPE=%s", googleCloudStorageArtifactServiceType),
 										fmt.Sprintf("--build-arg=MODEL_DEPENDENCIES_URL=gs%s", modelDependenciesURLSuffix),
 										fmt.Sprintf("--build-arg=MODEL_ARTIFACTS_URL=gs%s/model", testArtifactURISuffix),
 										fmt.Sprintf("--destination=%s", fmt.Sprintf("%s/%s-%s:%s", configWithGCRPushRegistry.DockerRegistry, project.Name, model.Name, modelVersionWithGCSArtifact.ID)),
@@ -1629,7 +1629,7 @@ func TestBuildImage(t *testing.T) {
 			wantDeleteJobName:        "",
 			wantImageRef:             fmt.Sprintf("%s/%s-%s:%s", configWithGCRPushRegistry.DockerRegistry, project.Name, model.Name, modelVersionWithGCSArtifact.ID),
 			config:                   configWithGCRPushRegistry,
-			artifactServiceType:      "gcs",
+			artifactServiceType:      googleCloudStorageArtifactServiceType,
 			artifactServiceURLScheme: "gs",
 		},
 	}