Add local-queue annotation to use created local queue by default in g…

…iven namespace
project-codeflare · Jul 31, 2024 · 9c0fc41 · 9c0fc41
1 parent b41e108
commit 9c0fc41
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 2 deletions.
diff --git a/support/environment.go b/support/environment.go
@@ -54,10 +54,12 @@ const (
 
 	// Storage bucket credentials
 	storageDefaultEndpoint = "AWS_DEFAULT_ENDPOINT"
+	storageDefaultRegion   = "AWS_DEFAULT_REGION"
 	storageAccessKeyId     = "AWS_ACCESS_KEY_ID"
 	storageSecretKey       = "AWS_SECRET_ACCESS_KEY"
 	storageBucketName      = "AWS_STORAGE_BUCKET"
 	storageBucketMnistDir  = "AWS_STORAGE_BUCKET_MNIST_DIR"
+	storage
 )
 
 type ClusterType string
@@ -133,6 +135,11 @@ func GetStorageBucketDefaultEndpoint() (string, bool) {
 	return storage_endpoint, exists
 }
 
+func GetStorageBucketDefaultRegion() (string, bool) {
+	storage_default_region, exists := os.LookupEnv(storageDefaultRegion)
+	return storage_default_region, exists
+}
+
 func GetStorageBucketAccessKeyId() (string, bool) {
 	storage_access_key_id, exists := os.LookupEnv(storageAccessKeyId)
 	return storage_access_key_id, exists

diff --git a/support/kueue.go b/support/kueue.go
@@ -65,7 +65,7 @@ func CreateKueueClusterQueue(t Test, clusterQueueSpec kueuev1beta1.ClusterQueueS
 	return clusterQueue
 }
 
-func CreateKueueLocalQueue(t Test, namespace, clusterQueueName string) *kueuev1beta1.LocalQueue {
+func CreateKueueLocalQueue(t Test, namespace, clusterQueueName string, annotations map[string]string) *kueuev1beta1.LocalQueue {
 	t.T().Helper()
 
 	localQueue := &kueuev1beta1.LocalQueue{
@@ -76,6 +76,7 @@ func CreateKueueLocalQueue(t Test, namespace, clusterQueueName string) *kueuev1b
 		ObjectMeta: metav1.ObjectMeta{
 			GenerateName: "lq-",
 			Namespace:    namespace,
+			Annotations:  annotations,
 		},
 		Spec: kueuev1beta1.LocalQueueSpec{
 			ClusterQueue: kueuev1beta1.ClusterQueueReference(clusterQueueName),

diff --git a/support/kueue_test.go b/support/kueue_test.go
@@ -52,7 +52,10 @@ func TestCreateKueueClusterQueue(t *testing.T) {
 func TestCreateKueueLocalQueue(t *testing.T) {
 	test := NewTest(t)
 
-	lq := CreateKueueLocalQueue(test, "ns-1", "cq-1")
+	annotations := map[string]string{
+		"kueue.x-k8s.io/default-queue": "true",
+	}
+	lq := CreateKueueLocalQueue(test, "ns-1", "cq-1", annotations)
 
 	test.Expect(lq).To(gomega.Not(gomega.BeNil()))
 	test.Expect(lq.GenerateName).To(gomega.Equal("lq-"))