feat(gpu): add the shared GPU placement and budget model

The binder re-checks what admission checked, and cozystack-api shows
the remaining budget a later order is checked against; computing these
in three places would let them disagree. The model is pure, so each
caller builds a snapshot from its own cache.

GPU budgets cap the whole subtree of the tenant that declares them,
walked over the namespace ancestry and independent of the tenant quota
pools. A raw GPU key blocks a class until the same tenant also declares
a GPU key, and then caps devices and raw-name pods together.

A node takes binds only when it is schedulable, Ready and carries no
NoSchedule or NoExecute taint the pod does not tolerate; the
node.kubernetes.io taints are left to the first two checks, and a
taint never counts as a cordon.

Assisted-by: LLM
Signed-off-by: Aleksei Sviridkin <f@lex.la>
This commit is contained in:
Aleksei Sviridkin committed 2026-10-20 15:36:10 +03:00
1 parent a11737a2dd
commit 5d90af051b
15 files changed
+1856

No files matched your search

+83
View File
@@ -0,0 +1,83 @@
/*
Copyright 2026 The Cozystack Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package model
import (
"sort"
"strings"
gpuv1alpha1 "github.com/cozystack/cozystack/api/gpu/v1alpha1"
)
// EffectiveMaxCount returns how many GPUs of the class one VM may take,
// and false when nothing beyond one node's GPUs limits it.
func EffectiveMaxCount(spec *gpuv1alpha1.GPUClassSpec) (int32, bool) {
if spec.MaxCount != nil {
return *spec.MaxCount, true
}
if spec.Mode == gpuv1alpha1.GPUModeVGPU {
return 1, true
}
return 0, false
}
// ResourceNames is the effective GPU resource-name set, sorted: the
// classes' resource names, the hoster's list and the names node agents
// report.
func ResourceNames(classes []gpuv1alpha1.GPUClass, configured []string, nodes []gpuv1alpha1.GPUNode) []string {
set := map[string]struct{}{}
add := func(n string) {
if n != "" {
set[n] = struct{}{}
}
}
for i := range classes {
if dp := classes[i].Spec.Source.DevicePlugin; dp != nil {
add(dp.ResourceName)
}
}
for _, n := range configured {
add(n)
}
for i := range nodes {
for _, n := range nodes[i].Status.ResourceNames {
add(n)
}
}
out := make([]string, 0, len(set))
for n := range set {
out = append(out, n)
}
sort.Strings(out)
return out
}
// FormatResourceNames renders names for the settings ConfigMap.
func FormatResourceNames(names []string) string {
return strings.Join(names, ",")
}
// ParseResourceNames reads the settings ConfigMap value.
func ParseResourceNames(s string) []string {
var out []string
for _, n := range strings.Split(s, ",") {
if n = strings.TrimSpace(n); n != "" {
out = append(out, n)
}
}
return out
}
+79
View File
@@ -0,0 +1,79 @@
/*
Copyright 2026 The Cozystack Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package model
import (
"reflect"
"testing"
"k8s.io/utils/ptr"
gpuv1alpha1 "github.com/cozystack/cozystack/api/gpu/v1alpha1"
)
func TestEffectiveMaxCount(t *testing.T) {
tests := []struct {
name string
spec gpuv1alpha1.GPUClassSpec
want int32
wantLimit bool
}{
{"vgpu unset defaults to 1", gpuv1alpha1.GPUClassSpec{Mode: gpuv1alpha1.GPUModeVGPU}, 1, true},
{"vgpu set", gpuv1alpha1.GPUClassSpec{Mode: gpuv1alpha1.GPUModeVGPU, MaxCount: ptr.To[int32](2)}, 2, true},
{"passthrough unset is no limit", gpuv1alpha1.GPUClassSpec{Mode: gpuv1alpha1.GPUModePassthrough}, 0, false},
{"passthrough set", gpuv1alpha1.GPUClassSpec{Mode: gpuv1alpha1.GPUModePassthrough, MaxCount: ptr.To[int32](4)}, 4, true},
}
for _, tc := range tests {
got, limited := EffectiveMaxCount(&tc.spec)
if got != tc.want || limited != tc.wantLimit {
t.Errorf("%s: EffectiveMaxCount = (%d, %v), want (%d, %v)", tc.name, got, limited, tc.want, tc.wantLimit)
}
}
}
func TestResourceNames(t *testing.T) {
classes := []gpuv1alpha1.GPUClass{
{Spec: gpuv1alpha1.GPUClassSpec{Source: gpuv1alpha1.GPUClassSource{DevicePlugin: &gpuv1alpha1.DevicePluginSource{ResourceName: "nvidia.com/NVIDIA_L40S-24C"}}}},
{Spec: gpuv1alpha1.GPUClassSpec{Source: gpuv1alpha1.GPUClassSource{DevicePlugin: &gpuv1alpha1.DevicePluginSource{ResourceName: "nvidia.com/GH100_H100"}}}},
{Spec: gpuv1alpha1.GPUClassSpec{}},
}
nodes := []gpuv1alpha1.GPUNode{
{Status: gpuv1alpha1.GPUNodeStatus{ResourceNames: []string{"nvidia.com/gpu", "nvidia.com/GH100_H100"}}},
}
got := ResourceNames(classes, []string{"nvidia.com/AD102GL_L40S", "", "nvidia.com/GH100_H100"}, nodes)
want := []string{"nvidia.com/AD102GL_L40S", "nvidia.com/GH100_H100", "nvidia.com/NVIDIA_L40S-24C", "nvidia.com/gpu"}
if !reflect.DeepEqual(got, want) {
t.Errorf("ResourceNames = %v, want %v", got, want)
}
}
func TestResourceNamesFormatRoundTrip(t *testing.T) {
names := []string{"nvidia.com/AD102GL_L40S", "nvidia.com/gpu"}
s := FormatResourceNames(names)
if s != "nvidia.com/AD102GL_L40S,nvidia.com/gpu" {
t.Errorf("FormatResourceNames = %q", s)
}
if got := ParseResourceNames(s); !reflect.DeepEqual(got, names) {
t.Errorf("ParseResourceNames = %v, want %v", got, names)
}
if got := ParseResourceNames(""); len(got) != 0 {
t.Errorf("ParseResourceNames(\"\") = %v, want empty", got)
}
if got := ParseResourceNames(" a/b , ,c/d"); !reflect.DeepEqual(got, []string{"a/b", "c/d"}) {
t.Errorf("ParseResourceNames tolerates spaces and empties: got %v", got)
}
}
+78
View File
@@ -0,0 +1,78 @@
/*
Copyright 2026 The Cozystack Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package model
import (
"strings"
"github.com/go-logr/logr"
corev1 "k8s.io/api/core/v1"
)
// systemTaintPrefix marks the taints Kubernetes sets for a cordon or a
// node that is not Ready, which Cordoned and Ready already decide.
const systemTaintPrefix = "node.kubernetes.io/"
// Cordoned reports whether the node is cordoned. Only spec.unschedulable
// counts: a taint, such as the one Fence Agents Remediation sets, is not a
// cordon and does not trigger the class's onCordon policy.
func Cordoned(n *corev1.Node) bool {
return n.Spec.Unschedulable
}
// Ready reports whether the node's Ready condition is True.
func Ready(n *corev1.Node) bool {
for _, c := range n.Status.Conditions {
if c.Type == corev1.NodeReady {
return c.Status == corev1.ConditionTrue
}
}
return false
}
// UntoleratedTaint returns the first NoSchedule or NoExecute taint outside
// node.kubernetes.io/ that tolerations do not tolerate, or nil.
func UntoleratedTaint(n *corev1.Node, tolerations []corev1.Toleration) *corev1.Taint {
for i := range n.Spec.Taints {
taint := &n.Spec.Taints[i]
if taint.Effect != corev1.TaintEffectNoSchedule && taint.Effect != corev1.TaintEffectNoExecute {
continue
}
if strings.HasPrefix(taint.Key, systemTaintPrefix) {
continue
}
tolerated := false
for j := range tolerations {
if tolerations[j].ToleratesTaint(logr.Discard(), taint, true) {
tolerated = true
break
}
}
if !tolerated {
return taint
}
}
return nil
}
// Eligible reports whether new binds, moves, rebinds and gathers may use
// the node, and whether its units count in a class's capacity.
// tolerations are those of the virt-launcher pod, nil when no VMI is
// known.
func Eligible(n *corev1.Node, tolerations []corev1.Toleration) bool {
return !Cordoned(n) && Ready(n) && UntoleratedTaint(n, tolerations) == nil
}
+104
View File
@@ -0,0 +1,104 @@
/*
Copyright 2026 The Cozystack Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package model
import (
"testing"
corev1 "k8s.io/api/core/v1"
)
func readyNode(mutate func(*corev1.Node)) *corev1.Node {
n := &corev1.Node{Status: corev1.NodeStatus{Conditions: []corev1.NodeCondition{{Type: corev1.NodeReady, Status: corev1.ConditionTrue}}}}
n.Name = "gpu-node-3"
if mutate != nil {
mutate(n)
}
return n
}
func withTaint(key string, effect corev1.TaintEffect) func(*corev1.Node) {
return func(n *corev1.Node) {
n.Spec.Taints = append(n.Spec.Taints, corev1.Taint{Key: key, Effect: effect})
}
}
const farTaint = "medik8s.io/fence-agents-remediation"
func TestCordoned(t *testing.T) {
tests := []struct {
name string
node *corev1.Node
want bool
}{
{"unschedulable", readyNode(func(n *corev1.Node) { n.Spec.Unschedulable = true }), true},
{"schedulable", readyNode(nil), false},
{"Fence Agents Remediation taint alone", readyNode(withTaint(farTaint, corev1.TaintEffectNoSchedule)), false},
{"unschedulable taint without the field", readyNode(withTaint(corev1.TaintNodeUnschedulable, corev1.TaintEffectNoSchedule)), false},
}
for _, tc := range tests {
if got := Cordoned(tc.node); got != tc.want {
t.Errorf("%s: Cordoned = %v, want %v", tc.name, got, tc.want)
}
}
}
func TestReady(t *testing.T) {
tests := []struct {
name string
node *corev1.Node
want bool
}{
{"ready", readyNode(nil), true},
{"not ready", readyNode(func(n *corev1.Node) { n.Status.Conditions[0].Status = corev1.ConditionFalse }), false},
{"unknown", readyNode(func(n *corev1.Node) { n.Status.Conditions[0].Status = corev1.ConditionUnknown }), false},
{"no condition", readyNode(func(n *corev1.Node) { n.Status.Conditions = nil }), false},
}
for _, tc := range tests {
if got := Ready(tc.node); got != tc.want {
t.Errorf("%s: Ready = %v, want %v", tc.name, got, tc.want)
}
}
}
func TestEligible(t *testing.T) {
tolerateFAR := []corev1.Toleration{{Key: farTaint, Operator: corev1.TolerationOpExists}}
tests := []struct {
name string
node *corev1.Node
tolerations []corev1.Toleration
want bool
}{
{name: "ready, schedulable, untainted", node: readyNode(nil), want: true},
{name: "cordoned", node: readyNode(func(n *corev1.Node) { n.Spec.Unschedulable = true }), want: false},
{name: "not ready", node: readyNode(func(n *corev1.Node) { n.Status.Conditions[0].Status = corev1.ConditionFalse }), want: false},
{name: "untolerated NoSchedule", node: readyNode(withTaint(farTaint, corev1.TaintEffectNoSchedule)), want: false},
{name: "untolerated NoExecute", node: readyNode(withTaint("example.com/maintenance", corev1.TaintEffectNoExecute)), want: false},
{name: "tolerated NoSchedule", node: readyNode(withTaint(farTaint, corev1.TaintEffectNoSchedule)), tolerations: tolerateFAR, want: true},
{name: "PreferNoSchedule does not count", node: readyNode(withTaint("example.com/soft", corev1.TaintEffectPreferNoSchedule)), want: true},
// The node.kubernetes.io taints mirror cordon and readiness, which
// the field and the condition already decide.
{name: "node.kubernetes.io/unschedulable taint on a schedulable node", node: readyNode(withTaint(corev1.TaintNodeUnschedulable, corev1.TaintEffectNoSchedule)), want: true},
{name: "node.kubernetes.io/not-ready taint on a ready node", node: readyNode(withTaint(corev1.TaintNodeNotReady, corev1.TaintEffectNoExecute)), want: true},
{name: "toleration of another effect", node: readyNode(withTaint(farTaint, corev1.TaintEffectNoExecute)), tolerations: []corev1.Toleration{{Key: farTaint, Operator: corev1.TolerationOpExists, Effect: corev1.TaintEffectNoSchedule}}, want: false},
}
for _, tc := range tests {
if got := Eligible(tc.node, tc.tolerations); got != tc.want {
t.Errorf("%s: Eligible = %v, want %v", tc.name, got, tc.want)
}
}
}
+40
View File
@@ -0,0 +1,40 @@
/*
Copyright 2026 The Cozystack Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
// Package model is the GPU logic shared by the binder, the admission
// webhooks and cozystack-api: the phase table, quota keys, the tenant
// ancestor walk, budgets, entitlement, capacity and node eligibility. It
// holds no client; callers build a Snapshot from their own caches.
package model
import gpuv1alpha1 "github.com/cozystack/cozystack/api/gpu/v1alpha1"
// HoldsGPUs reports whether a device in phase p holds units on a node.
func HoldsGPUs(p gpuv1alpha1.GPUDevicePhase) bool {
return p == gpuv1alpha1.GPUDeviceBound || p == gpuv1alpha1.GPUDeviceAttached
}
// CountsAgainstBudget reports whether a device in phase p counts against
// GPU budgets. A Lost device counts so that its rebind never fails a
// budget check.
func CountsAgainstBudget(p gpuv1alpha1.GPUDevicePhase) bool {
return HoldsGPUs(p) || p == gpuv1alpha1.GPUDeviceLost
}
// InUsageRecords reports whether a device in phase p has a usage record.
func InUsageRecords(p gpuv1alpha1.GPUDevicePhase) bool {
return HoldsGPUs(p)
}
+49
View File
@@ -0,0 +1,49 @@
/*
Copyright 2026 The Cozystack Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package model
import (
"testing"
gpuv1alpha1 "github.com/cozystack/cozystack/api/gpu/v1alpha1"
)
func TestPhaseTable(t *testing.T) {
tests := []struct {
phase gpuv1alpha1.GPUDevicePhase
holds, counts, records bool
}{
{phase: "", holds: false, counts: false, records: false},
{phase: gpuv1alpha1.GPUDevicePending, holds: false, counts: false, records: false},
{phase: gpuv1alpha1.GPUDeviceBound, holds: true, counts: true, records: true},
{phase: gpuv1alpha1.GPUDeviceAttached, holds: true, counts: true, records: true},
{phase: gpuv1alpha1.GPUDeviceLost, holds: false, counts: true, records: false},
}
for _, tc := range tests {
t.Run(string(tc.phase), func(t *testing.T) {
if got := HoldsGPUs(tc.phase); got != tc.holds {
t.Errorf("HoldsGPUs = %v, want %v", got, tc.holds)
}
if got := CountsAgainstBudget(tc.phase); got != tc.counts {
t.Errorf("CountsAgainstBudget = %v, want %v", got, tc.counts)
}
if got := InUsageRecords(tc.phase); got != tc.records {
t.Errorf("InUsageRecords = %v, want %v", got, tc.records)
}
})
}
}
+139
View File
@@ -0,0 +1,139 @@
/*
Copyright 2026 The Cozystack Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package model
import (
"reflect"
"testing"
corev1 "k8s.io/api/core/v1"
"k8s.io/apimachinery/pkg/api/resource"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/apimachinery/pkg/util/sets"
gpuv1alpha1 "github.com/cozystack/cozystack/api/gpu/v1alpha1"
)
func gpuContainer(limits, requests corev1.ResourceList) corev1.Container {
return corev1.Container{Resources: corev1.ResourceRequirements{Limits: limits, Requests: requests}}
}
func TestPodFrom(t *testing.T) {
names := sets.New(h100Name, l40sName)
one := resource.MustParse("1")
two := resource.MustParse("2")
tests := []struct {
name string
pod corev1.Pod
covered map[corev1.ResourceName]int64
want map[corev1.ResourceName]int64
wantOK bool
}{
{
name: "raw-name pod",
pod: corev1.Pod{Spec: corev1.PodSpec{Containers: []corev1.Container{gpuContainer(corev1.ResourceList{h100Name: two}, nil)}}},
want: map[corev1.ResourceName]int64{h100Name: 2},
wantOK: true,
},
{
name: "requests only",
pod: corev1.Pod{Spec: corev1.PodSpec{Containers: []corev1.Container{gpuContainer(nil, corev1.ResourceList{h100Name: one})}}},
want: map[corev1.ResourceName]int64{h100Name: 1},
wantOK: true,
},
{
name: "fully covered by devices",
pod: corev1.Pod{Spec: corev1.PodSpec{Containers: []corev1.Container{gpuContainer(corev1.ResourceList{h100Name: two}, nil)}}},
covered: map[corev1.ResourceName]int64{h100Name: 2},
wantOK: false,
},
{
name: "partly covered",
pod: corev1.Pod{Spec: corev1.PodSpec{Containers: []corev1.Container{gpuContainer(corev1.ResourceList{h100Name: two, l40sName: one}, nil)}}},
covered: map[corev1.ResourceName]int64{h100Name: 1},
want: map[corev1.ResourceName]int64{h100Name: 1, l40sName: 1},
wantOK: true,
},
{
name: "non-GPU resource ignored",
pod: corev1.Pod{Spec: corev1.PodSpec{Containers: []corev1.Container{gpuContainer(corev1.ResourceList{"nvidia.com/hostdev": one}, nil)}}},
wantOK: false,
},
{
name: "init container larger than the sum",
pod: corev1.Pod{Spec: corev1.PodSpec{
InitContainers: []corev1.Container{gpuContainer(corev1.ResourceList{h100Name: two}, nil)},
Containers: []corev1.Container{gpuContainer(corev1.ResourceList{h100Name: one}, nil)},
}},
want: map[corev1.ResourceName]int64{h100Name: 2},
wantOK: true,
},
{
name: "finished pod holds nothing",
pod: corev1.Pod{
Spec: corev1.PodSpec{Containers: []corev1.Container{gpuContainer(corev1.ResourceList{h100Name: two}, nil)}},
Status: corev1.PodStatus{Phase: corev1.PodFailed},
},
wantOK: false,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
tc.pod.Namespace, tc.pod.Spec.NodeName = "tenant-acme", "n1"
got, ok := PodFrom(&tc.pod, names, tc.covered)
if ok != tc.wantOK {
t.Fatalf("ok = %v, want %v", ok, tc.wantOK)
}
if ok && (!reflect.DeepEqual(got.Units, tc.want) || got.Namespace != "tenant-acme" || got.NodeName != "n1") {
t.Errorf("PodFrom = %+v, want units %v", got, tc.want)
}
})
}
}
func TestProjections(t *testing.T) {
class := &gpuv1alpha1.GPUClass{
ObjectMeta: metav1.ObjectMeta{Name: "h100"},
Spec: gpuv1alpha1.GPUClassSpec{
Mode: gpuv1alpha1.GPUModePassthrough,
NodeSelector: map[string]string{"gpu": "h100"},
Source: gpuv1alpha1.GPUClassSource{DevicePlugin: &gpuv1alpha1.DevicePluginSource{ResourceName: string(h100Name)}},
},
}
if got, want := ClassFrom(class), (Class{Name: "h100", Mode: gpuv1alpha1.GPUModePassthrough, ResourceName: h100Name, NodeSelector: map[string]string{"gpu": "h100"}}); !reflect.DeepEqual(got, want) {
t.Errorf("ClassFrom = %+v, want %+v", got, want)
}
device := &gpuv1alpha1.GPUDevice{
ObjectMeta: metav1.ObjectMeta{Name: "train-0", Namespace: "tenant-acme"},
Spec: gpuv1alpha1.GPUDeviceSpec{ClassName: "h100", Count: 2},
Status: gpuv1alpha1.GPUDeviceStatus{Phase: gpuv1alpha1.GPUDeviceBound, NodeName: "n1"},
}
if got, want := DeviceFrom(device), (Device{Namespace: "tenant-acme", Name: "train-0", Class: "h100", Count: 2, Phase: gpuv1alpha1.GPUDeviceBound, NodeName: "n1"}); got != want {
t.Errorf("DeviceFrom = %+v, want %+v", got, want)
}
n := readyNode(func(n *corev1.Node) {
n.Labels = map[string]string{"gpu": "h100"}
n.Status.Allocatable = corev1.ResourceList{h100Name: resource.MustParse("8")}
n.Spec.Unschedulable = true
})
got := NodeFrom(n, nil)
if got.Name != "gpu-node-3" || got.Allocatable[h100Name] != 8 || got.Eligible || got.Labels["gpu"] != "h100" {
t.Errorf("NodeFrom = %+v", got)
}
}
+89
View File
@@ -0,0 +1,89 @@
/*
Copyright 2026 The Cozystack Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package model
import (
"strings"
corev1 "k8s.io/api/core/v1"
"github.com/cozystack/cozystack/internal/gpu"
)
// QuotaKey is a parsed gpu.cozystack.io budget key.
type QuotaKey struct {
// Class is the class the key limits; empty when Total is set.
Class string
Total bool
}
// ParseQuotaKey parses a resourceQuotas key of the form
// gpu.cozystack.io/<class> or gpu.cozystack.io/total.
func ParseQuotaKey(key corev1.ResourceName) (QuotaKey, bool) {
rest, ok := strings.CutPrefix(string(key), gpu.QuotaKeyPrefix)
if !ok || rest == "" || strings.Contains(rest, "/") {
return QuotaKey{}, false
}
if rest == gpu.QuotaTotal {
return QuotaKey{Total: true}, true
}
return QuotaKey{Class: rest}, true
}
// ClassQuotaKey is the budget key of one class.
func ClassQuotaKey(class string) corev1.ResourceName {
return corev1.ResourceName(gpu.QuotaKeyPrefix + class)
}
// IsGPUQuotaKey reports whether key is in the gpu.cozystack.io namespace,
// well-formed or not. Such keys are budgets of GPUDevice, never pod
// quota, so they stay out of the namespace ResourceQuota, the declaration
// gate and the quota pool model.
func IsGPUQuotaKey(key corev1.ResourceName) bool {
return strings.HasPrefix(string(key), gpu.QuotaKeyPrefix)
}
// WithoutGPUQuotaKeys returns a copy of quota without GPU budget keys.
func WithoutGPUQuotaKeys(quota corev1.ResourceList) corev1.ResourceList {
out := corev1.ResourceList{}
for k, v := range quota {
if !IsGPUQuotaKey(k) {
out[k] = v
}
}
return out
}
// RawQuotaLimit returns the limit a tenant's resourceQuotas puts on the
// GPU resource name directly (a raw key), in any of the forms a quota
// accepts for an extended resource. When more than one form is present
// the smallest wins, as each caps the same pods.
func RawQuotaLimit(quota corev1.ResourceList, name corev1.ResourceName) (int64, bool) {
var limit int64
found := false
for _, k := range []corev1.ResourceName{name, "requests." + name, "limits." + name} {
q, ok := quota[k]
if !ok {
continue
}
if v := q.Value(); !found || v < limit {
limit = v
}
found = true
}
return limit, found
}
+116
View File
@@ -0,0 +1,116 @@
/*
Copyright 2026 The Cozystack Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package model
import (
"reflect"
"testing"
corev1 "k8s.io/api/core/v1"
"k8s.io/apimachinery/pkg/api/resource"
)
func TestParseQuotaKey(t *testing.T) {
tests := []struct {
key corev1.ResourceName
wantClass string
wantTotal bool
wantOK bool
}{
{key: "gpu.cozystack.io/l40s-24c", wantClass: "l40s-24c", wantOK: true},
{key: "gpu.cozystack.io/total", wantTotal: true, wantOK: true},
{key: "gpu.cozystack.io/", wantOK: false},
{key: "gpu.cozystack.io/a/b", wantOK: false},
{key: "nvidia.com/GH100_H100", wantOK: false},
{key: "cpu", wantOK: false},
{key: "requests.gpu.cozystack.io/h100", wantOK: false},
}
for _, tc := range tests {
t.Run(string(tc.key), func(t *testing.T) {
k, ok := ParseQuotaKey(tc.key)
if ok != tc.wantOK || k.Class != tc.wantClass || k.Total != tc.wantTotal {
t.Errorf("ParseQuotaKey = (%+v, %v), want ({Class:%s Total:%v}, %v)", k, ok, tc.wantClass, tc.wantTotal, tc.wantOK)
}
})
}
}
func TestClassQuotaKeyRoundTrip(t *testing.T) {
k, ok := ParseQuotaKey(ClassQuotaKey("h100"))
if !ok || k.Class != "h100" || k.Total {
t.Errorf("round trip gave (%+v, %v)", k, ok)
}
}
func TestIsGPUQuotaKey(t *testing.T) {
for key, want := range map[corev1.ResourceName]bool{
"gpu.cozystack.io/h100": true,
"gpu.cozystack.io/total": true,
"gpu.cozystack.io/": true,
"nvidia.com/GH100_H100": false,
"cpu": false,
} {
if got := IsGPUQuotaKey(key); got != want {
t.Errorf("IsGPUQuotaKey(%s) = %v, want %v", key, got, want)
}
}
}
func TestWithoutGPUQuotaKeys(t *testing.T) {
in := corev1.ResourceList{
"cpu": resource.MustParse("4"),
"gpu.cozystack.io/h100": resource.MustParse("2"),
"gpu.cozystack.io/total": resource.MustParse("4"),
"nvidia.com/GH100_H100": resource.MustParse("1"),
}
got := WithoutGPUQuotaKeys(in)
want := corev1.ResourceList{"cpu": resource.MustParse("4"), "nvidia.com/GH100_H100": resource.MustParse("1")}
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, want %v", got, want)
}
if len(in) != 4 {
t.Error("input was modified")
}
if got := WithoutGPUQuotaKeys(corev1.ResourceList{"gpu.cozystack.io/h100": resource.MustParse("1")}); len(got) != 0 {
t.Errorf("GPU-only list left %v", got)
}
}
func TestRawQuotaLimit(t *testing.T) {
const name = corev1.ResourceName("nvidia.com/GH100_H100")
tests := []struct {
name string
quota corev1.ResourceList
want int64
wantOK bool
}{
{name: "bare key", quota: corev1.ResourceList{name: resource.MustParse("2")}, want: 2, wantOK: true},
{name: "limits key", quota: corev1.ResourceList{"limits." + name: resource.MustParse("3")}, want: 3, wantOK: true},
{name: "requests key", quota: corev1.ResourceList{"requests." + name: resource.MustParse("1")}, want: 1, wantOK: true},
{name: "two forms take the smaller", quota: corev1.ResourceList{name: resource.MustParse("2"), "requests." + name: resource.MustParse("1")}, want: 1, wantOK: true},
{name: "other name", quota: corev1.ResourceList{"nvidia.com/other": resource.MustParse("2")}, wantOK: false},
{name: "none", quota: nil, wantOK: false},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
got, ok := RawQuotaLimit(tc.quota, name)
if got != tc.want || ok != tc.wantOK {
t.Errorf("RawQuotaLimit = (%d, %v), want (%d, %v)", got, ok, tc.want, tc.wantOK)
}
})
}
}
+86
View File
@@ -0,0 +1,86 @@
/*
Copyright 2026 The Cozystack Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package model
import (
"bytes"
"encoding/json"
"errors"
"fmt"
"strconv"
)
// DeviceRef is one entry of the gpu.AnnotationDevices list on a VMI:
// exactly one of Device and DeviceTemplate is set.
type DeviceRef struct {
// Device names a GPUDevice in the VMI's namespace.
Device string `json:"device,omitempty"`
DeviceTemplate *DeviceTemplate `json:"deviceTemplate,omitempty"`
}
// DeviceTemplate asks for a GPUDevice created at VMI start and deleted
// with the VMI.
type DeviceTemplate struct {
ClassName string `json:"className"`
Count int32 `json:"count"`
}
// ParseDeviceRefs reads and validates the gpu.AnnotationDevices value.
func ParseDeviceRefs(s string) ([]DeviceRef, error) {
dec := json.NewDecoder(bytes.NewReader([]byte(s)))
dec.DisallowUnknownFields()
var refs []DeviceRef
if err := dec.Decode(&refs); err != nil {
return nil, fmt.Errorf("parse device references: %w", err)
}
if len(refs) == 0 {
return nil, errors.New("device references: empty list")
}
seen := map[string]bool{}
for i, r := range refs {
switch {
case (r.Device == "") == (r.DeviceTemplate == nil):
return nil, fmt.Errorf("device reference %d: set exactly one of device and deviceTemplate", i)
case r.DeviceTemplate != nil && r.DeviceTemplate.ClassName == "":
return nil, fmt.Errorf("device reference %d: deviceTemplate needs className", i)
case r.DeviceTemplate != nil && r.DeviceTemplate.Count < 1:
return nil, fmt.Errorf("device reference %d: deviceTemplate count must be at least 1", i)
case seen[r.Device]:
return nil, fmt.Errorf("device reference %d: device %s is referenced twice", i, r.Device)
}
if r.Device != "" {
seen[r.Device] = true
}
}
return refs, nil
}
// FormatDeviceRefs renders refs as the gpu.AnnotationDevices value.
func FormatDeviceRefs(refs []DeviceRef) (string, error) {
b, err := json.Marshal(refs)
if err != nil {
return "", fmt.Errorf("format device references: %w", err)
}
return string(b), nil
}
// TemplateDeviceName names the GPUDevice created for the deviceTemplate
// entry at index of a VMI's references, so a repeated admission or a
// recreated VMI finds the same device.
func TemplateDeviceName(vmi string, index int) string {
return vmi + "-gpu-" + strconv.Itoa(index)
}
+74
View File
@@ -0,0 +1,74 @@
/*
Copyright 2026 The Cozystack Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package model
import (
"reflect"
"testing"
)
func TestParseDeviceRefs(t *testing.T) {
tests := []struct {
name string
in string
want []DeviceRef
wantErr bool
}{
{
name: "device and template",
in: `[{"device":"train-0"},{"deviceTemplate":{"className":"l40s-24c","count":2}}]`,
want: []DeviceRef{{Device: "train-0"}, {DeviceTemplate: &DeviceTemplate{ClassName: "l40s-24c", Count: 2}}},
},
{name: "empty list", in: `[]`, wantErr: true},
{name: "not JSON", in: `train-0`, wantErr: true},
{name: "both set", in: `[{"device":"a","deviceTemplate":{"className":"h100","count":1}}]`, wantErr: true},
{name: "neither set", in: `[{}]`, wantErr: true},
{name: "template count 0", in: `[{"deviceTemplate":{"className":"h100","count":0}}]`, wantErr: true},
{name: "template without class", in: `[{"deviceTemplate":{"count":1}}]`, wantErr: true},
{name: "unknown field", in: `[{"name":"nvidia.com/gpu"}]`, wantErr: true},
{name: "same device twice", in: `[{"device":"a"},{"device":"a"}]`, wantErr: true},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
got, err := ParseDeviceRefs(tc.in)
if (err != nil) != tc.wantErr {
t.Fatalf("err = %v, wantErr %v", err, tc.wantErr)
}
if !tc.wantErr && !reflect.DeepEqual(got, tc.want) {
t.Errorf("got %+v, want %+v", got, tc.want)
}
})
}
}
func TestFormatDeviceRefsRoundTrip(t *testing.T) {
refs := []DeviceRef{{Device: "train-0"}, {DeviceTemplate: &DeviceTemplate{ClassName: "h100", Count: 1}}}
s, err := FormatDeviceRefs(refs)
if err != nil {
t.Fatal(err)
}
got, err := ParseDeviceRefs(s)
if err != nil || !reflect.DeepEqual(got, refs) {
t.Errorf("round trip: %+v, %v", got, err)
}
}
func TestTemplateDeviceName(t *testing.T) {
if got := TemplateDeviceName("vm-instance-trainer", 1); got != "vm-instance-trainer-gpu-1" {
t.Errorf("TemplateDeviceName = %q", got)
}
}
+374
View File
@@ -0,0 +1,374 @@
/*
Copyright 2026 The Cozystack Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package model
import (
"fmt"
"sort"
corev1 "k8s.io/api/core/v1"
"k8s.io/apimachinery/pkg/labels"
"k8s.io/apimachinery/pkg/util/sets"
gpuv1alpha1 "github.com/cozystack/cozystack/api/gpu/v1alpha1"
"github.com/cozystack/cozystack/internal/gpu"
)
// Class is the part of a GPUClass that placement and budgets read.
type Class struct {
Name string
Mode gpuv1alpha1.GPUMode
ResourceName corev1.ResourceName
NodeSelector map[string]string
}
// ClassFrom projects a GPUClass.
func ClassFrom(c *gpuv1alpha1.GPUClass) Class {
out := Class{Name: c.Name, Mode: c.Spec.Mode, NodeSelector: c.Spec.NodeSelector}
if dp := c.Spec.Source.DevicePlugin; dp != nil {
out.ResourceName = corev1.ResourceName(dp.ResourceName)
}
return out
}
// Device is the part of a GPUDevice that placement and budgets read.
type Device struct {
Namespace string
Name string
Class string
Count int64
Phase gpuv1alpha1.GPUDevicePhase
NodeName string
}
// DeviceFrom projects a GPUDevice.
func DeviceFrom(d *gpuv1alpha1.GPUDevice) Device {
return Device{
Namespace: d.Namespace,
Name: d.Name,
Class: d.Spec.ClassName,
Count: int64(d.Spec.Count),
Phase: d.Status.Phase,
NodeName: d.Status.NodeName,
}
}
// Node is the part of a Node that placement reads.
type Node struct {
Name string
Labels map[string]string
Allocatable map[corev1.ResourceName]int64
// Eligible is Eligible() of the node, evaluated by the caller with the
// tolerations that apply.
Eligible bool
}
// NodeFrom projects a Node; tolerations are passed to Eligible.
func NodeFrom(n *corev1.Node, tolerations []corev1.Toleration) Node {
alloc := make(map[corev1.ResourceName]int64, len(n.Status.Allocatable))
for k, v := range n.Status.Allocatable {
alloc[k] = v.Value()
}
return Node{Name: n.Name, Labels: n.Labels, Allocatable: alloc, Eligible: Eligible(n, tolerations)}
}
// Pod is the GPU usage of a pod beyond what its attached devices cover:
// the units of VMs on raw resource names.
type Pod struct {
Namespace string
// NodeName is empty for a pod not yet scheduled; such a pod counts
// against raw quota keys but not against any node's capacity.
NodeName string
Units map[corev1.ResourceName]int64
}
// PodFrom returns the GPU units p requests under gpuNames beyond covered,
// the units its attached devices account for, and false when nothing is
// left or p is finished.
func PodFrom(p *corev1.Pod, gpuNames sets.Set[corev1.ResourceName], covered map[corev1.ResourceName]int64) (Pod, bool) {
if p.Status.Phase == corev1.PodSucceeded || p.Status.Phase == corev1.PodFailed {
return Pod{}, false
}
units := map[corev1.ResourceName]int64{}
for name := range gpuNames {
var sum, initMax int64
for i := range p.Spec.Containers {
sum += containerUnits(&p.Spec.Containers[i], name)
}
for i := range p.Spec.InitContainers {
initMax = max(initMax, containerUnits(&p.Spec.InitContainers[i], name))
}
if left := max(sum, initMax) - covered[name]; left > 0 {
units[name] = left
}
}
if len(units) == 0 {
return Pod{}, false
}
return Pod{Namespace: p.Namespace, NodeName: p.Spec.NodeName, Units: units}, true
}
// containerUnits reads the limit first: an extended resource may be given
// as a limit alone, and when both are set they are equal.
func containerUnits(c *corev1.Container, name corev1.ResourceName) int64 {
if q, ok := c.Resources.Limits[name]; ok {
return q.Value()
}
if q, ok := c.Resources.Requests[name]; ok {
return q.Value()
}
return 0
}
// Snapshot is one consistent view of what budgets and capacity depend on.
type Snapshot struct {
Classes map[string]Class
Devices []Device
Nodes []Node
Pods []Pod
// Quotas maps a tenant namespace to its declared resourceQuotas, all
// keys as written.
Quotas map[string]corev1.ResourceList
}
// RefusalReason says which limit refused an order.
type RefusalReason string
const (
ReasonClassNotFound RefusalReason = "ClassNotFound"
ReasonNotEntitled RefusalReason = "NotEntitled"
// ReasonBudget is a GPU key, or a raw key, that the order would exceed.
ReasonBudget RefusalReason = "Budget"
// ReasonRawKey is a raw key on a tenant that declares neither the
// class key nor total; it blocks the class until one is declared.
ReasonRawKey RefusalReason = "RawKey"
)
// Refusal is an order a budget or the entitlement refuses.
type Refusal struct {
Reason RefusalReason
Class string
Count int64
// Namespace is the tenant whose key refused the order.
Namespace string
Key corev1.ResourceName
Limit int64
Used int64
}
func (r *Refusal) Error() string {
switch r.Reason {
case ReasonClassNotFound:
return fmt.Sprintf("GPUClass %s does not exist", r.Class)
case ReasonNotEntitled:
return fmt.Sprintf("cannot order GPUClass %s: the tenant is not entitled to vGPU", r.Class)
case ReasonRawKey:
return fmt.Sprintf("cannot order %d of GPUClass %s: tenant %s limits %s directly and needs %s or %s first",
r.Count, r.Class, r.Namespace, r.Key, ClassQuotaKey(r.Class), gpu.QuotaKeyTotal)
default:
return fmt.Sprintf("cannot order %d of GPUClass %s: budget %s of tenant %s is %d and %d is in use",
r.Count, r.Class, r.Key, r.Namespace, r.Limit, r.Used)
}
}
type limit struct {
namespace string
key corev1.ResourceName
limit int64
used int64
blocked bool
}
// limits lists every key on ns or an ancestor that applies to orders of c,
// nearest tenant first.
func (s *Snapshot) limits(ns string, c Class) []limit {
var out []limit
for _, a := range Ancestors(ns) {
q := s.Quotas[a]
classQ, hasClass := q[ClassQuotaKey(c.Name)]
totalQ, hasTotal := q[gpu.QuotaKeyTotal]
if hasClass {
out = append(out, limit{namespace: a, key: ClassQuotaKey(c.Name), limit: classQ.Value(),
used: s.devicesUsed(a, func(d Device) bool { return d.Class == c.Name })})
}
if hasTotal {
out = append(out, limit{namespace: a, key: gpu.QuotaKeyTotal, limit: totalQ.Value(),
used: s.devicesUsed(a, func(Device) bool { return true })})
}
if c.ResourceName == "" {
continue
}
raw, ok := RawQuotaLimit(q, c.ResourceName)
if !ok {
continue
}
used := s.devicesUsed(a, func(d Device) bool { return s.Classes[d.Class].ResourceName == c.ResourceName }) +
s.podUnits(a, c.ResourceName)
out = append(out, limit{namespace: a, key: c.ResourceName, limit: raw, used: used, blocked: !hasClass && !hasTotal})
}
return out
}
func (s *Snapshot) devicesUsed(root string, match func(Device) bool) int64 {
var n int64
for _, d := range s.Devices {
if CountsAgainstBudget(d.Phase) && InSubtree(d.Namespace, root) && match(d) {
n += d.Count
}
}
return n
}
func (s *Snapshot) podUnits(root string, name corev1.ResourceName) int64 {
var n int64
for _, p := range s.Pods {
if InSubtree(p.Namespace, root) {
n += p.Units[name]
}
}
return n
}
// CheckOrder returns a *Refusal when an order of count GPUs of class in
// tenant namespace ns would exceed a budget of ns or an ancestor, or needs
// a vGPU entitlement ns lacks. Pending devices in the snapshot, the order
// itself included, do not count.
func (s *Snapshot) CheckOrder(ns, class string, count int64, entitled bool) error {
c, ok := s.Classes[class]
if !ok {
return &Refusal{Reason: ReasonClassNotFound, Class: class, Count: count}
}
if c.Mode == gpuv1alpha1.GPUModeVGPU && !entitled {
return &Refusal{Reason: ReasonNotEntitled, Class: class, Count: count}
}
for _, l := range s.limits(ns, c) {
r := &Refusal{Class: class, Count: count, Namespace: l.namespace, Key: l.key, Limit: l.limit, Used: l.used}
switch {
case l.blocked:
r.Reason = ReasonRawKey
return r
case l.used+count > l.limit:
r.Reason = ReasonBudget
return r
}
}
return nil
}
// RemainingBudget returns how many GPUs of class tenant ns may still
// order, and false when no budget limits it. It is 0 for a vGPU class
// without the entitlement.
func (s *Snapshot) RemainingBudget(ns, class string, entitled bool) (int64, bool) {
c, ok := s.Classes[class]
if !ok || (c.Mode == gpuv1alpha1.GPUModeVGPU && !entitled) {
return 0, true
}
var left int64
limited := false
for _, l := range s.limits(ns, c) {
v := l.limit - l.used
if l.blocked {
v = 0
}
if !limited || v < left {
left = v
}
limited = true
}
return max(left, 0), limited
}
func (s *Snapshot) node(name string) *Node {
for i := range s.Nodes {
if s.Nodes[i].Name == name {
return &s.Nodes[i]
}
}
return nil
}
func matches(c Class, n *Node) bool {
return labels.SelectorFromSet(c.NodeSelector).Matches(labels.Set(n.Labels))
}
// Conflicts returns, sorted, the classes that select the node and share a
// resource name with another class that selects it. Such a node serves
// none of them.
func (s *Snapshot) Conflicts(nodeName string) []string {
n := s.node(nodeName)
if n == nil {
return nil
}
byName := map[corev1.ResourceName][]string{}
for _, c := range s.Classes {
if c.ResourceName != "" && matches(c, n) {
byName[c.ResourceName] = append(byName[c.ResourceName], c.Name)
}
}
var out []string
for _, names := range byName {
if len(names) > 1 {
out = append(out, names...)
}
}
sort.Strings(out)
return out
}
// Serves reports whether the node is eligible, selected by the class and
// free of a resource-name conflict for it.
func (s *Snapshot) Serves(class, nodeName string) bool {
c, ok := s.Classes[class]
n := s.node(nodeName)
if !ok || n == nil || c.ResourceName == "" || !n.Eligible || !matches(c, n) {
return false
}
return !sets.New(s.Conflicts(nodeName)...).Has(class)
}
// FreeOnNode returns the free units of class on the node: allocatable,
// minus units held by devices there, minus units pods hold beyond their
// devices. It is 0 on a node that does not serve the class.
func (s *Snapshot) FreeOnNode(class, nodeName string) int64 {
if !s.Serves(class, nodeName) {
return 0
}
name := s.Classes[class].ResourceName
free := s.node(nodeName).Allocatable[name]
for _, d := range s.Devices {
if d.NodeName == nodeName && HoldsGPUs(d.Phase) && s.Classes[d.Class].ResourceName == name {
free -= d.Count
}
}
for _, p := range s.Pods {
if p.NodeName == nodeName {
free -= p.Units[name]
}
}
return max(free, 0)
}
// ClassCapacity returns the class's free units over the nodes that serve
// it and the most free units on one of them.
func (s *Snapshot) ClassCapacity(class string) (free, maxFreePerNode int64) {
for i := range s.Nodes {
f := s.FreeOnNode(class, s.Nodes[i].Name)
free += f
maxFreePerNode = max(maxFreePerNode, f)
}
return free, maxFreePerNode
}
+385
View File
@@ -0,0 +1,385 @@
/*
Copyright 2026 The Cozystack Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package model
import (
"errors"
"strings"
"testing"
corev1 "k8s.io/api/core/v1"
"k8s.io/apimachinery/pkg/api/resource"
gpuv1alpha1 "github.com/cozystack/cozystack/api/gpu/v1alpha1"
)
const (
h100Name = corev1.ResourceName("nvidia.com/GH100_H100")
l40sName = corev1.ResourceName("nvidia.com/NVIDIA_L40S-24C")
)
func classes() map[string]Class {
return map[string]Class{
"h100": {Name: "h100", Mode: gpuv1alpha1.GPUModePassthrough, ResourceName: h100Name, NodeSelector: map[string]string{"gpu": "h100"}},
"h100-alt": {Name: "h100-alt", Mode: gpuv1alpha1.GPUModePassthrough, ResourceName: h100Name, NodeSelector: map[string]string{"gpu": "h100-alt"}},
"l40s-24c": {Name: "l40s-24c", Mode: gpuv1alpha1.GPUModeVGPU, ResourceName: l40sName, NodeSelector: map[string]string{"gpu": "l40s"}},
}
}
func quota(kv ...any) corev1.ResourceList {
out := corev1.ResourceList{}
for i := 0; i < len(kv); i += 2 {
out[corev1.ResourceName(kv[i].(string))] = resource.MustParse(kv[i+1].(string))
}
return out
}
func dev(ns, class string, count int64, phase gpuv1alpha1.GPUDevicePhase, node string) Device {
return Device{Namespace: ns, Name: ns + "-" + class, Class: class, Count: count, Phase: phase, NodeName: node}
}
func refusalReason(err error) RefusalReason {
var r *Refusal
if errors.As(err, &r) {
return r.Reason
}
return ""
}
func TestCheckOrderBudgets(t *testing.T) {
tests := []struct {
name string
quotas map[string]corev1.ResourceList
devices []Device
pods []Pod
ns, class string
count int64
entitled bool
wantReason RefusalReason
wantKeyNS string
}{
{
name: "no GPU key anywhere is limited only by capacity",
ns: "tenant-acme-dev", class: "h100", count: 8,
},
{
name: "class key on the tenant fits",
quotas: map[string]corev1.ResourceList{"tenant-acme": quota("gpu.cozystack.io/h100", "2")},
ns: "tenant-acme", class: "h100", count: 2,
},
{
name: "class key on the tenant exceeded",
quotas: map[string]corev1.ResourceList{"tenant-acme": quota("gpu.cozystack.io/h100", "2")},
devices: []Device{dev("tenant-acme", "h100", 1, gpuv1alpha1.GPUDeviceBound, "n1")},
ns: "tenant-acme", class: "h100", count: 2,
wantReason: ReasonBudget, wantKeyNS: "tenant-acme",
},
{
name: "ancestor total counts the whole subtree",
quotas: map[string]corev1.ResourceList{"tenant-acme": quota("gpu.cozystack.io/total", "4")},
devices: []Device{dev("tenant-acme-a", "h100", 2, gpuv1alpha1.GPUDeviceAttached, "n1"), dev("tenant-acme-b-c", "l40s-24c", 1, gpuv1alpha1.GPUDeviceLost, "")},
ns: "tenant-acme-d", class: "h100", count: 2,
wantReason: ReasonBudget, wantKeyNS: "tenant-acme",
},
{
name: "pending devices do not count",
quotas: map[string]corev1.ResourceList{"tenant-acme": quota("gpu.cozystack.io/total", "4")},
devices: []Device{dev("tenant-acme-a", "h100", 3, gpuv1alpha1.GPUDevicePending, "")},
ns: "tenant-acme-d", class: "h100", count: 4,
},
{
name: "devices outside the subtree do not count",
quotas: map[string]corev1.ResourceList{"tenant-acme": quota("gpu.cozystack.io/total", "4")},
devices: []Device{dev("tenant-other", "h100", 4, gpuv1alpha1.GPUDeviceBound, "n1"), dev("tenant-acmex", "h100", 4, gpuv1alpha1.GPUDeviceBound, "n1")},
ns: "tenant-acme", class: "h100", count: 4,
},
{
name: "a child's larger key does not lift the ancestor's",
quotas: map[string]corev1.ResourceList{"tenant-acme": quota("gpu.cozystack.io/h100", "1"), "tenant-acme-dev": quota("gpu.cozystack.io/h100", "10")},
ns: "tenant-acme-dev", class: "h100", count: 2,
wantReason: ReasonBudget, wantKeyNS: "tenant-acme",
},
{
name: "a child key adds a limit",
quotas: map[string]corev1.ResourceList{"tenant-acme": quota("gpu.cozystack.io/h100", "10"), "tenant-acme-dev": quota("gpu.cozystack.io/h100", "1")},
ns: "tenant-acme-dev", class: "h100", count: 2,
wantReason: ReasonBudget, wantKeyNS: "tenant-acme-dev",
},
{
name: "sibling keys may overlap",
quotas: map[string]corev1.ResourceList{"tenant-acme": quota("gpu.cozystack.io/h100", "4"), "tenant-acme-a": quota("gpu.cozystack.io/h100", "4"), "tenant-acme-b": quota("gpu.cozystack.io/h100", "4")},
devices: []Device{dev("tenant-acme-b", "h100", 2, gpuv1alpha1.GPUDeviceBound, "n1")},
ns: "tenant-acme-a", class: "h100", count: 2,
},
{
name: "a class key does not limit another class",
quotas: map[string]corev1.ResourceList{"tenant-acme": quota("gpu.cozystack.io/l40s-24c", "1")},
devices: []Device{dev("tenant-acme", "l40s-24c", 1, gpuv1alpha1.GPUDeviceBound, "n1")},
ns: "tenant-acme", class: "h100", count: 3,
},
{
name: "an ancestor GPU key caps a child with its own cpu quota",
quotas: map[string]corev1.ResourceList{"tenant-acme": quota("gpu.cozystack.io/total", "1"), "tenant-acme-dev": quota("cpu", "8")},
devices: []Device{dev("tenant-acme", "h100", 1, gpuv1alpha1.GPUDeviceBound, "n1")},
ns: "tenant-acme-dev", class: "h100", count: 1,
wantReason: ReasonBudget, wantKeyNS: "tenant-acme",
},
{
name: "root keys cap every tenant",
quotas: map[string]corev1.ResourceList{"tenant-root": quota("gpu.cozystack.io/total", "0")},
ns: "tenant-acme", class: "h100", count: 1,
wantReason: ReasonBudget, wantKeyNS: "tenant-root",
},
{
name: "raw key without a GPU key blocks the class",
quotas: map[string]corev1.ResourceList{"tenant-acme": quota(string(h100Name), "4")},
ns: "tenant-acme-dev", class: "h100", count: 1,
wantReason: ReasonRawKey, wantKeyNS: "tenant-acme",
},
{
name: "raw key does not block a class of another resource name",
quotas: map[string]corev1.ResourceList{"tenant-acme": quota(string(h100Name), "4")},
ns: "tenant-acme", class: "l40s-24c", count: 1, entitled: true,
},
{
name: "raw key with total caps devices and raw-name pods together",
quotas: map[string]corev1.ResourceList{"tenant-acme": quota(string(h100Name), "3", "gpu.cozystack.io/total", "10")},
devices: []Device{dev("tenant-acme-a", "h100-alt", 1, gpuv1alpha1.GPUDeviceBound, "n2")},
pods: []Pod{{Namespace: "tenant-acme-b", NodeName: "n1", Units: map[corev1.ResourceName]int64{h100Name: 1}}},
ns: "tenant-acme", class: "h100", count: 2,
wantReason: ReasonBudget, wantKeyNS: "tenant-acme",
},
{
name: "raw key with class key fits",
quotas: map[string]corev1.ResourceList{"tenant-acme": quota(string(h100Name), "3", "gpu.cozystack.io/h100", "3")},
pods: []Pod{{Namespace: "tenant-acme", NodeName: "n1", Units: map[corev1.ResourceName]int64{h100Name: 1}}},
ns: "tenant-acme", class: "h100", count: 2,
},
{
name: "raw key on an ancestor blocks until that ancestor has a GPU key",
quotas: map[string]corev1.ResourceList{"tenant-acme": quota(string(h100Name), "3"), "tenant-acme-dev": quota("gpu.cozystack.io/total", "3")},
ns: "tenant-acme-dev", class: "h100", count: 1,
wantReason: ReasonRawKey, wantKeyNS: "tenant-acme",
},
{
name: "vgpu class without entitlement",
ns: "tenant-acme", class: "l40s-24c", count: 1,
wantReason: ReasonNotEntitled,
},
{
name: "vgpu class with entitlement and no key",
ns: "tenant-acme", class: "l40s-24c", count: 1, entitled: true,
},
{
name: "total caps vgpu",
quotas: map[string]corev1.ResourceList{"tenant-acme": quota("gpu.cozystack.io/total", "1")},
devices: []Device{dev("tenant-acme", "h100", 1, gpuv1alpha1.GPUDeviceBound, "n1")},
ns: "tenant-acme", class: "l40s-24c", count: 1, entitled: true,
wantReason: ReasonBudget, wantKeyNS: "tenant-acme",
},
{
name: "unknown class",
ns: "tenant-acme", class: "a100", count: 1,
wantReason: ReasonClassNotFound,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
s := &Snapshot{Classes: classes(), Quotas: tc.quotas, Devices: tc.devices, Pods: tc.pods}
err := s.CheckOrder(tc.ns, tc.class, tc.count, tc.entitled)
if got := refusalReason(err); got != tc.wantReason {
t.Fatalf("reason = %q, want %q (err %v)", got, tc.wantReason, err)
}
if err == nil {
return
}
var r *Refusal
errors.As(err, &r)
if r.Namespace != tc.wantKeyNS {
t.Errorf("refusing tenant = %q, want %q", r.Namespace, tc.wantKeyNS)
}
if msg := err.Error(); !strings.Contains(msg, tc.class) {
t.Errorf("message %q does not name the class", msg)
}
})
}
}
func TestRefusalMessageNamesClassCountAndLimit(t *testing.T) {
s := &Snapshot{Classes: classes(), Quotas: map[string]corev1.ResourceList{"tenant-acme": quota("gpu.cozystack.io/h100", "2")}}
err := s.CheckOrder("tenant-acme", "h100", 3, false)
if err == nil {
t.Fatal("order above the budget was accepted")
}
for _, want := range []string{"h100", "3", "gpu.cozystack.io/h100", "2", "tenant-acme"} {
if !strings.Contains(err.Error(), want) {
t.Errorf("message %q lacks %q", err.Error(), want)
}
}
}
func TestRemainingBudget(t *testing.T) {
tests := []struct {
name string
quotas map[string]corev1.ResourceList
devices []Device
pods []Pod
ns, class string
entitled bool
want int64
wantLimited bool
}{
{name: "no key is unlimited", ns: "tenant-acme", class: "h100"},
{
name: "smallest over the walk",
quotas: map[string]corev1.ResourceList{"tenant-acme": quota("gpu.cozystack.io/total", "4"), "tenant-acme-dev": quota("gpu.cozystack.io/h100", "3")},
devices: []Device{dev("tenant-acme-other", "l40s-24c", 2, gpuv1alpha1.GPUDeviceBound, "n1")},
ns: "tenant-acme-dev", class: "h100", want: 2, wantLimited: true,
},
{
name: "over budget clamps at zero",
quotas: map[string]corev1.ResourceList{"tenant-acme": quota("gpu.cozystack.io/total", "1")},
devices: []Device{dev("tenant-acme", "h100", 2, gpuv1alpha1.GPUDeviceLost, "")},
ns: "tenant-acme", class: "h100", want: 0, wantLimited: true,
},
{name: "vgpu without entitlement is zero", ns: "tenant-acme", class: "l40s-24c", want: 0, wantLimited: true},
{name: "vgpu with entitlement and no key is unlimited", ns: "tenant-acme", class: "l40s-24c", entitled: true},
{name: "unknown class is zero", ns: "tenant-acme", class: "a100", want: 0, wantLimited: true},
{
name: "raw key alone blocks",
quotas: map[string]corev1.ResourceList{"tenant-acme": quota(string(h100Name), "4")},
ns: "tenant-acme", class: "h100", want: 0, wantLimited: true,
},
{
name: "raw key with total counts raw-name pods",
quotas: map[string]corev1.ResourceList{"tenant-acme": quota(string(h100Name), "4", "gpu.cozystack.io/total", "10")},
pods: []Pod{{Namespace: "tenant-acme-x", Units: map[corev1.ResourceName]int64{h100Name: 3}}},
ns: "tenant-acme", class: "h100", want: 1, wantLimited: true,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
s := &Snapshot{Classes: classes(), Quotas: tc.quotas, Devices: tc.devices, Pods: tc.pods}
got, limited := s.RemainingBudget(tc.ns, tc.class, tc.entitled)
if got != tc.want || limited != tc.wantLimited {
t.Errorf("RemainingBudget = (%d, %v), want (%d, %v)", got, limited, tc.want, tc.wantLimited)
}
})
}
}
// CheckOrder and RemainingBudget answer the same question; an order of
// exactly the remaining budget fits and one more does not.
func TestRemainingBudgetAgreesWithCheckOrder(t *testing.T) {
s := &Snapshot{
Classes: classes(),
Quotas: map[string]corev1.ResourceList{"tenant-acme": quota("gpu.cozystack.io/total", "5", string(h100Name), "4"), "tenant-acme-dev": quota("gpu.cozystack.io/h100", "3")},
Devices: []Device{dev("tenant-acme", "l40s-24c", 1, gpuv1alpha1.GPUDeviceBound, "n1"), dev("tenant-acme-dev", "h100", 1, gpuv1alpha1.GPUDeviceAttached, "n1")},
Pods: []Pod{{Namespace: "tenant-acme-x", Units: map[corev1.ResourceName]int64{h100Name: 1}}},
}
left, limited := s.RemainingBudget("tenant-acme-dev", "h100", false)
if !limited || left != 2 {
t.Fatalf("RemainingBudget = (%d, %v), want (2, true)", left, limited)
}
if err := s.CheckOrder("tenant-acme-dev", "h100", left, false); err != nil {
t.Errorf("order of the remaining budget refused: %v", err)
}
if err := s.CheckOrder("tenant-acme-dev", "h100", left+1, false); err == nil {
t.Error("order above the remaining budget accepted")
}
}
func node(name string, labels map[string]string, eligible bool, alloc map[corev1.ResourceName]int64) Node {
return Node{Name: name, Labels: labels, Eligible: eligible, Allocatable: alloc}
}
func TestCapacity(t *testing.T) {
h100 := map[string]string{"gpu": "h100"}
s := &Snapshot{
Classes: classes(),
Nodes: []Node{
node("n1", h100, true, map[corev1.ResourceName]int64{h100Name: 8}),
node("n2", h100, true, map[corev1.ResourceName]int64{h100Name: 4}),
node("n3", h100, false, map[corev1.ResourceName]int64{h100Name: 8}),
node("n4", map[string]string{"gpu": "l40s"}, true, map[corev1.ResourceName]int64{l40sName: 16, h100Name: 8}),
},
Devices: []Device{
dev("tenant-a", "h100", 2, gpuv1alpha1.GPUDeviceBound, "n1"),
dev("tenant-b", "h100", 1, gpuv1alpha1.GPUDeviceAttached, "n1"),
dev("tenant-c", "h100", 4, gpuv1alpha1.GPUDeviceLost, "n1"),
dev("tenant-d", "h100", 4, gpuv1alpha1.GPUDevicePending, ""),
dev("tenant-e", "h100", 1, gpuv1alpha1.GPUDeviceBound, "n2"),
},
Pods: []Pod{
{Namespace: "tenant-f", NodeName: "n1", Units: map[corev1.ResourceName]int64{h100Name: 1}},
{Namespace: "tenant-g", NodeName: "n2", Units: map[corev1.ResourceName]int64{h100Name: 5}},
},
}
if got := s.FreeOnNode("h100", "n1"); got != 4 {
t.Errorf("FreeOnNode(n1) = %d, want 4 (8 - 3 bound - 1 raw pod)", got)
}
if got := s.FreeOnNode("h100", "n2"); got != 0 {
t.Errorf("FreeOnNode(n2) = %d, want 0 (clamped)", got)
}
free, maxFree := s.ClassCapacity("h100")
if free != 4 || maxFree != 4 {
t.Errorf("ClassCapacity(h100) = (%d, %d), want (4, 4): ineligible n3 and unmatched n4 must not count", free, maxFree)
}
if got := s.FreeOnNode("h100", "n3"); got != 0 {
t.Errorf("FreeOnNode(ineligible n3) = %d, want 0", got)
}
if got := s.FreeOnNode("h100", "n4"); got != 0 {
t.Errorf("FreeOnNode(unmatched n4) = %d, want 0", got)
}
}
func TestResourceNameConflictTakesTheNodeOutOfBothClasses(t *testing.T) {
s := &Snapshot{
Classes: classes(),
Nodes: []Node{
node("both", map[string]string{"gpu": "h100"}, true, map[corev1.ResourceName]int64{h100Name: 8}),
node("alt", map[string]string{"gpu": "h100-alt"}, true, map[corev1.ResourceName]int64{h100Name: 2}),
},
}
if got := s.Conflicts("both"); len(got) != 0 {
t.Fatalf("Conflicts(both) = %v before the overlap", got)
}
s.Classes["h100-alt"] = Class{Name: "h100-alt", Mode: gpuv1alpha1.GPUModePassthrough, ResourceName: h100Name, NodeSelector: map[string]string{}}
got := s.Conflicts("both")
if len(got) != 2 || got[0] != "h100" || got[1] != "h100-alt" {
t.Errorf("Conflicts(both) = %v, want [h100 h100-alt]", got)
}
if s.Serves("h100", "both") || s.Serves("h100-alt", "both") {
t.Error("a conflicted node still serves a class")
}
if free, _ := s.ClassCapacity("h100"); free != 0 {
t.Errorf("h100 free = %d, want 0", free)
}
if free, _ := s.ClassCapacity("h100-alt"); free != 2 {
t.Errorf("h100-alt free = %d, want 2 from the unconflicted node", free)
}
}
func TestEmptyNodeSelectorMatchesEveryNode(t *testing.T) {
s := &Snapshot{
Classes: map[string]Class{"any": {Name: "any", ResourceName: h100Name}},
Nodes: []Node{node("n1", map[string]string{"x": "y"}, true, map[corev1.ResourceName]int64{h100Name: 1})},
}
if !s.Serves("any", "n1") {
t.Error("a class without nodeSelector does not serve n1")
}
}
+71
View File
@@ -0,0 +1,71 @@
/*
Copyright 2026 The Cozystack Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package model
import (
"strings"
"github.com/cozystack/cozystack/internal/gpu"
)
const (
rootTenantNamespace = "tenant-root"
tenantNamespacePrefix = "tenant-"
)
// ParentNamespace returns the namespace of the parent of the tenant that
// owns ns, or "" for the root tenant and non-tenant namespaces. The tenant
// chart encodes the hierarchy in the name (a child x of namespace P owns
// P-x, or tenant-x directly under root), so the parent is ns without its
// last "-" segment.
func ParentNamespace(ns string) string {
if ns == rootTenantNamespace || !strings.HasPrefix(ns, tenantNamespacePrefix) {
return ""
}
segments := strings.Split(ns, "-")
if len(segments) == 2 {
return rootTenantNamespace
}
return strings.Join(segments[:len(segments)-1], "-")
}
// Ancestors returns ns followed by each ancestor tenant namespace up to
// the root.
func Ancestors(ns string) []string {
var out []string
for cur := ns; cur != ""; cur = ParentNamespace(cur) {
out = append(out, cur)
}
return out
}
// InSubtree reports whether ns is root or a descendant of it.
func InSubtree(ns, root string) bool {
for cur := ns; cur != ""; cur = ParentNamespace(cur) {
if cur == root {
return true
}
}
return false
}
// Entitled reports whether a tenant namespace's labels grant vGPU. The
// tenant chart writes the effective value, parent included, so the label
// alone decides.
func Entitled(namespaceLabels map[string]string) bool {
return namespaceLabels[gpu.LabelNamespaceVGPU] == gpu.LabelValueTrue
}
+89
View File
@@ -0,0 +1,89 @@
/*
Copyright 2026 The Cozystack Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package model
import (
"reflect"
"testing"
)
func TestParentNamespace(t *testing.T) {
for ns, want := range map[string]string{
"tenant-root": "",
"tenant-acme": "tenant-root",
"tenant-acme-dev": "tenant-acme",
"tenant-acme-dev-alice": "tenant-acme-dev",
"cozy-system": "",
"default": "",
} {
if got := ParentNamespace(ns); got != want {
t.Errorf("ParentNamespace(%s) = %q, want %q", ns, got, want)
}
}
}
func TestAncestors(t *testing.T) {
tests := map[string][]string{
"tenant-acme-dev": {"tenant-acme-dev", "tenant-acme", "tenant-root"},
"tenant-root": {"tenant-root"},
"cozy-system": {"cozy-system"},
}
for ns, want := range tests {
if got := Ancestors(ns); !reflect.DeepEqual(got, want) {
t.Errorf("Ancestors(%s) = %v, want %v", ns, got, want)
}
}
}
func TestInSubtree(t *testing.T) {
tests := []struct {
ns, root string
want bool
}{
{"tenant-acme-dev", "tenant-acme", true},
{"tenant-acme", "tenant-acme", true},
{"tenant-acme", "tenant-root", true},
{"tenant-acme-dev", "tenant-root", true},
{"tenant-acmex", "tenant-acme", false},
{"tenant-root", "tenant-acme", false},
{"tenant-other", "tenant-acme", false},
}
for _, tc := range tests {
if got := InSubtree(tc.ns, tc.root); got != tc.want {
t.Errorf("InSubtree(%s, %s) = %v, want %v", tc.ns, tc.root, got, tc.want)
}
}
}
func TestEntitled(t *testing.T) {
tests := []struct {
name string
labels map[string]string
want bool
}{
{"label true", map[string]string{"namespace.cozystack.io/vgpu": "true"}, true},
{"label false", map[string]string{"namespace.cozystack.io/vgpu": "false"}, false},
{"label True is not true", map[string]string{"namespace.cozystack.io/vgpu": "True"}, false},
{"no label", map[string]string{"other": "true"}, false},
{"nil labels", nil, false},
}
for _, tc := range tests {
if got := Entitled(tc.labels); got != tc.want {
t.Errorf("%s: Entitled = %v, want %v", tc.name, got, tc.want)
}
}
}