cupid/session.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 logging
import threading
import warnings

from odps.compat import six
from odps.utils import with_wait_argument

from .config import options
from .errors import CupidMasterTimeoutError
from .rpc import CupidTaskServiceRpcChannel
from .utils import get_environ

try:
    from .proto import cupidtaskparam_pb2 as task_param_pb
except TypeError:
    warnings.warn('Cannot import protos from pycupid: '
        'consider upgrading your protobuf python package.', ImportWarning)
    raise ImportError

logger = logging.getLogger(__name__)

_CUPID_CONF_PREFIXES = 'odps cupid'.split()


class CupidSession(object):
    def __init__(self, odps=None, project=None):
        from .runtime import context
        from odps import ODPS

        self._context = context()

        if odps is not None:
            self.odps = odps
        else:
            self.odps = ODPS.from_global()
        self.project = project if project is not None else self.odps.project

        self.lookup_name = get_environ('META_LOOKUP_NAME', '')
        self.running = False
        self.save_id = None
        self.job_running_event = threading.Event()

        self._kube_url = None

    def check_running(self, timeout=None):
        timeout = timeout or options.cupid.wait_am_start_time
        if not self.running:
            self.job_running_event.wait(timeout)
        return self.running

    @property
    def account(self):
        # todo add support for bearer token
        return self.odps.account

    def job_conf(self, conf=None):
        conf_items = {
            'odps.task.major.version': options.cupid.major_task_version,
            'odps.access.id': getattr(self.odps.account, 'access_id', None),
            'odps.access.key': getattr(self.odps.account, 'secret_access_key', None),
            'odps.end.point': self.odps.endpoint,
            'odps.project.name': self.project,
            'odps.moye.am.cores': '400',
            'odps.cupid.proxy.end.point': options.cupid.proxy_endpoint,
        }

        if conf:
            conf_items.update(conf)
        for k, v in six.iteritems(options.cupid.settings or {}):
            if any(k.startswith(pf) for pf in _CUPID_CONF_PREFIXES):
                conf_items[k] = v
        conf_obj_items = [task_param_pb.JobConfItem(key=k, value=str(v))
                          for k, v in six.iteritems(conf_items) if v is not None]
        return task_param_pb.JobConf(jobconfitem=conf_obj_items)

    def get_proxy_token(self, instance, app_name, expired_in_hours):
        if hasattr(instance, 'id'):
            instance = instance.id

        task_operator = task_param_pb.CupidTaskOperator(moperator='GetCupidProxyToken')
        proxy_token_request = task_param_pb.CupidProxyTokenRequest(
            instanceId=instance,
            appName=app_name,
            expiredInHours=expired_in_hours,
        )
        task_param = task_param_pb.CupidTaskParam(
            mcupidtaskoperator=task_operator,
            cupidProxyTokenRequest=proxy_token_request,
        )
        channel = CupidTaskServiceRpcChannel(self)
        inst = channel.submit_job(task_param, 'eAsyncNotFuxiJob', session=self)
        return channel.wait_cupid_instance(inst)

    def get_proxied_url(self, instance, app_name, expired_in_hours=None):
        expired_in_hours = expired_in_hours or 30 * 24
        return 'http://%s.%s' % (self.get_proxy_token(instance, app_name, expired_in_hours),
                                 options.cupid.proxy_endpoint)

    @with_wait_argument
    def start_kubernetes(self, async_=False, priority=None, running_cluster=None,
                         proxy_endpoint=None, major_task_version=None,
                         app_command=None, app_image=None, resources=None, **kw):
        priority = priority if priority is not None else options.priority
        if priority is None and options.get_priority is not None:
            priority = options.get_priority(self.odps)
        menginetype = options.cupid.engine_running_type

        if proxy_endpoint is not None:
            options.cupid.proxy_endpoint = proxy_endpoint
        if major_task_version is not None:
            options.cupid.major_task_version = major_task_version

        runtime_endpoint = kw.pop('runtime_endpoint', None)
        task_operator = task_param_pb.CupidTaskOperator(moperator='startam', menginetype=menginetype)
        task_name = kw.pop('task_name', None)

        if len(kw) > 0:
            raise ValueError('Got unexpected arguments: {}'.format(list(kw.keys())[0]))

        kub_conf = {
            'odps.cupid.kube.master.mode': options.cupid.kube.master_mode,
            'odps.cupid.master.type': options.cupid.master_type,
            'odps.cupid.engine.running.type': menginetype,
            'odps.cupid.job.capability.duration.hours': options.cupid.job_duration_hours,
            'odps.cupid.channel.init.timeout.seconds': options.cupid.channel_init_timeout_seconds,
            'odps.moye.runtime.type': options.cupid.application_type,
            'odps.runtime.end.point': runtime_endpoint or options.cupid.runtime.endpoint,
            'odps.cupid.resources': ','.join(resources or []),
            'odps.cupid.kube.appmaster.cmd': app_command,
            'odps.cupid.kube.appmaster.image': app_image
        }
        if running_cluster:
            kub_conf['odps.cupid.task.running.cluster'] = running_cluster
        if options.cupid.container_node_label is not None:
            kub_conf['odps.cupid.container.node.label'] = options.cupid.container_node_label
        if options.cupid.master.virtual_resource is not None:
            kub_conf['odps.cupid.master.virtual.resource'] = options.cupid.master.virtual_resource

        task_param = task_param_pb.CupidTaskParam(
            jobconf=self.job_conf(conf=kub_conf),
            mcupidtaskoperator=task_operator,
        )

        retrial = 0

        while True:
            retrial += 1
            channel = CupidTaskServiceRpcChannel(self)
            inst = channel.submit_job(task_param, 'eHasFuxiJob', with_resource=True, priority=priority,
                                    running_cluster=running_cluster, session=self, task_name=task_name)
            if async_:
                return inst
            else:
                try:
                    return self.get_instance_kube(inst)
                except CupidMasterTimeoutError:
                    if retrial == 3:
                        raise

    def get_instance_kube_api(self, inst, expired_in_hours=None):
        if self._kube_url is None:
            CupidTaskServiceRpcChannel.wait_cupid_instance(inst, 'running')
            self._kube_url = self.get_proxied_url(inst, '', expired_in_hours or 30 * 24)
        return self._kube_url

    def get_instance_kube(self, inst, expired_in_hours=None):
        from kubernetes import client
        config = client.Configuration()
        config.host = self.get_instance_kube_api(inst, expired_in_hours)
        return client.ApiClient(config)


from .io import *

__all__ = ['CupidSession', 'CupidTableUploadSession', 'CupidTableDownloadSession']