mirror of
https://github.com/cozystack/cozystack.git
synced 2026-11-03 10:37:21 +09:00
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:
15 files changed
+1856
No files matched your search
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user