diff --git a/.gitignore b/.gitignore index 4f3add4d1..8e94b11d0 100644 --- a/.gitignore +++ b/.gitignore @@ -9,5 +9,6 @@ libvgpu.so vendor license vgpuvalidator -_output/coverage/coverage_pkg.txt +_output/ +coverage.out .DS_Store diff --git a/Makefile b/Makefile index 3dd3709b2..148b9ed8b 100644 --- a/Makefile +++ b/Makefile @@ -43,7 +43,7 @@ clean: $(GO) clean -r -x ./cmd/... -rm -rf $(OUTPUT_DIR) -.PHONY: all build docker clean $(CMDS) +.PHONY: all build docker clean test $(CMDS) test: mkdir -p ./_output/coverage/ diff --git a/go.mod b/go.mod index dacc9cb0c..53a98b907 100644 --- a/go.mod +++ b/go.mod @@ -18,6 +18,7 @@ require ( github.com/prometheus/client_golang v1.18.0 github.com/sirupsen/logrus v1.9.3 github.com/spf13/cobra v1.8.1 + github.com/spf13/pflag v1.0.5 github.com/stretchr/testify v1.9.0 github.com/urfave/cli/v2 v2.27.1 golang.org/x/net v0.35.0 @@ -26,6 +27,7 @@ require ( google.golang.org/grpc v1.63.2 google.golang.org/protobuf v1.33.0 gopkg.in/yaml.v2 v2.4.0 + gopkg.in/yaml.v3 v3.0.1 gotest.tools/v3 v3.5.1 k8s.io/api v0.29.3 k8s.io/apimachinery v0.29.3 @@ -72,7 +74,6 @@ require ( github.com/prometheus/common v0.48.0 // indirect github.com/prometheus/procfs v0.13.0 // indirect github.com/russross/blackfriday/v2 v2.1.0 // indirect - github.com/spf13/pflag v1.0.5 // indirect github.com/syndtr/gocapability v0.0.0-20200815063812-42c35b437635 // indirect github.com/xrash/smetrics v0.0.0-20201216005158-039620a65673 // indirect golang.org/x/mod v0.17.0 // indirect @@ -84,7 +85,6 @@ require ( google.golang.org/appengine v1.6.8 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20240227224415-6ceb2ff114de // indirect gopkg.in/inf.v0 v0.9.1 // indirect - gopkg.in/yaml.v3 v3.0.1 // indirect k8s.io/kube-openapi v0.0.0-20240227032403-f107216b40e2 // indirect k8s.io/utils v0.0.0-20240102154912-e7106e64919e // indirect sigs.k8s.io/json v0.0.0-20221116044647-bc3834ca7abd // indirect diff --git a/go.sum b/go.sum index de9e20bbc..b0381f8eb 100644 --- a/go.sum +++ b/go.sum @@ -177,10 +177,6 @@ golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLL golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg= golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c= -golang.org/x/net v0.26.0 h1:soB7SVo0PWrY4vPW/+ay0jKDNScG2X9wFeYlXIvJsOQ= -golang.org/x/net v0.26.0/go.mod h1:5YKkiSynbBIh3p6iOc/vibscux0x38BZDkn8sCUPxHE= -golang.org/x/net v0.33.0 h1:74SYHlV8BIgHIFC/LrYkOGIwL19eTYXQ5wc6TBuO36I= -golang.org/x/net v0.33.0/go.mod h1:HXLR5J+9DxmrqMwG9qjGCxZ+zKXxBru04zlTvWlWuN4= golang.org/x/net v0.35.0 h1:T5GQRQb2y08kTAByq9L4/bz8cipCdA8FbRTXewonqY8= golang.org/x/net v0.35.0/go.mod h1:EglIi67kWsHKlRzzVMUD93VMSWGFOMSZgxFjparz1Qk= golang.org/x/oauth2 v0.17.0 h1:6m3ZPmLEFdVxKKWnKq4VqZ60gutO35zm+zrAHVmHyDQ= @@ -189,10 +185,8 @@ golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJ golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= -golang.org/x/sync v0.7.0 h1:YsImfSBoP9QPYL0xyKJPq0gcaJdG3rInoqxTWbfQu9M= -golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= -golang.org/x/sync v0.10.0 h1:3NQrjDixjgGwUOCaF8w2+VYHv0Ve/vGYSbdkTa98gmQ= golang.org/x/sync v0.11.0 h1:GGz8+XQP4FvTTrjZPzNKTMFtSXH80RAzG+5ghFPgK9w= +golang.org/x/sync v0.11.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= @@ -204,28 +198,16 @@ golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBc golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.21.0 h1:rF+pYz3DAGSQAxAu1CbC7catZg4ebC4UIeIhKxBZvws= -golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= -golang.org/x/sys v0.28.0 h1:Fksou7UEQUWlKvIdsqzJmUmCX3cZuD2+P3XyyzwMhlA= -golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc= golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= -golang.org/x/term v0.21.0 h1:WVXCp+/EBEHOj53Rvu+7KiT/iElMrO8ACK16SMZ3jaA= -golang.org/x/term v0.21.0/go.mod h1:ooXLefLobQVslOqselCNF4SxFAaoS6KujMbsGzSDmX0= -golang.org/x/term v0.27.0 h1:WP60Sv1nlK1T6SupCHbXzSaN0b9wUmsPoRS9b61A23Q= -golang.org/x/term v0.27.0/go.mod h1:iMsnZpn0cago0GOrHO2+Y7u7JPn5AylBrcoWkElMTSM= golang.org/x/term v0.29.0 h1:L6pJp37ocefwRRtYPKSWOWzOtWSxVajvz2ldH/xi3iU= golang.org/x/term v0.29.0/go.mod h1:6bl4lRlvVuDgSf3179VpIxBF0o10JUpXWOnI7nErv7s= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ= golang.org/x/text v0.3.8/go.mod h1:E6s5w1FMmriuDzIBO73fBruAKo1PCIq6d2Q6DHfQ8WQ= -golang.org/x/text v0.16.0 h1:a94ExnEXNtEwYLGJSIUxnWoxoRz/ZcCsV63ROupILh4= -golang.org/x/text v0.16.0/go.mod h1:GhwF1Be+LQoKShO3cGOHzqOgRrGaYc9AvblQOmPVHnI= -golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo= -golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ= golang.org/x/text v0.22.0 h1:bofq7m3/HAFvbF51jz3Q9wLg3jkvSPuiZu/pD1XwgtM= golang.org/x/text v0.22.0/go.mod h1:YRoo4H8PVmsu+E3Ou7cqLVH8oXWIHVoX0jqUWALQhfY= golang.org/x/time v0.5.0 h1:o7cqy6amK/52YcAKIPlM3a+Fpj35zvRj2TP+e1xFSfk= diff --git a/pkg/device/cambricon/device.go b/pkg/device/cambricon/device.go index 8e2b48475..e2b33b3c4 100644 --- a/pkg/device/cambricon/device.go +++ b/pkg/device/cambricon/device.go @@ -155,22 +155,18 @@ func (dev *CambriconDevices) ReleaseNodeLock(n *corev1.Node, p *corev1.Pod) erro return nil } - patchData := []byte(`[ - { - "op": "remove", - "path": "/metadata/annotations/cambricon.com~1dsmlu.lock" - } - ]`) - - _, err := client.GetClient().CoreV1().Nodes().Patch(context.TODO(), n.Name, types.JSONPatchType, patchData, metav1.PatchOptions{}) + newNode := n.DeepCopy() + delete(newNode.ObjectMeta.Annotations, DsmluLockTime) + _, err := client.GetClient().CoreV1().Nodes().Update(context.Background(), newNode, metav1.UpdateOptions{}) for i := 0; i < retry && err != nil; i++ { klog.ErrorS(err, "Failed to patch node annotation", "node", n.Name, "retry", i) time.Sleep(time.Duration(rand.Intn(i+1)) * 10 * time.Millisecond) - _, err = client.GetClient().CoreV1().Nodes().Patch(context.TODO(), n.Name, types.JSONPatchType, patchData, metav1.PatchOptions{}) + _, err = client.GetClient().CoreV1().Nodes().Update(context.Background(), newNode, metav1.UpdateOptions{}) } if err != nil { return fmt.Errorf("releaseNodeLock exceeds retry count %d", retry) } + delete(n.ObjectMeta.Annotations, DsmluLockTime) klog.InfoS("Node lock released", "node", n.Name) return nil } diff --git a/pkg/device/cambricon/device_test.go b/pkg/device/cambricon/device_test.go index 19d217d5a..11d627925 100644 --- a/pkg/device/cambricon/device_test.go +++ b/pkg/device/cambricon/device_test.go @@ -480,17 +480,9 @@ func Test_setNodeLock(t *testing.T) { }, expectErr: false, }, - { - name: "no node name", - node: corev1.Node{ - ObjectMeta: metav1.ObjectMeta{ - Annotations: map[string]string{}, - }, - }, - expectErr: false, - }, } + client.KubeClient = fake.NewSimpleClientset() k8sClient := client.GetClient() if k8sClient != nil { @@ -624,6 +616,7 @@ func Test_LockNode(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { node, pod, teardown, clientset := setupTest(t) + client.KubeClient = clientset defer teardown() // Set up the node with the specified annotations. diff --git a/pkg/device/devices_test.go b/pkg/device/devices_test.go index 18c574800..819d93cad 100644 --- a/pkg/device/devices_test.go +++ b/pkg/device/devices_test.go @@ -264,6 +264,21 @@ func createVNPUConfigs() []ascend.VNPUConfig { {Name: "vir16", Memory: 17476, AICore: 16}, }, }, + { + ChipName: "910B2", + CommonWord: "Ascend910B2", + ResourceName: "huawei.com/Ascend910B2", + ResourceMemoryName: "huawei.com/Ascend910B2-memory", + MemoryAllocatable: 65536, + MemoryCapacity: 65536, + AICore: 24, + AICPU: 6, + Templates: []ascend.Template{ + {Name: "vir03_1c_8g", Memory: 8192, AICore: 3, AICPU: 1}, + {Name: "vir06_1c_16g", Memory: 16384, AICore: 6, AICPU: 1}, + {Name: "vir12_3c_32g", Memory: 32768, AICore: 12, AICPU: 3}, + }, + }, { ChipName: "910B3", CommonWord: "Ascend910B", @@ -329,6 +344,7 @@ func setupTest(t *testing.T) (map[string]string, map[string]Devices) { iluvatar.IluvatarGPUDevice: iluvatar.IluvatarGPUCommonWord, mthreads.MthreadsGPUDevice: mthreads.MthreadsGPUCommonWord, metax.MetaxGPUDevice: metax.MetaxGPUCommonWord, + metax.MetaxSGPUDevice: metax.MetaxSGPUCommonWord, } return expectedDevices, devicesMap diff --git a/pkg/device/hygon/device_test.go b/pkg/device/hygon/device_test.go index 7d2f81cbd..81d20129a 100644 --- a/pkg/device/hygon/device_test.go +++ b/pkg/device/hygon/device_test.go @@ -28,6 +28,7 @@ import ( "gotest.tools/v3/assert" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/client-go/kubernetes/fake" "k8s.io/apimachinery/pkg/api/resource" "k8s.io/klog/v2" @@ -673,6 +674,7 @@ func Test_GenerateResourceRequests(t *testing.T) { } func Test_NodeCleanUp(t *testing.T) { + client.KubeClient = fake.NewSimpleClientset() tests := []struct { name string args string diff --git a/pkg/util/client/client.go b/pkg/util/client/client.go index 82a4137f9..730c0c251 100644 --- a/pkg/util/client/client.go +++ b/pkg/util/client/client.go @@ -33,12 +33,18 @@ var ( once sync.Once ) +func init() { + KubeClient = nil +} + func GetClient() kubernetes.Interface { once.Do(func() { var err error - KubeClient, err = newClient() - if err != nil { - klog.Fatalf("Failed to create Kubernetes client: %v", err) + if KubeClient == nil { + KubeClient, err = newClient() + if err != nil { + klog.Fatalf("Failed to create Kubernetes client: %v", err) + } } }) return KubeClient diff --git a/pkg/util/client/factory.go b/pkg/util/client/factory.go new file mode 100644 index 000000000..6fcbc3f32 --- /dev/null +++ b/pkg/util/client/factory.go @@ -0,0 +1,88 @@ +/* +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 client + +import ( + "sync" + + "k8s.io/client-go/kubernetes/fake" +) + +// KubeClientFactory is a factory class (singleton pattern). +type KubeClientFactory struct { + client KubeInterface +} + +var factoryOnce sync.Once + +// GetFactory gets the singleton factory object. +// Uses closure to maintain instance as a local variable. +var GetFactory = (func() func() *KubeClientFactory { + var instance *KubeClientFactory + + return func() *KubeClientFactory { + factoryOnce.Do(func() { + instance = &KubeClientFactory{} + instance.SetReal() // Use the real client by default. + }) + return instance + } +})() + +// NewInstance directly gets the Kubernetes client instance. +func NewInstance() KubeInterface { + return GetFactory().GetClient() +} + +func (f *KubeClientFactory) GetClient() KubeInterface { + //TODO: For compatibility with other direct "KubeClient" assignment call points, this line needs to be removed after replacement. + if KubeClient == nil { + f.client = &K8sClient{ + client: GetK8sClient().client, + } + KubeClient = GetK8sClient().client + } else { + f.client = &K8sClient{ + client: KubeClient, + } + } + return f.client +} + +func (f *KubeClientFactory) SetFake() *KubeClientFactory { + f.client = &K8sClient{ + client: fake.NewSimpleClientset(), + } + //TODO: For compatibility with other direct "KubeClient" assignment call points, this line needs to be removed after replacement. + KubeClient = fake.NewSimpleClientset() + return f +} + +func (f *KubeClientFactory) SetReal() *KubeClientFactory { + //TODO: For compatibility with other direct "KubeClient" assignment call points, this line needs to be removed after replacement. + if KubeClient == nil { + f.client = &K8sClient{ + client: GetK8sClient().client, + } + KubeClient = GetK8sClient().client + } else { + f.client = &K8sClient{ + client: KubeClient, + } + } + return f +} diff --git a/pkg/util/client/interface.go b/pkg/util/client/interface.go new file mode 100644 index 000000000..ad703b122 --- /dev/null +++ b/pkg/util/client/interface.go @@ -0,0 +1,34 @@ +/* +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 client + +import ( + "context" + + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/types" +) + +type KubeInterface interface { + GetNode(ctx context.Context, name string, opts metav1.GetOptions) (*corev1.Node, error) + ListPods(ctx context.Context, namespace string, opts metav1.ListOptions) (*corev1.PodList, error) + GetPod(ctx context.Context, namespace, name string, opts metav1.GetOptions) (*corev1.Pod, error) + PatchNode(ctx context.Context, name string, pt types.PatchType, data []byte, opts metav1.PatchOptions) (*corev1.Node, error) + PatchPod(ctx context.Context, namespace string, name string, pt types.PatchType, data []byte, opts metav1.PatchOptions) (*corev1.Pod, error) + CreateNode(ctx context.Context, node *corev1.Node, opts metav1.CreateOptions) (*corev1.Node, error) +} diff --git a/pkg/util/client/k8s_client.go b/pkg/util/client/k8s_client.go new file mode 100644 index 000000000..d69390b2d --- /dev/null +++ b/pkg/util/client/k8s_client.go @@ -0,0 +1,158 @@ +/* +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 client + +import ( + "context" + "fmt" + "os" + "path/filepath" + "sync" + + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/types" + "k8s.io/client-go/kubernetes" + "k8s.io/client-go/rest" + "k8s.io/client-go/tools/clientcmd" + "k8s.io/klog/v2" +) + +var ClientOnce sync.Once + +// GetK8sClient creates a real Kubernetes client (singleton pattern). +// Uses closure to maintain client as a local variable. +var GetK8sClient = (func() func() *K8sClient { + var client *K8sClient + + return func() *K8sClient { + var err error + ClientOnce.Do(func() { + client, err = createRealClient() + if err != nil { + klog.Fatalf("Failed to create Kubernetes client: %v", err) + } + }) + return client + } +})() + +type K8sClient struct { + client kubernetes.Interface +} + +var _ KubeInterface = (*K8sClient)(nil) + +func createRealClient() (*K8sClient, error) { + kubeConfigPath := os.Getenv("KUBECONFIG") + if kubeConfigPath == "" { + kubeConfigPath = filepath.Join(os.Getenv("HOME"), ".kube", "config") + } + + config, err := clientcmd.BuildConfigFromFlags("", kubeConfigPath) + if err != nil { + klog.ErrorS(err, "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) + } + } + + clientset, err := kubernetes.NewForConfig(config) + if err != nil { + return nil, fmt.Errorf("failed to create kubernetes client: %w", err) + } + + return &K8sClient{client: clientset}, nil +} + +func (c *K8sClient) GetNode(ctx context.Context, name string, opts metav1.GetOptions) (*corev1.Node, error) { + klog.V(4).InfoS("Retrieving node", "node", name) + node, err := c.client.CoreV1().Nodes().Get(ctx, name, opts) + if err != nil { + klog.ErrorS(err, "Failed to get node", "node", name) + return nil, fmt.Errorf("failed to get node %s: %w", name, err) + } + klog.V(4).InfoS("Successfully retrieved node", "node", name) + return node, nil +} + +func (c *K8sClient) GetPod(ctx context.Context, namespace, name string, opts metav1.GetOptions) (*corev1.Pod, error) { + klog.V(4).InfoS("Retrieving pod", "namespace", namespace, "pod", name) + pod, err := c.client.CoreV1().Pods(namespace).Get(ctx, name, opts) + if err != nil { + klog.ErrorS(err, "Failed to get pod", "namespace", namespace, "pod", name) + return nil, fmt.Errorf("failed to get pod %s/%s: %w", namespace, name, err) + } + klog.V(4).InfoS("Successfully retrieved pod", "namespace", namespace, "pod", name) + return pod, nil +} + +func (c *K8sClient) ListPods(ctx context.Context, namespace string, opts metav1.ListOptions) (*corev1.PodList, error) { + klog.V(4).InfoS("Listing pods", "namespace", namespace, "options", opts) + pods, err := c.client.CoreV1().Pods(namespace).List(ctx, opts) + if err != nil { + klog.ErrorS(err, "Failed to list pods", "namespace", namespace) + return nil, fmt.Errorf("failed to list pods in namespace %s: %w", namespace, err) + } + klog.V(4).InfoS("Successfully listed pods", "namespace", namespace, "count", len(pods.Items)) + return pods, nil +} + +func (c *K8sClient) PatchNode(ctx context.Context, name string, pt types.PatchType, data []byte, opts metav1.PatchOptions) (*corev1.Node, error) { + klog.V(4).InfoS("Patching node", "node", name, "patchType", pt) + node, err := c.client.CoreV1().Nodes().Patch(ctx, name, pt, data, opts) + if err != nil { + klog.ErrorS(err, "Failed to patch node", "node", name) + return nil, fmt.Errorf("failed to patch node %s: %w", name, err) + } + klog.V(4).InfoS("Successfully patched node", "node", name) + return node, nil +} + +func (c *K8sClient) PatchPod(ctx context.Context, namespace string, name string, pt types.PatchType, data []byte, opts metav1.PatchOptions) (*corev1.Pod, error) { + klog.V(4).InfoS("Patching pod", "namespace", namespace, "pod", name, "patchType", pt) + pod, err := c.client.CoreV1().Pods(namespace).Patch(ctx, name, pt, data, opts) + if err != nil { + klog.ErrorS(err, "Failed to patch pod", "namespace", namespace, "pod", name) + return nil, fmt.Errorf("failed to patch pod %s/%s: %w", namespace, name, err) + } + klog.V(4).InfoS("Successfully patched pod", "namespace", namespace, "pod", name) + return pod, nil +} + +func (c *K8sClient) CreateNode(ctx context.Context, node *corev1.Node, opts metav1.CreateOptions) (*corev1.Node, error) { + klog.V(4).InfoS("Creating node", "node", node.Name) + createdNode, err := c.client.CoreV1().Nodes().Create(ctx, node, opts) + if err != nil { + klog.ErrorS(err, "Failed to create node", "node", node.Name) + return nil, fmt.Errorf("failed to create node %s: %w", node.Name, err) + } + klog.V(4).InfoS("Successfully created node", "node", node.Name) + return createdNode, nil +} + +func (c *K8sClient) UpdateNode(ctx context.Context, node *corev1.Node, opts metav1.UpdateOptions) (*corev1.Node, error) { + klog.V(4).InfoS("Updating node", "node", node.Name) + updatedNode, err := c.client.CoreV1().Nodes().Update(ctx, node, opts) + if err != nil { + klog.ErrorS(err, "Failed to update node", "node", node.Name) + return nil, fmt.Errorf("failed to update node %s: %w", node.Name, err) + } + klog.V(4).InfoS("Successfully updated node", "node", node.Name) + return updatedNode, nil +} diff --git a/pkg/util/util.go b/pkg/util/util.go index 4a4bb2dba..5335ce2ce 100644 --- a/pkg/util/util.go +++ b/pkg/util/util.go @@ -63,7 +63,8 @@ func GetNode(nodename string) (*corev1.Node, error) { } klog.InfoS("Fetching node", "nodeName", nodename) - n, err := client.GetClient().CoreV1().Nodes().Get(context.Background(), nodename, metav1.GetOptions{}) + + n, err := client.NewInstance().GetNode(context.Background(), nodename, metav1.GetOptions{}) if err != nil { switch { case apierrors.IsNotFound(err): @@ -95,7 +96,7 @@ func GetPendingPod(ctx context.Context, node string) (*corev1.Pod, error) { podListOptions := metav1.ListOptions{ FieldSelector: selector, } - podlist, err := client.GetClient().CoreV1().Pods("").List(ctx, podListOptions) + podlist, err := client.NewInstance().ListPods(ctx, "", podListOptions) if err != nil { return nil, err } @@ -125,7 +126,7 @@ func GetPendingPod(ctx context.Context, node string) (*corev1.Pod, error) { } func GetAllocatePodByNode(ctx context.Context, nodeName string) (*corev1.Pod, error) { - node, err := client.GetClient().CoreV1().Nodes().Get(ctx, nodeName, metav1.GetOptions{}) + node, err := client.NewInstance().GetNode(ctx, nodeName, metav1.GetOptions{}) if err != nil { return nil, err } @@ -138,7 +139,7 @@ func GetAllocatePodByNode(ctx context.Context, nodeName string) (*corev1.Pod, er if ns == "" || name == "" { return nil, nil } - return client.GetClient().CoreV1().Pods(ns).Get(ctx, name, metav1.GetOptions{}) + return client.NewInstance().GetPod(ctx, ns, name, metav1.GetOptions{}) } return nil, nil } @@ -342,8 +343,7 @@ func PatchNodeAnnotations(node *corev1.Node, annotations map[string]string) erro if err != nil { return err } - _, err = client.GetClient().CoreV1().Nodes(). - Patch(context.Background(), node.Name, k8stypes.StrategicMergePatchType, bytes, metav1.PatchOptions{}) + _, err = client.NewInstance().PatchNode(context.Background(), node.Name, k8stypes.StrategicMergePatchType, bytes, metav1.PatchOptions{}) if err != nil { klog.Infoln("annotations=", annotations) klog.Infof("patch pod %v failed, %v", node.Name, err) @@ -374,8 +374,7 @@ func PatchPodAnnotations(pod *corev1.Pod, annotations map[string]string) error { return err } klog.V(5).Infof("patch pod %s/%s annotation content is %s", pod.Namespace, pod.Name, string(bytes)) - _, err = client.GetClient().CoreV1().Pods(pod.Namespace). - Patch(context.Background(), pod.Name, k8stypes.StrategicMergePatchType, bytes, metav1.PatchOptions{}) + _, err = client.NewInstance().PatchPod(context.Background(), pod.Namespace, pod.Name, k8stypes.StrategicMergePatchType, bytes, metav1.PatchOptions{}) if err != nil { klog.Infof("patch pod %v failed, %v", pod.Name, err) } diff --git a/pkg/util/util_test.go b/pkg/util/util_test.go index f8d4e782b..d04e70b98 100644 --- a/pkg/util/util_test.go +++ b/pkg/util/util_test.go @@ -26,7 +26,6 @@ import ( "gotest.tools/v3/assert" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" - "k8s.io/client-go/kubernetes/fake" "github.com/Project-HAMi/HAMi/pkg/util/client" ) @@ -38,7 +37,7 @@ func init() { inRequestDevices["NVIDIA"] = "hami.io/vgpu-devices-to-allocate" } -func TestExtractMigTemplatesFromUUID(t *testing.T) { +func Test_ExtractMigTemplatesFromUUID(t *testing.T) { testCases := []struct { name string uuid string @@ -110,7 +109,7 @@ func TestExtractMigTemplatesFromUUID(t *testing.T) { } } -func TestEmptyContainerDevicesCoding(t *testing.T) { +func Test_EmptyContainerDevicesCoding(t *testing.T) { cd1 := ContainerDevices{} s := EncodeContainerDevices(cd1) fmt.Println(s) @@ -118,7 +117,7 @@ func TestEmptyContainerDevicesCoding(t *testing.T) { assert.DeepEqual(t, cd1, cd2) } -func TestEmptyPodDeviceCoding(t *testing.T) { +func Test_EmptyPodDeviceCoding(t *testing.T) { pd1 := PodDevices{} s := EncodePodDevices(inRequestDevices, pd1) fmt.Println(s) @@ -126,7 +125,7 @@ func TestEmptyPodDeviceCoding(t *testing.T) { assert.DeepEqual(t, pd1, pd2) } -func TestPodDevicesCoding(t *testing.T) { +func Test_PodDevicesCoding(t *testing.T) { tests := []struct { name string args PodDevices @@ -252,7 +251,7 @@ func Test_DecodePodDevices(t *testing.T) { } } -func TestMarshalNodeDevices(t *testing.T) { +func Test_MarshalNodeDevices(t *testing.T) { type args struct { dlist []*DeviceInfo } @@ -332,7 +331,7 @@ func TestMarshalNodeDevices(t *testing.T) { } } -func TestUnMarshalNodeDevices(t *testing.T) { +func Test_UnMarshalNodeDevices(t *testing.T) { type args struct { str string } @@ -625,11 +624,17 @@ func Test_CheckHealth(t *testing.T) { } } -func TestMarkAnnotationsToDelete(t *testing.T) { - client.KubeClient = fake.NewSimpleClientset() - client.KubeClient.CoreV1().Nodes().Create(context.TODO(), &corev1.Node{ - ObjectMeta: metav1.ObjectMeta{Name: "node-worker2"}, - }, metav1.CreateOptions{}) +func Test_MarkAnnotationsToDelete(t *testing.T) { + client.GetFactory().SetFake().GetClient().CreateNode( + context.TODO(), + &corev1.Node{ + ObjectMeta: metav1.ObjectMeta{ + Name: "node-worker2", + }, + }, + metav1.CreateOptions{}, + ) + type args struct { devType string nn string