feat(backend): Add the Kubernetes native pipeline store (#11881)

* Add the Kubernetes native pipeline store

This also improves cache update race conditions in the webhooks.

Co-authored-by: Matt Prahl <[email protected]>
Signed-off-by: Ricardo M. Oliveira <[email protected]>

* Use controller-runtime for the non-caching client

Signed-off-by: mprahl <[email protected]>

* Put the Kubernetes native CI manifests under the Argo manifests

Signed-off-by: mprahl <[email protected]>

* Fix the flaky Kubernetes pipeline store tests

Some tests set the viper configuration of POD_NAMESPACE while others
didn't and so the order of the tests mattered. This now sets and resets
the viper configuration for each test.

Signed-off-by: mprahl <[email protected]>

* Modify the suggested pipeline version name to be a valid K8s name

This is more important for the Kubernetes pipeline store.

Signed-off-by: mprahl <[email protected]>

* Fix the ml-pipeline Service ports in K8s native mode

The KFP UI automatically uses the first port listed in the ml-pipeline
Service to communicate with the KFP API. Using a JSON patch to add the
webhook port ensures it doesn't change the order.

Signed-off-by: mprahl <[email protected]>

---------

Signed-off-by: Ricardo M. Oliveira <[email protected]>
Signed-off-by: mprahl <[email protected]>
Co-authored-by: Ricardo M. Oliveira <[email protected]>

Author: mprahl

.github/actions/kfp-cluster/action.yml

@@ -5,6 +5,10 @@ inputs:
   k8s_version:
     description: "The Kubernetes version to use for the Kind cluster"
     required: true
+  pipeline_store:
+    description: "Flag to deploy KFP with K8s Native API"
+    default: 'database'
+    required: false
   proxy:
     description: "If KFP should be deployed with proxy configuration"
     required: false
@@ -39,8 +43,12 @@ runs:
     - name: Deploy KFP
       shell: bash
       run: |
+        ARGS=""
+
         if [ "${{ inputs.proxy }}" = "true" ]; then
-          ./.github/resources/scripts/deploy-kfp.sh --proxy
-        else
-          ./.github/resources/scripts/deploy-kfp.sh
+          ARGS="${ARGS} --proxy"
+        elif [ "${{inputs.pipeline_store }}" = "kubernetes" ]; then
+          ARGS="${ARGS}  --deploy-k8s-native"
         fi
+        
+        ./.github/resources/scripts/deploy-kfp.sh $ARGS

.github/resources/manifests/kubernetes-native/overlays/no-proxy/apiserver-env.yaml renamed to .github/resources/manifests/argo/overlays/kubernetes-native/apiserver-env.yaml

@@ -24,43 +24,55 @@ C_DIR="${BASH_SOURCE%/*}"
 if [[ ! -d "$C_DIR" ]]; then C_DIR="$PWD"; fi
 source "${C_DIR}/helper-functions.sh"
 
+TEST_MANIFESTS=".github/resources/manifests/argo"
+PIPELINES_STORE="database"
 USE_PROXY=false
 
-while getopts ":p-:" OPT; do
-    case $OPT in
-        -) [ "$OPTARG" = "proxy" ] && USE_PROXY=true || { echo "Unknown option --$OPTARG"; exit 1; };;
-        \?) echo "Invalid option: -$OPTARG" >&2; exit 1;;
-    esac
+# Loop over script arguments passed. This uses a single switch-case
+# block with default value in case we want to make alternative deployments
+# in the future.
+while [ "$#" -gt 0 ]; do
+  case "$1" in
+    --deploy-k8s-native)
+      PIPELINES_STORE="kubernetes"
+      shift
+      ;;
+    --proxy)
+      USE_PROXY=true
+      shift
+      ;;
+  esac
 done
 
-shift $((OPTIND-1))
+if [ "${USE_PROXY}" == "true" && "${PIPELINES_STORE}" == "kubernetes" ]; then
+  echo "ERROR: Kubernetes Pipeline store cannot be deployed with proxy support."
+  exit 1
+fi
 
 kubectl apply -k "manifests/kustomize/cluster-scoped-resources/"
-kubectl apply -k "manifests/kustomize/base/crds"
 kubectl wait crd/applications.app.k8s.io --for condition=established --timeout=60s || EXIT_CODE=$?
 if [[ $EXIT_CODE -ne 0 ]]
 then
   echo "Failed to deploy cluster-scoped resources."
   exit $EXIT_CODE
 fi
 
-#Install cert-manager
-make -C ./backend install-cert-manager || EXIT_CODE=$?
-if [[ $EXIT_CODE -ne 0 ]]
-then
-  echo "Failed to deploy cert-manager."
-  exit $EXIT_CODE
-fi
-
-# Deploy manifest
-TEST_MANIFESTS=".github/resources/manifests/argo"
-
-if [[ "$PIPELINE_STORE" == "kubernetes" ]]; then
-  TEST_MANIFESTS=".github/resources/manifests/kubernetes-native"
+# If pipelines store is set to 'kubernetes', cert-manager must be deployed
+if [ "${PIPELINES_STORE}" == "kubernetes" ]; then
+  #Install cert-manager
+  make -C ./backend install-cert-manager || EXIT_CODE=$?
+  if [[ $EXIT_CODE -ne 0 ]]
+  then
+    echo "Failed to deploy cert-manager."
+    exit $EXIT_CODE
+  fi
 fi
 
+# Manifests will be deployed according to the flag provided
 if $USE_PROXY; then
   TEST_MANIFESTS="${TEST_MANIFESTS}/overlays/proxy"
+elif [ "${PIPELINES_STORE}" == "kubernetes" ]; then
+  TEST_MANIFESTS="${TEST_MANIFESTS}/overlays/kubernetes-native"
 else
   TEST_MANIFESTS="${TEST_MANIFESTS}/overlays/no-proxy"
 fi

.github/resources/manifests/kubernetes-native/overlays/no-proxy/kustomization.yaml renamed to .github/resources/manifests/argo/overlays/kubernetes-native/kustomization.yaml

@@ -152,8 +152,9 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
+        pipeline_store: [ "database", "kubernetes" ]
         k8s_version: [ "v1.29.2", "v1.30.2", "v1.31.0" ]
-    name: API integration tests v2 - K8s ${{ matrix.k8s_version }}
+    name: API integration tests v2 - K8s with ${{ matrix.pipeline_store }} ${{ matrix.k8s_version }}
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -167,6 +168,7 @@ jobs:
         uses: ./.github/actions/kfp-cluster
         with:
           k8s_version: ${{ matrix.k8s_version }}
+          pipeline_store: ${{ matrix.pipeline_store }}
 
       - name: Forward API port
         run: ./.github/resources/scripts/forward-port.sh "kubeflow" "ml-pipeline" 8888 8888
@@ -187,7 +189,7 @@ jobs:
         if: always()
         uses: actions/upload-artifact@v4
         with:
-          name: kfp-api-integration-tests-v2-artifacts-k8s-${{ matrix.k8s_version }}
+          name: kfp-api-integration-tests-v2-artifacts-k8s-${{ matrix.k8s_version }}-${{ matrix.pipeline_store }}
           path: /tmp/tmp*/*
 
   api-integration-tests-v2-with-proxy:

.github/resources/manifests/kubernetes-native/overlays/proxy/kustomization.yaml

@@ -29,8 +29,7 @@ jobs:
       uses: ./.github/actions/kfp-cluster
       with:
         k8s_version: ${{ matrix.k8s_version }}
-      env:
-        PIPELINE_STORE: kubernetes
+        pipeline_store: kubernetes
       continue-on-error: true
 
     - name: Run Webhook Integration Tests

.github/resources/manifests/kubernetes-native/overlays/proxy/proxy-env.yaml

@@ -15,6 +15,8 @@ on:
       - 'api/**/*.go'
       - 'kubernetes_platform/**/*.proto'
       - 'kubernetes_platform/**/*.go'
+      - 'backend/src/crd/kubernetes/**/*.go'
+      - 'manifests/kustomize/base/crds/*.yaml'
       - '!**/*.md'
       - '!**/OWNERS'
 
@@ -49,7 +51,10 @@ jobs:
       - name: Generate kfp-kubernetes proto files from source
         working-directory: ./kubernetes_platform
         run: make clean all
-      
+
+      - name: Generate K8s Native API CRDs
+        working-directory: ./backend/src/crd/kubernetes
+        run: make generate manifests
+
       - name: Check for Changes
         run: make check-diff
-          

.github/resources/scripts/deploy-kfp.sh

@@ -15,10 +15,12 @@
 package clientmanager
 
 import (
+	"context"
 	"database/sql"
 	"fmt"
 	"os"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/cenkalti/backoff"
@@ -36,6 +38,7 @@ import (
 	k8sapi "github.com/kubeflow/pipelines/backend/src/crd/kubernetes/v2beta1"
 	"github.com/minio/minio-go/v6"
 	"k8s.io/apimachinery/pkg/runtime"
+	"sigs.k8s.io/controller-runtime/pkg/cache"
 	ctrlclient "sigs.k8s.io/controller-runtime/pkg/client"
 )
 
@@ -108,12 +111,28 @@ type ClientManager struct {
 	uuid                      util.UUIDGeneratorInterface
 	authenticators            []auth.Authenticator
 	controllerClient          ctrlclient.Client
+	controllerClientNoCache   ctrlclient.Client
+}
+
+// Options to pass to Client Manager initialization
+type Options struct {
+	UsePipelineKubernetesStorage bool
+	Context                      context.Context
+	WaitGroup                    *sync.WaitGroup
 }
 
 func (c *ClientManager) TaskStore() storage.TaskStoreInterface {
 	return c.taskStore
 }
 
+func (c *ClientManager) ControllerClient(cacheEnabled bool) ctrlclient.Client {
+	if cacheEnabled {
+		return c.controllerClient
+	}
+
+	return c.controllerClientNoCache
+}
+
 func (c *ClientManager) ExperimentStore() storage.ExperimentStoreInterface {
 	return c.experimentStore
 }
@@ -182,48 +201,86 @@ func (c *ClientManager) Authenticators() []auth.Authenticator {
 	return c.authenticators
 }
 
-func (c *ClientManager) ControllerClient() ctrlclient.Client {
-	return c.controllerClient
-}
+func (c *ClientManager) init(options *Options) error {
+	// time
+	c.time = util.NewRealTime()
 
-func (c *ClientManager) init() error {
-	glog.Info("Initializing controller client...")
+	// UUID generator
+	c.uuid = util.NewUUIDGenerator()
 
-	restConfig, err := util.GetKubernetesConfig()
-	if err != nil {
-		return fmt.Errorf("failed to initialize the RestConfig: %w", err)
-	}
+	var pipelineStoreForRef storage.PipelineStoreInterface
 
-	controllerClient, err := ctrlclient.New(
-		restConfig, ctrlclient.Options{Scheme: scheme},
-	)
-	if err != nil {
-		return fmt.Errorf("failed to initialize the controller client: %w", err)
-	}
-	c.controllerClient = controllerClient
-	glog.Info("Controller client initialized successfully.")
+	if options.UsePipelineKubernetesStorage || common.IsOnlyKubernetesWebhookMode() {
+		glog.Info("Initializing controller client...")
+		restConfig, err := util.GetKubernetesConfig()
+		if err != nil {
+			return err
+		}
 
-	if common.IsOnlyKubernetesWebhookMode() {
-		return nil
+		var cacheConfig map[string]cache.Config
+
+		if !common.IsMultiUserMode() && common.GetPodNamespace() != "" && !common.IsOnlyKubernetesWebhookMode() {
+			cacheConfig = map[string]cache.Config{common.GetPodNamespace(): {}}
+		}
+
+		k8sAPICache, err := cache.New(restConfig,
+			cache.Options{
+				DefaultNamespaces: cacheConfig,
+				Scheme:            scheme,
+			},
+		)
+		if err != nil {
+			return err
+		}
+
+		options.WaitGroup.Add(1)
+		go func() {
+			defer options.WaitGroup.Done()
+
+			err := k8sAPICache.Start(options.Context)
+			if err != nil {
+				panic(fmt.Sprintf("Failed to start the cache to the cluster: %v", err))
+			}
+		}()
+
+		controllerClient, err := ctrlclient.New(
+			restConfig, ctrlclient.Options{Scheme: scheme, Cache: &ctrlclient.CacheOptions{Reader: k8sAPICache}},
+		)
+		if err != nil {
+			return fmt.Errorf("failed to initialize the controller client: %w", err)
+		}
+
+		controllerClientNoCache, err := ctrlclient.New(restConfig, ctrlclient.Options{Scheme: scheme})
+		if err != nil {
+			return fmt.Errorf("failed to initialize the no cache controller client: %w", err)
+		}
+
+		glog.Info("Controller client initialized successfully.")
+
+		c.controllerClient = controllerClient
+		c.controllerClientNoCache = controllerClientNoCache
+		if common.IsOnlyKubernetesWebhookMode() {
+			return nil
+		}
+
+		c.pipelineStore = storage.NewPipelineStoreKubernetes(controllerClient, controllerClientNoCache)
+		pipelineStoreForRef = c.pipelineStore
 	}
 
 	glog.Info("Initializing client manager")
 	glog.Info("Initializing DB client...")
 	db := InitDBClient(common.GetDurationConfig(initConnectionTimeout))
 	db.SetConnMaxLifetime(common.GetDurationConfig(dbConMaxLifeTime))
 	glog.Info("DB client initialized successfully")
-	// time
-	c.time = util.NewRealTime()
-
-	// UUID generator
-	c.uuid = util.NewUUIDGenerator()
 
 	c.db = db
+	if !options.UsePipelineKubernetesStorage {
+		c.pipelineStore = storage.NewPipelineStore(db, c.time, c.uuid)
+	}
 	c.experimentStore = storage.NewExperimentStore(db, c.time, c.uuid)
-	c.pipelineStore = storage.NewPipelineStore(db, c.time, c.uuid)
-	c.jobStore = storage.NewJobStore(db, c.time)
+	c.jobStore = storage.NewJobStore(db, c.time, pipelineStoreForRef)
 	c.taskStore = storage.NewTaskStore(db, c.time, c.uuid)
-	c.resourceReferenceStore = storage.NewResourceReferenceStore(db)
+	c.resourceReferenceStore = storage.NewResourceReferenceStore(db, pipelineStoreForRef)
 	c.dBStatusStore = storage.NewDBStatusStore(db)
 	c.defaultExperimentStore = storage.NewDefaultExperimentStore(db)
 	glog.Info("Initializing Object store client...")
@@ -573,11 +630,11 @@ func initLogArchive() (logArchive archive.LogArchiveInterface) {
 }
 
 // NewClientManager creates and Init a new instance of ClientManager.
-func NewClientManager() (ClientManager, error) {
-	clientManager := ClientManager{}
-	err := clientManager.init()
+func NewClientManager(options *Options) (*ClientManager, error) {
+	clientManager := &ClientManager{}
+	err := clientManager.init(options)
 	if err != nil {
-		return ClientManager{}, err
+		return nil, err
 	}
 
 	return clientManager, nil