Add feature gate

Signed-off-by: Heba Elayoty <hebaelayoty@gmail.com>
kaito-project · May 4, 2024 · 26f3f79 · 26f3f79
1 parent 754d983
commit 26f3f79
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 21 deletions.
diff --git a/pkg/controllers/workspace_controller.go b/pkg/controllers/workspace_controller.go
@@ -9,6 +9,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/azure/kaito/pkg/featuregates"
 	"github.com/azure/kaito/pkg/nodeclaim"
 	"github.com/azure/kaito/pkg/tuning"
 	"github.com/azure/kaito/pkg/utils/consts"
@@ -225,11 +226,14 @@ func (c *WorkspaceReconciler) applyWorkspaceResource(ctx context.Context, wObj *
 	if err := machine.WaitForPendingMachines(ctx, wObj, c.Client); err != nil {
 		return err
 	}
-	// Check nodeClaims only if the k8s version is greater than 1.29. As nodeClaim is supported from 1.29.
-	if c.KubernetesVersion.Major >= "29" {
-		// Wait for pending nodeClaims if any before we decide whether to create new node or not.
-		if err := nodeclaim.WaitForPendingNodeClaims(ctx, wObj, c.Client); err != nil {
-			return err
+
+	if featuregates.FeatureGates[consts.FeatureFlagKarpenter] {
+		// Check nodeClaims only if the k8s version is greater than 1.29. As nodeClaim is supported from 1.29.
+		if c.KubernetesVersion.Major >= consts.KubernetesSupportedMajorVersion {
+			// Wait for pending nodeClaims if any before we decide whether to create new node or not.
+			if err := nodeclaim.WaitForPendingNodeClaims(ctx, wObj, c.Client); err != nil {
+				return err
+			}
 		}
 	}
 
@@ -245,7 +249,8 @@ func (c *WorkspaceReconciler) applyWorkspaceResource(ctx context.Context, wObj *
 
 	if newNodesCount > 0 {
 		klog.InfoS("need to create more nodes", "NodeCount", newNodesCount)
-		if c.KubernetesVersion.Major < "29" {
+		if featuregates.FeatureGates[consts.FeatureFlagKarpenter] &&
+			c.KubernetesVersion.Major < consts.KubernetesSupportedMajorVersion {
 			if err := c.updateStatusConditionIfNotMatch(ctx, wObj, kaitov1alpha1.WorkspaceConditionTypeMachineStatus, metav1.ConditionUnknown,
 				"CreateMachinePending", fmt.Sprintf("creating %d machines", newNodesCount)); err != nil {
 				klog.ErrorS(err, "failed to update workspace status", "workspace", klog.KObj(wObj))
@@ -288,7 +293,8 @@ func (c *WorkspaceReconciler) applyWorkspaceResource(ctx context.Context, wObj *
 		}
 	}
 
-	if c.KubernetesVersion.Major < "29" {
+	if featuregates.FeatureGates[consts.FeatureFlagKarpenter] &&
+		c.KubernetesVersion.Major < consts.KubernetesSupportedMajorVersion {
 		if err = c.updateStatusConditionIfNotMatch(ctx, wObj, kaitov1alpha1.WorkspaceConditionTypeMachineStatus, metav1.ConditionTrue,
 			"installNodePluginsSuccess", "machines plugins have been installed successfully"); err != nil {
 			klog.ErrorS(err, "failed to update workspace status", "workspace", klog.KObj(wObj))
@@ -643,7 +649,7 @@ func (c *WorkspaceReconciler) SetupWithManager(mgr ctrl.Manager) error {
 		For(&kaitov1alpha1.Workspace{}).
 		Owns(&appsv1.Deployment{}).
 		Owns(&appsv1.StatefulSet{}).
-		Watches(&v1alpha5.Machine{}, c.watchMachines()). // watches for machine with labels indicating workspace name.
+		Watches(&v1alpha5.Machine{}, c.watchMachines()).    // watches for machine with labels indicating workspace name.
 		Watches(&v1beta1.NodeClaim{}, c.watchNodeClaims()). // watches for nodeClaim with labels indicating workspace name.
 		WithOptions(controller.Options{MaxConcurrentReconciles: 5}).
 		Complete(c)

diff --git a/pkg/controllers/workspace_controller_test.go b/pkg/controllers/workspace_controller_test.go
@@ -13,8 +13,10 @@ import (
 
 	"github.com/aws/karpenter-core/pkg/apis/v1alpha5"
 	"github.com/azure/kaito/api/v1alpha1"
+	"github.com/azure/kaito/pkg/featuregates"
 	"github.com/azure/kaito/pkg/machine"
 	"github.com/azure/kaito/pkg/nodeclaim"
+	"github.com/azure/kaito/pkg/utils/consts"
 	"github.com/azure/kaito/pkg/utils/test"
 	"github.com/stretchr/testify/mock"
 	"gotest.tools/assert"
@@ -717,10 +719,11 @@ func TestGetAllQualifiedNodes(t *testing.T) {
 func TestApplyWorkspaceResource(t *testing.T) {
 	test.RegisterTestModel()
 	testcases := map[string]struct {
-		callMocks     func(c *test.MockClient)
-		k8sVersion    *version.Info
-		expectedError error
-		workspace     v1alpha1.Workspace
+		callMocks                   func(c *test.MockClient)
+		k8sVersion                  *version.Info
+		karpenterFeatureGateEnabled bool
+		expectedError               error
+		workspace                   v1alpha1.Workspace
 	}{
 		"Fail to apply workspace because associated machines cannot be retrieved": {
 			callMocks: func(c *test.MockClient) {
@@ -759,9 +762,10 @@ func TestApplyWorkspaceResource(t *testing.T) {
 				c.On("List", mock.IsType(context.Background()), mock.IsType(&v1beta1.NodeClaimList{}), mock.Anything).Return(errors.New("failed to retrieve nodeClaims"))
 
 			},
-			k8sVersion:    &version.Info{Major: "29"},
-			workspace:     *test.MockWorkspaceDistributedModel,
-			expectedError: errors.New("failed to retrieve nodeClaims"),
+			k8sVersion:                  &version.Info{Major: "29"},
+			karpenterFeatureGateEnabled: true,
+			workspace:                   *test.MockWorkspaceDistributedModel,
+			expectedError:               errors.New("failed to retrieve nodeClaims"),
 		},
 		"Fail to apply workspace with nodeClaims because can't get qualified nodes": {
 			callMocks: func(c *test.MockClient) {
@@ -782,9 +786,10 @@ func TestApplyWorkspaceResource(t *testing.T) {
 
 				c.On("List", mock.IsType(context.Background()), mock.IsType(&corev1.NodeList{}), mock.Anything).Return(errors.New("failed to list nodes"))
 			},
-			k8sVersion:    &version.Info{Major: "29"},
-			workspace:     *test.MockWorkspaceDistributedModel,
-			expectedError: errors.New("failed to list nodes"),
+			k8sVersion:                  &version.Info{Major: "29"},
+			karpenterFeatureGateEnabled: true,
+			workspace:                   *test.MockWorkspaceDistributedModel,
+			expectedError:               errors.New("failed to list nodes"),
 		},
 		"Successfully apply workspace resource with machine": {
 			callMocks: func(c *test.MockClient) {
@@ -835,9 +840,10 @@ func TestApplyWorkspaceResource(t *testing.T) {
 				c.StatusMock.On("Update", mock.IsType(context.Background()), mock.IsType(&v1alpha1.Workspace{}), mock.Anything).Return(nil)
 
 			},
-			k8sVersion:    &version.Info{Major: "29"},
-			workspace:     *test.MockWorkspaceDistributedModel,
-			expectedError: nil,
+			k8sVersion:                  &version.Info{Major: "29"},
+			karpenterFeatureGateEnabled: true,
+			workspace:                   *test.MockWorkspaceDistributedModel,
+			expectedError:               nil,
 		},
 	}
 
@@ -874,6 +880,7 @@ func TestApplyWorkspaceResource(t *testing.T) {
 				Scheme:            test.NewTestScheme(),
 				KubernetesVersion: tc.k8sVersion,
 			}
+			featuregates.FeatureGates[consts.FeatureFlagKarpenter] = tc.karpenterFeatureGateEnabled
 			ctx := context.Background()
 
 			err := reconciler.applyWorkspaceResource(ctx, &tc.workspace)

diff --git a/pkg/utils/consts/consts.go b/pkg/utils/consts/consts.go
@@ -8,4 +8,7 @@ const (
 	WorkspaceFinalizer            = "workspace.finalizer.kaito.sh"
 	DefaultReleaseNamespaceEnvVar = "RELEASE_NAMESPACE"
 	FeatureFlagKarpenter          = "Karpenter"
+
+	//	KubernetesSupportedMajorVersion is the latest major version of Kubernetes that Kaito supports.
+	KubernetesSupportedMajorVersion = "29"
 )