fix(scheduler): honour HAMI_NODELOCK_EXPIRE when the flag is not set

cmd/scheduler/main.go

@@ -54,7 +54,7 @@ var (
 		Short: "kubernetes vgpu scheduler",
 		RunE: func(cmd *cobra.Command, args []string) error {
 			flag.PrintPFlags(cmd.Flags())
-			return start()
+			return start(cmd.Flags().Changed("node-lock-timeout"))
 		},
 	}
 )
@@ -109,9 +109,18 @@ func injectProfilingRoute(router *httprouter.Router) {
 	})
 }
 
-func start() error {
-	// Initialize node lock timeout from config
-	nodelock.NodeLockTimeout = config.NodeLockTimeout
+// applyNodeLockTimeout applies the flag only when it was set explicitly, so it
+// does not overwrite the value nodelock's init took from HAMI_NODELOCK_EXPIRE.
+func applyNodeLockTimeout(flagSet bool) {
+	if flagSet {
+		nodelock.NodeLockTimeout = config.NodeLockTimeout
+		return
+	}
+	config.NodeLockTimeout = nodelock.NodeLockTimeout
+}
+
+func start(nodeLockTimeoutFlagSet bool) error {
+	applyNodeLockTimeout(nodeLockTimeoutFlagSet)
 	klog.InfoS("Set node lock timeout", "timeout", nodelock.NodeLockTimeout)
 	client.InitGlobalClient(
 		client.WithBurst(config.Burst),

cmd/scheduler/nodelock_timeout_test.go

@@ -0,0 +1,59 @@
+/*
+Copyright 2024 The HAMi Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package main
+
+import (
+	"testing"
+	"time"
+
+	"github.com/Project-HAMi/HAMi/pkg/scheduler/config"
+	"github.com/Project-HAMi/HAMi/pkg/util/nodelock"
+)
+
+func TestApplyNodeLockTimeout(t *testing.T) {
+	const defaultTimeout = 5 * time.Minute
+
+	tests := []struct {
+		name    string
+		flagSet bool
+		flag    time.Duration
+		env     time.Duration
+		want    time.Duration
+	}{
+		{name: "flag unset keeps env value", flagSet: false, flag: defaultTimeout, env: 5 * time.Second, want: 5 * time.Second},
+		{name: "flag set overrides env", flagSet: true, flag: 5 * time.Second, env: defaultTimeout, want: 5 * time.Second},
+		{name: "neither set keeps default", flagSet: false, flag: defaultTimeout, env: defaultTimeout, want: defaultTimeout},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			origFlag, origEnv := config.NodeLockTimeout, nodelock.NodeLockTimeout
+			t.Cleanup(func() {
+				config.NodeLockTimeout, nodelock.NodeLockTimeout = origFlag, origEnv
+			})
+			config.NodeLockTimeout, nodelock.NodeLockTimeout = tc.flag, tc.env
+
+			applyNodeLockTimeout(tc.flagSet)
+
+			if nodelock.NodeLockTimeout != tc.want {
+				t.Errorf("nodelock.NodeLockTimeout = %v, want %v", nodelock.NodeLockTimeout, tc.want)
+			}
+			if config.NodeLockTimeout != tc.want {
+				t.Errorf("config.NodeLockTimeout = %v, want %v", config.NodeLockTimeout, tc.want)
+			}
+		})
+	}
+}