From 941e1e322436a42121ae4ad6193d09d564b435ea Mon Sep 17 00:00:00 2001
From: ChristianZaccaria <christian.zaccaria.cz@gmail.com>
Date: Thu, 21 Nov 2024 12:10:06 +0000
Subject: [PATCH] Add label to RayCluster if owned by AppWrapper

---
 src/codeflare_sdk/ray/cluster/build_ray_cluster.py           | 5 ++++-
 .../test_cluster_yamls/appwrapper/unit-test-all-params.yaml  | 1 +
 tests/test_cluster_yamls/kueue/aw_kueue.yaml                 | 1 +
 tests/test_cluster_yamls/kueue/ray_cluster_kueue.yaml        | 1 +
 tests/test_cluster_yamls/ray/default-appwrapper.yaml         | 1 +
 5 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/codeflare_sdk/ray/cluster/build_ray_cluster.py b/src/codeflare_sdk/ray/cluster/build_ray_cluster.py
index e590d4834..923d263c4 100644
--- a/src/codeflare_sdk/ray/cluster/build_ray_cluster.py
+++ b/src/codeflare_sdk/ray/cluster/build_ray_cluster.py
@@ -182,7 +182,7 @@ def build_ray_cluster(cluster: "codeflare_sdk.ray.cluster.Cluster"):
 # Metadata related functions
 def get_metadata(cluster: "codeflare_sdk.ray.cluster.Cluster"):
     """
-    The get_metadata() function builds and returns a V1ObjectMeta Object using cluster configurtation parameters
+    The get_metadata() function builds and returns a V1ObjectMeta Object using cluster configuration parameters
     """
     object_meta = V1ObjectMeta(
         name=cluster.config.name,
@@ -207,6 +207,9 @@ def get_labels(cluster: "codeflare_sdk.ray.cluster.Cluster"):
     if cluster.config.labels != {}:
         labels.update(cluster.config.labels)
 
+    if cluster.config.appwrapper is True:
+        labels.update({"resource.owner": "appwrapper"})
+
     if cluster.config.appwrapper is False:
         add_queue_label(cluster, labels)
 
diff --git a/tests/test_cluster_yamls/appwrapper/unit-test-all-params.yaml b/tests/test_cluster_yamls/appwrapper/unit-test-all-params.yaml
index 6d2c5440a..93877fddb 100644
--- a/tests/test_cluster_yamls/appwrapper/unit-test-all-params.yaml
+++ b/tests/test_cluster_yamls/appwrapper/unit-test-all-params.yaml
@@ -17,6 +17,7 @@ spec:
           controller-tools.k8s.io: '1.0'
           key1: value1
           key2: value2
+          resource.owner: appwrapper
         name: aw-all-params
         namespace: ns
       spec:
diff --git a/tests/test_cluster_yamls/kueue/aw_kueue.yaml b/tests/test_cluster_yamls/kueue/aw_kueue.yaml
index 402ffb6a6..dd167087d 100644
--- a/tests/test_cluster_yamls/kueue/aw_kueue.yaml
+++ b/tests/test_cluster_yamls/kueue/aw_kueue.yaml
@@ -13,6 +13,7 @@ spec:
       metadata:
         labels:
           controller-tools.k8s.io: '1.0'
+          resource.owner: appwrapper
         name: unit-test-aw-kueue
         namespace: ns
       spec:
diff --git a/tests/test_cluster_yamls/kueue/ray_cluster_kueue.yaml b/tests/test_cluster_yamls/kueue/ray_cluster_kueue.yaml
index a5cb36164..8af1f8d53 100644
--- a/tests/test_cluster_yamls/kueue/ray_cluster_kueue.yaml
+++ b/tests/test_cluster_yamls/kueue/ray_cluster_kueue.yaml
@@ -13,6 +13,7 @@ spec:
       metadata:
         labels:
           controller-tools.k8s.io: '1.0'
+          resource.owner: appwrapper
         name: unit-test-cluster-kueue
         namespace: ns
       spec:
diff --git a/tests/test_cluster_yamls/ray/default-appwrapper.yaml b/tests/test_cluster_yamls/ray/default-appwrapper.yaml
index 3e97474d4..7f4d585d9 100644
--- a/tests/test_cluster_yamls/ray/default-appwrapper.yaml
+++ b/tests/test_cluster_yamls/ray/default-appwrapper.yaml
@@ -11,6 +11,7 @@ spec:
       metadata:
         labels:
           controller-tools.k8s.io: '1.0'
+          resource.owner: appwrapper
         name: default-appwrapper
         namespace: ns
       spec: