From 47ba388c60021189f65a203d51cc9d2d8b80b794 Mon Sep 17 00:00:00 2001
From: Ewe Zi Yi <36802364+deadlycoconuts@users.noreply.github.com>
Date: Thu, 2 Jan 2025 16:12:57 +0700
Subject: [PATCH] feat(image-builder): Add mechanism to pass api server env
 vars to kaniko build jobs (#398)

* Add mechanism to pass api server env vars to kaniko build jobs

* Fix lint comments

* Fix inline comments
---
 api/turing/config/config.go             | 7 ++++++-
 api/turing/imagebuilder/imagebuilder.go | 6 ++++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/api/turing/config/config.go b/api/turing/config/config.go
index 8dc3596ff..5477b332e 100644
--- a/api/turing/config/config.go
+++ b/api/turing/config/config.go
@@ -208,6 +208,8 @@ type KanikoConfig struct {
 	ImageVersion string `validate:"required"`
 	// AdditionalArgs allows platform-level additional arguments to be configured for Kaniko jobs
 	AdditionalArgs []string
+	// APIServerEnvVars allows extra API-server environment variables to be passed to Kaniko jobs
+	APIServerEnvVars []string
 	// Kaniko kubernetes service account
 	ServiceAccount string
 	// ResourceRequestsLimits is the resources required by Kaniko executor.
@@ -439,8 +441,11 @@ type MlflowConfig struct {
 	// 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 KanikoConfig.AdditionalArgs e.g.
+	// 1) additional arguments in the config KanikoConfig.AdditionalArgs e.g.
 	// --build-arg=[AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY/AWS_DEFAULT_REGION/AWS_ENDPOINT_URL]=xxx
+	// OR
+	// 2) additional arguments in the config KanikoConfig.APIServerEnvVars, which will pass the specified environment
+	// variables PRESENT within the Turing API server's container to the image builder as build arguments
 	ArtifactServiceType string `validate:"required,oneof=nop gcs s3"`
 }
 
diff --git a/api/turing/imagebuilder/imagebuilder.go b/api/turing/imagebuilder/imagebuilder.go
index 5a89335ce..73d42e83c 100644
--- a/api/turing/imagebuilder/imagebuilder.go
+++ b/api/turing/imagebuilder/imagebuilder.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"os"
 	"sort"
 	"strings"
 	"time"
@@ -320,6 +321,11 @@ func (ib *imageBuilder) createKanikoJob(
 	volumes, volumeMounts = ib.configureVolumesAndVolumeMountsToAddCredentials(volumes, volumeMounts)
 	envVars = ib.configureEnvVarsToAddCredentials(envVars)
 
+	// Add all other env vars that are propagated from the API server as build args
+	for _, envVar := range ib.imageBuildingConfig.KanikoConfig.APIServerEnvVars {
+		kanikoArgs = append(kanikoArgs, fmt.Sprintf("--build-arg=%s=%s", envVar, os.Getenv(envVar)))
+	}
+
 	job := cluster.Job{
 		Name:                    kanikoJobName,
 		Namespace:               ib.imageBuildingConfig.BuildNamespace,