quick return for edge case

Signed-off-by: kerthcet <[email protected]>
NVIDIA · Dec 12, 2023 · 4b24c1b · 4b24c1b
1 parent bf96ab0
commit 4b24c1b
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/gpuallocator/besteffort_policy.go b/gpuallocator/besteffort_policy.go
@@ -46,6 +46,10 @@ func (p *bestEffortPolicy) Allocate(available []*Device, required []*Device, siz
 		return []*Device{}
 	}
 
+	if len(required) > len(available) {
+		return []*Device{}
+	}
+
 	// Find the highest scoring GPU partition with sets of of size 'size'.
 	// Don't consider partitions that don't have at least one set that contains
 	// all of the GPUs 'required' by the allocation.

diff --git a/gpuallocator/besteffort_test.go b/gpuallocator/besteffort_test.go
@@ -110,6 +110,14 @@ func TestBestEffortAllocate(t *testing.T) {
 			4,
 			[]int{},
 		},
+		{
+			"Required too many devices than available",
+			devices,
+			[]int{0, 1, 2, 3, 4, 5},
+			[]int{1, 2, 3, 4, 5, 6},
+			1,
+			[]int{},
+		},
 	}
 
 	RunPolicyAllocTests(t, policy, tests)