Add autoscaling support for Dask and Ray clusters

Author: polyaxon-ci

cli/polyaxon/_flow/run/dask/dask.py

@@ -1,7 +1,8 @@
 from typing import Optional, Union
 from typing_extensions import Literal
 
-from clipped.types.ref_or_obj import RefField
+from clipped.compact.pydantic import Field
+from clipped.types.ref_or_obj import IntOrRef, RefField
 
 from polyaxon._flow.run.base import BaseRun
 from polyaxon._flow.run.dask.replica import V1DaskReplica
@@ -22,6 +23,8 @@ class V1DaskCluster(BaseRun, DestinationImageMixin):
         kind: str, should be equal `daskcluster`
         worker: V1DaskReplica, optional - Worker replica specification
         scheduler: V1DaskReplica, optional - Scheduler replica specification
+        min_replicas: int, optional - Minimum number of workers for autoscaling
+        max_replicas: int, optional - Maximum number of workers for autoscaling
 
     ## YAML usage
 
@@ -44,6 +47,25 @@ class V1DaskCluster(BaseRun, DestinationImageMixin):
     >>>           memory: 1Gi
     >>>           cpu: 1
     ```
+
+    ## Autoscaling
+
+    Enable autoscaling by setting `minReplicas` and `maxReplicas`.
+    When both are set, a DaskAutoscaler resource will be created.
+
+    ```yaml
+    >>> run:
+    >>>   kind: daskcluster
+    >>>   minReplicas: 1
+    >>>   maxReplicas: 10
+    >>>   worker:
+    >>>     replicas: 2
+    >>>     container:
+    >>>       image: daskdev/dask:latest
+    >>>   scheduler:
+    >>>     container:
+    >>>       image: daskdev/dask:latest
+    ```
     """
 
     _IDENTIFIER = V1RunKind.DASKCLUSTER
@@ -52,6 +74,8 @@ class V1DaskCluster(BaseRun, DestinationImageMixin):
     kind: Literal[_IDENTIFIER] = _IDENTIFIER
     worker: Optional[Union[V1DaskReplica, RefField]] = None
     scheduler: Optional[Union[V1DaskReplica, RefField]] = None
+    min_replicas: Optional[IntOrRef] = Field(alias="minReplicas", default=None)
+    max_replicas: Optional[IntOrRef] = Field(alias="maxReplicas", default=None)
 
     def apply_image_destination(self, image: str):
         if self.worker:

cli/polyaxon/_flow/run/ray/__init__.py

@@ -1,2 +1,3 @@
+from polyaxon._flow.run.ray.autoscaler import V1RayAutoscalerOptions
 from polyaxon._flow.run.ray.ray import V1RayCluster
 from polyaxon._flow.run.ray.replica import V1RayReplica

cli/polyaxon/_flow/run/ray/autoscaler.py

@@ -0,0 +1,83 @@
+from typing import Optional, Union
+
+from clipped.compact.pydantic import Field
+from clipped.types.ref_or_obj import RefField
+
+from polyaxon._k8s import k8s_schemas
+from polyaxon._schemas.base import BaseSchemaModel
+
+
+class V1RayAutoscalerOptions(BaseSchemaModel):
+    """Ray autoscaler options for configuring automatic scaling behavior.
+
+    Args:
+        upscaling_mode: str, optional - "Conservative", "Default", or "Aggressive"
+        image_pull_policy: str, optional - "IfNotPresent", "Always", or "Never"
+        resources: V1ResourceRequirements, optional - Resources for the autoscaler container
+
+    ## YAML usage
+
+    ```yaml
+    >>> run:
+    >>>   kind: raycluster
+    >>>   autoscalerOptions:
+    >>>     upscalingMode: Default
+    >>>     imagePullPolicy: IfNotPresent
+    ```
+
+    ## Python usage
+
+    ```python
+    >>> from polyaxon.schemas import V1RayAutoscalerOptions
+    >>> autoscaler_options = V1RayAutoscalerOptions(
+    >>>     upscaling_mode="Default",
+    >>>     image_pull_policy="IfNotPresent",
+    >>> )
+    ```
+
+    ## Fields
+
+    ### upscalingMode
+
+    Controls how aggressively the autoscaler scales up workers.
+
+    - `Conservative`: Upscaling is rate-limited; the number of pending worker pods
+      is at most the size of the Ray cluster.
+    - `Default`: Upscaling is not rate-limited.
+    - `Aggressive`: An alias for Default; upscaling is not rate-limited.
+
+    ```yaml
+    >>> autoscalerOptions:
+    >>>   upscalingMode: Default
+    ```
+
+    ### imagePullPolicy
+
+    Image pull policy for the autoscaler container.
+
+    ```yaml
+    >>> autoscalerOptions:
+    >>>   imagePullPolicy: IfNotPresent
+    ```
+
+    ### resources
+
+    Resource requirements for the autoscaler container.
+
+    ```yaml
+    >>> autoscalerOptions:
+    >>>   resources:
+    >>>     limits:
+    >>>       cpu: "500m"
+    >>>       memory: "512Mi"
+    >>>     requests:
+    >>>       cpu: "250m"
+    >>>       memory: "256Mi"
+    ```
+    """
+
+    _IDENTIFIER = "ray_autoscaler_options"
+
+    upscaling_mode: Optional[str] = Field(alias="upscalingMode", default=None)
+    image_pull_policy: Optional[str] = Field(alias="imagePullPolicy", default=None)
+    resources: Optional[Union[k8s_schemas.V1ResourceRequirements, RefField]] = None

cli/polyaxon/_flow/run/ray/ray.py

@@ -2,10 +2,11 @@
 from typing_extensions import Literal
 
 from clipped.compact.pydantic import Field
-from clipped.types.ref_or_obj import RefField
+from clipped.types.ref_or_obj import BoolOrRef, RefField
 
 from polyaxon._flow.run.base import BaseRun
 from polyaxon._flow.run.enums import V1RunKind
+from polyaxon._flow.run.ray.autoscaler import V1RayAutoscalerOptions
 from polyaxon._flow.run.ray.replica import V1RayReplica
 from polyaxon._flow.run.resources import V1RunResources
 from polyaxon._flow.run.utils import DestinationImageMixin
@@ -27,6 +28,8 @@ class V1RayCluster(BaseRun, DestinationImageMixin):
         ray_version: str, optional
         head: [V1RayReplica](/docs/experimentation/distributed/ray-replica/), optional
         workers: Dict[str, [V1RayReplica](/docs/experimentation/distributed/ray-replica/)], optional
+        enable_in_tree_autoscaling: bool, optional
+        autoscaler_options: [V1RayAutoscalerOptions](/docs/experimentation/distributed/ray-autoscaler/), optional
 
 
     ## YAML usage
@@ -40,6 +43,8 @@ class V1RayCluster(BaseRun, DestinationImageMixin):
     >>>   rayVersion:
     >>>   head:
     >>>   workers:
+    >>>   enableInTreeAutoscaling:
+    >>>   autoscalerOptions:
     ```
 
     ## Python usage
@@ -52,6 +57,7 @@ class V1RayCluster(BaseRun, DestinationImageMixin):
     >>>     ray_version="2.5.0",
     >>>     head=V1RayReplica(...),
     >>>     worker=V1RayReplica(...),
+    >>>     enable_in_tree_autoscaling=True,
     >>> )
     ```
 
@@ -132,10 +138,42 @@ class V1RayCluster(BaseRun, DestinationImageMixin):
     >>>     ...
     >>>   ...
     ```
+
+    ### enableInTreeAutoscaling
+
+    Enable the KubeRay in-tree autoscaler for automatic worker scaling.
+    When enabled, the Ray autoscaler will automatically scale workers based on workload demands.
+
+    > **Note**: The `idleTimeoutSeconds` for the autoscaler is derived from `termination.culling.timeout`.
+
+    ```yaml
+    >>> run:
+    >>>   kind: raycluster
+    >>>   enableInTreeAutoscaling: true
+    >>>   workers:
+    >>>     gpu-workers:
+    >>>       minReplicas: 0
+    >>>       maxReplicas: 10
+    >>>   ...
+    ```
+
+    ### autoscalerOptions
+
+    Optional configuration for the Ray autoscaler behavior.
+
+    ```yaml
+    >>> run:
+    >>>   kind: raycluster
+    >>>   enableInTreeAutoscaling: true
+    >>>   autoscalerOptions:
+    >>>     upscalingMode: Default
+    >>>     imagePullPolicy: IfNotPresent
+    >>>   ...
+    ```
     """
 
     _IDENTIFIER = V1RunKind.RAYCLUSTER
-    _CUSTOM_DUMP_FIELDS = {"head", "workers"}
+    _CUSTOM_DUMP_FIELDS = {"head", "workers", "autoscaler_options"}
     _FIELDS_DICT_PATCH = ["workers"]
 
     kind: Literal[_IDENTIFIER] = _IDENTIFIER
@@ -147,6 +185,12 @@ class V1RayCluster(BaseRun, DestinationImageMixin):
     ray_version: Optional[str] = Field(alias="rayVersion", default=None)
     head: Optional[Union[V1RayReplica, RefField]] = None
     workers: Optional[Dict[str, Union[V1RayReplica, RefField]]] = Field(default=None)
+    enable_in_tree_autoscaling: Optional[BoolOrRef] = Field(
+        alias="enableInTreeAutoscaling", default=None
+    )
+    autoscaler_options: Optional[Union[V1RayAutoscalerOptions, RefField]] = Field(
+        alias="autoscalerOptions", default=None
+    )
 
     def apply_image_destination(self, image: str):
         if self.head:

cli/polyaxon/_k8s/converter/converters/dask_cluster.py

@@ -64,4 +64,6 @@ def _get_replica(replica: Optional[V1DaskReplica]) -> Optional[ReplicaSpec]:
             notifications=plugins.notifications,
             labels=labels,
             annotations=self.annotations,
+            min_replicas=cluster.min_replicas,
+            max_replicas=cluster.max_replicas,
         )

cli/polyaxon/_k8s/converter/converters/ray_cluster.py

@@ -77,4 +77,6 @@ def _get_replica(replica: Optional[V1RayReplica]) -> Optional[ReplicaSpec]:
             notifications=plugins.notifications,
             labels=labels,
             annotations=self.annotations,
+            enable_in_tree_autoscaling=cluster.enable_in_tree_autoscaling,
+            autoscaler_options=cluster.autoscaler_options,
         )

cli/polyaxon/_k8s/custom_resources/dask_cluster.py

@@ -85,6 +85,8 @@ def get_dask_cluster_custom_resource(
     notifications: List[V1Notification],
     labels: Dict[str, str],
     annotations: Dict[str, str],
+    min_replicas: Optional[int] = None,
+    max_replicas: Optional[int] = None,
 ) -> Dict:
     template_spec = {}
     get_dask_replicas_template(
@@ -160,6 +162,10 @@ def get_dask_cluster_custom_resource(
         ],
     )
     template_spec = {"replicaSpecs": template_spec, "service": service}
+    # Add autoscaling configuration if both min and max replicas are set
+    if min_replicas is not None and max_replicas is not None:
+        template_spec["minReplicas"] = min_replicas
+        template_spec["maxReplicas"] = max_replicas
     custom_object = {"daskClusterSpec": template_spec}
     custom_object = set_termination(
         custom_object=custom_object, termination=termination

cli/polyaxon/_k8s/custom_resources/ray_cluster.py

@@ -1,6 +1,7 @@
 from typing import Dict, List, Optional
 
 from polyaxon._flow import V1Notification, V1Termination
+from polyaxon._flow.run.ray.autoscaler import V1RayAutoscalerOptions
 from polyaxon._k8s import k8s_schemas
 from polyaxon._k8s.converter.pod.spec import get_pod_spec, get_pod_template_spec
 from polyaxon._k8s.custom_resources.operation import (
@@ -138,6 +139,24 @@ def get_ray_worker_replicas_template(
         template_spec["workers"] = workers
 
 
+def _get_autoscaler_options(
+    autoscaler_options: Optional[V1RayAutoscalerOptions],
+) -> Optional[Dict]:
+    """Build autoscaler options dict from user-provided options."""
+    if not autoscaler_options:
+        return None
+
+    result = {}
+    if autoscaler_options.upscaling_mode:
+        result["upscalingMode"] = autoscaler_options.upscaling_mode
+    if autoscaler_options.image_pull_policy:
+        result["imagePullPolicy"] = autoscaler_options.image_pull_policy
+    if autoscaler_options.resources:
+        result["resources"] = autoscaler_options.resources
+
+    return result if result else None
+
+
 def get_ray_cluster_custom_resource(
     resource_name: str,
     namespace: str,
@@ -153,6 +172,8 @@ def get_ray_cluster_custom_resource(
     ray_version: Optional[str],
     labels: Dict[str, str],
     annotations: Dict[str, str],
+    enable_in_tree_autoscaling: Optional[bool] = None,
+    autoscaler_options: Optional[V1RayAutoscalerOptions] = None,
 ) -> Dict:
     template_spec = {}
 
@@ -181,6 +202,14 @@ def get_ray_cluster_custom_resource(
     if ray_version:
         template_spec["rayVersion"] = ray_version
 
+    # Add autoscaling configuration
+    if enable_in_tree_autoscaling is not None:
+        template_spec["enableInTreeAutoscaling"] = enable_in_tree_autoscaling
+
+    autoscaler_opts = _get_autoscaler_options(autoscaler_options)
+    if autoscaler_opts:
+        template_spec["autoscalerOptions"] = autoscaler_opts
+
     custom_object = {"rayClusterSpec": template_spec}
     custom_object = set_termination(
         custom_object=custom_object, termination=termination