Ray cluster CRD and example CR + multi-ray-cluster operator (#12098)

python/ray/operator/operator.py

@@ -0,0 +1,154 @@
+"""
+Ray operator for Kubernetes.
+
+Reads ray cluster config from a k8s ConfigMap, starts a ray head node pod using
+create_or_update_cluster(), then runs an autoscaling loop in the operator pod
+executing this script. Writes autoscaling logs to the directory
+/root/ray-operator-logs.
+
+In this setup, the ray head node does not run an autoscaler. It is important
+NOT to supply an --autoscaling-config argument to head node's ray start command
+in the cluster config when using this operator.
+
+To run, first create a ConfigMap named ray-operator-configmap from a ray
+cluster config. Then apply the manifest at python/ray/autoscaler/kubernetes/operator_configs/operator_config.yaml
+
+For example:
+kubectl create namespace raytest
+kubectl -n raytest create configmap ray-operator-configmap --from-file=python/ray/autoscaler/kubernetes/operator_configs/test_cluster_config.yaml
+kubectl -n raytest apply -f python/ray/autoscaler/kubernetes/operator_configs/operator_config.yaml
+""" # noqa
+import logging
+import multiprocessing as mp
+import os
+from typing import Any, Callable, Dict, Optional
+
+from kubernetes.client.exceptions import ApiException
+import yaml
+
+from ray._private import services
+from ray.autoscaler._private import commands
+from ray import monitor
+from ray.operator import operator_utils
+from ray import ray_constants
+
+
+class RayCluster():
+    def __init__(self, config: Dict[str, Any]):
+        self.config = config
+        self.name = self.config["cluster_name"]
+        self.config_path = operator_utils.config_path(self.name)
+
+        self.setup_logging()
+
+        self.subprocess = None  # type: Optional[mp.Process]
+
+    def do_in_subprocess(self,
+                         f: Callable[[], None],
+                         wait_to_finish: bool = False) -> None:
+        # First stop the subprocess if it's alive
+        self.clean_up_subprocess()
+        # Reinstantiate process with f as target and start.
+        self.subprocess = mp.Process(name=self.name, target=f)
+        # Kill subprocess if monitor dies
+        self.subprocess.daemon = True
+        self.subprocess.start()
+        if wait_to_finish:
+            self.subprocess.join()
+
+    def clean_up_subprocess(self):
+        if self.subprocess and self.subprocess.is_alive():
+            self.subprocess.terminate()
+            self.subprocess.join()
+
+    def create_or_update(self) -> None:
+        self.do_in_subprocess(self._create_or_update)
+
+    def _create_or_update(self) -> None:
+        self.start_head()
+        self.start_monitor()
+
+    def start_head(self) -> None:
+        self.write_config()
+        self.config = commands.create_or_update_cluster(
+            self.config_path,
+            override_min_workers=None,
+            override_max_workers=None,
+            no_restart=False,
+            restart_only=False,
+            yes=True,
+            no_config_cache=True)
+        self.write_config()
+
+    def start_monitor(self) -> None:
+        ray_head_pod_ip = commands.get_head_node_ip(self.config_path)
+        # TODO: Add support for user-specified redis port and password
+        redis_address = services.address(ray_head_pod_ip,
+                                         ray_constants.DEFAULT_PORT)
+        self.mtr = monitor.Monitor(
+            redis_address=redis_address,
+            autoscaling_config=self.config_path,
+            redis_password=ray_constants.REDIS_DEFAULT_PASSWORD,
+            prefix_cluster_info=True)
+        self.mtr.run()
+
+    def clean_up(self) -> None:
+        self.clean_up_subprocess()
+        self.clean_up_logging()
+        self.delete_config()
+
+    def setup_logging(self) -> None:
+        self.handler = logging.StreamHandler()
+        self.handler.addFilter(lambda rec: rec.processName == self.name)
+        logging_format = ":".join([self.name, ray_constants.LOGGER_FORMAT])
+        self.handler.setFormatter(logging.Formatter(logging_format))
+        operator_utils.root_logger.addHandler(self.handler)
+
+    def clean_up_logging(self) -> None:
+        operator_utils.root_logger.removeHandler(self.handler)
+
+    def write_config(self) -> None:
+        with open(self.config_path, "w") as file:
+            yaml.dump(self.config, file)
+
+    def delete_config(self) -> None:
+        os.remove(self.config_path)
+
+
+ray_clusters = {}
+
+
+def cluster_action(cluster_config: Dict[str, Any], event_type: str) -> None:
+    cluster_name = cluster_config["cluster_name"]
+    if event_type == "ADDED":
+        ray_clusters[cluster_name] = RayCluster(cluster_config)
+        ray_clusters[cluster_name].create_or_update()
+    elif event_type == "MODIFIED":
+        ray_clusters[cluster_name].create_or_update()
+    elif event_type == "DELETED":
+        ray_clusters[cluster_name].clean_up()
+        del ray_clusters[cluster_name]
+
+
+def main() -> None:
+    # Make directory for ray cluster configs
+    if not os.path.isdir(operator_utils.RAY_CONFIG_DIR):
+        os.mkdir(operator_utils.RAY_CONFIG_DIR)
+    # Control loop
+    cluster_cr_stream = operator_utils.cluster_cr_stream()
+    try:
+        for event in cluster_cr_stream:
+            cluster_cr = event["object"]
+            event_type = event["type"]
+            cluster_config = operator_utils.cr_to_config(cluster_cr)
+            cluster_action(cluster_config, event_type)
+    except ApiException as e:
+        if e.status == 404:
+            raise Exception(
+                "Caught a 404 error. Has the RayCluster CRD been created?")
+        else:
+            raise
+
+
+if __name__ == "__main__":
+    main()

python/ray/operator/operator_utils.py

@@ -0,0 +1,114 @@
+import copy
+import logging
+import os
+from typing import Any, Dict, Iterator, List
+
+from kubernetes.watch import Watch
+
+from ray.autoscaler._private.kubernetes import custom_objects_api
+
+RAY_NAMESPACE = os.environ.get("RAY_OPERATOR_POD_NAMESPACE")
+
+RAY_CONFIG_DIR = os.path.expanduser("~/ray_cluster_configs")
+CONFIG_SUFFIX = "_config.yaml"
+
+CONFIG_FIELDS = {
+    "maxWorkers": "max_workers",
+    "upscalingSpeed": "upscaling_speed",
+    "idleTimeoutMinutes": "idle_timeout_minutes",
+    "headPodType": "head_node_type",
+    "workerDefaultPodType": "worker_default_node_type",
+    "workerStartRayCommands": "worker_start_ray_commands",
+    "headStartRayCommands": "head_start_ray_commands",
+    "podTypes": "available_node_types"
+}
+
+NODE_TYPE_FIELDS = {
+    "minWorkers": "min_workers",
+    "maxWorkers": "max_workers",
+    "podConfig": "node_config",
+    "rayResources": "resources",
+    "setupCommands": "worker_setup_commands"
+}
+
+PROVIDER_CONFIG = {
+    "type": "kubernetes",
+    "use_internal_ips": True,
+    "namespace": RAY_NAMESPACE
+}
+
+root_logger = logging.getLogger("ray")
+root_logger.setLevel(logging.getLevelName("DEBUG"))
+"""
+ownerReferences:
+  - apiVersion: apps/v1
+    controller: true
+    blockOwnerDeletion: true
+    kind: ReplicaSet
+    name: my-repset
+    uid: d9607e19-f88f-11e6-a518-42010a800195
+"""
+
+
+def config_path(cluster_name: str) -> str:
+    file_name = cluster_name + CONFIG_SUFFIX
+    return os.path.join(RAY_CONFIG_DIR, file_name)
+
+
+def cluster_cr_stream() -> Iterator:
+    w = Watch()
+    return w.stream(
+        custom_objects_api().list_namespaced_custom_object,
+        namespace=RAY_NAMESPACE,
+        group="cluster.ray.io",
+        version="v1",
+        plural="rayclusters")
+
+
+def cr_to_config(cluster_resource: Dict[str, Any]) -> Dict[str, Any]:
+    """Convert RayCluster custom resource to a ray cluster config for use by the
+    autoscaler."""
+    cr_spec = cluster_resource["spec"]
+    cr_meta = cluster_resource["metadata"]
+    config = translate(cr_spec, dictionary=CONFIG_FIELDS)
+    pod_types = cr_spec["podTypes"]
+    config["available_node_types"] = get_node_types(
+        pod_types, cluster_name=cr_meta["name"], cluster_uid=cr_meta["uid"])
+    config["cluster_name"] = cr_meta["name"]
+    config["provider"] = PROVIDER_CONFIG
+    return config
+
+
+def get_node_types(pod_types: List[Dict[str, Any]], cluster_name: str,
+                   cluster_uid: str) -> Dict[str, Any]:
+    cluster_owner_reference = get_cluster_owner_reference(
+        cluster_name, cluster_uid)
+    node_types = {}
+    for pod_type in pod_types:
+        name = pod_type["name"]
+        pod_type_copy = copy.deepcopy(pod_type)
+        pod_type_copy.pop("name")
+        node_types[name] = translate(
+            pod_type_copy, dictionary=NODE_TYPE_FIELDS)
+        # Deleting a RayCluster CR will also delete the associated pods.
+        node_types[name]["node_config"]["metadata"].update({
+            "ownerReferences": [cluster_owner_reference]
+        })
+    return node_types
+
+
+def get_cluster_owner_reference(cluster_name: str,
+                                cluster_uid: str) -> Dict[str, Any]:
+    return {
+        "apiVersion": "apps/v1",
+        "controller": True,
+        "blockOwnerDeletion": True,
+        "kind": "RayCluster",
+        "name": cluster_name,
+        "uid": cluster_uid
+    }
+
+
+def translate(configuration: Dict[str, Any],
+              dictionary: Dict[str, str]) -> Dict[str, Any]:
+    return {dictionary[field]: configuration[field] for field in configuration}