diff --git a/go.mod b/go.mod index 195697f90..1da17922b 100644 --- a/go.mod +++ b/go.mod @@ -23,7 +23,7 @@ go 1.22.0 toolchain go1.22.5 require ( - github.com/apache/yunikorn-core v0.0.0-20241017135039-079a02dbdfa7 + github.com/apache/yunikorn-core v0.0.0-20241104184802-6ef347b31c1c github.com/apache/yunikorn-scheduler-interface v0.0.0-20241016105739-f0e241aa0146 github.com/google/go-cmp v0.6.0 github.com/google/uuid v1.6.0 diff --git a/go.sum b/go.sum index 5cbb64964..4674b3ff1 100644 --- a/go.sum +++ b/go.sum @@ -8,8 +8,8 @@ github.com/NYTimes/gziphandler v1.1.1 h1:ZUDjpQae29j0ryrS0u/B8HZfJBtBQHjqw2rQ2cq github.com/NYTimes/gziphandler v1.1.1/go.mod h1:n/CVRwUEOgIxrgPvAQhUUr9oeUtvrhMomdKFjzJNB0c= github.com/antlr4-go/antlr/v4 v4.13.0 h1:lxCg3LAv+EUK6t1i0y1V6/SLeUi0eKEKdhQAlS8TVTI= github.com/antlr4-go/antlr/v4 v4.13.0/go.mod h1:pfChB/xh/Unjila75QW7+VU4TSnWnnk9UTnmpPaOR2g= -github.com/apache/yunikorn-core v0.0.0-20241017135039-079a02dbdfa7 h1:PY3kIiQYxsNcs42DK+8b7NxfTvMF0Z6eIuK+aJNWl18= -github.com/apache/yunikorn-core v0.0.0-20241017135039-079a02dbdfa7/go.mod h1:JA8Uee+D+T9v3p+YznGiGM9cLk5tzX+EM+YYr1TdFYo= +github.com/apache/yunikorn-core v0.0.0-20241104184802-6ef347b31c1c h1:c0+cVnKSAOiJHC6lNUKEl+tt7lZLIEfqv0cPaTI//4U= +github.com/apache/yunikorn-core v0.0.0-20241104184802-6ef347b31c1c/go.mod h1:JA8Uee+D+T9v3p+YznGiGM9cLk5tzX+EM+YYr1TdFYo= github.com/apache/yunikorn-scheduler-interface v0.0.0-20241016105739-f0e241aa0146 h1:CZ4U7y19YSxNJVBNox3DahhuoxDL++naBl/kj+kqVFc= github.com/apache/yunikorn-scheduler-interface v0.0.0-20241016105739-f0e241aa0146/go.mod h1:co3uU98sj1CUTPNTM13lTyi+CY0DOgDndDW2KiUjktU= github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5 h1:0CwZNZbxp69SHPdPJAN/hZIm0C4OItdklCFmMRWYpio= diff --git a/pkg/shim/scheduler_mock_test.go b/pkg/shim/scheduler_mock_test.go index a220ff267..a397c39a7 100644 --- a/pkg/shim/scheduler_mock_test.go +++ b/pkg/shim/scheduler_mock_test.go @@ -336,6 +336,25 @@ func (fc *MockScheduler) waitAndAssertForeignAllocationInCore(partition, allocat }, time.Second, 5*time.Second) } +func (fc *MockScheduler) waitAndAssertForeignAllocationResources(partition, allocationID, nodeID string, expected *si.Resource) error { + return utils.WaitForCondition(func() bool { + node := fc.coreContext.Scheduler.GetClusterContext().GetNode(nodeID, partition) + if node == nil { + log.Log(log.Test).Warn("Node not found", zap.String("node ID", nodeID)) + return false + } + allocs := node.GetForeignAllocations() + for _, alloc := range allocs { + if alloc.GetAllocationKey() == allocationID { + current := alloc.GetAllocatedResource().ToProto() + return common.Equals(expected, current) + } + } + + return false + }, time.Second, 5*time.Second) +} + func (fc *MockScheduler) getApplicationFromCore(appID, partition string) *objects.Application { return fc.coreContext.Scheduler.GetClusterContext().GetApplication(appID, partition) } diff --git a/pkg/shim/scheduler_test.go b/pkg/shim/scheduler_test.go index 835b9572a..1babc11aa 100644 --- a/pkg/shim/scheduler_test.go +++ b/pkg/shim/scheduler_test.go @@ -350,6 +350,21 @@ partitions: err = cluster.waitAndAssertForeignAllocationInCore(partitionName, "foreign-2", "node-1", true) assert.NilError(t, err) + // update pod resources + pod1Copy := pod1.DeepCopy() + pod1Copy.Spec.Containers[0].Resources.Requests[siCommon.Memory] = *resource.NewQuantity(500, resource.DecimalSI) + pod1Copy.Spec.Containers[0].Resources.Requests[siCommon.CPU] = *resource.NewMilliQuantity(2000, resource.DecimalSI) + + cluster.UpdatePod(pod1, pod1Copy) + expectedUsage := common.NewResourceBuilder(). + AddResource(siCommon.Memory, 500). + AddResource(siCommon.CPU, 2). + AddResource("pods", 1). + Build() + err = cluster.waitAndAssertForeignAllocationResources(partitionName, "foreign-1", "node-1", expectedUsage) + assert.NilError(t, err) + + // delete pods cluster.DeletePod(pod1) cluster.DeletePod(pod2)