diff --git a/internal/pkg/api/device/ascend/device.go b/internal/pkg/api/device/ascend/device.go index dc3edf2..4da7ea8 100644 --- a/internal/pkg/api/device/ascend/device.go +++ b/internal/pkg/api/device/ascend/device.go @@ -41,14 +41,22 @@ type VNPUConfig struct { ChipName string `yaml:"chipName"` ResourceName string `yaml:"resourceName"` ResourceMemoryName string `yaml:"resourceMemoryName"` + ResourceCoreName string `yaml:"resourceCoreName"` MemoryAllocatable int64 `yaml:"memoryAllocatable"` MemoryCapacity int64 `yaml:"memoryCapacity"` MemoryFactor int32 `yaml:"memoryFactor"` AICore int32 `yaml:"aiCore"` AICPU int32 `yaml:"aiCPU"` + RuntimeClassName string `yaml:"runtimeClassName"` + OverwriteEnv bool `yaml:"overwriteEnv"` Templates []Template `yaml:"templates"` } +type VNPUs struct { + HamiVnpuCore bool `yaml:"hamiVnpuCore"` + Configs []VNPUConfig `yaml:"configs"` +} + type Devices struct { config VNPUConfig nodeRegisterAnno string diff --git a/internal/pkg/api/device/device.go b/internal/pkg/api/device/device.go index 60a61a6..860d41e 100644 --- a/internal/pkg/api/device/device.go +++ b/internal/pkg/api/device/device.go @@ -52,12 +52,14 @@ type DeviceInfo struct { type MigTemplate struct { Name string `yaml:"name"` + Core int32 `yaml:"core"` Memory int32 `yaml:"memory"` Count int32 `yaml:"count"` } type MigTemplateUsage struct { Name string `json:"name,omitempty"` + Core int32 `json:"core,omitempty"` Memory int32 `json:"memory,omitempty"` InUse bool `json:"inuse,omitempty"` } diff --git a/internal/pkg/config/config.go b/internal/pkg/config/config.go index d7a7dfc..998b6c4 100644 --- a/internal/pkg/config/config.go +++ b/internal/pkg/config/config.go @@ -48,7 +48,7 @@ type Config struct { KunlunConfig kunlun.KunlunConfig `yaml:"kunlun"` AWSNeuronConfig awsneuron.AWSNeuronConfig `yaml:"awsneuron"` AMDGPUConfig amd.AMDConfig `yaml:"amd"` - VNPUs []ascend.VNPUConfig `yaml:"vnpus"` + VNPUs ascend.VNPUs `yaml:"vnpus"` } var ( @@ -76,7 +76,7 @@ func InitDevicesWithConfig(config *Config) error { if amdDevice != nil { device.DevicesMap[amdDevice.CommonWord()] = amdDevice }*/ - for _, dev := range ascend.InitDevices(config.VNPUs) { + for _, dev := range ascend.InitDevices(config.VNPUs.Configs) { commonWord := dev.CommonWord() device.DevicesMap[commonWord] = dev klog.Infof("Ascend device %s initialized", commonWord) diff --git a/internal/pkg/config/config_test.go b/internal/pkg/config/config_test.go new file mode 100644 index 0000000..9752682 --- /dev/null +++ b/internal/pkg/config/config_test.go @@ -0,0 +1,138 @@ +/* +Copyright 2025 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 config + +import ( + "os" + "path/filepath" + "testing" + + deviceapi "github.com/HAMi/mock-device-plugin/internal/pkg/api/device" + "github.com/HAMi/mock-device-plugin/internal/pkg/api/device/ascend" +) + +func writeTempConfig(t *testing.T, body string) string { + t.Helper() + + path := filepath.Join(t.TempDir(), "device-config.yaml") + if err := os.WriteFile(path, []byte(body), 0o600); err != nil { + t.Fatalf("write config: %v", err) + } + return path +} + +func TestLoadConfigLoadsLatestWrappedVNPUs(t *testing.T) { + path := writeTempConfig(t, `vnpus: + hamiVnpuCore: true + configs: + - commonWord: Ascend910A + chipName: 910A + resourceName: huawei.com/Ascend910A + resourceMemoryName: huawei.com/Ascend910A-memory + resourceCoreName: huawei.com/Ascend910A-core + memoryAllocatable: 32768 + memoryCapacity: 32768 + memoryFactor: 1 + aiCore: 20 + aiCPU: 7 + runtimeClassName: hami-ascend + overwriteEnv: true + templates: + - name: vir05_1c_8g + memory: 8192 + aiCore: 5 + aiCPU: 1 +`) + + cfg, err := LoadConfig(path) + if err != nil { + t.Fatalf("LoadConfig returned error: %v", err) + } + + if !cfg.VNPUs.HamiVnpuCore { + t.Fatalf("expected HamiVnpuCore=true") + } + if len(cfg.VNPUs.Configs) != 1 { + t.Fatalf("expected 1 VNPU config, got %d", len(cfg.VNPUs.Configs)) + } + + got := cfg.VNPUs.Configs[0] + if got.ResourceCoreName != "huawei.com/Ascend910A-core" { + t.Fatalf("unexpected ResourceCoreName: %s", got.ResourceCoreName) + } + if got.RuntimeClassName != "hami-ascend" { + t.Fatalf("unexpected RuntimeClassName: %s", got.RuntimeClassName) + } + if !got.OverwriteEnv { + t.Fatalf("expected OverwriteEnv=true") + } +} + +func TestLoadConfigRejectsFlatVNPUsShape(t *testing.T) { + path := writeTempConfig(t, `vnpus: + - commonWord: Ascend910A + chipName: 910A + resourceName: huawei.com/Ascend910A + resourceMemoryName: huawei.com/Ascend910A-memory + memoryAllocatable: 32768 + memoryCapacity: 32768 + memoryFactor: 1 +`) + + _, err := LoadConfig(path) + if err == nil { + t.Fatalf("expected LoadConfig to reject flat vnpus shape") + } +} + +func TestInitDevicesWithConfigUsesWrappedConfigs(t *testing.T) { + t.Cleanup(func() { + deviceapi.DevicesMap = nil + }) + + cfg := &Config{ + VNPUs: ascend.VNPUs{ + Configs: []ascend.VNPUConfig{ + { + CommonWord: "Ascend910A", + ChipName: "910A", + ResourceName: "huawei.com/Ascend910A", + ResourceMemoryName: "huawei.com/Ascend910A-memory", + ResourceCoreName: "huawei.com/Ascend910A-core", + MemoryAllocatable: 32768, + MemoryCapacity: 32768, + MemoryFactor: 1, + AICore: 20, + AICPU: 7, + RuntimeClassName: "hami-ascend", + OverwriteEnv: true, + Templates: []ascend.Template{ + {Name: "vir05_1c_8g", Memory: 8192, AICore: 5, AICPU: 1}, + }, + }, + }, + }, + } + + if err := InitDevicesWithConfig(cfg); err != nil { + t.Fatalf("InitDevicesWithConfig returned error: %v", err) + } + + if _, ok := deviceapi.DevicesMap["Ascend910A"]; !ok { + t.Fatalf("expected Ascend910A device to be registered from wrapped configs") + } +}