wip tests for provider controller

sky-uk · Dec 17, 2024 · 4cf3376 · 4cf3376
1 parent 79dc7b4
commit 4cf3376
Show file tree

Hide file tree

Showing 9 changed files with 7,588 additions and 46 deletions.
diff --git a/apis/config/v1alpha6/zz_generated.deepcopy.go b/apis/config/v1alpha6/zz_generated.deepcopy.go
diff --git a/config/crd/bases/config.kubeflow.org_kfpcontrollerconfigs.yaml b/config/crd/bases/config.kubeflow.org_kfpcontrollerconfigs.yaml
diff --git a/config/crd/bases/pipelines.kubeflow.org_experiments.yaml b/config/crd/bases/pipelines.kubeflow.org_experiments.yaml
@@ -247,6 +247,8 @@ spec:
                   name:
                     type: string
                 type: object
+              serving:
+                type: string
               synchronizationState:
                 type: string
               version:

diff --git a/config/crd/bases/pipelines.kubeflow.org_pipelines.yaml b/config/crd/bases/pipelines.kubeflow.org_pipelines.yaml
@@ -304,6 +304,8 @@ spec:
                   name:
                     type: string
                 type: object
+              serving:
+                type: string
               synchronizationState:
                 type: string
               version:

diff --git a/config/crd/bases/pipelines.kubeflow.org_providers.yaml b/config/crd/bases/pipelines.kubeflow.org_providers.yaml
@@ -177,6 +177,8 @@ spec:
             type: object
           spec:
             properties:
+              cliImage:
+                type: string
               defaultBeamArgs:
                 items:
                   properties:
@@ -205,6 +207,7 @@ spec:
               type:
                 type: string
             required:
+            - cliImage
             - executionMode
             - image
             - pipelineRootStorage
@@ -291,6 +294,8 @@ spec:
                   name:
                     type: string
                 type: object
+              serving:
+                type: string
               synchronizationState:
                 type: string
               version:

diff --git a/config/crd/bases/pipelines.kubeflow.org_runs.yaml b/config/crd/bases/pipelines.kubeflow.org_runs.yaml
@@ -395,6 +395,8 @@ spec:
                   name:
                     type: string
                 type: object
+              serving:
+                type: string
               synchronizationState:
                 type: string
               version:

diff --git a/config/crd/bases/pipelines.kubeflow.org_runschedules.yaml b/config/crd/bases/pipelines.kubeflow.org_runschedules.yaml
@@ -316,6 +316,8 @@ spec:
                   name:
                     type: string
                 type: object
+              serving:
+                type: string
               synchronizationState:
                 type: string
               version:

diff --git a/controllers/pipelines/provider_controller.go b/controllers/pipelines/provider_controller.go
@@ -24,14 +24,14 @@ type ProviderReconciler struct {
 	Scheme *runtime.Scheme
 }
 
-func NewProviderReconciler(ec K8sExecutionContext, scheme *runtime.Scheme, config config.KfpControllerConfigSpec) *ProviderReconciler {
+func NewProviderReconciler(ec K8sExecutionContext, config config.KfpControllerConfigSpec) *ProviderReconciler {
 	return &ProviderReconciler{
 		StateHandler: StateHandler[*pipelinesv1.Provider]{},
 		ResourceReconciler: ResourceReconciler[*pipelinesv1.Provider]{
 			EC:     ec,
 			Config: config,
 		},
-		Scheme: scheme,
+		Scheme: ec.Scheme,
 	}
 }
 
@@ -51,28 +51,18 @@ func (r *ProviderReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
 		return ctrl.Result{}, client.IgnoreNotFound(err)
 	}
 
-	logger.V(2).Info("found provider", "resource", provider)
+	logger.Info("found provider", "resource", provider)
 
-	existingProviderDeployment, err := r.fetchProviderServiceDeployment(ctx, provider, req.Namespace)
+	desiredProviderDeployment, err := r.constructProviderServiceDeployment(provider, req.Namespace, r.Config)
 	if err != nil {
-		logger.Error(err, "unable to fetch provider service deployment")
+		logger.Error(err, "unable to construct provider service deployment")
 		return ctrl.Result{}, err
 	}
+	logger.Info("created provider deployment", "deployment", desiredProviderDeployment)
 
-	if existingProviderDeployment == nil {
-		providerDeployment, err := r.constructProviderServiceDeployment(provider, req.Namespace, r.Config)
-		if err != nil {
-			logger.Error(err, "unable to construct provider service deployment")
-			return ctrl.Result{}, err
-		}
-		if err := r.EC.Client.Create(ctx, providerDeployment); err != nil {
-			logger.Error(err, "unable to create provider service deployment")
-			return ctrl.Result{}, err
-		}
-		logger.Info("created provider service deployment", "resource", providerDeployment)
-	} else {
-		logger.V(2).Info("found provider service deployment", "resource", existingProviderDeployment)
-		// ensure that deployment matches the desired state
+	if err := r.EC.Client.Patch(ctx, desiredProviderDeployment, client.Apply, client.FieldOwner("provider-controller-provider")); err != nil {
+		logger.Error(err, "unable to update provider service deployment")
+		return ctrl.Result{}, err
 	}
 
 	duration := time.Since(startTime)
@@ -89,26 +79,6 @@ func (r *ProviderReconciler) SetupWithManager(mgr ctrl.Manager) error {
 		Complete(r)
 }
 
-func (r *ProviderReconciler) fetchProviderServiceDeployment(ctx context.Context, provider *pipelinesv1.Provider, namespace string) (*appsv1.Deployment, error) {
-	deployments := &appsv1.DeploymentList{}
-	if err := r.EC.Client.NonCached.List(ctx, deployments, client.InNamespace(provider.Namespace)); err != nil {
-		return nil, err
-	}
-	var providerDeployment *appsv1.Deployment
-	for _, deployment := range deployments.Items {
-		if metav1.IsControlledBy(&deployment, provider) {
-			providerDeployment = &deployment
-			break // TODO: handle multiple deployments
-		}
-	}
-
-	if providerDeployment == nil {
-		// Do we need to return an error here?
-		return nil, fmt.Errorf("provider service deployment not found")
-	}
-	return providerDeployment, nil
-}
-
 func (r *ProviderReconciler) constructProviderServiceDeployment(provider *pipelinesv1.Provider, namespace string, config config.KfpControllerConfigSpec) (*appsv1.Deployment, error) {
 	replicas := int32(config.DefaultProviderValues.Replicas)
 
@@ -125,12 +95,8 @@ func (r *ProviderReconciler) constructProviderServiceDeployment(provider *pipeli
 		}
 	}
 
-	const providerLabelKey = "provider"
-	podTemplate.ObjectMeta.Labels[providerLabelKey] = provider.Name
 	podTemplate.Spec.ServiceAccountName = provider.Spec.ServiceAccount
-
 	labels := config.DefaultProviderValues.Labels
-	labels[providerLabelKey] = provider.Name
 
 	deployment := &appsv1.Deployment{
 		ObjectMeta: metav1.ObjectMeta{

diff --git a/main.go b/main.go
@@ -137,7 +137,7 @@ func main() {
 		os.Exit(1)
 	}
 
-	if err = pipelinescontrollers.NewProviderReconciler(ec, scheme, ctrlConfig.Spec).SetupWithManager(mgr); err != nil {
+	if err = pipelinescontrollers.NewProviderReconciler(ec, ctrlConfig.Spec).SetupWithManager(mgr); err != nil {
 		setupLog.Error(err, "unable to create controller", "controller", "Provider")
 		os.Exit(1)
 	}