feat(api): Refactor API server to introduce support for Docker image registries and s3 bucket support

.github/workflows/merlin.yml

@@ -137,13 +137,6 @@ jobs:
           E2E_USE_GOOGLE_OAUTH: false
         working-directory: ./python/sdk
         run: make unit-test
-      - uses: codecov/codecov-action@v4
-        with:
-          flags: sdk-test-${{ matrix.python-version }}
-          name: sdk-test-${{ matrix.python-version }}
-          token: ${{ secrets.CODECOV_TOKEN }}
-          working-directory: ./python/sdk
-          codecov_yml_path: ../../.github/workflows/codecov-config/codecov.yml
 
   lint-api:
     runs-on: ubuntu-latest
@@ -188,13 +181,6 @@ jobs:
           POSTGRES_USER: ${{ secrets.DB_USERNAME }}
           POSTGRES_PASSWORD: ${{ secrets.DB_PASSWORD }}
         run: make it-test-api-ci
-      - uses: codecov/codecov-action@v4
-        with:
-          flags: api-test
-          name: api-test
-          token: ${{ secrets.CODECOV_TOKEN }}
-          working-directory: ./api
-          codecov_yml_path: ../.github/workflows/codecov-config/codecov.yml
 
   test-observation-publisher:
     runs-on: ubuntu-latest

api/cmd/api/setup.go

@@ -6,7 +6,6 @@ import (
 	"net/http"
 	"time"
 
-	gcs "cloud.google.com/go/storage"
 	"github.com/GoogleCloudPlatform/spark-on-k8s-operator/pkg/client/clientset/versioned"
 	"github.com/caraml-dev/merlin/webhook"
 	"github.com/caraml-dev/mlp/api/pkg/artifact"
@@ -105,55 +104,65 @@ func initImageBuilder(cfg *config.Config) (webserviceBuilder imagebuilder.ImageB
 	}
 
 	var artifactService artifact.Service
-	if cfg.ImageBuilderConfig.ArtifactServiceType == "gcs" {
-		gcsClient, err := gcs.NewClient(context.Background())
-		if err != nil {
-			log.Panicf("%s,failed initializing gcs for mlflow delete package", err.Error())
-		}
-		artifactService = artifact.NewGcsArtifactClient(gcsClient)
-	} else if cfg.ImageBuilderConfig.ArtifactServiceType == "nop" {
+	if cfg.MlflowConfig.ArtifactServiceType == "gcs" {
+		artifactService, err = artifact.NewGcsArtifactClient()
+	} else if cfg.MlflowConfig.ArtifactServiceType == "s3" {
+		artifactService, err = artifact.NewS3ArtifactClient()
+	} else if cfg.MlflowConfig.ArtifactServiceType == "nop" {
 		artifactService = artifact.NewNopArtifactClient()
 	} else {
-		log.Panicf("invalid artifact service type %s", cfg.ImageBuilderConfig.ArtifactServiceType)
+		log.Panicf("invalid artifact service type %s", cfg.MlflowConfig.ArtifactServiceType)
+	}
+	if err != nil {
+		log.Panicf("%s,failed initializing mlflow artifact service", err.Error())
+	}
+
+	if cfg.ImageBuilderConfig.KanikoPushRegistryType != "gcr" &&
+		cfg.ImageBuilderConfig.KanikoPushRegistryType != "docker" {
+		log.Panicf("invalid kaniko push registry type %s", cfg.ImageBuilderConfig.KanikoPushRegistryType)
 	}
 
 	webServiceConfig := imagebuilder.Config{
-		BaseImage:               cfg.ImageBuilderConfig.BaseImage,
-		BuildNamespace:          cfg.ImageBuilderConfig.BuildNamespace,
-		DockerRegistry:          cfg.ImageBuilderConfig.DockerRegistry,
-		BuildTimeoutDuration:    timeout,
-		KanikoImage:             cfg.ImageBuilderConfig.KanikoImage,
-		KanikoServiceAccount:    cfg.ImageBuilderConfig.KanikoServiceAccount,
-		KanikoAdditionalArgs:    cfg.ImageBuilderConfig.KanikoAdditionalArgs,
-		DefaultResources:        cfg.ImageBuilderConfig.DefaultResources,
-		Tolerations:             cfg.ImageBuilderConfig.Tolerations,
-		NodeSelectors:           cfg.ImageBuilderConfig.NodeSelectors,
-		MaximumRetry:            cfg.ImageBuilderConfig.MaximumRetry,
-		SafeToEvict:             cfg.ImageBuilderConfig.SafeToEvict,
-		ClusterName:             cfg.ImageBuilderConfig.ClusterName,
-		GcpProject:              cfg.ImageBuilderConfig.GcpProject,
-		Environment:             cfg.Environment,
-		SupportedPythonVersions: cfg.ImageBuilderConfig.SupportedPythonVersions,
+		BaseImage:                        cfg.ImageBuilderConfig.BaseImage,
+		BuildNamespace:                   cfg.ImageBuilderConfig.BuildNamespace,
+		DockerRegistry:                   cfg.ImageBuilderConfig.DockerRegistry,
+		BuildTimeoutDuration:             timeout,
+		KanikoImage:                      cfg.ImageBuilderConfig.KanikoImage,
+		KanikoPushRegistryType:           cfg.ImageBuilderConfig.KanikoPushRegistryType,
+		KanikoDockerCredentialSecretName: cfg.ImageBuilderConfig.KanikoDockerCredentialSecretName,
+		KanikoServiceAccount:             cfg.ImageBuilderConfig.KanikoServiceAccount,
+		KanikoAdditionalArgs:             cfg.ImageBuilderConfig.KanikoAdditionalArgs,
+		DefaultResources:                 cfg.ImageBuilderConfig.DefaultResources,
+		Tolerations:                      cfg.ImageBuilderConfig.Tolerations,
+		NodeSelectors:                    cfg.ImageBuilderConfig.NodeSelectors,
+		MaximumRetry:                     cfg.ImageBuilderConfig.MaximumRetry,
+		SafeToEvict:                      cfg.ImageBuilderConfig.SafeToEvict,
+		ClusterName:                      cfg.ImageBuilderConfig.ClusterName,
+		GcpProject:                       cfg.ImageBuilderConfig.GcpProject,
+		Environment:                      cfg.Environment,
+		SupportedPythonVersions:          cfg.ImageBuilderConfig.SupportedPythonVersions,
 	}
 	webserviceBuilder = imagebuilder.NewModelServiceImageBuilder(kubeClient, webServiceConfig, artifactService)
 
 	predJobConfig := imagebuilder.Config{
-		BaseImage:               cfg.ImageBuilderConfig.PredictionJobBaseImage,
-		BuildNamespace:          cfg.ImageBuilderConfig.BuildNamespace,
-		DockerRegistry:          cfg.ImageBuilderConfig.DockerRegistry,
-		BuildTimeoutDuration:    timeout,
-		KanikoImage:             cfg.ImageBuilderConfig.KanikoImage,
-		KanikoServiceAccount:    cfg.ImageBuilderConfig.KanikoServiceAccount,
-		KanikoAdditionalArgs:    cfg.ImageBuilderConfig.KanikoAdditionalArgs,
-		DefaultResources:        cfg.ImageBuilderConfig.DefaultResources,
-		Tolerations:             cfg.ImageBuilderConfig.Tolerations,
-		NodeSelectors:           cfg.ImageBuilderConfig.NodeSelectors,
-		MaximumRetry:            cfg.ImageBuilderConfig.MaximumRetry,
-		SafeToEvict:             cfg.ImageBuilderConfig.SafeToEvict,
-		ClusterName:             cfg.ImageBuilderConfig.ClusterName,
-		GcpProject:              cfg.ImageBuilderConfig.GcpProject,
-		Environment:             cfg.Environment,
-		SupportedPythonVersions: cfg.ImageBuilderConfig.SupportedPythonVersions,
+		BaseImage:                        cfg.ImageBuilderConfig.PredictionJobBaseImage,
+		BuildNamespace:                   cfg.ImageBuilderConfig.BuildNamespace,
+		DockerRegistry:                   cfg.ImageBuilderConfig.DockerRegistry,
+		BuildTimeoutDuration:             timeout,
+		KanikoImage:                      cfg.ImageBuilderConfig.KanikoImage,
+		KanikoPushRegistryType:           cfg.ImageBuilderConfig.KanikoPushRegistryType,
+		KanikoDockerCredentialSecretName: cfg.ImageBuilderConfig.KanikoDockerCredentialSecretName,
+		KanikoServiceAccount:             cfg.ImageBuilderConfig.KanikoServiceAccount,
+		KanikoAdditionalArgs:             cfg.ImageBuilderConfig.KanikoAdditionalArgs,
+		DefaultResources:                 cfg.ImageBuilderConfig.DefaultResources,
+		Tolerations:                      cfg.ImageBuilderConfig.Tolerations,
+		NodeSelectors:                    cfg.ImageBuilderConfig.NodeSelectors,
+		MaximumRetry:                     cfg.ImageBuilderConfig.MaximumRetry,
+		SafeToEvict:                      cfg.ImageBuilderConfig.SafeToEvict,
+		ClusterName:                      cfg.ImageBuilderConfig.ClusterName,
+		GcpProject:                       cfg.ImageBuilderConfig.GcpProject,
+		Environment:                      cfg.Environment,
+		SupportedPythonVersions:          cfg.ImageBuilderConfig.SupportedPythonVersions,
 	}
 	predJobBuilder = imagebuilder.NewPredictionJobImageBuilder(kubeClient, predJobConfig, artifactService)
 

api/config/config.go

@@ -104,7 +104,7 @@ type BaseImageConfig struct {
 	ImageName string `validate:"required" json:"imageName"`
 	// Dockerfile Path within the build context
 	DockerfilePath string `validate:"required" json:"dockerfilePath"`
-	// GCS URL Containing build context
+	// GCS/S3 URL Containing build context
 	BuildContextURI string `validate:"required" json:"buildContextURI"`
 	// Path to sub folder which is intended to build instead of using root folder
 	BuildContextSubPath string `json:"buildContextSubPath"`
@@ -209,18 +209,19 @@ type ClusterConfig struct {
 }
 
 type ImageBuilderConfig struct {
-	ClusterName            string `validate:"required"`
-	GcpProject             string
-	ArtifactServiceType    string
-	BaseImage              BaseImageConfig `validate:"required"`
-	PredictionJobBaseImage BaseImageConfig `validate:"required"`
-	BuildNamespace         string          `validate:"required" default:"mlp"`
-	DockerRegistry         string          `validate:"required"`
-	BuildTimeout           string          `validate:"required" default:"10m"`
-	KanikoImage            string          `validate:"required" default:"gcr.io/kaniko-project/executor:v1.6.0"`
-	KanikoServiceAccount   string
-	KanikoAdditionalArgs   []string
-	DefaultResources       ResourceRequestsLimits `validate:"required"`
+	ClusterName                      string `validate:"required"`
+	GcpProject                       string
+	BaseImage                        BaseImageConfig `validate:"required"`
+	PredictionJobBaseImage           BaseImageConfig `validate:"required"`
+	BuildNamespace                   string          `validate:"required" default:"mlp"`
+	DockerRegistry                   string          `validate:"required"`
+	BuildTimeout                     string          `validate:"required" default:"10m"`
+	KanikoImage                      string          `validate:"required" default:"gcr.io/kaniko-project/executor:v1.6.0"`
+	KanikoServiceAccount             string
+	KanikoPushRegistryType           string `validate:"required,oneof=docker gcr" default:"docker"`
+	KanikoDockerCredentialSecretName string
+	KanikoAdditionalArgs             []string
+	DefaultResources                 ResourceRequestsLimits `validate:"required"`
 	// How long to keep the image building job resource in the Kubernetes cluster. Default: 2 days (48 hours).
 	Retention     time.Duration `validate:"required" default:"48h"`
 	Tolerations   Tolerations
@@ -453,8 +454,13 @@ type JaegerConfig struct {
 }
 
 type MlflowConfig struct {
-	TrackingURL         string `validate:"required"`
-	ArtifactServiceType string `validate:"required"`
+	TrackingURL string `validate:"required_if=ArtifactServiceType gcs ArtifactServiceType s3"`
+	// Note that the Kaniko image builder needs to be configured correctly to have the necessary credentials to download
+	// the artifacts from the blob storage tool depending on the artifact service type selected (gcs/s3). For gcs, the
+	// credentials can be provided via a k8s service account or a secret but for s3, the credentials can be provided via
+	// additional arguments in the config KanikoAdditionalArgs e.g.
+	// --build-arg=[AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY/AWS_DEFAULT_REGION/AWS_ENDPOINT_URL]=xxx
+	ArtifactServiceType string `validate:"required,oneof=nop gcs s3"`
 }
 
 func (cfg *Config) Validate() error {

api/config/config_test.go

@@ -387,9 +387,8 @@ func TestLoad(t *testing.T) {
 					EnvironmentConfigs:    []*EnvironmentConfig{},
 				},
 				ImageBuilderConfig: ImageBuilderConfig{
-					ClusterName:         "test-cluster",
-					GcpProject:          "test-project",
-					ArtifactServiceType: "gcs",
+					ClusterName: "test-cluster",
+					GcpProject:  "test-project",
 					BaseImage: BaseImageConfig{
 						ImageName:           "ghcr.io/caraml-dev/merlin/merlin-pyfunc-base:0.0.0",
 						DockerfilePath:      "pyfunc-server/docker/Dockerfile",
@@ -403,12 +402,13 @@ func TestLoad(t *testing.T) {
 						BuildContextSubPath: "python",
 						MainAppPath:         "/home/spark/merlin-spark-app/main.py",
 					},
-					BuildNamespace:       "caraml",
-					DockerRegistry:       "test-docker.pkg.dev/test/caraml-registry",
-					BuildTimeout:         "30m",
-					KanikoImage:          "gcr.io/kaniko-project/executor:v1.21.0",
-					KanikoServiceAccount: "kaniko-merlin",
-					KanikoAdditionalArgs: []string{"--test=true", "--no-logs=false"},
+					BuildNamespace:         "caraml",
+					DockerRegistry:         "test-docker.pkg.dev/test/caraml-registry",
+					BuildTimeout:           "30m",
+					KanikoImage:            "gcr.io/kaniko-project/executor:v1.21.0",
+					KanikoServiceAccount:   "kaniko-merlin",
+					KanikoPushRegistryType: "docker",
+					KanikoAdditionalArgs:   []string{"--test=true", "--no-logs=false"},
 					DefaultResources: ResourceRequestsLimits{
 						Requests: Resource{
 							CPU:    "1",

api/go.mod

@@ -4,15 +4,14 @@ go 1.22
 
 require (
 	cloud.google.com/go/bigtable v1.11.0
-	cloud.google.com/go/storage v1.39.0
 	github.com/GoogleCloudPlatform/spark-on-k8s-operator v0.0.0-20221025152940-c261df66a006
 	github.com/afex/hystrix-go v0.0.0-20180502004556-fa1af6a1f4f5
 	github.com/antihax/optional v1.0.0
 	github.com/antonmedv/expr v1.12.5
 	github.com/bboughton/gcp-helpers v0.1.0
 	github.com/buger/jsonparser v1.1.1
 	github.com/caraml-dev/merlin-pyspark-app v0.0.3
-	github.com/caraml-dev/mlp v1.13.2-rc2
+	github.com/caraml-dev/mlp v1.13.2
 	github.com/caraml-dev/protopath v0.1.0
 	github.com/caraml-dev/universal-prediction-interface v1.0.0
 	github.com/cenkalti/backoff/v4 v4.2.1
@@ -246,11 +245,31 @@ require (
 )
 
 require (
+	cloud.google.com/go/storage v1.39.0 // indirect
 	github.com/avast/retry-go/v4 v4.6.0 // indirect
+	github.com/aws/aws-sdk-go-v2 v1.30.6-0.20240906182417-827d25db0048 // indirect
+	github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.4 // indirect
+	github.com/aws/aws-sdk-go-v2/config v1.17.8 // indirect
+	github.com/aws/aws-sdk-go-v2/credentials v1.12.21 // indirect
+	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.17 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.17 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.17 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/ini v1.3.24 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.17 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.4 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.3.19 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.19 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.17.17 // indirect
+	github.com/aws/aws-sdk-go-v2/service/s3 v1.61.2 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sso v1.11.23 // indirect
+	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.13.6 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sts v1.16.19 // indirect
+	github.com/aws/smithy-go v1.20.4 // indirect
 	golang.org/x/time v0.5.0 // indirect
 	k8s.io/klog/v2 v2.120.1 // indirect
 	k8s.io/kube-openapi v0.0.0-20231113174909-778a5567bc1e // indirect
 	sigs.k8s.io/structured-merge-diff/v4 v4.4.1 // indirect
+
 )
 
 replace (