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

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

Author: chipspeak

src/codeflare_sdk/common/utils/unit_test_support.py

@@ -496,7 +496,10 @@ 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={
+            "key1": "value1",
+            "key2": "value2",
+        },
         volumes=volumes,
         volume_mounts=volume_mounts,
     )

src/codeflare_sdk/ray/cluster/build_ray_cluster.py

@@ -138,13 +138,16 @@ def build_ray_cluster(cluster: "codeflare_sdk.ray.cluster.Cluster"):
                     "num-gpus": str(head_gpu_count),
                     "resources": head_resources,
                 },
-                "template": {
-                    "spec": get_pod_spec(
+                "template": V1PodTemplateSpec(
+                    metadata=V1ObjectMeta(cluster.config.annotations)
+                    if cluster.config.annotations
+                    else None,
+                    spec=get_pod_spec(
                         cluster,
                         [get_head_container_spec(cluster)],
                         cluster.config.head_tolerations,
-                    )
-                },
+                    ),
+                ),
             },
             "workerGroupSpecs": [
                 {
@@ -158,11 +161,14 @@ def build_ray_cluster(cluster: "codeflare_sdk.ray.cluster.Cluster"):
                         "resources": worker_resources,
                     },
                     "template": V1PodTemplateSpec(
+                        metadata=V1ObjectMeta(cluster.config.annotations)
+                        if cluster.config.annotations
+                        else None,
                         spec=get_pod_spec(
                             cluster,
                             [get_worker_container_spec(cluster)],
                             cluster.config.worker_tolerations,
-                        )
+                        ),
                     ),
                 }
             ],
@@ -215,7 +221,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,
@@ -227,6 +233,7 @@ def get_metadata(cluster: "codeflare_sdk.ray.cluster.Cluster"):
     annotations = with_nb_annotations(cluster.config.annotations)
     if annotations != {}:
         object_meta.annotations = annotations  # As annotations are not a guarantee they are appended to the metadata after creation.
+
     return object_meta
 
 

src/codeflare_sdk/ray/cluster/test_config.py

@@ -41,6 +41,7 @@ def test_default_cluster_creation(mocker):
         f"{expected_clusters_dir}/ray/default-ray-cluster.yaml",
         get_template_variables(),
     )
+
     assert cluster.resource_yaml == expected_rc
 
 
@@ -114,6 +115,7 @@ def test_config_creation_all_parameters(mocker):
 @pytest.mark.filterwarnings("ignore::UserWarning")
 def test_all_config_params_aw(mocker):
     create_cluster_all_config_params(mocker, "aw-all-params", True)
+
     assert filecmp.cmp(
         f"{aw_dir}aw-all-params.yaml",
         f"{expected_clusters_dir}/appwrapper/unit-test-all-params.yaml",

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/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: