diff --git a/cmd/workspace/main.go b/cmd/workspace/main.go index fdb48684c..8a0e24f04 100644 --- a/cmd/workspace/main.go +++ b/cmd/workspace/main.go @@ -32,6 +32,7 @@ import ( "github.com/kaito-project/kaito/pkg/featuregates" "github.com/kaito-project/kaito/pkg/k8sclient" kaitoutils "github.com/kaito-project/kaito/pkg/utils" + "github.com/kaito-project/kaito/pkg/utils/consts" "github.com/kaito-project/kaito/pkg/utils/nodeclaim" "github.com/kaito-project/kaito/pkg/workspace/controllers" "github.com/kaito-project/kaito/pkg/workspace/webhooks" @@ -161,9 +162,11 @@ func main() { exitWithErrorFunc() } - err = nodeclaim.CheckNodeClass(ctx, kClient) - if err != nil { - exitWithErrorFunc() + if featuregates.FeatureGates[consts.FeatureFlagEnsureNodeClass] { + err := nodeclaim.CheckNodeClass(ctx, kClient) + if err != nil { + exitWithErrorFunc() + } } klog.InfoS("starting manager") diff --git a/pkg/featuregates/featuregates.go b/pkg/featuregates/featuregates.go index 015a2f072..9b5ddbf8c 100644 --- a/pkg/featuregates/featuregates.go +++ b/pkg/featuregates/featuregates.go @@ -13,9 +13,10 @@ import ( ) var ( - // FeatureGates is a map that holds the feature gates and their default values for Kaito. + // FeatureGates is a map that holds the feature gate names and their default values for Kaito. FeatureGates = map[string]bool{ - consts.FeatureFlagVLLM: true, + consts.FeatureFlagVLLM: true, + consts.FeatureFlagEnsureNodeClass: false, // Add more feature gates here } ) diff --git a/pkg/featuregates/featuregates_test.go b/pkg/featuregates/featuregates_test.go index 0fde9ac85..3b7499914 100644 --- a/pkg/featuregates/featuregates_test.go +++ b/pkg/featuregates/featuregates_test.go @@ -13,19 +13,22 @@ func TestParseFeatureGates(t *testing.T) { name string featureGates string expectedError bool - expectedValue string + targetFeature string + expectedValue bool }{ { - name: "WithValidEnableFeatureGates", + name: "WithValidEnableFeatureGates-vLLM", featureGates: "vLLM=true", expectedError: false, - expectedValue: "true", + targetFeature: "vLLM", + expectedValue: true, }, { - name: "WithDuplicateFeatureGates", + name: "WithDuplicateFeatureGates-vLLM", featureGates: "vLLM=false,vLLM=true", expectedError: false, - expectedValue: "true", // Apply the last value. + targetFeature: "vLLM", + expectedValue: true, // Apply the last value. }, { name: "WithInvalidFeatureGates", @@ -38,10 +41,25 @@ func TestParseFeatureGates(t *testing.T) { expectedError: true, }, { - name: "WithValidDisableFeatureGates", + name: "WithValidDisableFeatureGates-vLLM", featureGates: "vLLM=false", expectedError: false, - expectedValue: "false", + targetFeature: "vLLM", + expectedValue: false, + }, + { + name: "WithValidEnableFeatureGates-ensureNodeClass", + featureGates: "ensureNodeClass=true", + expectedError: false, + targetFeature: "ensureNodeClass", + expectedValue: true, + }, + { + name: "WithValidDisableFeatureGates-ensureNodeClass", + featureGates: "ensureNodeClass=false", + expectedError: false, + targetFeature: "ensureNodeClass", + expectedValue: false, }, { name: "WithEmptyFeatureGates", @@ -57,6 +75,9 @@ func TestParseFeatureGates(t *testing.T) { assert.Check(t, err != nil, "expected error but got nil") } else { assert.NilError(t, err) + if tt.targetFeature != "" && FeatureGates[tt.targetFeature] != tt.expectedValue { + t.Errorf("feature gate test %s fails", tt.name) + } } }) } diff --git a/pkg/utils/consts/consts.go b/pkg/utils/consts/consts.go index 68aab1cc7..3c8551828 100644 --- a/pkg/utils/consts/consts.go +++ b/pkg/utils/consts/consts.go @@ -20,7 +20,8 @@ const ( NvidiaGPU = "nvidia.com/gpu" // Feature flags - FeatureFlagVLLM = "vLLM" + FeatureFlagVLLM = "vLLM" + FeatureFlagEnsureNodeClass = "ensureNodeClass" // Nodeclaim related consts KaitoNodePoolName = "kaito"