fix v2.8.1 monitor for new version of cache file

pkg/monitor/nvidia/api/api.go

@@ -0,0 +1,75 @@
+/*
+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 api
+
+import "sync"
+
+type Header struct {
+	InitializedFlag int32
+	MajorVersion    int32
+	MinorVersion    int32
+}
+
+type UsageInfo interface {
+	DeviceMax() int
+	DeviceNum() int
+	DeviceMemoryContextSize(idx int) uint64
+	DeviceMemoryModuleSize(idx int) uint64
+	DeviceMemoryBufferSize(idx int) uint64
+	DeviceMemoryOffset(idx int) uint64
+	DeviceMemoryTotal(idx int) uint64
+	DeviceSmUtil(idx int) uint64
+	SetDeviceSmLimit(l uint64)
+	IsValidUUID(idx int) bool
+	DeviceUUID(idx int) string
+	DeviceMemoryLimit(idx int) uint64
+	SetDeviceMemoryLimit(l uint64)
+	LastKernelTime() int64
+	GetPriority() int
+	GetRecentKernel() int32
+	SetRecentKernel(v int32)
+	GetUtilizationSwitch() int32
+	SetUtilizationSwitch(v int32)
+}
+
+type CacheFactory interface {
+	Match(header *Header, fileSize int64) bool
+	Cast(data []byte) UsageInfo
+	Name() string
+}
+
+var (
+	factories   []CacheFactory
+	factoriesMu sync.RWMutex
+)
+
+func RegisterFactory(f CacheFactory) {
+	factoriesMu.Lock()
+	defer factoriesMu.Unlock()
+	factories = append(factories, f)
+}
+
+func FindFactory(header *Header, fileSize int64) CacheFactory {
+	factoriesMu.RLock()
+	defer factoriesMu.RUnlock()
+	for _, f := range factories {
+		if f.Match(header, fileSize) {
+			return f
+		}
+	}
+	return nil
+}

pkg/monitor/nvidia/cudevshr.go

@@ -27,8 +27,7 @@ import (
 	"time"
 	"unsafe"
 
-	v0 "github.com/Project-HAMi/HAMi/pkg/monitor/nvidia/v0"
-	v1 "github.com/Project-HAMi/HAMi/pkg/monitor/nvidia/v1"
+	"github.com/Project-HAMi/HAMi/pkg/monitor/nvidia/api"
 	"github.com/Project-HAMi/HAMi/pkg/util"
 
 	corev1 "k8s.io/api/core/v1"
@@ -44,34 +43,8 @@ import (
 
 const SharedRegionMagicFlag = 19920718
 
-type headerT struct {
-	initializedFlag int32
-	majorVersion    int32
-	minorVersion    int32
-}
-
-type UsageInfo interface {
-	DeviceMax() int
-	DeviceNum() int
-	DeviceMemoryContextSize(idx int) uint64
-	DeviceMemoryModuleSize(idx int) uint64
-	DeviceMemoryBufferSize(idx int) uint64
-	DeviceMemoryOffset(idx int) uint64
-	DeviceMemoryTotal(idx int) uint64
-	DeviceSmUtil(idx int) uint64
-	SetDeviceSmLimit(l uint64)
-	IsValidUUID(idx int) bool
-	DeviceUUID(idx int) string
-	DeviceMemoryLimit(idx int) uint64
-	SetDeviceMemoryLimit(l uint64)
-	LastKernelTime() int64
-	//UsedMemory(idx int) (uint64, error)
-	GetPriority() int
-	GetRecentKernel() int32
-	SetRecentKernel(v int32)
-	GetUtilizationSwitch() int32
-	SetUtilizationSwitch(v int32)
-}
+type HeaderT = api.Header
+type UsageInfo = api.UsageInfo
 
 type ContainerUsage struct {
 	PodUID        string
@@ -95,7 +68,7 @@ type ContainerLister struct {
 	stopCh          chan struct{}
 }
 
-var resyncInterval time.Duration = 5 * time.Minute
+var resyncInterval = 5 * time.Minute
 
 func init() {
 	if os.Getenv("HAMI_RESYNC_INTERVAL") != "" {
@@ -109,6 +82,7 @@ func init() {
 }
 
 func NewContainerLister() (*ContainerLister, error) {
+	ensureBuiltinsRegistered()
 	hookPath, ok := os.LookupEnv("HOOK_PATH")
 	if !ok {
 		return nil, fmt.Errorf("HOOK_PATH not set")
@@ -252,7 +226,7 @@ func loadCache(fpath string) (*ContainerUsage, error) {
 		klog.Errorf("Failed to stat cache file: %s, error: %v", cacheFile, err)
 		return nil, err
 	}
-	if info.Size() < int64(unsafe.Sizeof(headerT{})) {
+	if info.Size() < int64(unsafe.Sizeof(HeaderT{})) {
 		return nil, fmt.Errorf("cache file size %d too small", info.Size())
 	}
 	f, err := os.OpenFile(cacheFile, os.O_RDWR, 0666)
@@ -269,21 +243,20 @@ func loadCache(fpath string) (*ContainerUsage, error) {
 		klog.Errorf("Failed to mmap cache file: %s, error: %v", cacheFile, err)
 		return nil, err
 	}
-	head := (*headerT)(unsafe.Pointer(&usage.data[0]))
-	if head.initializedFlag != SharedRegionMagicFlag {
+	head := (*HeaderT)(unsafe.Pointer(&usage.data[0]))
+	if head.InitializedFlag != SharedRegionMagicFlag {
 		_ = syscall.Munmap(usage.data)
 		return nil, fmt.Errorf("cache file magic flag not matched")
 	}
-	if info.Size() == 1197897 {
-		klog.Infoln("casting......v0")
-		usage.Info = v0.CastSpec(usage.data)
-	} else if head.majorVersion == 1 {
-		klog.Infoln("casting......v1")
-		usage.Info = v1.CastSpec(usage.data)
-	} else {
+	factory := findFactory(head, info.Size())
+	if factory == nil {
+		majorVersion := head.MajorVersion
+		minorVersion := head.MinorVersion
 		_ = syscall.Munmap(usage.data)
-		return nil, fmt.Errorf("unknown cache file size %d version %d.%d", info.Size(), head.majorVersion, head.minorVersion)
+		return nil, fmt.Errorf("unknown cache file size %d version %d.%d", info.Size(), majorVersion, minorVersion)
 	}
+	klog.Infof("casting......%s", factory.Name())
+	usage.Info = factory.Cast(usage.data)
 	return usage, nil
 }
 

pkg/monitor/nvidia/factory.go

@@ -0,0 +1,38 @@
+/*
+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 nvidia
+
+import (
+	"sync"
+
+	"github.com/Project-HAMi/HAMi/pkg/monitor/nvidia/api"
+	nvidiav0 "github.com/Project-HAMi/HAMi/pkg/monitor/nvidia/v0"
+	nvidiav1 "github.com/Project-HAMi/HAMi/pkg/monitor/nvidia/v1"
+)
+
+var registerBuiltinsOnce sync.Once
+
+func ensureBuiltinsRegistered() {
+	registerBuiltinsOnce.Do(func() {
+		nvidiav0.Register()
+		nvidiav1.Register()
+	})
+}
+
+func findFactory(header *HeaderT, fileSize int64) api.CacheFactory {
+	return api.FindFactory(header, fileSize)
+}

pkg/monitor/nvidia/factory_test.go

@@ -0,0 +1,70 @@
+/*
+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 nvidia
+
+import "testing"
+
+func TestFindFactory_RoutesByVersion(t *testing.T) {
+	ensureBuiltinsRegistered()
+
+	tests := []struct {
+		name     string
+		header   HeaderT
+		size     int64
+		expected string
+	}{
+		{
+			name:     "v0 matches by known cache size",
+			header:   HeaderT{MajorVersion: 0, MinorVersion: 0},
+			size:     1197897,
+			expected: "v0",
+		},
+		{
+			name:     "v1 base matches major 1 minor <= 1",
+			header:   HeaderT{MajorVersion: 1, MinorVersion: 1},
+			size:     1024,
+			expected: "v1",
+		},
+		{
+			name:     "v1 sem matches major 1 minor >= 2",
+			header:   HeaderT{MajorVersion: 1, MinorVersion: 2},
+			size:     1024,
+			expected: "v1-sem",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			factory := findFactory(&tt.header, tt.size)
+			if factory == nil {
+				t.Fatalf("expected factory %q, got nil", tt.expected)
+			}
+			if factory.Name() != tt.expected {
+				t.Fatalf("expected factory %q, got %q", tt.expected, factory.Name())
+			}
+		})
+	}
+}
+
+func TestFindFactory_UnknownVersionReturnsNil(t *testing.T) {
+	ensureBuiltinsRegistered()
+
+	factory := findFactory(&HeaderT{MajorVersion: 9, MinorVersion: 9}, 1)
+	if factory != nil {
+		t.Fatalf("expected nil factory for unknown version, got %q", factory.Name())
+	}
+}

pkg/monitor/nvidia/v0/spec.go

@@ -16,7 +16,11 @@ limitations under the License.
 
 package v0
 
-import "unsafe"
+import (
+	"unsafe"
+
+	"github.com/Project-HAMi/HAMi/pkg/monitor/nvidia/api"
+)
 
 const maxDevices = 16
 
@@ -190,3 +194,13 @@ func (s Spec) GetUtilizationSwitch() int32 {
 func (s Spec) SetUtilizationSwitch(v int32) {
 	s.sr.utilizationSwitch = v
 }
+
+func Register() {
+	api.RegisterFactory(&v0Factory{})
+}
+
+type v0Factory struct{}
+
+func (v0Factory) Match(h *api.Header, size int64) bool { return size == 1197897 }
+func (v0Factory) Cast(data []byte) api.UsageInfo       { return CastSpec(data) }
+func (v0Factory) Name() string                         { return "v0" }