From 6279006513397a42a865f871ef84107eaf46352d Mon Sep 17 00:00:00 2001 From: Bobbins228 Date: Thu, 19 Dec 2024 12:03:44 +0000 Subject: [PATCH 1/3] feat: add annotations parameter to cluster configuration --- src/codeflare_sdk/ray/cluster/build_ray_cluster.py | 7 +++---- src/codeflare_sdk/ray/cluster/config.py | 3 +++ 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/codeflare_sdk/ray/cluster/build_ray_cluster.py b/src/codeflare_sdk/ray/cluster/build_ray_cluster.py index e590d4834..111713381 100644 --- a/src/codeflare_sdk/ray/cluster/build_ray_cluster.py +++ b/src/codeflare_sdk/ray/cluster/build_ray_cluster.py @@ -191,7 +191,7 @@ def get_metadata(cluster: "codeflare_sdk.ray.cluster.Cluster"): ) # Get the NB annotation if it exists - could be useful in future for a "annotations" parameter. - annotations = get_nb_annotations() + 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 @@ -213,11 +213,10 @@ def get_labels(cluster: "codeflare_sdk.ray.cluster.Cluster"): return labels -def get_nb_annotations(): +def with_nb_annotations(annotations: dict): """ - The get_nb_annotations() function generates the annotation for NB Prefix if the SDK is running in a notebook + The with_nb_annotations() function generates the annotation for NB Prefix if the SDK is running in a notebook and appends any user set annotations """ - annotations = {} # Notebook annotation nb_prefix = os.environ.get("NB_PREFIX") diff --git a/src/codeflare_sdk/ray/cluster/config.py b/src/codeflare_sdk/ray/cluster/config.py index b8b097f8c..e30690292 100644 --- a/src/codeflare_sdk/ray/cluster/config.py +++ b/src/codeflare_sdk/ray/cluster/config.py @@ -89,6 +89,8 @@ class ClusterConfiguration: A dictionary of custom resource mappings to map extended resource requests to RayCluster resource names overwrite_default_resource_mapping: A boolean indicating whether to overwrite the default resource mapping. + annotations: + A dictionary of annotations to apply to the cluster. """ name: str @@ -126,6 +128,7 @@ class ClusterConfiguration: extended_resource_mapping: Dict[str, str] = field(default_factory=dict) overwrite_default_resource_mapping: bool = False local_queue: Optional[str] = None + annotations: Dict[str, str] = field(default_factory=dict) def __post_init__(self): if not self.verify_tls: From 5c9ced01073d40aef47cc98d22f288e9ea463db5 Mon Sep 17 00:00:00 2001 From: Bobbins228 Date: Thu, 19 Dec 2024 12:05:16 +0000 Subject: [PATCH 2/3] test: add unit test for annotations parameter --- src/codeflare_sdk/common/utils/unit_test_support.py | 1 + src/codeflare_sdk/ray/cluster/test_config.py | 5 +++++ .../test_cluster_yamls/appwrapper/unit-test-all-params.yaml | 2 ++ tests/test_cluster_yamls/ray/unit-test-all-params.yaml | 2 ++ 4 files changed, 10 insertions(+) diff --git a/src/codeflare_sdk/common/utils/unit_test_support.py b/src/codeflare_sdk/common/utils/unit_test_support.py index 4182fc93d..388b6b4ee 100644 --- a/src/codeflare_sdk/common/utils/unit_test_support.py +++ b/src/codeflare_sdk/common/utils/unit_test_support.py @@ -442,5 +442,6 @@ 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"}, ) return Cluster(config) diff --git a/src/codeflare_sdk/ray/cluster/test_config.py b/src/codeflare_sdk/ray/cluster/test_config.py index b16cbaf28..aafb23dd9 100644 --- a/src/codeflare_sdk/ray/cluster/test_config.py +++ b/src/codeflare_sdk/ray/cluster/test_config.py @@ -93,6 +93,11 @@ def test_config_creation_all_parameters(mocker): ) assert cluster.config.overwrite_default_resource_mapping == True assert cluster.config.local_queue == "local-queue-default" + assert cluster.config.annotations == { + "app.kubernetes.io/managed-by": "test-prefix", + "key1": "value1", + "key2": "value2", + } assert filecmp.cmp( f"{aw_dir}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 6d2c5440a..36d186d10 100644 --- a/tests/test_cluster_yamls/appwrapper/unit-test-all-params.yaml +++ b/tests/test_cluster_yamls/appwrapper/unit-test-all-params.yaml @@ -13,6 +13,8 @@ spec: metadata: annotations: app.kubernetes.io/managed-by: test-prefix + key1: value1 + key2: value2 labels: controller-tools.k8s.io: '1.0' key1: value1 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 8426eede6..a6514bd80 100644 --- a/tests/test_cluster_yamls/ray/unit-test-all-params.yaml +++ b/tests/test_cluster_yamls/ray/unit-test-all-params.yaml @@ -3,6 +3,8 @@ kind: RayCluster metadata: annotations: app.kubernetes.io/managed-by: test-prefix + key1: value1 + key2: value2 labels: controller-tools.k8s.io: '1.0' key1: value1 From e1572ead15e77ebe5d8e45745799e5c16bd8f6a0 Mon Sep 17 00:00:00 2001 From: Bobbins228 Date: Thu, 19 Dec 2024 12:07:26 +0000 Subject: [PATCH 3/3] docs: update user docs with annotations parameter example --- docs/sphinx/user-docs/cluster-configuration.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/sphinx/user-docs/cluster-configuration.rst b/docs/sphinx/user-docs/cluster-configuration.rst index dc3f2cf42..411760f8b 100644 --- a/docs/sphinx/user-docs/cluster-configuration.rst +++ b/docs/sphinx/user-docs/cluster-configuration.rst @@ -26,6 +26,7 @@ requirements for creating the Ray Cluster. worker_memory_limits=2, # Default 2 # image="", # Optional Field labels={"exampleLabel": "example", "secondLabel": "example"}, + annotations={"key1":"value1", "key2":"value2"}, )) .. note::