From 52ea5cf2c3c1e52d63e36329ca660c26eba97d6f Mon Sep 17 00:00:00 2001 From: mhsong2 Date: Wed, 9 Sep 2026 17:28:50 +0800 Subject: [PATCH] feat: support xpu devices in runc sandboxes Add provider-owned GPU and Ascend discovery, runtime selection, exclusive leases, and validated OCI edits for native runc sandboxes. Keep NVIDIA runsc support and preserve application library paths when injecting driver directories. Consume Ascend integration through the standalone adapter protocol so sandboxd no longer contains the MindCluster source dependency or adapter build target. Signed-off-by: mhsong2 --- AGENTS.md | 4 + Makefile | 2 +- README.md | 21 +- api/runtime/v1/sandbox-api.pb.go | 2 +- api/runtime/v1/sandbox-api.proto | 2 +- config/config.go | 14 + configs/ascend/mounts.json | 14 + configs/sandboxd.toml | 7 + doc/runtime.md | 15 +- internal/server/server.go | 23 +- internal/server/server_test.go | 14 +- pkg/runtime/handler.go | 11 + pkg/runtime/library_paths_test.go | 70 +++ pkg/runtime/oci_loader.go | 92 ++++ pkg/runtime/oci_loader_test.go | 47 ++ pkg/xpumanager/ascend.go | 621 +++++++++++++++++++++++ pkg/xpumanager/ascend_test.go | 136 +++++ pkg/xpumanager/device_stat_linux.go | 34 ++ pkg/xpumanager/device_stat_linux_test.go | 36 ++ pkg/xpumanager/device_stat_other.go | 23 + pkg/xpumanager/lease.go | 98 ++++ pkg/xpumanager/manager.go | 286 +++-------- pkg/xpumanager/manager_test.go | 205 ++++++-- pkg/xpumanager/nvidia.go | 281 +++++++++- test/e2e/README.md | 4 +- test/e2e/gpu.Dockerfile | 3 + 26 files changed, 1760 insertions(+), 305 deletions(-) create mode 100644 configs/ascend/mounts.json create mode 100644 pkg/runtime/library_paths_test.go create mode 100644 pkg/xpumanager/ascend.go create mode 100644 pkg/xpumanager/ascend_test.go create mode 100644 pkg/xpumanager/device_stat_linux.go create mode 100644 pkg/xpumanager/device_stat_linux_test.go create mode 100644 pkg/xpumanager/device_stat_other.go create mode 100644 pkg/xpumanager/lease.go diff --git a/AGENTS.md b/AGENTS.md index c18d8b9..a57d77a 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -92,3 +92,7 @@ virtio-fs staging and restored live-memory files, and runtime state. Every virtio-fs export must remain read-only; never extend this path to writable host sharing. Bounded read-only regular-file injection remains a separate startup-metadata mechanism for files such as `resolv.conf`. + +# Accelerator Contract + +Ascend is an optional external adapter for physical 310P and 910 A2/A3 devices with runc. Keep driver-reported model normalization separate from runtime-family validation; do not add per-SKU allowlists. The SDK/scheduler matches the normalized reported model. Provider library directories use additive `SpecUpdates.PrependLibraryPaths`, preserving the application environment; device-visibility variables still use authoritative overrides. diff --git a/Makefile b/Makefile index f4061a6..bfec594 100644 --- a/Makefile +++ b/Makefile @@ -251,7 +251,7 @@ fmt: ## format Go code go fmt ./... check-fmt: ## verify Go code is gofmt-clean - @files="$$(gofmt -l .)" || exit $$?; \ + @files="$$(gofmt -l $$(git ls-files '*.go'))" || exit $$?; \ test -z "$$files" || { printf '%s\n' "$$files" >&2; exit 1; } vet: ## run go vet diff --git a/README.md b/README.md index f22f9c2..8ce24e7 100644 --- a/README.md +++ b/README.md @@ -33,13 +33,15 @@ The `sbox` binary is an administrative CLI for managing sandboxes. ### NVIDIA GPU sandboxes -GPU support is experimental and currently uses gVisor runsc with nvproxy. The -scheduler passes concrete node-local device IDs through +GPU support is experimental. It supports gVisor runsc through nvproxy and +native runc through the NVIDIA OCI prestart hook. The scheduler passes concrete +node-local device IDs through `StartRequest.xpu_allocations`; sandboxd resolves them to NVIDIA UUIDs and maintains a local exclusive lease: ```bash sbox start \ + --runtime runc \ --rootfs /path/to/directory-rootfs \ --xpu-allocation gpu:0,2 \ /bin/sleep 300 @@ -103,6 +105,15 @@ make networkacl-test make bpfnat-test ``` +The Ascend OCI adapter is an optional external process and is not linked into +the default sandboxd binary. Build or download it from +[openYuanrong-mirror/ascend-oci-adapter](https://github.com/openYuanrong-mirror/ascend-oci-adapter). +The default trusted read-only driver mount profile remains +`configs/ascend/mounts.json`; deployments install it alongside the adapter. +Physical Ascend 310P and 910 A2/A3 support uses driver-reported model names +without a per-SKU allowlist. See [Sandbox runtimes](doc/runtime.md) for model +selection and library-path handling. + `networkacl-test` runs one backend-neutral conformance suite against native iptables and TC eBPF enforcement in isolated network namespaces. It covers allow and deny precedence, exact and wildcard peers, peer and sandbox ports, @@ -162,11 +173,7 @@ tools/ pinned protobuf code-generation image ## Known limitations - Kata Containers and Firecracker require a usable `/dev/kvm`; nodes without KVM continue to support gVisor. Firecracker additionally requires a compatible guest kernel/initrd and the ext4 image tool. Its root filesystem may be an immutable EROFS image or a directory exported through virtio-fs; directory-backed OCI/Nydus roots require `virtiofs_enabled` and a compatible virtiofsd. -- NVIDIA GPU sandboxes require runsc, a directory/lisafs-backed rootfs, - `nvidia-container-cli`, accessible NVIDIA devices and userspace driver - libraries, and a host driver supported by the pinned runsc nvproxy. Kata, - Firecracker, runc, MIG, fractional GPUs, and regular-file/EROFS rootfs are - not supported. +- NVIDIA GPU sandboxes require `nvidia-container-cli`, the executable NVIDIA OCI runtime hook, accessible NVIDIA devices and userspace driver libraries. Runsc additionally requires a host driver supported by its pinned nvproxy compatibility gate and a directory/lisafs-backed rootfs. Runc accepts the normal directory or EROFS rootfs through its writable host overlay. Kata, Firecracker, MIG, and fractional GPUs are not supported. - sandboxd detects the local cgroup mode at startup. Legacy and hybrid hosts use cgroup v1; unified hosts use cgroup v2. The gRPC API and resource-cache behavior are identical in both modes. - `[plugin.resource].disable_cgroup = true` enables an experimental/debug compatibility mode for environments where sandboxd cannot write the diff --git a/api/runtime/v1/sandbox-api.pb.go b/api/runtime/v1/sandbox-api.pb.go index 55bef6a..07ede78 100644 --- a/api/runtime/v1/sandbox-api.pb.go +++ b/api/runtime/v1/sandbox-api.pb.go @@ -1260,7 +1260,7 @@ func (*Mount_ImageUrl) isMount_Source() {} // scheduler on the selected node. type XpuAllocation struct { state protoimpl.MessageState `protogen:"open.v1"` - // Type is the normalized accelerator type, such as gpu. + // Type is the normalized accelerator type, such as gpu or npu. Type string `protobuf:"bytes,1,opt,name=type,proto3" json:"type,omitempty"` // DeviceIDs contains node-local physical device IDs selected by the scheduler. DeviceIds []uint32 `protobuf:"varint,2,rep,packed,name=device_ids,json=deviceIds,proto3" json:"device_ids,omitempty"` diff --git a/api/runtime/v1/sandbox-api.proto b/api/runtime/v1/sandbox-api.proto index f6f45bb..2599843 100644 --- a/api/runtime/v1/sandbox-api.proto +++ b/api/runtime/v1/sandbox-api.proto @@ -207,7 +207,7 @@ message Mount { // XpuAllocation describes concrete accelerator devices assigned by the // scheduler on the selected node. message XpuAllocation { - // Type is the normalized accelerator type, such as gpu. + // Type is the normalized accelerator type, such as gpu or npu. string type = 1; // DeviceIDs contains node-local physical device IDs selected by the scheduler. repeated uint32 device_ids = 2; diff --git a/config/config.go b/config/config.go index c8d5e08..1a23fe6 100644 --- a/config/config.go +++ b/config/config.go @@ -53,9 +53,23 @@ type PluginConfig struct { NodeResourceConfig `toml:"node_resource" json:"nodeResource"` + XPUConfig `toml:"xpu" json:"xpu"` + ImageManagerConfig `toml:"image" json:"image"` } +// XPUConfig contains optional vendor accelerator providers. +type XPUConfig struct { + Ascend AscendConfig `toml:"ascend" json:"ascend"` +} + +// AscendConfig enables the external Ascend OCI adapter for runc sandboxes. +type AscendConfig struct { + Enabled bool `toml:"enabled" json:"enabled"` + Adapter string `toml:"adapter" json:"adapter"` + MountProfile string `toml:"mount_profile" json:"mountProfile"` +} + // ImageManagerConfig configures image and mount lifecycle management. type ImageManagerConfig struct { ImageManagerRoot string `toml:"root" json:"root"` diff --git a/configs/ascend/mounts.json b/configs/ascend/mounts.json new file mode 100644 index 0000000..bead0fe --- /dev/null +++ b/configs/ascend/mounts.json @@ -0,0 +1,14 @@ +{ + "default": [ + { + "path": [ + "/usr/local/Ascend/driver/lib64", + "/usr/local/Ascend/driver/include", + "/usr/local/dcmi", + "/usr/local/bin/npu-smi", + "/usr/sbin/dmidecode", + "/var/queue_schedule" + ] + } + ] +} diff --git a/configs/sandboxd.toml b/configs/sandboxd.toml index f45c42d..a569796 100644 --- a/configs/sandboxd.toml +++ b/configs/sandboxd.toml @@ -126,3 +126,10 @@ cgroup_memory_limit = "0" # [plugin.node_resource] # provider = "kubernetes" # or "cgroup" # sock_path = "/run/sandboxd/resource.sock" + +# Ascend is optional and fail-closed. Enabling it requires the runc runtime, +# a root-owned adapter, and a versioned read-only mount profile. +# [plugin.xpu.ascend] +# enabled = true +# adapter = "/usr/local/bin/ascend-oci-adapter" +# mount_profile = "/etc/akernel/ascend/mounts.json" diff --git a/doc/runtime.md b/doc/runtime.md index cd60520..492cf6d 100644 --- a/doc/runtime.md +++ b/doc/runtime.md @@ -18,13 +18,26 @@ binaries, boot artifacts, and host prerequisites pass validation. | Published-port DNAT | Supported | Supported | Supported | Supported | | Writable-layer quota | Supported | Not supported | Not supported | Supported | | Checkpoint and restore | Supported (systrap and KVM) | Not supported | Not supported | Supported | -| NVIDIA GPU | Experimental nvproxy support | Not supported | Not supported | Not supported | +| NVIDIA GPU | Experimental nvproxy support | Experimental NVIDIA OCI-hook support | Not supported | Not supported | +| Ascend NPU | Not supported | Experimental physical 310P and 910 A2/A3 support | Not supported | Not supported | | Cgroup-disabled mode | Experimental | Not supported | Not supported | Not supported | | KVM | Optional execution platform; not exposed to the sandbox | Optional guest exposure | Required by the runtime | Required by the runtime; nested KVM is not exposed | See [Checkpoint and restore](checkpoint-restore.md) for the API design, artifact ownership, failure semantics, and compatibility requirements. +The NVIDIA provider owns a single physical-device inventory and exclusive +lease table shared by runsc and runc. Runsc remains subject to its nvproxy +driver-compatibility gate, while an unsupported nvproxy driver does not disable +runc GPU support. Both runtimes consume the provider's UUIDs and OCI environment +through the NVIDIA prestart hook. + +The Ascend provider reports the normalized model name returned by the driver, without a per-model allowlist. Requests must match that reported name exactly; for example, a node reporting `ascend910` uses `npu:ascend910:1`, not an inferred SKU. The adapter independently resolves the driver device family to the supported physical 310P or 910 A2/A3 injection profile. Unsupported families and empty model names are rejected. Allocated devices must share a model and runtime family. + +Ascend driver library paths are prepended to the final image/request `LD_LIBRARY_PATH`, preserving application and CANN directories while removing duplicates and empty entries. Provider-owned device visibility variables remain authoritative. + +GPU capacity is not qualified by runtime in the scheduler. Deploy GPU nodes with consistent runtime capabilities and select `runtime="runc"` explicitly when their GPUs are available only through runc. + ## Selection and configuration A start request selects a runtime by name. Each adapter must have an entry diff --git a/internal/server/server.go b/internal/server/server.go index 0516d06..0178ea1 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -695,7 +695,9 @@ func NewSandboxService(root, configPath string) (result SandboxService, retErr e } xpuMgr := xpumanager.New( cfg.RuntimeConfig.RuntimeBinary[config.RuntimeNameRunsc], + cfg.RuntimeConfig.RuntimeBinary[config.RuntimeNameRunc] != "", sandboxRoot, + cfg.XPUConfig, ) // The optional node-resource module comes up first so its external resource @@ -1215,10 +1217,21 @@ func (h *sandboxService) Start(ctx context.Context, request *runtime.StartReques return &runtime.StartResponse{Code: -1, Message: err.Error()}, errord.ToGRPC(errord.ErrInvalidArgument) } - if len(startReq.XpuAllocations) > 0 && startReq.Runtime != config.RuntimeNameRunsc { - err := fmt.Errorf("XPU allocations require runtime %q", config.RuntimeNameRunsc) - return &runtime.StartResponse{Code: -1, Message: err.Error()}, - errord.ToGRPC(errord.ErrInvalidArgument) + if len(startReq.XpuAllocations) > 0 { + if startReq.Runtime != config.RuntimeNameRunsc && startReq.Runtime != config.RuntimeNameRunc { + err := fmt.Errorf("XPU allocations require runtime %q or %q", config.RuntimeNameRunsc, config.RuntimeNameRunc) + return &runtime.StartResponse{Code: -1, Message: err.Error()}, + errord.ToGRPC(errord.ErrInvalidArgument) + } + if h.xpuMgr == nil { + err := errors.New("XPU manager is not configured") + return &runtime.StartResponse{Code: -1, Message: err.Error()}, + errord.ToGRPC(errord.ErrFailedPrecondition) + } + if err := h.xpuMgr.ValidateRuntime(startReq.Runtime, startReq.XpuAllocations); err != nil { + return &runtime.StartResponse{Code: -1, Message: err.Error()}, + errord.ToGRPC(errord.ErrInvalidArgument) + } } if startReq.WritableLayerLimitBytes > 0 { if startReq.Runtime != config.RuntimeNameRunsc && @@ -1387,7 +1400,7 @@ func (h *sandboxService) Start(ctx context.Context, request *runtime.StartReques return &runtime.StartResponse{Code: -1, Message: err.Error()}, errord.ToGRPC(errord.ErrFailedPrecondition) } - specUpdates, err = h.xpuMgr.Acquire(sandboxID, startReq.XpuAllocations) + specUpdates, err = h.xpuMgr.Acquire(sandboxID, startReq.Runtime, startReq.XpuAllocations) if err != nil { return &runtime.StartResponse{ Code: -1, diff --git a/internal/server/server_test.go b/internal/server/server_test.go index 7bcb9f3..21fc139 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -284,7 +284,6 @@ func TestStartRejectsFirecrackerOCIImageBeforeFilesystemPrepare(t *testing.T) { func TestStartRejectsXPUForUnsupportedRuntimes(t *testing.T) { for _, runtimeName := range []string{ - config.RuntimeNameRunc, config.RuntimeNameKata, config.RuntimeNameFirecracker, } { @@ -303,6 +302,19 @@ func TestStartRejectsXPUForUnsupportedRuntimes(t *testing.T) { } } +func TestStartRejectsXPUWhenManagerIsUnavailable(t *testing.T) { + s := newTestService(t, map[string]svc.Handler{ + config.RuntimeNameRunc: svc.NewFakeRuntimeHandler(), + }) + response, err := s.Start(context.Background(), &runtime.StartRequest{ + Runtime: config.RuntimeNameRunc, + Rootfs: &runtime.RootfsConfig{}, + XpuAllocations: []*runtime.XpuAllocation{{Type: "npu"}}, + }) + assert.Equal(t, codes.FailedPrecondition, status.Code(err)) + assert.Contains(t, response.Message, "XPU manager is not configured") +} + func TestStartRejectsEnableKVMForRunsc(t *testing.T) { s := newTestService(t, map[string]svc.Handler{ config.RuntimeNameRunsc: svc.NewFakeRuntimeHandler(), diff --git a/pkg/runtime/handler.go b/pkg/runtime/handler.go index a51c875..b7b3eb6 100644 --- a/pkg/runtime/handler.go +++ b/pkg/runtime/handler.go @@ -107,9 +107,20 @@ type StartConfig struct { // this boundary so vendor-specific discovery and authorization do not leak // into the runsc client. type SpecUpdates struct { + // Envs carries provider-owned device visibility settings. Prestart is used + // by providers such as NVIDIA whose toolkit performs injection in an OCI + // hook; Ascend resolves the corresponding OCI objects before runc starts. Envs []*runtime.KeyValue Prestart []Hook Annotations map[string]string + // LinuxDevices and DeviceCgroupRules authorize provider-selected device + // nodes. Mounts carries the provider's read-only driver/runtime files. + LinuxDevices []LinuxDevice + DeviceCgroupRules []LinuxDeviceCgroup + Mounts []Mount + // PrependLibraryPaths adds driver library directories ahead of the final + // image/request LD_LIBRARY_PATH without replacing application directories. + PrependLibraryPaths []string // RequiresHostWritableRootfs requests a private writable rootfs view // before provider hooks execute. It is separate from the writable layer // visible to workloads after the sandbox starts. diff --git a/pkg/runtime/library_paths_test.go b/pkg/runtime/library_paths_test.go new file mode 100644 index 0000000..52fe5e5 --- /dev/null +++ b/pkg/runtime/library_paths_test.go @@ -0,0 +1,70 @@ +// Copyright (c) 2026 Ant Group Corporation. +// +// 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 runtime + +import ( + "reflect" + "testing" + + api "github.com/inclusionAI/sandboxd/api/runtime/v1" +) + +func TestGenerateOciPreservesApplicationLibrariesWithProviderPaths(t *testing.T) { + loader, err := NewBundleLoader("", t.TempDir()) + if err != nil { + t.Fatal(err) + } + loader.baseSpec.Process.Env = []string{"LD_LIBRARY_PATH=/image/lib"} + _, spec, err := loader.GenerateOci(OciLoadOptions{ + SandboxID: "sbox-libraries", CgroupPath: "/sandbox/libraries", + Config: StartConfig{ + Rootfs: t.TempDir(), Resources: &api.LinuxSandboxResources{}, + Envs: []*api.KeyValue{{Key: "LD_LIBRARY_PATH", Value: ":/opt/cann/lib64:/driver/lib:/application/lib::"}}, + SpecUpdates: &SpecUpdates{ + PrependLibraryPaths: []string{"/driver/lib", "", "/driver/lib"}, + }, + }, + }) + if err != nil { + t.Fatal(err) + } + want := "LD_LIBRARY_PATH=/driver/lib:/opt/cann/lib64:/application/lib" + if !containsString(spec.Process.Env, want) { + t.Fatalf("OCI env = %v, want %s", spec.Process.Env, want) + } + if containsString(spec.Process.Env, "LD_LIBRARY_PATH=/image/lib") { + t.Fatal("request must override image libraries before driver paths are merged") + } +} + +func TestPrependLibraryPathsPreservesOtherEnvironment(t *testing.T) { + envs := []string{"CUSTOM=a=b", "LD_LIBRARY_PATH=/application/lib"} + got := prependLibraryPaths(envs, []string{"/driver/lib"}) + want := []string{"CUSTOM=a=b", "LD_LIBRARY_PATH=/driver/lib:/application/lib"} + if !reflect.DeepEqual(got, want) { + t.Fatalf("env = %v, want %v", got, want) + } + if got := prependLibraryPaths(envs, nil); !reflect.DeepEqual(got, envs) { + t.Fatalf("provider without library paths changed env: %v", got) + } +} + +func TestPrependLibraryPathsWithoutApplicationPath(t *testing.T) { + got := prependLibraryPaths([]string{"PATH=/bin"}, []string{"", "/driver/lib", "/driver/lib", ""}) + want := []string{"PATH=/bin", "LD_LIBRARY_PATH=/driver/lib"} + if !reflect.DeepEqual(got, want) { + t.Fatalf("env = %v, want %v", got, want) + } +} diff --git a/pkg/runtime/oci_loader.go b/pkg/runtime/oci_loader.go index 53e01dc..2330d65 100644 --- a/pkg/runtime/oci_loader.go +++ b/pkg/runtime/oci_loader.go @@ -187,6 +187,7 @@ func (r *BundleLoader) GenerateOci(options OciLoadOptions) (string, *Spec, error if len(updates.Envs) > 0 { ociSpec.Process.Env = combineEnvs(ociSpec.Process.Env, updates.Envs) } + ociSpec.Process.Env = prependLibraryPaths(ociSpec.Process.Env, updates.PrependLibraryPaths) if len(updates.Prestart) > 0 { if ociSpec.Hooks == nil { ociSpec.Hooks = &Hooks{} @@ -194,6 +195,9 @@ func (r *BundleLoader) GenerateOci(options OciLoadOptions) (string, *Spec, error ociSpec.Hooks.Prestart = append(ociSpec.Hooks.Prestart, updates.Prestart...) } ociSpec.Annotations = combineAnnotations(ociSpec.Annotations, updates.Annotations) + if err := applyProviderDevicesAndMounts(ociSpec, updates); err != nil { + return "", ociSpec, err + } } ociFile := filepath.Join(bundleDir, config.SandboxSpecFile) @@ -211,6 +215,63 @@ func (r *BundleLoader) GenerateOci(options OciLoadOptions) (string, *Spec, error return bundleDir, ociSpec, os.WriteFile(ociFile, buf, 0644) } +func applyProviderDevicesAndMounts(spec *Spec, updates *SpecUpdates) error { + if spec.Linux == nil { + spec.Linux = &Linux{} + } + if spec.Linux.Resources == nil { + spec.Linux.Resources = &LinuxResources{} + } + devicePaths := make(map[string]struct{}, len(spec.Linux.Devices)+len(updates.LinuxDevices)) + for _, device := range spec.Linux.Devices { + devicePaths[device.Path] = struct{}{} + } + for _, device := range updates.LinuxDevices { + if device.Path == "" || !filepath.IsAbs(device.Path) || device.Type != "c" || + device.Major < 0 || device.Minor < 0 { + return fmt.Errorf("invalid provider device %q", device.Path) + } + if _, duplicate := devicePaths[device.Path]; duplicate { + return fmt.Errorf("provider device conflicts at %s", device.Path) + } + devicePaths[device.Path] = struct{}{} + spec.Linux.Devices = append(spec.Linux.Devices, device) + } + for _, rule := range updates.DeviceCgroupRules { + if !rule.Allow || rule.Type != "c" || rule.Major == nil || rule.Minor == nil || + *rule.Major < 0 || *rule.Minor < 0 || rule.Access != "rwm" { + return errors.New("invalid provider device cgroup rule") + } + spec.Linux.Resources.Devices = append(spec.Linux.Resources.Devices, rule) + } + mountTargets := make(map[string]struct{}, len(spec.Mounts)+len(updates.Mounts)) + for _, mount := range spec.Mounts { + mountTargets[filepath.Clean(mount.Destination)] = struct{}{} + } + for _, mount := range updates.Mounts { + destination := filepath.Clean(mount.Destination) + if !filepath.IsAbs(destination) || destination == "/" || !filepath.IsAbs(mount.Source) || + mount.Type != "bind" || !containsMountOption(mount.Options, "ro") { + return fmt.Errorf("invalid provider mount %s -> %s", mount.Source, mount.Destination) + } + if _, conflict := mountTargets[destination]; conflict { + return fmt.Errorf("provider mount conflicts at %s", destination) + } + mountTargets[destination] = struct{}{} + spec.Mounts = append(spec.Mounts, mount) + } + return nil +} + +func containsMountOption(options []string, expected string) bool { + for _, option := range options { + if option == expected { + return true + } + } + return false +} + func setNetworkNamespace(linux *Linux, path string) { for index := range linux.Namespaces { if linux.Namespaces[index].Type == NetworkNamespace { @@ -413,6 +474,37 @@ func combineAnnotations(annotations map[string]string, annoToAdd map[string]stri return annotations } +func prependLibraryPaths(envs, providerPaths []string) []string { + if len(providerPaths) == 0 { + return envs + } + const prefix = "LD_LIBRARY_PATH=" + var existing string + result := make([]string, 0, len(envs)+1) + for _, env := range envs { + if strings.HasPrefix(env, prefix) { + existing = strings.TrimPrefix(env, prefix) + } else { + result = append(result, env) + } + } + paths := make([]string, 0, len(providerPaths)) + seen := make(map[string]bool) + appendPaths := func(values []string) { + for _, value := range values { + // Empty entries search the process working directory, so never + // introduce or retain them in the merged driver search path. + if value != "" && !seen[value] { + paths = append(paths, value) + seen[value] = true + } + } + } + appendPaths(providerPaths) + appendPaths(strings.Split(existing, ":")) + return append(result, prefix+strings.Join(paths, ":")) +} + func combineEnvs(envs []string, overrides []*runtime.KeyValue) []string { envMap := map[string]string{} for _, env := range envs { diff --git a/pkg/runtime/oci_loader_test.go b/pkg/runtime/oci_loader_test.go index 94c5678..24136f0 100644 --- a/pkg/runtime/oci_loader_test.go +++ b/pkg/runtime/oci_loader_test.go @@ -204,6 +204,53 @@ func TestGenerateOciAppliesProviderUpdatesLast(t *testing.T) { } } +func TestGenerateOciAppliesAscendDevicesCgroupAndMounts(t *testing.T) { + loader, err := NewBundleLoader("", t.TempDir()) + if err != nil { + t.Fatal(err) + } + major, minor := int64(240), int64(3) + _, spec, err := loader.GenerateOci(OciLoadOptions{ + SandboxID: "sandbox-npu", CgroupPath: "/sandbox/npu", + Config: StartConfig{ + Rootfs: t.TempDir(), Resources: &runtime.LinuxSandboxResources{}, + SpecUpdates: &SpecUpdates{ + LinuxDevices: []LinuxDevice{{Path: "/dev/davinci3", Type: "c", Major: major, Minor: minor}}, + DeviceCgroupRules: []LinuxDeviceCgroup{{ + Allow: true, Type: "c", Major: &major, Minor: &minor, Access: "rwm", + }}, + Mounts: []Mount{{ + Destination: "/usr/local/Ascend/driver/lib64/driver", Type: "bind", + Source: "/usr/local/Ascend/driver/lib64/driver", Options: []string{"ro", "rbind"}, + }}, + }, + }, + }) + if err != nil { + t.Fatal(err) + } + if len(spec.Linux.Devices) != 1 || spec.Linux.Devices[0].Path != "/dev/davinci3" { + t.Fatalf("provider devices missing from %+v", spec.Linux.Devices) + } + for name, set := range map[string][]string{ + "bounding": spec.Process.Capabilities.Bounding, + "effective": spec.Process.Capabilities.Effective, + "inheritable": spec.Process.Capabilities.Inheritable, + "permitted": spec.Process.Capabilities.Permitted, + } { + if containsString(set, "CAP_DAC_OVERRIDE") { + t.Fatalf("provider unexpectedly added CAP_DAC_OVERRIDE to %s set: %v", name, set) + } + } + if len(spec.Linux.Resources.Devices) != 1 || spec.Linux.Resources.Devices[0].Access != "rwm" { + t.Fatalf("provider cgroup rule missing from %+v", spec.Linux.Resources.Devices) + } + if len(spec.Mounts) == 0 || spec.Mounts[len(spec.Mounts)-1].Destination != + "/usr/local/Ascend/driver/lib64/driver" { + t.Fatalf("provider mount missing from %+v", spec.Mounts) + } +} + func TestGenerateOciWithoutCgroup(t *testing.T) { loader, err := NewBundleLoader("", t.TempDir()) if err != nil { diff --git a/pkg/xpumanager/ascend.go b/pkg/xpumanager/ascend.go new file mode 100644 index 0000000..fbc9e75 --- /dev/null +++ b/pkg/xpumanager/ascend.go @@ -0,0 +1,621 @@ +// Copyright (c) 2026 Ant Group Corporation. +// +// 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 xpumanager + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "sort" + "strings" + "sync" + "syscall" + "time" + + api "github.com/inclusionAI/sandboxd/api/runtime/v1" + "github.com/inclusionAI/sandboxd/config" + svc "github.com/inclusionAI/sandboxd/pkg/runtime" +) + +const ( + ascendSchemaVersion = 1 + defaultAdapterTimeout = 10 * time.Second + maxAdapterInputBytes = 1 << 20 + maxAdapterOutputBytes = 4 << 20 + ascendRuntimeFamily = "Ascend910" + ascend310PRuntimeFamily = "Ascend310P" + ascendProviderName = "ascend-cdi" + ascendVisibleDevicesEnv = "ASCEND_VISIBLE_DEVICES" + ascendRTVisibleDevicesEnv = "ASCEND_RT_VISIBLE_DEVICES" +) + +var allowedAscendSharedDevices = map[string]struct{}{ + "/dev/davinci_manager": {}, + "/dev/devmm_svm": {}, + "/dev/hisi_hdc": {}, + "/dev/dvpp_cmdlist": {}, +} + +type adapterVersion struct { + SchemaVersion int `json:"schema_version"` + ProviderVersion string `json:"provider_version"` +} + +type ascendDevice struct { + SchedulerID uint32 `json:"scheduler_id"` + LogicID int32 `json:"logic_id"` + PhysicalID int32 `json:"physical_id"` + StableID string `json:"stable_id"` + ProductModel string `json:"product_model"` + Generation string `json:"generation"` + RuntimeFamily string `json:"runtime_family"` + ResourceFamily string `json:"resource_family"` + RawProduct string `json:"raw_product"` + Healthy bool `json:"healthy"` +} + +type adapterDiscovery struct { + SchemaVersion int `json:"schema_version"` + ProviderVersion string `json:"provider_version"` + Devices []ascendDevice `json:"devices"` +} + +type adapterEditsRequest struct { + SchemaVersion int `json:"schema_version"` + LogicIDs []int32 `json:"logic_ids"` + RuntimeFamily string `json:"runtime_family"` + PhysicalOnly bool `json:"physical_only"` + MountProfile string `json:"mount_profile"` +} + +type adapterDevice struct { + HostPath string `json:"host_path"` + ContainerPath string `json:"container_path"` + Type string `json:"type"` + Major int64 `json:"major"` + Minor int64 `json:"minor"` + Permissions string `json:"permissions"` +} + +type adapterMount struct { + Source string `json:"source"` + Destination string `json:"destination"` + Type string `json:"type"` + Options []string `json:"options"` +} + +type adapterEdits struct { + SchemaVersion int `json:"schema_version"` + ProviderVersion string `json:"provider_version"` + Devices []adapterDevice `json:"devices"` + SharedDevices []adapterDevice `json:"shared_devices"` + Mounts []adapterMount `json:"mounts"` + Env map[string]string `json:"env"` +} + +type mountProfile map[string][]profileMountGroup + +type profileMountGroup struct { + Paths []string `json:"path"` + Type string `json:"type,omitempty"` +} + +type deviceStatFunc func(string) (deviceType string, major, minor int64, err error) + +type boundedBuffer struct { + bytes.Buffer + limit int + overflow bool +} + +func (b *boundedBuffer) Write(data []byte) (int, error) { + originalLength := len(data) + remaining := b.limit - b.Len() + if remaining <= 0 { + b.overflow = true + return originalLength, nil + } + if len(data) > remaining { + data = data[:remaining] + b.overflow = true + } + _, _ = b.Buffer.Write(data) + return originalLength, nil +} + +type ascendProvider struct { + mu sync.RWMutex + + config config.AscendConfig + sandboxRoot string + providerVersion string + profile mountProfile + devices map[uint32]ascendDevice + resources []Resource + leases map[string]string + healthy bool + reason error + statDevice deviceStatFunc +} + +func newAscendProvider(cfg config.AscendConfig, sandboxRoot string) *ascendProvider { + provider := &ascendProvider{ + config: cfg, + sandboxRoot: sandboxRoot, + devices: make(map[uint32]ascendDevice), + leases: make(map[string]string), + statDevice: statCharacterDevice, + } + if err := provider.initialize(); err != nil { + provider.reason = err + return provider + } + provider.healthy = true + return provider +} + +func (p *ascendProvider) initialize() error { + if err := validateRootOwnedExecutable(p.config.Adapter); err != nil { + return err + } + profileData, err := os.ReadFile(p.config.MountProfile) + if err != nil { + return fmt.Errorf("read Ascend mount profile: %w", err) + } + if err := json.Unmarshal(profileData, &p.profile); err != nil { + return fmt.Errorf("parse Ascend mount profile: %w", err) + } + if len(p.profile) == 0 { + return errors.New("Ascend mount profile must not be empty") + } + var version adapterVersion + if err := p.invoke(defaultAdapterTimeout, nil, &version, "version", "--output=json"); err != nil { + return fmt.Errorf("query Ascend adapter version: %w", err) + } + if version.SchemaVersion != ascendSchemaVersion || version.ProviderVersion == "" { + return errors.New("Ascend adapter returned an incompatible version") + } + var discovery adapterDiscovery + if err := p.invoke(defaultAdapterTimeout, nil, &discovery, "discover", "--output=json"); err != nil { + return fmt.Errorf("discover Ascend devices: %w", err) + } + if discovery.SchemaVersion != ascendSchemaVersion || discovery.ProviderVersion != version.ProviderVersion { + return errors.New("Ascend discovery schema or provider version mismatch") + } + p.providerVersion = version.ProviderVersion + if err := p.acceptDiscovery(discovery.Devices); err != nil { + return err + } + return p.restoreLeases() +} + +func validateRootOwnedExecutable(path string) error { + if path == "" { + return errors.New("Ascend adapter path is required") + } + info, err := os.Stat(path) + if err != nil { + return fmt.Errorf("stat Ascend adapter %s: %w", path, err) + } + if !info.Mode().IsRegular() || info.Mode().Perm()&0111 == 0 { + return errors.New("Ascend adapter must be an executable regular file") + } + if info.Mode().Perm()&0022 != 0 { + return errors.New("Ascend adapter must not be group/other writable") + } + if stat, ok := info.Sys().(*syscall.Stat_t); ok && stat.Uid != 0 && os.Geteuid() == 0 { + return errors.New("Ascend adapter must be owned by root") + } + return nil +} + +func (p *ascendProvider) invoke(timeout time.Duration, input any, output any, args ...string) error { + var stdin []byte + var err error + if input != nil { + stdin, err = json.Marshal(input) + if err != nil { + return err + } + if len(stdin) > maxAdapterInputBytes { + return errors.New("Ascend adapter input exceeds size limit") + } + } + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() + command := exec.CommandContext(ctx, p.config.Adapter, args...) + command.Stdin = bytes.NewReader(stdin) + stdout := &boundedBuffer{limit: maxAdapterOutputBytes} + stderr := &boundedBuffer{limit: 64 << 10} + command.Stdout = stdout + command.Stderr = stderr + err = command.Run() + if ctx.Err() != nil { + return fmt.Errorf("Ascend adapter timed out: %w", ctx.Err()) + } + if err != nil { + return fmt.Errorf("Ascend adapter failed: %w: %s", err, strings.TrimSpace(stderr.String())) + } + if stdout.overflow { + return errors.New("Ascend adapter output exceeds size limit") + } + if err := json.Unmarshal(stdout.Bytes(), output); err != nil { + return fmt.Errorf("parse Ascend adapter output: %w", err) + } + return nil +} + +func (p *ascendProvider) acceptDiscovery(devices []ascendDevice) error { + logicIDs := make(map[int32]struct{}) + stableIDs := make(map[string]struct{}) + for _, device := range devices { + if !device.Healthy { + continue + } + device.ProductModel = strings.ToLower(strings.TrimSpace(device.ProductModel)) + if device.ProductModel == "" { + return errors.New("Ascend product model must not be empty") + } + if !validAscendDeviceFamily(device) { + return fmt.Errorf("unsupported or inconsistent Ascend device family for product model %q", device.ProductModel) + } + if device.LogicID < 0 || device.StableID == "" { + return fmt.Errorf("invalid Ascend identity for scheduler ID %d", device.SchedulerID) + } + if _, duplicate := p.devices[device.SchedulerID]; duplicate { + return fmt.Errorf("duplicate Ascend scheduler ID %d", device.SchedulerID) + } + if _, duplicate := logicIDs[device.LogicID]; duplicate { + return fmt.Errorf("duplicate Ascend logic ID %d", device.LogicID) + } + if _, duplicate := stableIDs[device.StableID]; duplicate { + return fmt.Errorf("duplicate Ascend stable ID %q", device.StableID) + } + p.devices[device.SchedulerID] = device + logicIDs[device.LogicID] = struct{}{} + stableIDs[device.StableID] = struct{}{} + } + if len(p.devices) == 0 { + return errors.New("Ascend discovery returned no healthy physical NPU") + } + p.resources = buildAscendResources(p.devices) + return nil +} + +func validAscendDeviceFamily(device ascendDevice) bool { + switch device.Generation { + case "310P": + return device.RuntimeFamily == ascend310PRuntimeFamily && device.ResourceFamily == "huawei.com/Ascend310P" + case "A2", "A3": + return device.RuntimeFamily == ascendRuntimeFamily && device.ResourceFamily == "huawei.com/Ascend910" + default: + return false + } +} + +func buildAscendResources(devices map[uint32]ascendDevice) []Resource { + byModel := make(map[string][]uint32) + for id, device := range devices { + byModel[device.ProductModel] = append(byModel[device.ProductModel], id) + } + models := make([]string, 0, len(byModel)) + for model := range byModel { + models = append(models, model) + } + sort.Strings(models) + resources := make([]Resource, 0, len(models)) + for _, model := range models { + ids := byModel[model] + sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] }) + resources = append(resources, Resource{Type: TypeNPU, ProductModel: model, DeviceIDs: ids}) + } + return resources +} + +func (p *ascendProvider) Type() string { return TypeNPU } + +func (p *ascendProvider) SupportsRuntime(runtimeName string) bool { + return runtimeName == config.RuntimeNameRunc +} + +func (p *ascendProvider) Healthy() (bool, error) { + p.mu.RLock() + defer p.mu.RUnlock() + return p.healthy, p.reason +} + +func reservedAscendEnv(key string) bool { + switch key { + case ascendVisibleDevicesEnv, ascendRTVisibleDevicesEnv, "ASCEND_DOCKER_RUNTIME", + "ASCEND_RUNTIME_OPTIONS", "ASCEND_RUNTIME_MOUNTS", "ASCEND_VNPU_SPECS", + "ASCEND_ALLOW_LINK", "DISABLE_UB_MOUNT": + return true + default: + return false + } +} + +func (p *ascendProvider) Resources() []Resource { + p.mu.RLock() + defer p.mu.RUnlock() + if !p.healthy { + return []Resource{} + } + resources := make([]Resource, len(p.resources)) + copy(resources, p.resources) + for index := range resources { + resources[index].DeviceIDs = append([]uint32(nil), resources[index].DeviceIDs...) + } + return resources +} + +func (p *ascendProvider) Acquire( + sandboxID, runtimeName string, + allocation *api.XpuAllocation, +) (*svc.SpecUpdates, error) { + if sandboxID == "" || allocation == nil { + return nil, errors.New("sandbox ID and NPU allocation are required") + } + if !p.SupportsRuntime(runtimeName) { + return nil, fmt.Errorf("NPU allocations require runtime %q", config.RuntimeNameRunc) + } + if strings.ToLower(strings.TrimSpace(allocation.Type)) != TypeNPU { + return nil, errors.New("invalid NPU allocation") + } + if len(allocation.DeviceIds) == 0 { + return nil, errors.New("NPU device IDs must not be empty") + } + + p.mu.Lock() + defer p.mu.Unlock() + if !p.healthy { + return nil, fmt.Errorf("NPU support is unavailable: %w", p.reason) + } + seen := make(map[uint32]struct{}, len(allocation.DeviceIds)) + devices := make([]ascendDevice, 0, len(allocation.DeviceIds)) + for _, id := range allocation.DeviceIds { + if _, duplicate := seen[id]; duplicate { + return nil, fmt.Errorf("duplicate NPU device ID %d", id) + } + seen[id] = struct{}{} + device, ok := p.devices[id] + if !ok { + return nil, fmt.Errorf("NPU device ID %d is not in the node inventory", id) + } + if owner, leased := p.leases[device.StableID]; leased && owner != sandboxID { + return nil, fmt.Errorf("NPU device ID %d is already leased by sandbox %s", id, owner) + } + devices = append(devices, device) + } + for _, device := range devices { + p.leases[device.StableID] = sandboxID + } + rollback := func() { + for _, device := range devices { + if p.leases[device.StableID] == sandboxID { + delete(p.leases, device.StableID) + } + } + } + logicIDs := make([]int32, len(devices)) + stableIDs := make([]string, len(devices)) + model := devices[0].ProductModel + runtimeFamily := devices[0].RuntimeFamily + for index, device := range devices { + if device.ProductModel != model { + rollback() + return nil, errors.New("all NPU devices in one allocation must have the same product model") + } + if device.RuntimeFamily != runtimeFamily { + rollback() + return nil, errors.New("all NPU devices in one allocation must use the same runtime family") + } + logicIDs[index] = device.LogicID + stableIDs[index] = device.StableID + } + var edits adapterEdits + if err := p.invoke(defaultAdapterTimeout, adapterEditsRequest{ + SchemaVersion: ascendSchemaVersion, + LogicIDs: logicIDs, + RuntimeFamily: runtimeFamily, + PhysicalOnly: true, + MountProfile: p.config.MountProfile, + }, &edits, "edits", "--input=-", "--output=json"); err != nil { + rollback() + return nil, err + } + updates, err := p.validateEdits(logicIDs, runtimeFamily, edits) + if err != nil { + rollback() + return nil, err + } + recordJSON, err := encodeLease(leaseRecord{ + SandboxID: sandboxID, Type: TypeNPU, Runtime: runtimeName, + ProductModel: devices[0].ProductModel, + SchedulerIDs: append([]uint32(nil), allocation.DeviceIds...), + LogicIDs: append([]int32(nil), logicIDs...), StableIDs: stableIDs, + Provider: ascendProviderName, ProviderVersion: p.providerVersion, + }) + if err != nil { + rollback() + return nil, err + } + updates.Annotations = map[string]string{AllocationAnnotation: string(recordJSON)} + return updates, nil +} + +func (p *ascendProvider) validateEdits( + logicIDs []int32, runtimeFamily string, edits adapterEdits, +) (*svc.SpecUpdates, error) { + if edits.SchemaVersion != ascendSchemaVersion || edits.ProviderVersion != p.providerVersion { + return nil, errors.New("Ascend edits schema or provider version mismatch") + } + expectedDevices := make(map[string]struct{}, len(logicIDs)) + for _, id := range logicIDs { + expectedDevices[fmt.Sprintf("/dev/davinci%d", id)] = struct{}{} + } + updates := &svc.SpecUpdates{} + appendDevice := func(device adapterDevice, shared bool) error { + if device.Type != "c" || device.Permissions != "rwm" || device.Major < 0 || device.Minor < 0 { + return fmt.Errorf("invalid Ascend device edit for %s", device.ContainerPath) + } + if shared { + if _, allowed := allowedAscendSharedDevices[device.ContainerPath]; !allowed { + return fmt.Errorf("Ascend shared device %s is not allowed", device.ContainerPath) + } + if !sharedDeviceHostAllowed(device.HostPath, device.ContainerPath) { + return fmt.Errorf("Ascend shared device mapping %s -> %s is not allowed", device.HostPath, device.ContainerPath) + } + } else { + if _, expected := expectedDevices[device.ContainerPath]; !expected || device.HostPath != device.ContainerPath { + return fmt.Errorf("Ascend device %s is outside the current lease", device.ContainerPath) + } + delete(expectedDevices, device.ContainerPath) + } + deviceType, major, minor, err := p.statDevice(device.HostPath) + if err != nil || deviceType != device.Type || major != device.Major || minor != device.Minor { + return fmt.Errorf("Ascend device identity mismatch for %s", device.HostPath) + } + majorCopy, minorCopy := device.Major, device.Minor + updates.LinuxDevices = append(updates.LinuxDevices, svc.LinuxDevice{ + Path: device.ContainerPath, Type: device.Type, Major: device.Major, Minor: device.Minor, + }) + updates.DeviceCgroupRules = append(updates.DeviceCgroupRules, svc.LinuxDeviceCgroup{ + Allow: true, Type: device.Type, Major: &majorCopy, Minor: &minorCopy, Access: device.Permissions, + }) + return nil + } + for _, device := range edits.Devices { + if err := appendDevice(device, false); err != nil { + return nil, err + } + } + if len(expectedDevices) != 0 { + return nil, errors.New("Ascend edits did not include every leased device") + } + for _, device := range edits.SharedDevices { + if err := appendDevice(device, true); err != nil { + return nil, err + } + } + for _, mount := range edits.Mounts { + if !p.mountAllowed(runtimeFamily, mount.Source) || mount.Destination != mount.Source || mount.Type != "bind" || + !containsOption(mount.Options, "ro") { + return nil, fmt.Errorf("Ascend mount %s -> %s is not an allowed read-only bind", mount.Source, mount.Destination) + } + if !filepath.IsAbs(mount.Source) || !filepath.IsAbs(mount.Destination) { + return nil, errors.New("Ascend mount paths must be absolute") + } + updates.Mounts = append(updates.Mounts, svc.Mount{ + Destination: mount.Destination, Type: mount.Type, Source: mount.Source, + Options: append([]string(nil), mount.Options...), + }) + } + for key, value := range edits.Env { + if key == "LD_LIBRARY_PATH" { + updates.PrependLibraryPaths = strings.Split(value, ":") + continue + } + if !reservedAscendEnv(key) { + return nil, fmt.Errorf("Ascend adapter returned unowned environment variable %q", key) + } + updates.Envs = append(updates.Envs, &api.KeyValue{Key: key, Value: value}) + } + return updates, nil +} + +func sharedDeviceHostAllowed(hostPath, containerPath string) bool { + if hostPath == containerPath { + return true + } + return hostPath == "/dev/davinci_manager_docker" && containerPath == "/dev/davinci_manager" +} + +func (p *ascendProvider) mountAllowed(runtimeFamily, source string) bool { + groups := p.profile[runtimeFamily] + if len(groups) == 0 { + groups = p.profile["default"] + } + for _, group := range groups { + if strings.EqualFold(group.Type, "UB") { + continue + } + for _, pattern := range group.Paths { + if matched, err := filepath.Match(pattern, source); err == nil && matched { + return true + } + if pattern == source { + return true + } + } + } + return false +} + +func containsOption(options []string, expected string) bool { + for _, option := range options { + if option == expected { + return true + } + } + return false +} + +func (p *ascendProvider) Release(sandboxID string) { + if sandboxID == "" { + return + } + p.mu.Lock() + defer p.mu.Unlock() + for stableID, owner := range p.leases { + if owner == sandboxID { + delete(p.leases, stableID) + } + } +} + +func (p *ascendProvider) restoreLeases() error { + leases, err := readPersistedLeases(p.sandboxRoot) + if err != nil { + return err + } + for _, persisted := range leases { + record := persisted.record + if record.Type != TypeNPU { + continue + } + if record.SchemaVersion != leaseSchemaVersion || record.Runtime != config.RuntimeNameRunc || + record.Provider != ascendProviderName || record.ProviderVersion != p.providerVersion || + len(record.SchedulerIDs) == 0 || len(record.SchedulerIDs) != len(record.StableIDs) { + return fmt.Errorf("invalid NPU allocation annotation in %s", persisted.bundlePath) + } + for index, id := range record.SchedulerIDs { + device, ok := p.devices[id] + if !ok || device.StableID != record.StableIDs[index] || device.ProductModel != record.ProductModel { + return fmt.Errorf("NPU identity changed for device ID %d in %s", id, persisted.bundlePath) + } + if owner, duplicate := p.leases[device.StableID]; duplicate && owner != record.SandboxID { + return fmt.Errorf("NPU stable ID %s is assigned to both %s and %s", device.StableID, owner, record.SandboxID) + } + p.leases[device.StableID] = record.SandboxID + } + } + return nil +} diff --git a/pkg/xpumanager/ascend_test.go b/pkg/xpumanager/ascend_test.go new file mode 100644 index 0000000..87bcc18 --- /dev/null +++ b/pkg/xpumanager/ascend_test.go @@ -0,0 +1,136 @@ +// Copyright (c) 2026 Ant Group Corporation. +// +// 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 xpumanager + +import ( + "encoding/json" + "os" + "path/filepath" + "testing" + + api "github.com/inclusionAI/sandboxd/api/runtime/v1" + "github.com/inclusionAI/sandboxd/config" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestAscendProviderDiscoversAcquiresAndReleases(t *testing.T) { + tempDir := t.TempDir() + adapterPath := filepath.Join(tempDir, "fake-ascend-adapter") + profilePath := filepath.Join(tempDir, "mounts.json") + require.NoError(t, os.WriteFile(profilePath, []byte(`{"default":[]}`), 0600)) + require.NoError(t, os.WriteFile(adapterPath, []byte(`#!/bin/sh +case "$1" in + version) printf '%s' '{"schema_version":1,"provider_version":"fake-v1"}' ;; + discover) printf '%s' '{"schema_version":1,"provider_version":"fake-v1","devices":[{"scheduler_id":0,"logic_id":0,"physical_id":0,"stable_id":"die-0","product_model":"ascend310p3","generation":"310P","runtime_family":"Ascend310P","resource_family":"huawei.com/Ascend310P","raw_product":"310P3","healthy":true}]}' ;; + edits) + request=$(cat) + case "$request" in + *'"runtime_family":"Ascend310P"'*) ;; + *) exit 3 ;; + esac + printf '%s' '{"schema_version":1,"provider_version":"fake-v1","devices":[{"host_path":"/dev/davinci0","container_path":"/dev/davinci0","type":"c","major":1,"minor":3,"permissions":"rwm"}],"shared_devices":[],"mounts":[],"env":{"ASCEND_RT_VISIBLE_DEVICES":"0","LD_LIBRARY_PATH":"/driver/lib:/dcmi/lib"}}' + ;; + *) exit 2 ;; +esac +`), 0700)) + + provider := newAscendProvider(config.AscendConfig{ + Enabled: true, Adapter: adapterPath, MountProfile: profilePath, + }, tempDir) + healthy, reason := provider.Healthy() + require.True(t, healthy, reason) + assert.Equal(t, []Resource{{Type: TypeNPU, ProductModel: "ascend310p3", DeviceIDs: []uint32{0}}}, provider.Resources()) + provider.statDevice = func(path string) (string, int64, int64, error) { + assert.Equal(t, "/dev/davinci0", path) + return "c", 1, 3, nil + } + + updates, err := provider.Acquire("sbox-npu", config.RuntimeNameRunc, &api.XpuAllocation{ + Type: "npu", DeviceIds: []uint32{0}, + }) + require.NoError(t, err) + require.Len(t, updates.LinuxDevices, 1) + require.Len(t, updates.DeviceCgroupRules, 1) + assert.Equal(t, []string{"/driver/lib", "/dcmi/lib"}, updates.PrependLibraryPaths) + assert.Contains(t, updates.Annotations, AllocationAnnotation) + var record leaseRecord + require.NoError(t, json.Unmarshal([]byte(updates.Annotations[AllocationAnnotation]), &record)) + assert.Equal(t, "ascend310p3", record.ProductModel) + + _, err = provider.Acquire("sbox-other", config.RuntimeNameRunc, &api.XpuAllocation{ + Type: "npu", DeviceIds: []uint32{0}, + }) + require.ErrorContains(t, err, "already leased") + provider.Release("sbox-npu") + _, err = provider.Acquire("sbox-other", config.RuntimeNameRunc, &api.XpuAllocation{ + Type: "npu", DeviceIds: []uint32{0}, + }) + require.NoError(t, err) +} + +func TestAscendProviderGroupsMixedModelsAndRejectsMixedAllocation(t *testing.T) { + provider := &ascendProvider{devices: make(map[uint32]ascendDevice)} + err := provider.acceptDiscovery([]ascendDevice{ + {SchedulerID: 0, LogicID: 0, StableID: "a", ProductModel: "ascend910", Generation: "A3", RuntimeFamily: ascendRuntimeFamily, ResourceFamily: "huawei.com/Ascend910", Healthy: true}, + {SchedulerID: 1, LogicID: 1, StableID: "b", ProductModel: "ascend910b4", Generation: "A2", RuntimeFamily: ascendRuntimeFamily, ResourceFamily: "huawei.com/Ascend910", Healthy: true}, + }) + require.NoError(t, err) + provider.healthy = true + assert.Equal(t, []Resource{ + {Type: TypeNPU, ProductModel: "ascend910", DeviceIDs: []uint32{0}}, + {Type: TypeNPU, ProductModel: "ascend910b4", DeviceIDs: []uint32{1}}, + }, provider.Resources()) + provider.leases = make(map[string]string) + _, err = provider.Acquire("mixed", config.RuntimeNameRunc, &api.XpuAllocation{Type: "npu", DeviceIds: []uint32{0, 1}}) + require.ErrorContains(t, err, "same product model") + assert.Empty(t, provider.leases) +} + +func TestAscendProviderValidatesDeviceFamilyAndUnauthorizedEdits(t *testing.T) { + provider := &ascendProvider{devices: make(map[uint32]ascendDevice)} + err := provider.acceptDiscovery([]ascendDevice{ + {SchedulerID: 0, LogicID: 0, StableID: "a", ProductModel: "custom-driver-model", Generation: "A3", RuntimeFamily: ascendRuntimeFamily, ResourceFamily: "huawei.com/Ascend910", Healthy: true}, + }) + require.NoError(t, err) + + provider = &ascendProvider{devices: make(map[uint32]ascendDevice)} + err = provider.acceptDiscovery([]ascendDevice{ + {SchedulerID: 0, LogicID: 0, StableID: "a", ProductModel: "ascend910", Generation: "A1", RuntimeFamily: ascendRuntimeFamily, ResourceFamily: "huawei.com/Ascend910", Healthy: true}, + }) + require.ErrorContains(t, err, "unsupported or inconsistent Ascend device family") + + provider = &ascendProvider{devices: make(map[uint32]ascendDevice)} + err = provider.acceptDiscovery([]ascendDevice{ + {SchedulerID: 0, LogicID: 0, StableID: "a", ProductModel: "", Generation: "A3", RuntimeFamily: ascendRuntimeFamily, ResourceFamily: "huawei.com/Ascend910", Healthy: true}, + }) + require.ErrorContains(t, err, "must not be empty") + + provider = &ascendProvider{providerVersion: "v1", statDevice: func(string) (string, int64, int64, error) { + return "c", 1, 2, nil + }} + _, err = provider.validateEdits([]int32{0}, ascendRuntimeFamily, adapterEdits{ + SchemaVersion: 1, ProviderVersion: "v1", + Devices: []adapterDevice{{HostPath: "/dev/davinci9", ContainerPath: "/dev/davinci9", Type: "c", Major: 1, Minor: 2, Permissions: "rwm"}}, + }) + require.ErrorContains(t, err, "outside the current lease") + + _, err = provider.validateEdits([]int32{0}, ascendRuntimeFamily, adapterEdits{ + SchemaVersion: 1, ProviderVersion: "v1", + Devices: []adapterDevice{{HostPath: "/dev/davinci0", ContainerPath: "/dev/davinci0", Type: "c", Major: 1, Minor: 2, Permissions: "rwm"}}, + SharedDevices: []adapterDevice{{HostPath: "/dev/random", ContainerPath: "/dev/davinci_manager", Type: "c", Major: 1, Minor: 2, Permissions: "rwm"}}, + }) + require.ErrorContains(t, err, "mapping") +} diff --git a/pkg/xpumanager/device_stat_linux.go b/pkg/xpumanager/device_stat_linux.go new file mode 100644 index 0000000..5537235 --- /dev/null +++ b/pkg/xpumanager/device_stat_linux.go @@ -0,0 +1,34 @@ +//go:build linux + +// Copyright (c) 2026 Ant Group Corporation. +// +// 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 xpumanager + +import ( + "fmt" + + "golang.org/x/sys/unix" +) + +func statCharacterDevice(path string) (string, int64, int64, error) { + var stat unix.Stat_t + if err := unix.Stat(path, &stat); err != nil { + return "", 0, 0, err + } + if stat.Mode&unix.S_IFMT != unix.S_IFCHR { + return "", 0, 0, fmt.Errorf("%s is not a character device", path) + } + return "c", int64(unix.Major(uint64(stat.Rdev))), int64(unix.Minor(uint64(stat.Rdev))), nil +} diff --git a/pkg/xpumanager/device_stat_linux_test.go b/pkg/xpumanager/device_stat_linux_test.go new file mode 100644 index 0000000..7a85d36 --- /dev/null +++ b/pkg/xpumanager/device_stat_linux_test.go @@ -0,0 +1,36 @@ +//go:build linux + +// Copyright (c) 2026 Ant Group Corporation. +// +// 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 xpumanager + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestStatCharacterDevice(t *testing.T) { + deviceType, major, minor, err := statCharacterDevice("/dev/null") + require.NoError(t, err) + require.Equal(t, "c", deviceType) + require.Equal(t, int64(1), major) + require.Equal(t, int64(3), minor) +} + +func TestStatCharacterDeviceRejectsRegularFile(t *testing.T) { + _, _, _, err := statCharacterDevice(t.TempDir()) + require.Error(t, err) +} diff --git a/pkg/xpumanager/device_stat_other.go b/pkg/xpumanager/device_stat_other.go new file mode 100644 index 0000000..6137df8 --- /dev/null +++ b/pkg/xpumanager/device_stat_other.go @@ -0,0 +1,23 @@ +//go:build !linux + +// Copyright (c) 2026 Ant Group Corporation. +// +// 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 xpumanager + +import "errors" + +func statCharacterDevice(string) (string, int64, int64, error) { + return "", 0, 0, errors.New("Ascend device validation requires Linux") +} diff --git a/pkg/xpumanager/lease.go b/pkg/xpumanager/lease.go new file mode 100644 index 0000000..147bf04 --- /dev/null +++ b/pkg/xpumanager/lease.go @@ -0,0 +1,98 @@ +// Copyright (c) 2026 Ant Group Corporation. +// +// 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 xpumanager + +import ( + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/inclusionAI/sandboxd/config" +) + +const leaseSchemaVersion = 2 + +type leaseRecord struct { + SchemaVersion int `json:"schema_version,omitempty"` + SandboxID string `json:"sandbox_id"` + Type string `json:"type"` + Runtime string `json:"runtime,omitempty"` + ProductModel string `json:"product_model,omitempty"` + SchedulerIDs []uint32 `json:"scheduler_ids,omitempty"` + LogicIDs []int32 `json:"logic_ids,omitempty"` + StableIDs []string `json:"stable_ids,omitempty"` + Provider string `json:"provider,omitempty"` + ProviderVersion string `json:"provider_version,omitempty"` + + // Legacy GPU fields are accepted for rolling upgrades from schema v1. + DeviceIDs []uint32 `json:"device_ids,omitempty"` + DeviceUUID []string `json:"device_uuids,omitempty"` +} + +type persistedLease struct { + bundlePath string + record leaseRecord +} + +func encodeLease(record leaseRecord) ([]byte, error) { + record.SchemaVersion = leaseSchemaVersion + return json.Marshal(record) +} + +func readPersistedLeases(sandboxRoot string) ([]persistedLease, error) { + entries, err := os.ReadDir(sandboxRoot) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return nil, nil + } + return nil, fmt.Errorf("read sandbox root %s: %w", sandboxRoot, err) + } + var leases []persistedLease + for _, entry := range entries { + if !entry.IsDir() || !strings.HasPrefix(entry.Name(), config.SandboxIDPrefix) { + continue + } + configPath := filepath.Join(sandboxRoot, entry.Name(), config.SandboxSpecFile) + data, err := os.ReadFile(configPath) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + continue + } + return nil, fmt.Errorf("read XPU lease from %s: %w", configPath, err) + } + var spec struct { + Annotations map[string]string `json:"annotations"` + } + if err := json.Unmarshal(data, &spec); err != nil { + return nil, fmt.Errorf("parse XPU lease from %s: %w", configPath, err) + } + raw := spec.Annotations[AllocationAnnotation] + if raw == "" { + continue + } + var record leaseRecord + if err := json.Unmarshal([]byte(raw), &record); err != nil { + return nil, fmt.Errorf("parse XPU allocation annotation in %s: %w", configPath, err) + } + if record.SandboxID != entry.Name() || record.Type == "" { + return nil, fmt.Errorf("invalid XPU allocation annotation in %s", configPath) + } + leases = append(leases, persistedLease{bundlePath: configPath, record: record}) + } + return leases, nil +} diff --git a/pkg/xpumanager/manager.go b/pkg/xpumanager/manager.go index 118c659..09df3d9 100644 --- a/pkg/xpumanager/manager.go +++ b/pkg/xpumanager/manager.go @@ -17,14 +17,10 @@ package xpumanager import ( - "encoding/json" "errors" "fmt" - "os" - "path/filepath" "sort" "strings" - "sync" api "github.com/inclusionAI/sandboxd/api/runtime/v1" "github.com/inclusionAI/sandboxd/config" @@ -34,6 +30,7 @@ import ( const ( TypeGPU = "gpu" + TypeNPU = "npu" AllocationAnnotation = "sandbox.akernel.dev/xpu-allocation" ) @@ -45,246 +42,107 @@ type Resource struct { DeviceIDs []uint32 `json:"device_ids"` } -// Device contains provider-private identity for one scheduler-visible ID. -type Device struct { - ID uint32 - UUID string - ProductModel string +// Provider is the vendor boundary below the node-local XPU coordinator. +type Provider interface { + Type() string + Resources() []Resource + SupportsRuntime(string) bool + Acquire(sandboxID, runtimeName string, allocation *api.XpuAllocation) (*svc.SpecUpdates, error) + Release(sandboxID string) + Healthy() (bool, error) } -type leaseRecord struct { - SandboxID string `json:"sandbox_id"` - Type string `json:"type"` - DeviceIDs []uint32 `json:"device_ids"` - DeviceUUID []string `json:"device_uuids"` -} - -// Manager owns an immutable discovery snapshot and UUID-keyed leases. +// Manager is the XPU coordinator and provider registry. type Manager struct { - mu sync.RWMutex - - runscBinary string - sandboxRoot string - run commandRunner - stat statFunc - - devices map[uint32]Device - resources []Resource - leases map[string]string - healthy bool - reason error + providers map[string]Provider } -// New discovers the local NVIDIA inventory. Discovery failure is intentionally -// non-fatal for sandboxd: CPU-only nodes stay usable and advertise no XPU. -func New(runscBinary, sandboxRoot string) *Manager { - manager := &Manager{ - runscBinary: runscBinary, - sandboxRoot: sandboxRoot, - run: runCommand, - stat: os.Stat, - devices: make(map[uint32]Device), - leases: make(map[string]string), +// New constructs independent NVIDIA and optional Ascend providers. A provider +// discovery failure is non-fatal for sandboxd and does not affect other types. +func New(runscBinary string, runcConfigured bool, sandboxRoot string, xpuConfig config.XPUConfig) *Manager { + manager := &Manager{providers: make(map[string]Provider)} + manager.register(newNVIDIAProvider(runscBinary, runcConfigured, sandboxRoot)) + if xpuConfig.Ascend.Enabled { + manager.register(newAscendProvider(xpuConfig.Ascend, sandboxRoot)) } - if err := manager.discoverNVIDIA(); err != nil { - manager.reason = err - logrus.Infof("xpumanager: NVIDIA GPU support unavailable: %v", err) - return manager - } - if err := manager.restoreLeases(); err != nil { - manager.healthy = false - manager.resources = nil - manager.reason = err - logrus.Errorf("xpumanager: refusing GPU allocations after lease recovery failure: %v", err) - return manager - } - manager.healthy = true - logrus.Infof("xpumanager: discovered %d schedulable NVIDIA GPU(s)", len(manager.devices)) return manager } -// Resources returns a deep copy of the stable capacity inventory. Active -// leases never alter this list. -func (m *Manager) Resources() []Resource { - m.mu.RLock() - defer m.mu.RUnlock() - if !m.healthy { - return []Resource{} +func (m *Manager) register(provider Provider) { + if provider == nil { + return } - resources := make([]Resource, len(m.resources)) - for index := range m.resources { - resources[index] = m.resources[index] - resources[index].DeviceIDs = append([]uint32(nil), m.resources[index].DeviceIDs...) + m.providers[provider.Type()] = provider + if healthy, reason := provider.Healthy(); !healthy { + logrus.Infof("xpumanager: %s provider unavailable: %v", provider.Type(), reason) } - return resources } -// ReservedAnnotation reports whether key stores provider-owned allocation -// state. Callers must not be able to forge recovery metadata through labels. -func ReservedAnnotation(key string) bool { - return key == AllocationAnnotation +// Resources returns a deterministic deep copy of every healthy provider's +// stable inventory. Active leases never alter this list. +func (m *Manager) Resources() []Resource { + resources := make([]Resource, 0) + for _, provider := range m.providers { + if healthy, _ := provider.Healthy(); !healthy { + continue + } + resources = append(resources, provider.Resources()...) + } + sort.Slice(resources, func(i, j int) bool { + if resources[i].Type == resources[j].Type { + return resources[i].ProductModel < resources[j].ProductModel + } + return resources[i].Type < resources[j].Type + }) + return resources } -// Acquire validates and atomically leases all requested devices. -func (m *Manager) Acquire(sandboxID string, allocations []*api.XpuAllocation) (*svc.SpecUpdates, error) { +// ValidateRuntime rejects unsupported type/runtime combinations before node +// resources and filesystems are prepared. +func (m *Manager) ValidateRuntime(runtimeName string, allocations []*api.XpuAllocation) error { if len(allocations) == 0 { - return nil, nil - } - if sandboxID == "" { - return nil, errors.New("sandbox ID is required for XPU allocation") + return nil } if len(allocations) != 1 || allocations[0] == nil { - return nil, errors.New("exactly one XPU allocation is supported") - } - allocation := allocations[0] - if strings.ToLower(strings.TrimSpace(allocation.Type)) != TypeGPU { - return nil, fmt.Errorf("unsupported XPU type %q", allocation.Type) - } - if len(allocation.DeviceIds) == 0 { - return nil, errors.New("XPU device IDs must not be empty") + return errors.New("exactly one XPU allocation is supported") } - - m.mu.Lock() - defer m.mu.Unlock() - if !m.healthy { - if m.reason != nil { - return nil, fmt.Errorf("GPU support is unavailable: %w", m.reason) - } - return nil, errors.New("GPU support is unavailable") - } - - seen := make(map[uint32]struct{}, len(allocation.DeviceIds)) - devices := make([]Device, 0, len(allocation.DeviceIds)) - for _, id := range allocation.DeviceIds { - if _, duplicate := seen[id]; duplicate { - return nil, fmt.Errorf("duplicate GPU device ID %d", id) - } - seen[id] = struct{}{} - device, ok := m.devices[id] - if !ok { - return nil, fmt.Errorf("GPU device ID %d is not in the node inventory", id) - } - if owner, leased := m.leases[device.UUID]; leased && owner != sandboxID { - return nil, fmt.Errorf("GPU device ID %d is already leased by sandbox %s", id, owner) - } - devices = append(devices, device) - } - model := devices[0].ProductModel - for _, device := range devices[1:] { - if device.ProductModel != model { - return nil, errors.New("all GPU devices in one allocation must have the same product model") - } + typeName := strings.ToLower(strings.TrimSpace(allocations[0].Type)) + provider, ok := m.providers[typeName] + if !ok { + return fmt.Errorf("unsupported XPU type %q", allocations[0].Type) } - for _, device := range devices { - m.leases[device.UUID] = sandboxID + if !provider.SupportsRuntime(runtimeName) { + return fmt.Errorf("XPU type %q does not support runtime %q", typeName, runtimeName) } + return nil +} - uuids := make([]string, len(devices)) - for index, device := range devices { - uuids[index] = device.UUID +// Acquire routes a trusted scheduler allocation to the matching provider. +func (m *Manager) Acquire(sandboxID, runtimeName string, allocations []*api.XpuAllocation) (*svc.SpecUpdates, error) { + if err := m.ValidateRuntime(runtimeName, allocations); err != nil { + return nil, err } - record := leaseRecord{ - SandboxID: sandboxID, - Type: TypeGPU, - DeviceIDs: append([]uint32(nil), allocation.DeviceIds...), - DeviceUUID: append([]string(nil), uuids...), - } - recordJSON, err := json.Marshal(record) - if err != nil { - for _, uuid := range uuids { - delete(m.leases, uuid) - } - return nil, fmt.Errorf("encode GPU lease: %w", err) + if len(allocations) == 0 { + return nil, nil } - - return nvidiaSpecUpdates(uuids, recordJSON), nil + allocation := allocations[0] + provider := m.providers[strings.ToLower(strings.TrimSpace(allocation.Type))] + return provider.Acquire(sandboxID, runtimeName, allocation) } -// Release releases all UUID leases owned by sandboxID. It is idempotent. +// Release releases every provider lease owned by sandboxID. It is idempotent. func (m *Manager) Release(sandboxID string) { - if sandboxID == "" { - return - } - m.mu.Lock() - defer m.mu.Unlock() - for uuid, owner := range m.leases { - if owner == sandboxID { - delete(m.leases, uuid) - } + for _, provider := range m.providers { + provider.Release(sandboxID) } } -func (m *Manager) restoreLeases() error { - entries, err := os.ReadDir(m.sandboxRoot) - if err != nil { - if errors.Is(err, os.ErrNotExist) { - return nil - } - return fmt.Errorf("read sandbox root %s: %w", m.sandboxRoot, err) - } - for _, entry := range entries { - if !entry.IsDir() || !strings.HasPrefix(entry.Name(), config.SandboxIDPrefix) { - continue - } - configPath := filepath.Join(m.sandboxRoot, entry.Name(), config.SandboxSpecFile) - data, err := os.ReadFile(configPath) - if err != nil { - if errors.Is(err, os.ErrNotExist) { - continue - } - return fmt.Errorf("read XPU lease from %s: %w", configPath, err) - } - var spec struct { - Annotations map[string]string `json:"annotations"` - } - if err := json.Unmarshal(data, &spec); err != nil { - return fmt.Errorf("parse XPU lease from %s: %w", configPath, err) - } - raw := spec.Annotations[AllocationAnnotation] - if raw == "" { - continue - } - var record leaseRecord - if err := json.Unmarshal([]byte(raw), &record); err != nil { - return fmt.Errorf("parse XPU allocation annotation in %s: %w", configPath, err) - } - if record.SandboxID != entry.Name() || record.Type != TypeGPU || - len(record.DeviceIDs) == 0 || len(record.DeviceIDs) != len(record.DeviceUUID) { - return fmt.Errorf("invalid XPU allocation annotation in %s", configPath) - } - for index, id := range record.DeviceIDs { - device, ok := m.devices[id] - if !ok || device.UUID != record.DeviceUUID[index] { - return fmt.Errorf("GPU identity changed for device ID %d in %s", id, configPath) - } - if owner, duplicate := m.leases[device.UUID]; duplicate && owner != record.SandboxID { - return fmt.Errorf("GPU UUID %s is assigned to both %s and %s", device.UUID, owner, record.SandboxID) - } - m.leases[device.UUID] = record.SandboxID - } - } - return nil +// ReservedEnv reports whether any accelerator provider owns key. +func ReservedEnv(key string) bool { + return reservedNVIDIAEnv(key) || reservedAscendEnv(key) } -func buildResources(devices map[uint32]Device) []Resource { - byModel := make(map[string][]uint32) - for id, device := range devices { - byModel[device.ProductModel] = append(byModel[device.ProductModel], id) - } - models := make([]string, 0, len(byModel)) - for model := range byModel { - models = append(models, model) - } - sort.Strings(models) - resources := make([]Resource, 0, len(models)) - for _, model := range models { - ids := byModel[model] - sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] }) - resources = append(resources, Resource{ - Type: TypeGPU, - ProductModel: model, - DeviceIDs: ids, - }) - } - return resources +// ReservedAnnotation reports whether key stores provider-owned recovery state. +func ReservedAnnotation(key string) bool { + return key == AllocationAnnotation } diff --git a/pkg/xpumanager/manager_test.go b/pkg/xpumanager/manager_test.go index 4121e5c..03b18a7 100644 --- a/pkg/xpumanager/manager_test.go +++ b/pkg/xpumanager/manager_test.go @@ -15,6 +15,7 @@ package xpumanager import ( + "context" "encoding/json" "os" "path/filepath" @@ -43,18 +44,46 @@ Brand: Tesla GPU UUID: GPU-uuid-0 ` -func testManager(t *testing.T) *Manager { +func testNVIDIAProvider(t *testing.T) *nvidiaProvider { t.Helper() _, devices, err := parseNVIDIAInfo(sampleNvidiaInfo) require.NoError(t, err) - return &Manager{ - devices: devices, - resources: buildResources(devices), - leases: make(map[string]string), - healthy: true, + return &nvidiaProvider{ + devices: devices, + resources: buildResources(devices), + leases: make(map[string]string), + healthy: true, + runscReady: true, + runcEnabled: true, } } +func TestNVIDIARuntimeSupportIsIndependent(t *testing.T) { + provider := &nvidiaProvider{ + runscBinary: "/usr/local/bin/runsc", + runcEnabled: true, + run: func(context.Context, string, ...string) ([]byte, error) { + return []byte("570.195.03\n"), nil + }, + } + require.NoError(t, provider.configureRuntimeSupport(context.Background(), "470.223.02")) + assert.False(t, provider.SupportsRuntime(config.RuntimeNameRunsc)) + assert.True(t, provider.SupportsRuntime(config.RuntimeNameRunc)) + + provider.runcEnabled = false + require.ErrorContains( + t, + provider.configureRuntimeSupport(context.Background(), "470.223.02"), + "not supported", + ) + + provider.run = func(context.Context, string, ...string) ([]byte, error) { + return []byte("470.223.02\n"), nil + } + require.NoError(t, provider.configureRuntimeSupport(context.Background(), "470.223.02")) + assert.True(t, provider.SupportsRuntime(config.RuntimeNameRunsc)) +} + func TestParseNVIDIAInfoAndResources(t *testing.T) { driver, devices, err := parseNVIDIAInfo(sampleNvidiaInfo) require.NoError(t, err) @@ -68,125 +97,191 @@ func TestParseNVIDIAInfoAndResources(t *testing.T) { }}, buildResources(devices)) } +func TestValidateNVIDIARuntimeHook(t *testing.T) { + path := filepath.Join(t.TempDir(), "nvidia-container-runtime-hook") + require.NoError(t, os.WriteFile(path, []byte("#!/bin/sh\n"), 0644)) + require.ErrorContains(t, validateNVIDIARuntimeHook(path, os.Stat), "executable regular file") + require.NoError(t, os.Chmod(path, 0755)) + require.NoError(t, validateNVIDIARuntimeHook(path, os.Stat)) +} + func TestAcquireMultipleGPUs(t *testing.T) { - manager := testManager(t) - updates, err := manager.Acquire("sbox-gpu", []*api.XpuAllocation{{ - Type: "gpu", - DeviceIds: []uint32{0, 2}, - }}) - require.NoError(t, err) - require.Len(t, updates.Prestart, 1) - assert.True(t, updates.RequiresHostWritableRootfs) - assert.Equal(t, nvidiaRuntimeHookPath, updates.Prestart[0].Path) - assert.Equal(t, "GPU-uuid-0,GPU-uuid-2", updates.Envs[0].Value) - assert.Equal(t, "compute,utility", updates.Envs[1].Value) - assert.Equal(t, "0,1", updates.Envs[2].Value) + for _, runtimeName := range []string{config.RuntimeNameRunsc, config.RuntimeNameRunc} { + t.Run(runtimeName, func(t *testing.T) { + provider := testNVIDIAProvider(t) + updates, err := provider.Acquire("sbox-gpu", runtimeName, &api.XpuAllocation{ + Type: "gpu", + DeviceIds: []uint32{0, 2}, + }) + require.NoError(t, err) + require.Len(t, updates.Prestart, 1) + assert.Equal(t, runtimeName == config.RuntimeNameRunsc, updates.RequiresHostWritableRootfs) + assert.Equal(t, nvidiaRuntimeHookPath, updates.Prestart[0].Path) + assert.Equal(t, "GPU-uuid-0,GPU-uuid-2", updates.Envs[0].Value) + assert.Equal(t, "compute,utility", updates.Envs[1].Value) + assert.Equal(t, "0,1", updates.Envs[2].Value) - var record leaseRecord - require.NoError(t, json.Unmarshal([]byte(updates.Annotations[AllocationAnnotation]), &record)) - assert.Equal(t, []uint32{0, 2}, record.DeviceIDs) - assert.Equal(t, []string{"GPU-uuid-0", "GPU-uuid-2"}, record.DeviceUUID) + var record leaseRecord + require.NoError(t, json.Unmarshal([]byte(updates.Annotations[AllocationAnnotation]), &record)) + assert.Equal(t, runtimeName, record.Runtime) + assert.Equal(t, []uint32{0, 2}, record.SchedulerIDs) + assert.Equal(t, []string{"GPU-uuid-0", "GPU-uuid-2"}, record.StableIDs) + }) + } } func TestAcquireIsAtomicAndReleaseIsIdempotent(t *testing.T) { - manager := testManager(t) - _, err := manager.Acquire("sbox-owner", []*api.XpuAllocation{{ + provider := testNVIDIAProvider(t) + _, err := provider.Acquire("sbox-owner", config.RuntimeNameRunsc, &api.XpuAllocation{ Type: "gpu", DeviceIds: []uint32{0}, - }}) + }) require.NoError(t, err) - _, err = manager.Acquire("sbox-other", []*api.XpuAllocation{{ + _, err = provider.Acquire("sbox-other", config.RuntimeNameRunsc, &api.XpuAllocation{ Type: "gpu", DeviceIds: []uint32{2, 0}, - }}) + }) require.ErrorContains(t, err, "already leased") - assert.NotContains(t, manager.leases, "GPU-uuid-2") + assert.NotContains(t, provider.leases, "GPU-uuid-2") - manager.Release("sbox-owner") - manager.Release("sbox-owner") - _, err = manager.Acquire("sbox-other", []*api.XpuAllocation{{ + provider.Release("sbox-owner") + provider.Release("sbox-owner") + _, err = provider.Acquire("sbox-other", config.RuntimeNameRunsc, &api.XpuAllocation{ Type: "gpu", DeviceIds: []uint32{2, 0}, - }}) + }) require.NoError(t, err) } func TestAcquireRejectsInvalidAllocations(t *testing.T) { tests := []struct { name string - allocation []*api.XpuAllocation + allocation *api.XpuAllocation errorText string }{ - {name: "empty", allocation: []*api.XpuAllocation{{Type: "gpu"}}, errorText: "must not be empty"}, - {name: "duplicate", allocation: []*api.XpuAllocation{{Type: "gpu", DeviceIds: []uint32{0, 0}}}, errorText: "duplicate"}, - {name: "unknown ID", allocation: []*api.XpuAllocation{{Type: "gpu", DeviceIds: []uint32{1}}}, errorText: "not in the node inventory"}, - {name: "unknown type", allocation: []*api.XpuAllocation{{Type: "npu", DeviceIds: []uint32{0}}}, errorText: "unsupported"}, - {name: "multiple", allocation: []*api.XpuAllocation{{Type: "gpu", DeviceIds: []uint32{0}}, {Type: "gpu", DeviceIds: []uint32{2}}}, errorText: "exactly one"}, + {name: "empty", allocation: &api.XpuAllocation{Type: "gpu"}, errorText: "must not be empty"}, + {name: "duplicate", allocation: &api.XpuAllocation{Type: "gpu", DeviceIds: []uint32{0, 0}}, errorText: "duplicate"}, + {name: "unknown ID", allocation: &api.XpuAllocation{Type: "gpu", DeviceIds: []uint32{1}}, errorText: "not in the node inventory"}, + {name: "unknown type", allocation: &api.XpuAllocation{Type: "npu", DeviceIds: []uint32{0}}, errorText: "invalid GPU"}, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { - _, err := testManager(t).Acquire("sbox-test", test.allocation) + _, err := testNVIDIAProvider(t).Acquire("sbox-test", config.RuntimeNameRunsc, test.allocation) require.ErrorContains(t, err, test.errorText) }) } } func TestResourcesAreStableAcrossLeases(t *testing.T) { - manager := testManager(t) - before := manager.Resources() - _, err := manager.Acquire("sbox-test", []*api.XpuAllocation{{ + provider := testNVIDIAProvider(t) + before := provider.Resources() + _, err := provider.Acquire("sbox-test", config.RuntimeNameRunsc, &api.XpuAllocation{ Type: "gpu", DeviceIds: []uint32{0}, - }}) + }) require.NoError(t, err) - assert.Equal(t, before, manager.Resources()) + assert.Equal(t, before, provider.Resources()) } func TestReservedEnv(t *testing.T) { assert.True(t, ReservedEnv("NVIDIA_VISIBLE_DEVICES")) assert.True(t, ReservedEnv("NVIDIA_DRIVER_CAPABILITIES")) assert.True(t, ReservedEnv("CUDA_VISIBLE_DEVICES")) + assert.True(t, ReservedEnv("ASCEND_VISIBLE_DEVICES")) + assert.True(t, ReservedEnv("ASCEND_RT_VISIBLE_DEVICES")) assert.False(t, ReservedEnv("CUDA_VERSION")) } +func TestManagerRejectsMixedOrUnsupportedRuntime(t *testing.T) { + manager := &Manager{providers: map[string]Provider{TypeGPU: testNVIDIAProvider(t)}} + require.NoError(t, manager.ValidateRuntime(config.RuntimeNameRunc, []*api.XpuAllocation{{ + Type: TypeGPU, DeviceIds: []uint32{0}, + }})) + require.ErrorContains(t, manager.ValidateRuntime(config.RuntimeNameKata, []*api.XpuAllocation{{ + Type: TypeGPU, DeviceIds: []uint32{0}, + }}), "does not support runtime") + require.ErrorContains(t, manager.ValidateRuntime(config.RuntimeNameRunsc, []*api.XpuAllocation{ + {Type: TypeGPU, DeviceIds: []uint32{0}}, + {Type: TypeNPU, DeviceIds: []uint32{1}}, + }), "exactly one") + require.ErrorContains(t, manager.ValidateRuntime(config.RuntimeNameRunc, []*api.XpuAllocation{{ + Type: TypeNPU, DeviceIds: []uint32{1}, + }}), "unsupported XPU type") +} + func TestReservedAnnotation(t *testing.T) { assert.True(t, ReservedAnnotation(AllocationAnnotation)) assert.False(t, ReservedAnnotation("sandbox.akernel.dev/env-id")) } func TestRestoreLeases(t *testing.T) { - manager := testManager(t) - manager.sandboxRoot = t.TempDir() - writeLeaseSpec(t, manager.sandboxRoot, leaseRecord{ + provider := testNVIDIAProvider(t) + provider.sandboxRoot = t.TempDir() + writeLeaseSpec(t, provider.sandboxRoot, leaseRecord{ SandboxID: "sbox-recovered", Type: TypeGPU, DeviceIDs: []uint32{0, 2}, DeviceUUID: []string{"GPU-uuid-0", "GPU-uuid-2"}, }) - require.NoError(t, manager.restoreLeases()) - assert.Equal(t, "sbox-recovered", manager.leases["GPU-uuid-0"]) - assert.Equal(t, "sbox-recovered", manager.leases["GPU-uuid-2"]) + require.NoError(t, provider.restoreLeases()) + assert.Equal(t, "sbox-recovered", provider.leases["GPU-uuid-0"]) + assert.Equal(t, "sbox-recovered", provider.leases["GPU-uuid-2"]) +} + +func TestRestoreRuncLease(t *testing.T) { + provider := testNVIDIAProvider(t) + provider.sandboxRoot = t.TempDir() + writeLeaseSpec(t, provider.sandboxRoot, leaseRecord{ + SchemaVersion: leaseSchemaVersion, + SandboxID: "sbox-runc-recovered", + Type: TypeGPU, + Runtime: config.RuntimeNameRunc, + ProductModel: "l20", + SchedulerIDs: []uint32{0}, + StableIDs: []string{"GPU-uuid-0"}, + Provider: "nvidia", + }) + + require.NoError(t, provider.restoreLeases()) + assert.Equal(t, "sbox-runc-recovered", provider.leases["GPU-uuid-0"]) +} + +func TestRestoreLeaseRejectsUnsupportedRuntime(t *testing.T) { + provider := testNVIDIAProvider(t) + provider.sandboxRoot = t.TempDir() + writeLeaseSpec(t, provider.sandboxRoot, leaseRecord{ + SchemaVersion: leaseSchemaVersion, + SandboxID: "sbox-kata", + Type: TypeGPU, + Runtime: config.RuntimeNameKata, + ProductModel: "l20", + SchedulerIDs: []uint32{0}, + StableIDs: []string{"GPU-uuid-0"}, + Provider: "nvidia", + }) + + require.ErrorContains(t, provider.restoreLeases(), "invalid GPU runtime") } func TestRestoreLeasesFailsClosedOnDuplicateUUID(t *testing.T) { - manager := testManager(t) - manager.sandboxRoot = t.TempDir() - writeLeaseSpec(t, manager.sandboxRoot, leaseRecord{ + provider := testNVIDIAProvider(t) + provider.sandboxRoot = t.TempDir() + writeLeaseSpec(t, provider.sandboxRoot, leaseRecord{ SandboxID: "sbox-first", Type: TypeGPU, DeviceIDs: []uint32{0}, DeviceUUID: []string{"GPU-uuid-0"}, }) - writeLeaseSpec(t, manager.sandboxRoot, leaseRecord{ + writeLeaseSpec(t, provider.sandboxRoot, leaseRecord{ SandboxID: "sbox-second", Type: TypeGPU, DeviceIDs: []uint32{0}, DeviceUUID: []string{"GPU-uuid-0"}, }) - require.ErrorContains(t, manager.restoreLeases(), "assigned to both") + require.ErrorContains(t, provider.restoreLeases(), "assigned to both") } func writeLeaseSpec(t *testing.T, sandboxRoot string, record leaseRecord) { diff --git a/pkg/xpumanager/nvidia.go b/pkg/xpumanager/nvidia.go index 08be9ee..21e4e65 100644 --- a/pkg/xpumanager/nvidia.go +++ b/pkg/xpumanager/nvidia.go @@ -22,12 +22,16 @@ import ( "os" "os/exec" "regexp" + "sort" "strconv" "strings" + "sync" "time" api "github.com/inclusionAI/sandboxd/api/runtime/v1" + "github.com/inclusionAI/sandboxd/config" svc "github.com/inclusionAI/sandboxd/pkg/runtime" + "github.com/sirupsen/logrus" ) const ( @@ -48,6 +52,54 @@ var nvidiaModelSeparator = regexp.MustCompile(`[^a-z0-9._-]+`) type commandRunner func(context.Context, string, ...string) ([]byte, error) type statFunc func(string) (os.FileInfo, error) +// Device contains NVIDIA-private identity for one scheduler-visible ID. +type Device struct { + ID uint32 + UUID string + ProductModel string +} + +type nvidiaProvider struct { + mu sync.RWMutex + + runscBinary string + runcEnabled bool + sandboxRoot string + run commandRunner + stat statFunc + runscReady bool + + devices map[uint32]Device + resources []Resource + leases map[string]string + healthy bool + reason error +} + +func newNVIDIAProvider(runscBinary string, runcEnabled bool, sandboxRoot string) *nvidiaProvider { + provider := &nvidiaProvider{ + runscBinary: runscBinary, + runcEnabled: runcEnabled, + sandboxRoot: sandboxRoot, + run: runCommand, + stat: os.Stat, + devices: make(map[uint32]Device), + leases: make(map[string]string), + } + if err := provider.discoverNVIDIA(); err != nil { + provider.reason = err + return provider + } + if err := provider.restoreLeases(); err != nil { + provider.reason = err + provider.resources = nil + return provider + } + provider.healthy = true + logrus.Infof("xpumanager: discovered %d schedulable NVIDIA GPU(s)", len(provider.devices)) + return provider +} + func runCommand(ctx context.Context, binary string, args ...string) ([]byte, error) { command := exec.CommandContext(ctx, binary, args...) output, err := command.CombinedOutput() @@ -63,14 +115,17 @@ func runCommand(ctx context.Context, binary string, args ...string) ([]byte, err return output, nil } -func (m *Manager) discoverNVIDIA() error { - if m.runscBinary == "" { - return errors.New("runsc runtime is not configured") +func (m *nvidiaProvider) discoverNVIDIA() error { + if m.runscBinary == "" && !m.runcEnabled { + return errors.New("neither runsc nor runc runtime is configured") } cliPath, err := exec.LookPath(nvidiaContainerCLI) if err != nil { return fmt.Errorf("locate %s: %w", nvidiaContainerCLI, err) } + if err := validateNVIDIARuntimeHook(nvidiaRuntimeHookPath, m.stat); err != nil { + return err + } ctx, cancel := context.WithTimeout(context.Background(), nvidiaDiscoveryTimeout) defer cancel() @@ -82,30 +137,54 @@ func (m *Manager) discoverNVIDIA() error { if err != nil { return err } - supportedOutput, err := m.run(ctx, m.runscBinary, "nvproxy", "list-supported-drivers") - if err != nil { - return fmt.Errorf("list runsc nvproxy drivers: %w", err) - } - if !nvidiaDriverSupported(driverVersion, string(supportedOutput)) { - return fmt.Errorf( - "NVIDIA driver %s is not supported by %s nvproxy", - driverVersion, - m.runscBinary, - ) - } for _, path := range []string{nvidiaControlDevice, nvidiaUVMDevice} { if _, err := m.stat(path); err != nil { return fmt.Errorf("required NVIDIA device %s is unavailable: %w", path, err) } } + if err := m.configureRuntimeSupport(ctx, driverVersion); err != nil { + return err + } m.devices = devices m.resources = buildResources(devices) return nil } +func (m *nvidiaProvider) configureRuntimeSupport(ctx context.Context, driverVersion string) error { + if m.runscBinary == "" { + return nil + } + supportedOutput, err := m.run(ctx, m.runscBinary, "nvproxy", "list-supported-drivers") + if err == nil && nvidiaDriverSupported(driverVersion, string(supportedOutput)) { + m.runscReady = true + return nil + } + if err != nil { + err = fmt.Errorf("list runsc nvproxy drivers: %w", err) + } else { + err = fmt.Errorf("NVIDIA driver %s is not supported by %s nvproxy", driverVersion, m.runscBinary) + } + if !m.runcEnabled { + return err + } + logrus.Warnf("xpumanager: runsc GPU support unavailable; runc remains enabled: %v", err) + return nil +} + +func validateNVIDIARuntimeHook(path string, stat statFunc) error { + hookInfo, err := stat(path) + if err != nil { + return fmt.Errorf("stat NVIDIA runtime hook %s: %w", path, err) + } + if !hookInfo.Mode().IsRegular() || hookInfo.Mode().Perm()&0111 == 0 { + return fmt.Errorf("NVIDIA runtime hook %s must be an executable regular file", path) + } + return nil +} + // ReservedEnv reports whether key is controlled by the NVIDIA provider. -func ReservedEnv(key string) bool { +func reservedNVIDIAEnv(key string) bool { switch key { case nvidiaVisibleDevicesEnv, nvidiaDriverCapabilities, cudaVisibleDevicesEnv: return true @@ -114,7 +193,7 @@ func ReservedEnv(key string) bool { } } -func nvidiaSpecUpdates(uuids []string, recordJSON []byte) *svc.SpecUpdates { +func nvidiaSpecUpdates(runtimeName string, uuids []string, recordJSON []byte) *svc.SpecUpdates { logicalIDs := make([]string, len(uuids)) for index := range uuids { logicalIDs[index] = strconv.Itoa(index) @@ -132,8 +211,176 @@ func nvidiaSpecUpdates(uuids []string, recordJSON []byte) *svc.SpecUpdates { Annotations: map[string]string{ AllocationAnnotation: string(recordJSON), }, - RequiresHostWritableRootfs: true, + RequiresHostWritableRootfs: runtimeName == config.RuntimeNameRunsc, + } +} + +func (m *nvidiaProvider) Type() string { return TypeGPU } + +func (m *nvidiaProvider) SupportsRuntime(runtimeName string) bool { + return runtimeName == config.RuntimeNameRunsc && m.runscReady || + runtimeName == config.RuntimeNameRunc && m.runcEnabled +} + +func (m *nvidiaProvider) Healthy() (bool, error) { + m.mu.RLock() + defer m.mu.RUnlock() + return m.healthy, m.reason +} + +func (m *nvidiaProvider) Resources() []Resource { + m.mu.RLock() + defer m.mu.RUnlock() + if !m.healthy { + return []Resource{} + } + resources := make([]Resource, len(m.resources)) + for index := range m.resources { + resources[index] = m.resources[index] + resources[index].DeviceIDs = append([]uint32(nil), m.resources[index].DeviceIDs...) + } + return resources +} + +func (m *nvidiaProvider) Acquire( + sandboxID, runtimeName string, + allocation *api.XpuAllocation, +) (*svc.SpecUpdates, error) { + if sandboxID == "" { + return nil, errors.New("sandbox ID is required for XPU allocation") + } + if !m.SupportsRuntime(runtimeName) { + return nil, fmt.Errorf("GPU allocations are unavailable for runtime %q", runtimeName) + } + if allocation == nil || strings.ToLower(strings.TrimSpace(allocation.Type)) != TypeGPU { + return nil, errors.New("invalid GPU allocation") + } + if len(allocation.DeviceIds) == 0 { + return nil, errors.New("XPU device IDs must not be empty") + } + + m.mu.Lock() + defer m.mu.Unlock() + if !m.healthy { + return nil, fmt.Errorf("GPU support is unavailable: %w", m.reason) + } + seen := make(map[uint32]struct{}, len(allocation.DeviceIds)) + devices := make([]Device, 0, len(allocation.DeviceIds)) + for _, id := range allocation.DeviceIds { + if _, duplicate := seen[id]; duplicate { + return nil, fmt.Errorf("duplicate GPU device ID %d", id) + } + seen[id] = struct{}{} + device, ok := m.devices[id] + if !ok { + return nil, fmt.Errorf("GPU device ID %d is not in the node inventory", id) + } + if owner, leased := m.leases[device.UUID]; leased && owner != sandboxID { + return nil, fmt.Errorf("GPU device ID %d is already leased by sandbox %s", id, owner) + } + devices = append(devices, device) + } + model := devices[0].ProductModel + for _, device := range devices[1:] { + if device.ProductModel != model { + return nil, errors.New("all GPU devices in one allocation must have the same product model") + } + } + uuids := make([]string, len(devices)) + for index, device := range devices { + m.leases[device.UUID] = sandboxID + uuids[index] = device.UUID + } + recordJSON, err := encodeLease(leaseRecord{ + SandboxID: sandboxID, + Type: TypeGPU, + Runtime: runtimeName, + ProductModel: model, + SchedulerIDs: append([]uint32(nil), allocation.DeviceIds...), + StableIDs: append([]string(nil), uuids...), + Provider: "nvidia", + ProviderVersion: "nvidia-container-cli", + }) + if err != nil { + for _, uuid := range uuids { + delete(m.leases, uuid) + } + return nil, fmt.Errorf("encode GPU lease: %w", err) + } + return nvidiaSpecUpdates(runtimeName, uuids, recordJSON), nil +} + +func (m *nvidiaProvider) Release(sandboxID string) { + if sandboxID == "" { + return + } + m.mu.Lock() + defer m.mu.Unlock() + for uuid, owner := range m.leases { + if owner == sandboxID { + delete(m.leases, uuid) + } + } +} + +func (m *nvidiaProvider) restoreLeases() error { + leases, err := readPersistedLeases(m.sandboxRoot) + if err != nil { + return err + } + for _, persisted := range leases { + record := persisted.record + if record.Type != TypeGPU { + continue + } + if record.SchemaVersion >= leaseSchemaVersion && + (record.Runtime != config.RuntimeNameRunsc && record.Runtime != config.RuntimeNameRunc) { + return fmt.Errorf("invalid GPU runtime %q in %s", record.Runtime, persisted.bundlePath) + } + ids := record.SchedulerIDs + stableIDs := record.StableIDs + if record.SchemaVersion <= 1 { + ids = record.DeviceIDs + stableIDs = record.DeviceUUID + } + if len(ids) == 0 || len(ids) != len(stableIDs) { + return fmt.Errorf("invalid GPU allocation annotation in %s", persisted.bundlePath) + } + for index, id := range ids { + device, ok := m.devices[id] + if !ok || device.UUID != stableIDs[index] { + return fmt.Errorf("GPU identity changed for device ID %d in %s", id, persisted.bundlePath) + } + if owner, duplicate := m.leases[device.UUID]; duplicate && owner != record.SandboxID { + return fmt.Errorf("GPU UUID %s is assigned to both %s and %s", device.UUID, owner, record.SandboxID) + } + m.leases[device.UUID] = record.SandboxID + } + } + return nil +} + +func buildResources(devices map[uint32]Device) []Resource { + byModel := make(map[string][]uint32) + for id, device := range devices { + byModel[device.ProductModel] = append(byModel[device.ProductModel], id) + } + models := make([]string, 0, len(byModel)) + for model := range byModel { + models = append(models, model) + } + sort.Strings(models) + resources := make([]Resource, 0, len(models)) + for _, model := range models { + ids := byModel[model] + sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] }) + resources = append(resources, Resource{ + Type: TypeGPU, + ProductModel: model, + DeviceIDs: ids, + }) } + return resources } func parseNVIDIAInfo(output string) (string, map[uint32]Device, error) { diff --git a/test/e2e/README.md b/test/e2e/README.md index ea70d93..ebff650 100644 --- a/test/e2e/README.md +++ b/test/e2e/README.md @@ -181,8 +181,8 @@ state cannot leak between runtimes. ## GPU debug image `gpu.Dockerfile` builds a standalone debug image with sandboxd, sbox, the -checksum-verified gVisor runsc release, `nvidia-container-cli` 1.19.1, and the -CUDA vectorAdd sample rootfs. It starts sandboxd in experimental +checksum-verified gVisor runsc release, NVIDIA Container Toolkit 1.19.1, and +the CUDA vectorAdd sample rootfs. It starts sandboxd in experimental cgroup-disabled mode: ```bash diff --git a/test/e2e/gpu.Dockerfile b/test/e2e/gpu.Dockerfile index 0da8a23..55d01ee 100644 --- a/test/e2e/gpu.Dockerfile +++ b/test/e2e/gpu.Dockerfile @@ -65,8 +65,11 @@ RUN set -eux; \ > /etc/apt/sources.list.d/nvidia-container-toolkit.list; \ apt-get update; \ apt-get install -y --no-install-recommends \ + "nvidia-container-toolkit=${LIBNVIDIA_CONTAINER_VERSION}" \ "libnvidia-container1=${LIBNVIDIA_CONTAINER_VERSION}" \ "libnvidia-container-tools=${LIBNVIDIA_CONTAINER_VERSION}"; \ + command -v nvidia-container-cli; \ + test -x /usr/bin/nvidia-container-runtime-hook; \ rm -rf /var/lib/apt/lists/* COPY third_party/runtime-versions.env /tmp/runtime-versions.env