Skip to content

Commit f59d768

Browse files
committed
feat: add annotations parameter to cluster configuration
1 parent c8e77fc commit f59d768

File tree

2 files changed

+6
-4
lines changed

2 files changed

+6
-4
lines changed

src/codeflare_sdk/ray/cluster/build_ray_cluster.py

+3-4
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,7 @@ def get_metadata(cluster: "codeflare_sdk.ray.cluster.Cluster"):
191191
)
192192

193193
# Get the NB annotation if it exists - could be useful in future for a "annotations" parameter.
194-
annotations = get_nb_annotations()
194+
annotations = get_annotations(cluster.config.annotations)
195195
if annotations != {}:
196196
object_meta.annotations = annotations # As annotations are not a guarantee they are appended to the metadata after creation.
197197
return object_meta
@@ -213,11 +213,10 @@ def get_labels(cluster: "codeflare_sdk.ray.cluster.Cluster"):
213213
return labels
214214

215215

216-
def get_nb_annotations():
216+
def get_annotations(annotations):
217217
"""
218-
The get_nb_annotations() function generates the annotation for NB Prefix if the SDK is running in a notebook
218+
The get_annotations() function generates the annotation for NB Prefix if the SDK is running in a notebook and appends any user set annotations
219219
"""
220-
annotations = {}
221220

222221
# Notebook annotation
223222
nb_prefix = os.environ.get("NB_PREFIX")

src/codeflare_sdk/ray/cluster/config.py

+3
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,8 @@ class ClusterConfiguration:
8989
A dictionary of custom resource mappings to map extended resource requests to RayCluster resource names
9090
overwrite_default_resource_mapping:
9191
A boolean indicating whether to overwrite the default resource mapping.
92+
annotations:
93+
A dictionary of annotations to apply to the cluster.
9294
"""
9395

9496
name: str
@@ -126,6 +128,7 @@ class ClusterConfiguration:
126128
extended_resource_mapping: Dict[str, str] = field(default_factory=dict)
127129
overwrite_default_resource_mapping: bool = False
128130
local_queue: Optional[str] = None
131+
annotations: Dict[str, str] = field(default_factory=dict)
129132

130133
def __post_init__(self):
131134
if not self.verify_tls:

0 commit comments

Comments
 (0)