Separate options from client to make the responsibility more clear.

cmd/scheduler/main.go

@@ -67,12 +67,12 @@ func init() {
 	rootCmd.Flags().StringVar(&config.GPUSchedulerPolicy, "gpu-scheduler-policy", util.GPUSchedulerPolicySpread.String(), "GPU scheduler policy")
 	rootCmd.Flags().StringVar(&config.MetricsBindAddress, "metrics-bind-address", ":9395", "The TCP address that the scheduler should bind to for serving prometheus metrics(e.g. 127.0.0.1:9395, :9395)")
 	rootCmd.Flags().StringToStringVar(&config.NodeLabelSelector, "node-label-selector", nil, "key=value pairs separated by commas")
-	// add QPS and Burst to the global flagset
-	// qps and burst settings for the client-go client
-	rootCmd.Flags().Float32Var(&config.QPS, "kube-qps", 5.0, "QPS to use while talking with kube-apiserver.")
-	rootCmd.Flags().IntVar(&config.Burst, "kube-burst", 10, "Burst to use while talking with kube-apiserver.")
-	// Add profiling related flags
+
+	rootCmd.Flags().Float32Var(&config.QPS, "kube-qps", client.DefaultQPS, "QPS to use while talking with kube-apiserver.")
+	rootCmd.Flags().IntVar(&config.Burst, "kube-burst", client.DefaultBurst, "Burst to use while talking with kube-apiserver.")
+	rootCmd.Flags().IntVar(&config.Timeout, "kube-timeout", client.DefaultTimeout, "Timeout to use while talking with kube-apiserver.")
 	rootCmd.Flags().BoolVar(&enableProfiling, "profiling", false, "Enable pprof profiling via HTTP server")
+
 	rootCmd.PersistentFlags().AddGoFlagSet(device.GlobalFlagSet())
 	rootCmd.AddCommand(version.VersionCmd)
 	rootCmd.Flags().AddGoFlagSet(util.InitKlogFlags())
@@ -99,7 +99,12 @@ func injectProfilingRoute(router *httprouter.Router) {
 }
 
 func start() error {
-	client.InitGlobalClient(client.WithBurst(config.Burst), client.WithQPS(config.QPS))
+	client.InitGlobalClient(
+		client.WithBurst(config.Burst),
+		client.WithQPS(config.QPS),
+		client.WithTimeout(config.Timeout),
+	)
+
 	device.InitDevices()
 	sher = scheduler.NewScheduler()
 	sher.Start()

pkg/device/cambricon/device_test.go

@@ -453,7 +453,11 @@ func Test_PatchAnnotations(t *testing.T) {
 }
 
 func Test_setNodeLock(t *testing.T) {
-	client.InitGlobalClient(client.WithBurst(10), client.WithQPS(5.0))
+	client.InitGlobalClient(
+		client.WithBurst(10),
+		client.WithQPS(5.0),
+		client.WithTimeout(60),
+	)
 	tests := []struct {
 		name      string
 		node      corev1.Node

pkg/scheduler/config/config.go

@@ -21,6 +21,7 @@ import "github.com/Project-HAMi/HAMi/pkg/util"
 var (
 	QPS                float32
 	Burst              int
+	Timeout            int
 	HTTPBind           string
 	SchedulerName      string
 	MetricsBindAddress string

pkg/util/client/client.go

@@ -28,6 +28,12 @@ import (
 	"k8s.io/klog/v2"
 )
 
+type Client struct {
+	// Embedded kubernetes.Interface to avoid name conflicts.
+	kubernetes.Interface
+	config *rest.Config
+}
+
 var (
 	KubeClient kubernetes.Interface
 	once       sync.Once
@@ -37,71 +43,59 @@ func init() {
 	KubeClient = nil
 }
 
+// GetClient returns the global Kubernetes client.
 func GetClient() kubernetes.Interface {
 	return KubeClient
 }
 
-// Client is a kubernetes client.
-type Client struct {
-	Client kubernetes.Interface
-	QPS    float32
-	Burst  int
-}
-
-// WithQPS sets the QPS of the client.
-func WithQPS(qps float32) func(*Client) {
-	return func(c *Client) {
-		c.QPS = qps
+// NewClient creates a new Kubernetes client with the given options.
+func NewClient(opts ...Option) (*Client, error) {
+	restConfig, err := loadKubeConfig()
+	if err != nil {
+		return nil, fmt.Errorf("failed to load kubeconfig: %w", err)
 	}
-}
 
-func WithBurst(burst int) func(*Client) {
-	return func(c *Client) {
-		c.Burst = burst
+	// Apply WithDefaults option first to set default values.
+	WithDefaults()(restConfig)
+
+	// Then apply user-provided options that will override defaults if specified.
+	for _, opt := range opts {
+		opt(restConfig)
 	}
-}
 
-// NewClientWithConfig creates a new client with a given config.
-func NewClientWithConfig(config *rest.Config, opts ...func(*Client)) (*Client, error) {
-	client, err := kubernetes.NewForConfig(config)
+	clientset, err := kubernetes.NewForConfig(restConfig)
 	if err != nil {
-		return nil, err
-	}
-	c := &Client{
-		Client: client,
-	}
-	for _, opt := range opts {
-		opt(c)
+		return nil, fmt.Errorf("failed to create kubernetes client: %w", err)
 	}
-	return c, nil
+
+	return &Client{
+		Interface: clientset,
+		config:    restConfig,
+	}, nil
+}
+
+// InitGlobalClient initializes the global Kubernetes client with the given options.
+func InitGlobalClient(opts ...Option) {
+	once.Do(func() {
+		client, err := NewClient(opts...)
+		if err != nil {
+			klog.Fatalf("Failed to initialize global client: %v", err)
+		}
+		KubeClient = client.Interface
+	})
 }
 
-// NewClient creates a new client.
-func NewClient(ops ...func(*Client)) (*Client, error) {
+// loadKubeConfig loads Kubernetes configuration from the environment or in-cluster.
+func loadKubeConfig() (*rest.Config, error) {
 	kubeConfigPath := os.Getenv("KUBECONFIG")
 	if kubeConfigPath == "" {
 		kubeConfigPath = filepath.Join(os.Getenv("HOME"), ".kube", "config")
 	}
+
 	config, err := clientcmd.BuildConfigFromFlags("", kubeConfigPath)
 	if err != nil {
 		klog.Infof("BuildConfigFromFlags failed for file %s: %v. Using in-cluster config.", kubeConfigPath, err)
-		config, err = rest.InClusterConfig()
-		if err != nil {
-			return nil, fmt.Errorf("failed to get in-cluster config: %w", err)
-		}
-	}
-	c, err := NewClientWithConfig(config, ops...)
-	if err != nil {
-		return nil, fmt.Errorf("failed to create kubernetes client: %w", err)
-	}
-	return c, err
-}
-
-// InitGlobalClient creates a new global client.
-func InitGlobalClient(ops ...func(*Client)) {
-	c, err := NewClient(ops...)
-	if err != nil {
-		klog.Fatalf("new client error %s", err.Error())
+		return rest.InClusterConfig()
 	}
-	KubeClient = c.Client
+	return config, nil
 }

pkg/util/client/client_test.go

@@ -17,37 +17,26 @@ limitations under the License.
 package client
 
 import (
+	"context"
 	"errors"
+	"fmt"
 	"os"
 	"path/filepath"
+	"sync"
 	"testing"
+	"time"
 
+	"gotest.tools/v3/assert"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/rest"
 	"k8s.io/client-go/tools/clientcmd"
-	clientcmdapi "k8s.io/client-go/tools/clientcmd/api"
 )
 
-// MockClientConfig is a mock implementation of clientcmd.ClientConfig.
-type MockClientConfig struct {
-	config *rest.Config
-	err    error
-}
-
-func (m *MockClientConfig) RawConfig() (clientcmdapi.Config, error) {
-	return clientcmdapi.Config{}, nil
-}
-
-func (m *MockClientConfig) ClientConfig() (*rest.Config, error) {
-	return m.config, m.err
-}
-
-func (m *MockClientConfig) Namespace() (string, bool, error) {
-	return "", false, nil
-}
-
-func (m *MockClientConfig) ConfigAccess() clientcmd.ConfigAccess {
-	return nil
-}
+// Mock functions for testing.
+var (
+	buildConfigFromFlags = clientcmd.BuildConfigFromFlags
+	inClusterConfig      = rest.InClusterConfig
+)
 
 // TestGetClient tests the GetClient function.
 func TestGetClient(t *testing.T) {
@@ -117,8 +106,152 @@ func TestGetClient(t *testing.T) {
 	}
 }
 
-// Mock functions for testing.
-var (
-	buildConfigFromFlags = clientcmd.BuildConfigFromFlags
-	inClusterConfig      = rest.InClusterConfig
-)
+// TestClientWithOptions tests client initialization with options.
+func TestClientWithOptions(t *testing.T) {
+	KubeClient = nil
+	once = sync.Once{}
+
+	timeout := 1
+	client, _ := NewClient(WithTimeout(timeout))
+
+	assert.Equal(t, client.config.Timeout, time.Duration(timeout)*time.Second)
+	assert.Equal(t, client.config.QPS, DefaultQPS)
+	assert.Equal(t, client.config.Burst, DefaultBurst)
+
+	KubeClient = nil
+	once = sync.Once{}
+
+	qps := float32(50.0)
+	client, _ = NewClient(WithQPS(qps))
+
+	assert.Equal(t, client.config.Timeout, time.Duration(DefaultTimeout)*time.Second)
+	assert.Equal(t, client.config.QPS, qps)
+	assert.Equal(t, client.config.Burst, DefaultBurst)
+
+	KubeClient = nil
+	once = sync.Once{}
+	burst := 100
+	client, _ = NewClient(WithBurst(burst))
+
+	assert.Equal(t, client.config.Timeout, time.Duration(DefaultTimeout)*time.Second)
+	assert.Equal(t, client.config.QPS, DefaultQPS)
+	assert.Equal(t, client.config.Burst, burst)
+
+	KubeClient = nil
+	once = sync.Once{}
+	timeout = 2
+	qps = 0.5
+	burst = 100
+	client, _ = NewClient(WithTimeout(timeout), WithQPS(qps), WithBurst(burst))
+
+	assert.Equal(t, client.config.Timeout, time.Duration(timeout)*time.Second)
+	assert.Equal(t, client.config.QPS, qps)
+	assert.Equal(t, client.config.Burst, burst)
+}
+
+// TestClientRealNodePerformance tests the performance with a real Kubernetes cluster if available.
+func TestClientRealNodePerformance(t *testing.T) {
+
+	skipRealClusterTest := true
+	// Skip this test by default as it requires a real Kubernetes cluster.
+	if skipRealClusterTest == true {
+		t.Skip("Skipping real cluster test. Set TEST_WITH_REAL_CLUSTER=true to run this test.")
+	}
+
+	tests := []struct {
+		name    string
+		qps     float32
+		burst   int
+		updates int
+		timeout int
+	}{
+		{
+			name:    "Real Cluster - Low QPS and Burst",
+			qps:     1,
+			burst:   1,
+			updates: 10,
+			timeout: 1,
+		},
+		{
+			name:    "Real Cluster - Standard Timeout",
+			qps:     5,
+			burst:   10,
+			updates: 10,
+			timeout: 5,
+		},
+		{
+			name:    "Real Cluster - High Timeout",
+			qps:     10,
+			burst:   20,
+			updates: 15,
+			timeout: 10,
+		},
+		{
+			name:    "Real Cluster - Very Short Timeout",
+			qps:     5,
+			burst:   5,
+			updates: 5,
+			timeout: 1,
+		},
+	}
+
+	labelKey := "test-performance-label"
+	var nodeName string
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			client, err := NewClient(WithQPS(tt.qps), WithBurst(tt.burst), WithTimeout(tt.timeout))
+			if err != nil {
+				t.Fatalf("Failed to create client: %v", err)
+			}
+
+			if nodeName == "" {
+				nodes, err := client.CoreV1().Nodes().List(context.TODO(), metav1.ListOptions{})
+				if err != nil {
+					t.Fatalf("Failed to list nodes: %v", err)
+				}
+				if len(nodes.Items) == 0 {
+					t.Fatal("No nodes found in the cluster")
+				}
+				nodeName = nodes.Items[0].Name
+				t.Logf("Using node %s for testing", nodeName)
+			}
+			start := time.Now()
+			for i := 0; i < tt.updates; i++ {
+				labelValue := fmt.Sprintf("perf-test-value-%d", i)
+				node, err := client.CoreV1().Nodes().Get(context.TODO(), nodeName, metav1.GetOptions{})
+				if err != nil {
+					t.Fatalf("Failed to get node: %v", err)
+				}
+				if node.Labels == nil {
+					node.Labels = make(map[string]string)
+				}
+				node.Labels[labelKey] = labelValue
+				_, err = client.CoreV1().Nodes().Update(context.TODO(), node, metav1.UpdateOptions{})
+				if err != nil {
+					t.Fatalf("Failed to update node: %v", err)
+				}
+			}
+
+			elapsed := time.Since(start)
+
+			node, err := client.CoreV1().Nodes().Get(context.TODO(), nodeName, metav1.GetOptions{})
+			if err != nil {
+				t.Fatalf("Failed to get node during cleanup: %v", err)
+			}
+			delete(node.Labels, labelKey)
+			_, err = client.CoreV1().Nodes().Update(context.TODO(), node, metav1.UpdateOptions{})
+			if err != nil {
+				t.Fatalf("Failed to cleanup test label: %v", err)
+			}
+
+			opsPerSecond := float64(tt.updates) / elapsed.Seconds()
+
+			t.Logf("Real cluster performance test results for %s:", tt.name)
+			t.Logf("  - QPS: %.1f, Burst: %d", tt.qps, tt.burst)
+			t.Logf("  - Updates performed: %d", tt.updates)
+			t.Logf("  - Total time: %v", elapsed)
+			t.Logf("  - Operations per second: %.2f", opsPerSecond)
+		})
+	}
+}