feat: add head tolerations & node labels for flexible ray cluster pods scheduling

CHANGELOG.md

@@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - updated mlflow version to 2.16.0 to support LLM tracing
 - remove CDK overhead from `mlflow-image` module
 - renamed mlflow manifests and updated README.MD
+- added head tolerations & node labels for flexible ray cluster pods scheduling
 
 ## v1.5.0
 

modules/eks/ray-cluster/README.md

@@ -32,6 +32,17 @@ This module creates a Ray cluster in AWS EKS Kubernetes cluster. It deploys a Ra
         cpu: "1"
         memory: "8G"
 ```
+- `head_tolerations` - List of head group tolerations. Empty by default (no tolerations). Example input:
+```yaml
+  - key: "nvidia.com/gpu"
+    value: "true"
+    # operator: "Equal"
+    effect: "NoSchedule"
+```
+- `head_labels` - Dictionary of head group labels. Empty by default (no labels). Example input:
+```yaml
+    usage: gpu
+```
 - `worker_replicas` - The requested number of the Pod replicas in worker group. Defaults to `1`.
 - `worker_min_replicas` - The minimum number of the Pod replicas in worker group. Defaults to `1`.
 - `worker_max_replicas` - The maximum number of the Pod replicas in worker group. Defaults to `10`.

modules/eks/ray-cluster/app.py

@@ -31,6 +31,8 @@
     enable_autoscaling=app_settings.parameters.enable_autoscaling,
     autoscaler_idle_timeout_seconds=app_settings.parameters.autoscaler_idle_timeout_seconds,
     head_resources=app_settings.parameters.head_resources,
+    head_tolerations=app_settings.parameters.head_tolerations,
+    head_labels=app_settings.parameters.head_labels,
     worker_replicas=app_settings.parameters.worker_replicas,
     worker_min_replicas=app_settings.parameters.worker_min_replicas,
     worker_max_replicas=app_settings.parameters.worker_max_replicas,

modules/eks/ray-cluster/ray_cluster_stack.py

@@ -32,6 +32,8 @@ def __init__(
         enable_autoscaling: bool,
         autoscaler_idle_timeout_seconds: int,
         head_resources: Dict[str, Dict[str, str]],
+        head_tolerations: List[Dict[str, str]],
+        head_labels: Dict[str, str],
         worker_replicas: int,
         worker_min_replicas: int,
         worker_max_replicas: int,
@@ -127,6 +129,8 @@ def __init__(
                     "resources": head_resources,
                     "volumes": volumes,
                     "volumeMounts": volume_mounts,
+                    "tolerations": head_tolerations,
+                    "nodeSelector": head_labels,
                 },
                 "worker": {
                     "replicas": worker_replicas,

modules/eks/ray-cluster/settings.py

@@ -51,6 +51,8 @@ class SeedFarmerParameters(CdkBaseSettings):
     enable_autoscaling: bool = Field(default=True)
     autoscaler_idle_timeout_seconds: int = Field(default=60)
     head_resources: Dict[str, Dict[str, str]] = Field(default=DEFAULT_POD_RESOURCES)
+    head_tolerations: List[Dict[str, str]] = Field(default=[])
+    head_labels: Dict[str, str] = Field(default={})
     worker_replicas: int = Field(default=1)
     worker_min_replicas: int = Field(default=1)
     worker_max_replicas: int = Field(default=10)

modules/eks/ray-cluster/tests/test_stack.py

@@ -45,6 +45,8 @@ def ray_cluster_stack(stack_defaults) -> cdk.Stack:
     worker_min_replicas = 4
     worker_max_replicas = 10
     head_resources = {"limits": {"cpu": "1", "memory": "8G"}, "requests": {"cpu": "1", "memory": "8G"}}
+    head_tolerations = []
+    head_labels = {}
     worker_resources = {"limits": {"cpu": "1", "memory": "8G"}, "requests": {"cpu": "1", "memory": "8G"}}
     worker_tolerations = []
     worker_labels = {}
@@ -68,6 +70,8 @@ def ray_cluster_stack(stack_defaults) -> cdk.Stack:
         enable_autoscaling=enable_autoscaling,
         autoscaler_idle_timeout_seconds=autoscaler_idle_timeout_seconds,
         head_resources=head_resources,
+        head_tolerations=head_tolerations,
+        head_labels=head_labels,
         worker_replicas=worker_replicas,
         worker_min_replicas=worker_min_replicas,
         worker_max_replicas=worker_max_replicas,