fix(RHOAIENG-20531): propagate annotations to ray pods

Signed-off-by: Pat O'Connor <[email protected]>

Author: chipspeak

demo-notebooks/guided-demos/0_basic_ray.ipynb

@@ -203,12 +203,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.18"
-  },
-  "vscode": {
-   "interpreter": {
-    "hash": "f9f85f796d01129d0dd105a088854619f454435301f6ffec2fea96ecbd9be4ac"
-   }
+   "version": "3.11.11"
   }
  },
  "nbformat": 4,

src/codeflare_sdk/common/utils/unit_test_support.py

@@ -40,6 +40,11 @@ def create_cluster_config(num_workers=2, write_to_file=False):
         worker_memory_limits=6,
         appwrapper=True,
         write_to_file=write_to_file,
+        annotations={
+            "app.kubernetes.io/managed-by": "test-prefix",
+            "key1": "value1",
+            "key2": "value2",
+        },
     )
     return config
 
@@ -498,7 +503,11 @@ def create_cluster_all_config_params(mocker, cluster_name, is_appwrapper) -> Clu
         extended_resource_mapping={"example.com/gpu": "GPU", "intel.com/gpu": "TPU"},
         overwrite_default_resource_mapping=True,
         local_queue="local-queue-default",
-        annotations={"key1": "value1", "key2": "value2"},
+        annotations={
+            "app.kubernetes.io/managed-by": "test-prefix",
+            "key1": "value1",
+            "key2": "value2",
+        },
         volumes=volumes,
         volume_mounts=volume_mounts,
     )

src/codeflare_sdk/ray/cluster/build_ray_cluster.py

@@ -140,11 +140,12 @@ def build_ray_cluster(cluster: "codeflare_sdk.ray.cluster.Cluster"):
                     "resources": head_resources,
                 },
                 "template": {
+                    "metadata": V1ObjectMeta(cluster.config.annotations),
                     "spec": get_pod_spec(
                         cluster,
                         [get_head_container_spec(cluster)],
                         cluster.config.head_tolerations,
-                    )
+                    ),
                 },
             },
             "workerGroupSpecs": [
@@ -159,11 +160,12 @@ def build_ray_cluster(cluster: "codeflare_sdk.ray.cluster.Cluster"):
                         "resources": worker_resources,
                     },
                     "template": V1PodTemplateSpec(
+                        metadata=V1ObjectMeta(cluster.config.annotations),
                         spec=get_pod_spec(
                             cluster,
                             [get_worker_container_spec(cluster)],
                             cluster.config.worker_tolerations,
-                        )
+                        ),
                     ),
                 }
             ],

src/codeflare_sdk/ray/cluster/test_config.py

@@ -35,7 +35,17 @@ def test_default_cluster_creation(mocker):
     mocker.patch("kubernetes.client.ApisApi.get_api_versions")
     mocker.patch("kubernetes.client.CustomObjectsApi.list_namespaced_custom_object")
 
-    cluster = Cluster(ClusterConfiguration(name="default-cluster", namespace="ns"))
+    cluster = Cluster(
+        ClusterConfiguration(
+            name="default-cluster",
+            namespace="ns",
+            annotations={
+                "app.kubernetes.io/managed-by": "test-prefix",
+                "key1": "value1",
+                "key2": "value2",
+            },
+        )
+    )
 
     expected_rc = apply_template(
         f"{expected_clusters_dir}/ray/default-ray-cluster.yaml",
@@ -50,7 +60,16 @@ def test_default_appwrapper_creation(mocker):
     mocker.patch("kubernetes.client.CustomObjectsApi.list_namespaced_custom_object")
 
     cluster = Cluster(
-        ClusterConfiguration(name="default-appwrapper", namespace="ns", appwrapper=True)
+        ClusterConfiguration(
+            name="default-appwrapper",
+            namespace="ns",
+            appwrapper=True,
+            annotations={
+                "app.kubernetes.io/managed-by": "test-prefix",
+                "key1": "value1",
+                "key2": "value2",
+            },
+        )
     )
 
     expected_aw = apply_template(

tests/test_cluster_yamls/appwrapper/unit-test-all-params.yaml

@@ -42,6 +42,11 @@ spec:
             resources: '"{\"TPU\": 2}"'
           serviceType: ClusterIP
           template:
+            metadata:
+              annotations:
+                app.kubernetes.io/managed-by: test-prefix
+                key1: value1
+                key2: value2
             spec:
               containers:
               - env:
@@ -142,6 +147,11 @@ spec:
             resources: '"{}"'
           replicas: 10
           template:
+            metadata:
+              annotations:
+                app.kubernetes.io/managed-by: test-prefix
+                key1: value1
+                key2: value2
             spec:
               containers:
               - env:

tests/test_cluster_yamls/kueue/aw_kueue.yaml

@@ -11,6 +11,10 @@ spec:
       apiVersion: ray.io/v1
       kind: RayCluster
       metadata:
+        annotations:
+          app.kubernetes.io/managed-by: test-prefix
+          key1: value1
+          key2: value2
         labels:
           controller-tools.k8s.io: '1.0'
         name: unit-test-aw-kueue
@@ -36,6 +40,11 @@ spec:
             resources: '"{}"'
           serviceType: ClusterIP
           template:
+            metadata:
+              annotations:
+                app.kubernetes.io/managed-by: test-prefix
+                key1: value1
+                key2: value2
             spec:
               containers:
               - image: "${image}"
@@ -101,6 +110,11 @@ spec:
             resources: '"{}"'
           replicas: 2
           template:
+            metadata:
+              annotations:
+                app.kubernetes.io/managed-by: test-prefix
+                key1: value1
+                key2: value2
             spec:
               containers:
               - image: "${image}"

tests/test_cluster_yamls/kueue/ray_cluster_kueue.yaml

@@ -11,6 +11,10 @@ spec:
       apiVersion: ray.io/v1
       kind: RayCluster
       metadata:
+        annotations:
+          app.kubernetes.io/managed-by: test-prefix
+          key1: value1
+          key2: value2
         labels:
           controller-tools.k8s.io: '1.0'
         name: unit-test-cluster-kueue
@@ -36,6 +40,11 @@ spec:
             resources: '"{}"'
           serviceType: ClusterIP
           template:
+            metadata:
+              annotations:
+                app.kubernetes.io/managed-by: test-prefix
+                key1: value1
+                key2: value2
             spec:
               containers:
               - image: "${image}"
@@ -101,6 +110,11 @@ spec:
             resources: '"{}"'
           replicas: 2
           template:
+            metadata:
+              annotations:
+                app.kubernetes.io/managed-by: test-prefix
+                key1: value1
+                key2: value2
             spec:
               containers:
               - image: "${image}"

tests/test_cluster_yamls/ray/default-appwrapper.yaml

@@ -9,6 +9,10 @@ spec:
       apiVersion: ray.io/v1
       kind: RayCluster
       metadata:
+        annotations:
+          app.kubernetes.io/managed-by: test-prefix
+          key1: value1
+          key2: value2
         labels:
           controller-tools.k8s.io: '1.0'
         name: default-appwrapper
@@ -34,6 +38,11 @@ spec:
             resources: '"{}"'
           serviceType: ClusterIP
           template:
+            metadata:
+              annotations:
+                app.kubernetes.io/managed-by: test-prefix
+                key1: value1
+                key2: value2
             spec:
               containers:
               - image: "${image}"
@@ -99,6 +108,11 @@ spec:
             resources: '"{}"'
           replicas: 1
           template:
+            metadata:
+              annotations:
+                app.kubernetes.io/managed-by: test-prefix
+                key1: value1
+                key2: value2
             spec:
               containers:
               - image: "${image}"

tests/test_cluster_yamls/ray/default-ray-cluster.yaml

@@ -1,6 +1,10 @@
 apiVersion: ray.io/v1
 kind: RayCluster
 metadata:
+  annotations:
+    app.kubernetes.io/managed-by: test-prefix
+    key1: value1
+    key2: value2
   labels:
     controller-tools.k8s.io: '1.0'
   name: default-cluster
@@ -26,6 +30,11 @@ spec:
       resources: '"{}"'
     serviceType: ClusterIP
     template:
+      metadata:
+        annotations:
+          app.kubernetes.io/managed-by: test-prefix
+          key1: value1
+          key2: value2
       spec:
         containers:
         - image: "${image}"
@@ -91,6 +100,11 @@ spec:
       resources: '"{}"'
     replicas: 1
     template:
+      metadata:
+        annotations:
+          app.kubernetes.io/managed-by: test-prefix
+          key1: value1
+          key2: value2
       spec:
         containers:
         - image: "${image}"

tests/test_cluster_yamls/ray/unit-test-all-params.yaml

@@ -33,6 +33,11 @@ spec:
       resources: '"{\"TPU\": 2}"'
     serviceType: ClusterIP
     template:
+      metadata:
+        annotations:
+          app.kubernetes.io/managed-by: test-prefix
+          key1: value1
+          key2: value2
       spec:
         containers:
         - env:
@@ -133,6 +138,11 @@ spec:
       resources: '"{}"'
     replicas: 10
     template:
+      metadata:
+        annotations:
+          app.kubernetes.io/managed-by: test-prefix
+          key1: value1
+          key2: value2
       spec:
         containers:
         - env: