diff --git a/src/codeflare_sdk/common/utils/unit_test_support.py b/src/codeflare_sdk/common/utils/unit_test_support.py index 53fbc696..79c7bf01 100644 --- a/src/codeflare_sdk/common/utils/unit_test_support.py +++ b/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, ) diff --git a/src/codeflare_sdk/ray/cluster/build_ray_cluster.py b/src/codeflare_sdk/ray/cluster/build_ray_cluster.py index 9c626200..317091a8 100644 --- a/src/codeflare_sdk/ray/cluster/build_ray_cluster.py +++ b/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 diff --git a/src/codeflare_sdk/ray/cluster/test_config.py b/src/codeflare_sdk/ray/cluster/test_config.py index e6bcb8ba..6007f60b 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 825318a0..f1b75410 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 16a703b8..d5d8059d 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: