diff --git a/pkg/monitor/nvidia/api/api.go b/pkg/monitor/nvidia/api/api.go new file mode 100644 index 000000000..bba3d5dbc --- /dev/null +++ b/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 +} diff --git a/pkg/monitor/nvidia/cudevshr.go b/pkg/monitor/nvidia/cudevshr.go index ad322d3d1..5592d51e0 100644 --- a/pkg/monitor/nvidia/cudevshr.go +++ b/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 } diff --git a/pkg/monitor/nvidia/factory.go b/pkg/monitor/nvidia/factory.go new file mode 100644 index 000000000..71c496868 --- /dev/null +++ b/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) +} diff --git a/pkg/monitor/nvidia/factory_test.go b/pkg/monitor/nvidia/factory_test.go new file mode 100644 index 000000000..02bc107ac --- /dev/null +++ b/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()) + } +} diff --git a/pkg/monitor/nvidia/v0/spec.go b/pkg/monitor/nvidia/v0/spec.go index 9163e7627..0c73df048 100644 --- a/pkg/monitor/nvidia/v0/spec.go +++ b/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" } diff --git a/pkg/monitor/nvidia/v1/spec.go b/pkg/monitor/nvidia/v1/spec.go index bf01ca411..1fce5abb0 100644 --- a/pkg/monitor/nvidia/v1/spec.go +++ b/pkg/monitor/nvidia/v1/spec.go @@ -16,7 +16,11 @@ limitations under the License. package v1 -import "unsafe" +import ( + "unsafe" + + "github.com/Project-HAMi/HAMi/pkg/monitor/nvidia/api" +) const maxDevices = 16 @@ -59,7 +63,7 @@ type sharedRegionT struct { majorVersion int32 minorVersion int32 smInitFlag int32 - ownerPid uint32 + ownerPid uint64 sem semT num uint64 uuids [16]uuid @@ -137,10 +141,8 @@ func (s Spec) DeviceSmUtil(idx int) uint64 { } func (s Spec) SetDeviceSmLimit(l uint64) { - idx := uint64(0) - for idx < s.sr.num { + for idx := range min(int(s.sr.num), maxDevices) { s.sr.smLimit[idx] = l - idx += 1 } } @@ -157,10 +159,8 @@ func (s Spec) DeviceMemoryLimit(idx int) uint64 { } func (s Spec) SetDeviceMemoryLimit(l uint64) { - idx := uint64(0) - for idx < s.sr.num { + for idx := range min(int(s.sr.num), maxDevices) { s.sr.limit[idx] = l - idx += 1 } } @@ -197,3 +197,173 @@ func (s Spec) GetUtilizationSwitch() int32 { func (s Spec) SetUtilizationSwitch(v int32) { s.sr.utilizationSwitch = v } + +// --- sem_postinit variant (C major=1, minor>=2, e.g. hami-core) --- + +type shrregProcSlotTWithSeqlock struct { + pid int32 + hostpid int32 + used [16]deviceMemory + monitorused [16]uint64 + deviceUtil [16]deviceUtilization + status int32 + seqlock uint64 + unused [2]uint64 +} + +type sharedRegionTWithSemPostinit struct { + initializedFlag int32 + majorVersion int32 + minorVersion int32 + smInitFlag int32 + ownerPid uint64 + sem semT + num uint64 + uuids [16]uuid + limit [16]uint64 + smLimit [16]uint64 + procs [1024]shrregProcSlotTWithSeqlock + procnum int32 + utilizationSwitch int32 + recentKernel int32 + priority int32 + lastKernelTime int64 + semPostinit semT +} + +type SpecWithSemPostinit struct { + sr *sharedRegionTWithSemPostinit +} + +func (s SpecWithSemPostinit) DeviceMax() int { + return maxDevices +} + +func (s SpecWithSemPostinit) DeviceNum() int { + return int(s.sr.num) +} + +func (s SpecWithSemPostinit) DeviceMemoryContextSize(idx int) uint64 { + v := uint64(0) + for _, p := range s.sr.procs[:int(s.sr.procnum)] { + v += p.used[idx].contextSize + } + return v +} + +func (s SpecWithSemPostinit) DeviceMemoryModuleSize(idx int) uint64 { + v := uint64(0) + for _, p := range s.sr.procs[:int(s.sr.procnum)] { + v += p.used[idx].moduleSize + } + return v +} + +func (s SpecWithSemPostinit) DeviceMemoryBufferSize(idx int) uint64 { + v := uint64(0) + for _, p := range s.sr.procs[:int(s.sr.procnum)] { + v += p.used[idx].bufferSize + } + return v +} + +func (s SpecWithSemPostinit) DeviceMemoryOffset(idx int) uint64 { + v := uint64(0) + for _, p := range s.sr.procs[:int(s.sr.procnum)] { + v += p.used[idx].offset + } + return v +} + +func (s SpecWithSemPostinit) DeviceMemoryTotal(idx int) uint64 { + v := uint64(0) + for _, p := range s.sr.procs[:int(s.sr.procnum)] { + v += p.used[idx].total + } + return v +} + +func (s SpecWithSemPostinit) DeviceSmUtil(idx int) uint64 { + v := uint64(0) + for _, p := range s.sr.procs[:int(s.sr.procnum)] { + v += p.deviceUtil[idx].smUtil + } + return v +} + +func (s SpecWithSemPostinit) SetDeviceSmLimit(l uint64) { + for idx := range min(int(s.sr.num), maxDevices) { + s.sr.smLimit[idx] = l + } +} + +func (s SpecWithSemPostinit) IsValidUUID(idx int) bool { + return s.sr.uuids[idx].uuid[0] != 0 +} + +func (s SpecWithSemPostinit) DeviceUUID(idx int) string { + return string(s.sr.uuids[idx].uuid[:]) +} + +func (s SpecWithSemPostinit) DeviceMemoryLimit(idx int) uint64 { + return s.sr.limit[idx] +} + +func (s SpecWithSemPostinit) SetDeviceMemoryLimit(l uint64) { + for idx := range min(int(s.sr.num), maxDevices) { + s.sr.limit[idx] = l + } +} + +func (s SpecWithSemPostinit) LastKernelTime() int64 { + return s.sr.lastKernelTime +} + +func (s SpecWithSemPostinit) GetPriority() int { + return int(s.sr.priority) +} + +func (s SpecWithSemPostinit) GetRecentKernel() int32 { + return s.sr.recentKernel +} + +func (s SpecWithSemPostinit) SetRecentKernel(v int32) { + s.sr.recentKernel = v +} + +func (s SpecWithSemPostinit) GetUtilizationSwitch() int32 { + return s.sr.utilizationSwitch +} + +func (s SpecWithSemPostinit) SetUtilizationSwitch(v int32) { + s.sr.utilizationSwitch = v +} + +func CastSpecWithSemPostinit(data []byte) SpecWithSemPostinit { + return SpecWithSemPostinit{ + sr: (*sharedRegionTWithSemPostinit)(unsafe.Pointer(&data[0])), + } +} + +// --- Factory registrations --- + +func Register() { + api.RegisterFactory(&v1SemFactory{}) // major=1, minor=2 + api.RegisterFactory(&v1BaseFactory{}) // major=1, minor=1 +} + +type v1SemFactory struct{} + +func (v1SemFactory) Match(h *api.Header, size int64) bool { + return h.MajorVersion == 1 && h.MinorVersion >= 2 +} +func (v1SemFactory) Cast(data []byte) api.UsageInfo { return CastSpecWithSemPostinit(data) } +func (v1SemFactory) Name() string { return "v1-sem" } + +type v1BaseFactory struct{} + +func (v1BaseFactory) Match(h *api.Header, size int64) bool { + return h.MajorVersion == 1 && h.MinorVersion <= 1 +} +func (v1BaseFactory) Cast(data []byte) api.UsageInfo { return CastSpec(data) } +func (v1BaseFactory) Name() string { return "v1" }