fix: sync moved packages with KAI-Scheduler main; add utilities/math

Carries the DRA GPU-count overflow fix (#1874) into utilities/resources,
which requires the new utilities/math SaturatingAdd helper. Also re-syncs
scheduling/v2alpha2 PodGroup types and podgrouper/constants with main so the
module reflects current upstream. Prevents consumers repointed at the api
module from silently losing these upstream fixes.

Author: SiorMeir

.changes/unreleased/Added-20260723-140000.yaml

@@ -0,0 +1,3 @@
+kind: Added
+body: |-
+  `kai.scheduler/v1` API group (`kai/v1`): `Config` and `SchedulingShard` types plus their component sub-packages, with generated CRD manifests. `usagedb` config structs (`UsageDBConfig`, `UsageParams`, `WindowType`), `podgrouper/constants`, and `utilities/math` (`SaturatingAdd`).

.changes/unreleased/Changed-20260723-140100.yaml

@@ -0,0 +1,3 @@
+kind: Changed
+body: |-
+  Re-synced `utilities/resources` (DRA GPU-count overflow fix), `scheduling/v2alpha2` PodGroup types, and `podgrouper/constants` with KAI-Scheduler main.

CHANGELOG.md

@@ -5,18 +5,6 @@ All notable changes to api will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-## [Unreleased]
-
-### Added
-
-- **`kai.scheduler/v1` API group** (`kai/v1`) — Config and SchedulingShard operator types plus their
-  component sub-packages (`admission`, `binder`, `common`, `node_scale_adjuster`,
-  `numa_placement_exporter`, `pod_group_controller`, `pod_grouper`, `prometheus`, `queue_controller`,
-  `scheduler`), with `config` and `schedulingshard` CRD manifests.
-- **`usagedb`** — usage-database config structs (`UsageDBConfig`, `UsageParams`, `WindowType`) consumed by
-  SchedulingShard and external API users.
-- **`podgrouper/constants`** — PodGroup labeling/annotation constants for external integrators.
-
 ## [v0.1.0]
 
 Initial release of the standalone `github.com/kai-scheduler/api` module, seeded from KAI-Scheduler `main`.

config/crd/scheduling.run.ai_podgroups.yaml

@@ -101,6 +101,12 @@ spec:
                   -1 and 1
                 format: int32
                 type: integer
+              stalenessGracePeriod:
+                description: |-
+                  StalenessGracePeriod is the minimum duration a stale PodGroup it allowed to remain in stale
+                  status before stale workloads may be evicted to make room. Negative values disable stale gang
+                  eviction for this PodGroup. Defaults to the scheduler's global staleness grace period.
+                type: string
               subGroups:
                 description: SubGroups defines finer-grained subsets of pods within
                   the PodGroup with individual scheduling constraints

podgrouper/constants/constants.go

@@ -10,7 +10,8 @@ const (
 	PreemptibilityLabelKey = "kai.scheduler/preemptibility"
 	UserLabelKey           = "user"
 
-	PreemptionDelayAnnotationKey = "kai.scheduler/preemption-delay"
+	PreemptionDelayAnnotationKey      = "kai.scheduler/preemption-delay"
+	StalenessGracePeriodAnnotationKey = "kai.scheduler/staleness-grace-period"
 
 	BuildPriorityClass     = "build"
 	TrainPriorityClass     = "train"

scheduling/v2alpha2/podgroup_types.go

@@ -82,6 +82,12 @@ type PodGroupSpec struct {
 	// allocation into free capacity, nor the PodGroup's own evictability.
 	// +optional
 	PreemptionDelay *metav1.Duration `json:"preemptionDelay,omitempty" protobuf:"bytes,9,opt,name=preemptionDelay"`
+
+	// StalenessGracePeriod is the minimum duration a stale PodGroup it allowed to remain in stale
+	// status before stale workloads may be evicted to make room. Negative values disable stale gang
+	// eviction for this PodGroup. Defaults to the scheduler's global staleness grace period.
+	// +optional
+	StalenessGracePeriod *metav1.Duration `json:"stalenessGracePeriod,omitempty" protobuf:"bytes,10,opt,name=stalenessGracePeriod"`
 }
 
 // Preemptibility defines whether this PodGroup can be preempted
@@ -128,6 +134,16 @@ func ParsePreemptionDelay(value string) (*metav1.Duration, error) {
 	return &metav1.Duration{Duration: delay}, nil
 }
 
+// ParseStalenessGracePeriod parses a staleness grace period string (e.g. "-10s" "30s", "5m").
+// Returns an error for invalid values.
+func ParseStalenessGracePeriod(value string) (*metav1.Duration, error) {
+	stale, err := time.ParseDuration(value)
+	if err != nil {
+		return nil, err
+	}
+	return &metav1.Duration{Duration: stale}, nil
+}
+
 type SubGroup struct {
 	// Name uniquely identifies the SubGroup within the PodGroup.
 	// Must be a valid DNS label (RFC 1123).

scheduling/v2alpha2/zz_generated.deepcopy.go

@@ -0,0 +1,23 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+// Package math holds dependency-light arithmetic helpers that the scheduler's
+// resource model can use without pulling in Kubernetes or DRA client packages.
+package math
+
+import stdmath "math"
+
+// SaturatingAdd returns a + b, clamped to [math.MinInt64, math.MaxInt64]
+// on overflow instead of wrapping around to the opposite sign.
+func SaturatingAdd(a, b int64) int64 {
+	sum := a + b
+	// Overflow can only happen when both operands share a sign and the sign of
+	// the result flips.
+	if a > 0 && b > 0 && sum < 0 {
+		return stdmath.MaxInt64
+	}
+	if a < 0 && b < 0 && sum >= 0 {
+		return stdmath.MinInt64
+	}
+	return sum
+}

utilities/math/saturating.go

@@ -0,0 +1,30 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package math
+
+import (
+	stdmath "math"
+	"testing"
+)
+
+func TestSaturatingAdd(t *testing.T) {
+	cases := []struct {
+		name string
+		a, b int64
+		want int64
+	}{
+		{"normal sum", 2, 3, 5},
+		{"positive overflow saturates to MaxInt64", stdmath.MaxInt64, stdmath.MaxInt64, stdmath.MaxInt64},
+		{"positive overflow by one", stdmath.MaxInt64, 1, stdmath.MaxInt64},
+		{"negative overflow saturates to MinInt64", stdmath.MinInt64, stdmath.MinInt64, stdmath.MinInt64},
+		{"mixed signs do not overflow", stdmath.MaxInt64, stdmath.MinInt64, -1},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			if got := SaturatingAdd(tc.a, tc.b); got != tc.want {
+				t.Fatalf("SaturatingAdd(%d, %d) = %d, want %d", tc.a, tc.b, got, tc.want)
+			}
+		})
+	}
+}

utilities/math/saturating_test.go

@@ -23,6 +23,8 @@ import (
 	draclient "k8s.io/dynamic-resource-allocation/client"
 	resourceinstall "k8s.io/kubernetes/pkg/apis/resource/install"
 	"sigs.k8s.io/controller-runtime/pkg/client"
+
+	commonmath "github.com/kai-scheduler/api/utilities/math"
 )
 
 func GetResourceClaimName(pod *v1.Pod, podClaim *v1.PodResourceClaim) (string, error) {
@@ -129,7 +131,7 @@ func ExtractDRAGPUResourcesFromClaims(podResourceClaims []*resourceapi.ResourceC
 			// Find the DeviceClassName for this claim
 			deviceClassName := getGPUDeviceClassNameFromClaim(claim)
 			if deviceClassName != "" {
-				deviceClassCounts[deviceClassName] += gpuCount
+				deviceClassCounts[deviceClassName] = commonmath.SaturatingAdd(deviceClassCounts[deviceClassName], gpuCount)
 			}
 		}
 	}
@@ -163,6 +165,10 @@ func getGPUDeviceClassNameFromClaim(claim *resourceapi.ResourceClaim) string {
 
 // countGPUDevicesFromClaim counts GPU devices from a ResourceClaim.
 // Returns the total count of GPU devices requested by this claim.
+//
+// Exactly.Count is user-controlled (the apiserver only enforces > 0), so counts
+// are summed with SaturatingAdd to keep a very large total from wrapping
+// into a negative GPU request in downstream quota accounting.
 func countGPUDevicesFromClaim(claim *resourceapi.ResourceClaim) int64 {
 	totalCount := int64(0)
 
@@ -178,16 +184,16 @@ func countGPUDevicesFromClaim(claim *resourceapi.ResourceClaim) int64 {
 		switch request.Exactly.AllocationMode {
 		case resourceapi.DeviceAllocationModeExactCount:
 			if request.Exactly.Count > 0 {
-				totalCount += request.Exactly.Count
+				totalCount = commonmath.SaturatingAdd(totalCount, request.Exactly.Count)
 			} else {
 				// Default to 1 if Count is not specified for ExactCount mode
-				totalCount += 1
+				totalCount = commonmath.SaturatingAdd(totalCount, 1)
 			}
 		case resourceapi.DeviceAllocationModeAll:
 			// For "All" mode, we can't determine the exact count without allocation info.
 			// For bookkeeping purposes, we'll treat it as requesting 1 device.
 			// This is a conservative estimate for queue resource tracking.
-			totalCount += 1
+			totalCount = commonmath.SaturatingAdd(totalCount, 1)
 		default:
 			// Unknown allocation mode, skip this request
 			continue