From f0f722bd9b9f74762ca4f4713145b529b319a906 Mon Sep 17 00:00:00 2001 From: Pat O'Connor Date: Thu, 6 Mar 2025 13:40:26 +0000 Subject: [PATCH] fix(RHOAIENG-20531): propagate annotations to ray pods Signed-off-by: Pat O'Connor --- .../common/utils/unit_test_support.py | 5 ++++- .../ray/cluster/build_ray_cluster.py | 19 +++++++++++++------ src/codeflare_sdk/ray/cluster/test_config.py | 2 ++ .../appwrapper/unit-test-all-params.yaml | 10 ++++++++++ .../ray/unit-test-all-params.yaml | 10 ++++++++++ 5 files changed, 39 insertions(+), 7 deletions(-) diff --git a/src/codeflare_sdk/common/utils/unit_test_support.py b/src/codeflare_sdk/common/utils/unit_test_support.py index 28a303813..a5d31ce5b 100644 --- a/src/codeflare_sdk/common/utils/unit_test_support.py +++ b/src/codeflare_sdk/common/utils/unit_test_support.py @@ -498,7 +498,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, ) diff --git a/src/codeflare_sdk/ray/cluster/build_ray_cluster.py b/src/codeflare_sdk/ray/cluster/build_ray_cluster.py index 215ac32e2..74c4fad38 100644 --- a/src/codeflare_sdk/ray/cluster/build_ray_cluster.py +++ b/src/codeflare_sdk/ray/cluster/build_ray_cluster.py @@ -139,13 +139,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": [ { @@ -159,11 +162,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, - ) + ), ), } ], @@ -191,7 +197,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, @@ -203,6 +209,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 diff --git a/src/codeflare_sdk/ray/cluster/test_config.py b/src/codeflare_sdk/ray/cluster/test_config.py index 34cc4237b..45dd1fc70 100644 --- a/src/codeflare_sdk/ray/cluster/test_config.py +++ b/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", 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 0977d659d..dd3275bd5 100644 --- a/tests/test_cluster_yamls/appwrapper/unit-test-all-params.yaml +++ b/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: diff --git a/tests/test_cluster_yamls/ray/unit-test-all-params.yaml b/tests/test_cluster_yamls/ray/unit-test-all-params.yaml index 188319ab1..4e018981a 100644 --- a/tests/test_cluster_yamls/ray/unit-test-all-params.yaml +++ b/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: