odps/mars_extension/legacy/deploy/app.py

# Copyright 1999-2022 Alibaba Group Holding Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import base64
import logging
import json
import os
import sys

from mars.deploy.kubernetes.client import KubernetesCluster, new_cluster
from mars.deploy.kubernetes.config import (
    VolumeConfig,
    MarsReplicationControllerConfig,
    MarsSchedulersConfig,
    MarsWorkersConfig,
    MarsWebsConfig,
)
from mars.deploy.utils import wait_services_ready
from mars.utils import readable_size

from ...utils import build_mars_image_name

logger = logging.getLogger(__name__)

UI_PORT = 50002
NOTEBOOK_PORT = 50003
GS_COORDINATOR_PORT = 63800
GS_COORDINATOR_GATEWAY_PORT = 64800


class DiskDriverVolumeConfig(VolumeConfig):
    def __init__(self, name, mount_path, device_size, renew_interval=3, timeout=120):
        super().__init__(name, mount_path)
        self._device_size = device_size
        self._timeout = timeout
        self._renew_interval = renew_interval

    def build(self):
        return {
            "name": self.name,
            "diskDriverEphemeralDevice": {
                "deviceSize": self._device_size,
                "timeout": self._timeout,
                "renewInterval": self._renew_interval,
            },
        }


class CupidMarsConfigMixin:
    def __init__(self, *args, **kwargs):
        self._container_port = kwargs.pop("container_port", None)
        super().__init__(*args, **kwargs)

    @property
    def stat_type(self):
        return "cgroup" if os.environ.get("VM_ENGINE_TYPE") != "hyper" else None

    @staticmethod
    def get_local_app_module(mod_name):
        from . import scheduler

        return scheduler.__name__.rsplit(".", 1)[0] + "." + mod_name

    def build_container(self):
        container_config = super().build_container()
        container_config.update({"enableChannel": True, "enableNativeLib": True})
        if self._container_port is not None:
            container_config["ports"] = {
                "containerPort": int(self._container_port),
            }
        return container_config

    def build_template_spec(self):
        spec_config = super().build_template_spec()
        if os.environ.get("VM_ENGINE_TYPE") == "hyper":
            spec_config.update({"restartPolicy": "Never"})
        else:
            spec_config.update({"hostNetwork": True, "restartPolicy": "Never"})
        return spec_config

    def config_readiness_probe(self):
        return None

    def build(self):
        result = super().build()
        with open("pod_%s.json" % self.rc_name, "w") as pod_file:
            pod_file.write(json.dumps(result, indent=2, sort_keys=True))
        return result

    def remove_env(self, name):
        try:
            super().remove_env(name)
        except AttributeError:
            self._envs.pop(name, None)


class CupidMarsSchedulersConfig(CupidMarsConfigMixin, MarsSchedulersConfig):
    def add_default_envs(self):
        super().add_default_envs()
        self.remove_env("MARS_USE_CGROUP_STAT")
        self.add_env("MARS_DISABLE_FAILOVER", "1")
        self.add_env("MARS_POD_ROLE", "scheduler")
        if self.stat_type == "cgroup":
            self.add_env("MARS_MEM_USE_CGROUP_STAT", "1")
            self.add_env("MARS_CPU_USE_PROCESS_STAT", "1")


class CupidMarsWorkersConfig(CupidMarsConfigMixin, MarsWorkersConfig):
    def __init__(self, *args, **kwargs):
        if os.environ.get("VM_ENGINE_TYPE") == "hyper":
            kwargs["mount_shm"] = kwargs.get("mount_shm", None) or False
        kwargs["min_cache_mem"] = kwargs.get("min_cache_mem", None) or "10%"
        super().__init__(*args, **kwargs)

    def add_default_envs(self):
        super().add_default_envs()
        self.add_env("MARS_LOCK_FREE_FILEIO", "1")
        self.add_env("MARS_DISABLE_PROC_RECOVER", "1")
        self.add_env("MARS_PLASMA_LIMIT", "95%")
        self.add_env("MARS_POD_ROLE", "worker")
        self.remove_env("MARS_USE_CGROUP_STAT")
        if self.stat_type == "cgroup":
            self.add_env("MARS_CPU_USE_PROCESS_STAT", "1")
            self.add_env("MARS_MEM_USE_CGROUP_STAT", "1")


class CupidMarsWebsConfig(CupidMarsConfigMixin, MarsWebsConfig):
    def __init__(self, *args, **kwargs):
        if os.environ.get("VM_ENGINE_TYPE") == "hyper":
            default_port = UI_PORT
        else:
            default_port = None
        kwargs["service_port"] = kwargs.get("service_port", None) or default_port
        super().__init__(*args, **kwargs)

    def add_default_envs(self):
        super().add_default_envs()
        self.add_env("MARS_POD_ROLE", "web")
        if os.environ.get("VM_ENGINE_TYPE") == "hyper":
            self.add_env("MARS_UI_PORT", str(UI_PORT))


class CupidMarsNotebooksConfig(CupidMarsConfigMixin, MarsReplicationControllerConfig):
    rc_name = "marsnotebook"

    def __init__(self, *args, **kwargs):
        if os.environ.get("VM_ENGINE_TYPE") == "hyper":
            default_port = NOTEBOOK_PORT
        else:
            default_port = None
        kwargs["service_port"] = kwargs.get("service_port", None) or default_port
        super().__init__(*args, **kwargs)

    def build_container_command(self):
        return [
            "/srv/entrypoint.sh",
            self.get_local_app_module("notebook"),
        ]

    def add_default_envs(self):
        super().add_default_envs()
        self.add_env("MARS_POD_ROLE", "notebook")
        if os.environ.get("VM_ENGINE_TYPE") == "hyper":
            self.add_env("MARS_NOTEBOOK_PORT", self._container_port)


class CupidGSCoordinatorConfig(CupidMarsConfigMixin, MarsReplicationControllerConfig):
    rc_name = "gscoordinator"

    def __init__(self, *args, **kwargs):
        if os.environ.get("VM_ENGINE_TYPE") == "hyper":
            default_port = GS_COORDINATOR_GATEWAY_PORT
        else:
            default_port = None
        kwargs["service_port"] = kwargs.get("service_port", None) or default_port
        self._worker_pod_name_list = kwargs.pop("worker_pod_name_list", None)
        self._worker_pod_ip_list = kwargs.pop("worker_pod_ip_list", None)
        self._port = GS_COORDINATOR_PORT
        self._gateway_port = kwargs["service_port"]
        self._worker_num = kwargs.pop("worker_num", 1)
        super().__init__(*args, **kwargs)

    def build_container_command(self):
        # replace command with coordinator
        coordinator_args = dict(
            port=self._port,
            gateway_port=self._gateway_port,
            worker_pod_name_list=self._worker_pod_name_list,
            worker_pod_ip_list=self._worker_pod_ip_list,
            num_workers=self._worker_num,
        )
        cmd = [
            "/srv/entrypoint.sh",
            self.get_local_app_module("gscoordinator"),
            base64.b64encode(json.dumps(coordinator_args).encode()).decode(),
        ]
        return cmd

    def add_default_envs(self):
        super().add_default_envs()
        self.add_env("MARS_POD_ROLE", "gscoordinator")
        if os.environ.get("VM_ENGINE_TYPE") == "hyper":
            self.add_env("GS_COORDINATOR_GATEWAY_PORT", self._container_port)
            self.add_env("GS_COORDINATOR_PORT", self._container_port)


class CupidKubernetesCluster(KubernetesCluster):
    _scheduler_config_cls = CupidMarsSchedulersConfig
    _worker_config_cls = CupidMarsWorkersConfig
    _web_config_cls = CupidMarsWebsConfig
    _default_service_port = None

    def __init__(self, *args, **kwargs):
        self._with_notebook = kwargs.pop("with_notebook", False)
        self._notebook_cpu = kwargs.pop("notebook_cpu", None)
        self._notebook_mem = kwargs.pop("notebook_mem", None)

        self._notebook_extra_env = kwargs.get("extra_env", None) or dict()
        self._notebook_extra_env.update(
            kwargs.pop("notebook_extra_env", None) or dict()
        )

        self._with_graphscope = kwargs.pop("with_graphscope", False)
        self._coordinator_extra_env = kwargs.get("extra_env", None) or dict()
        self._coordinator_extra_env.update(
            kwargs.pop("coordinator_extra_env", None) or dict()
        )
        self._coordinator_cpu = kwargs.pop("coordinator_cpu", None)
        self._coordinator_mem = kwargs.pop("coordinator_mem", None)

        self._node_blacklist = set(kwargs.pop("node_blacklist", None) or [])
        self._blacklisted_pods = set()

        kwargs["image"] = build_mars_image_name(kwargs.pop("image", None))
        super().__init__(*args, **kwargs)

    def _create_notebook(self):
        if self._with_notebook:
            notebook_config = CupidMarsNotebooksConfig(
                1,
                image=self._image,
                cpu=self._notebook_cpu,
                memory=self._notebook_mem,
                volumes=self._extra_volumes,
                pre_stop_command=self._pre_stop_command,
            )
            notebook_config.add_simple_envs(self._notebook_extra_env)
            self._core_api.create_namespaced_replication_controller(
                self._namespace, notebook_config.build()
            )

    def _create_graphscope(self):
        if self._with_graphscope:
            # check the worker services are ready
            limits = [self._worker_num]
            selectors = ["mars/service-type=" + CupidMarsWorkersConfig.rc_name]
            wait_services_ready(
                selectors,
                limits,
                lambda sel: self._get_ready_pod_count(sel),
                timeout=self._timeout,
            )
            # get the worker pod name and ip
            worker_pod_name_list, worker_pod_ip_list = self._get_pods_name_and_ip(
                selectors[0]
            )
            coordinator_config = CupidGSCoordinatorConfig(
                1,
                image=self._image,
                cpu=self._coordinator_cpu,
                memory=self._coordinator_mem,
                volumes=self._extra_volumes,
                pre_stop_command=self._pre_stop_command,
                worker_pod_name_list=worker_pod_name_list,
                worker_pod_ip_list=worker_pod_ip_list,
                worker_num=self._worker_num,
            )
            coordinator_config.add_simple_envs(self._coordinator_extra_env)
            self._core_api.create_namespaced_replication_controller(
                self._namespace, coordinator_config.build()
            )

    def _create_services(self):
        super()._create_services()
        self._create_graphscope()
        self._create_notebook()

    def _create_kube_service(self):
        # does not create k8s service as not supported in cupid
        pass

    def _get_ready_pod_count(self, label_selector):
        if self._node_blacklist:
            query = self._core_api.list_namespaced_pod(
                namespace=self._namespace, label_selector=label_selector
            ).to_dict()
            for el in query["items"]:
                node_name = el.get("spec", {}).get("node_name")
                pod_name = el["metadata"]["name"]

                if pod_name in self._blacklisted_pods:
                    continue
                if node_name and node_name in self._node_blacklist:
                    logger.warning(
                        "Found node %s in blacklist, will terminate pod %s",
                        node_name,
                        pod_name,
                    )
                    try:
                        self._core_api.delete_namespaced_pod(
                            name=pod_name, namespace=self._namespace
                        )
                        self._blacklisted_pods.add(pod_name)
                    except:  # noqa: E722
                        pass
        return super()._get_ready_pod_count(label_selector)

    def _get_pods_name_and_ip(self, label_selector):
        pod_name_list = []
        pod_ip_list = []
        query = self._core_api.list_namespaced_pod(
            namespace=self._namespace, label_selector=label_selector
        ).to_dict()
        for el in query["items"]:
            if el["status"]["reason"] == "CupidStarted":
                pod_name_list.append(el["metadata"]["name"])
                pod_ip_list.append(el["status"]["pod_ip"])
        return pod_name_list, pod_ip_list


class MarsCupidServer(object):
    def __init__(self):
        self.args = None
        self._instance_id = None
        self._kube_url = None
        self._kube_client = None

    def __call__(self, argv=None):
        if argv is None:
            argv = sys.argv[1:]
        return self._main(argv)

    def config_logging(self):
        import logging.config
        import mars

        log_conf = self.args.log_conf or "logging.conf"

        conf_file_paths = [
            "",
            os.path.abspath("."),
            os.path.dirname(os.path.dirname(mars.__file__)),
        ]
        log_configured = False
        for path in conf_file_paths:
            conf_path = os.path.join(path, log_conf) if path else log_conf
            if os.path.exists(conf_path):
                logging.config.fileConfig(conf_path, disable_existing_loggers=False)
                log_configured = True
                break

        if not log_configured:
            log_level = self.args.log_level
            level = getattr(logging, log_level.upper()) if log_level else logging.INFO
            logging.getLogger("mars").setLevel(level)
            logging.basicConfig(format=self.args.log_format)

    def _main(self, argv=None):
        parser = argparse.ArgumentParser(description="Mars Cupid Application")
        parser.add_argument("--log-level", help="log level")
        parser.add_argument("--log-format", help="log format")
        parser.add_argument(
            "--log-conf", help="log config file, logging.conf by default"
        )
        parser.add_argument("encoded_args")
        self.args = parser.parse_args(argv)
        self.config_logging()

        self._kube_url = os.environ["KUBE_API_ADDRESS"].strip('"')
        self._instance_id = os.environ["KUBE_NAMESPACE"].strip('"')

        args_dict = json.loads(base64.b64decode(self.args.encoded_args).decode())
        args_dict["namespace"] = self._instance_id

        if args_dict.get("worker_disk_num") is not None:
            disk_num = args_dict.pop("worker_disk_num")
            disk_size = args_dict.pop("worker_disk_size")
            args_dict["worker_spill_paths"] = [
                DiskDriverVolumeConfig(
                    name="diskdriver-volume%d" % i,
                    mount_path="/diskdriver%d" % i,
                    device_size=readable_size(disk_size, trunc=True).lower(),
                )
                for i in range(disk_num)
            ]

        scheduler_extra_env = args_dict.get("scheduler_extra_env") or dict()
        if args_dict.get("instance_idle_timeout") is not None:
            idle_timeout = args_dict.pop("instance_idle_timeout")
            scheduler_extra_env["MARS_INSTANCE_IDLE_TIMEOUT"] = str(idle_timeout)
        args_dict["scheduler_extra_env"] = scheduler_extra_env

        extra_env = args_dict.get("extra_env") or dict()
        extra_env["KUBE_API_ADDRESS"] = self._kube_url
        extra_env["KUBE_NAMESPACE"] = self._instance_id
        args_dict["extra_env"] = extra_env

        new_cluster(
            self.get_instance_kube(), cluster_cls=CupidKubernetesCluster, **args_dict
        )
        logger.info("Cluster creation finished")

    def get_instance_kube(self):
        from kubernetes import client

        config = client.Configuration()
        config.host = self._kube_url
        return client.ApiClient(config)


main = MarsCupidServer()

if __name__ == "__main__":
    main()