spark-instrumented-optimizer/python/pyspark/resource/profile.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.
#

from pyspark.resource.requests import TaskResourceRequest, TaskResourceRequests, \
    ExecutorResourceRequests, ExecutorResourceRequest


class ResourceProfile(object):

    """
    Resource profile to associate with an RDD. A :class:`pyspark.resource.ResourceProfile`
    allows the user to specify executor and task requirements for an RDD that will get
    applied during a stage. This allows the user to change the resource requirements between
    stages. This is meant to be immutable so user cannot change it after building.

    .. versionadded:: 3.1.0

    Notes
    -----
    This API is evolving.
    """

    def __init__(self, _java_resource_profile=None, _exec_req=None, _task_req=None):
        if _java_resource_profile is not None:
            self._java_resource_profile = _java_resource_profile
        else:
            self._java_resource_profile = None
            self._executor_resource_requests = _exec_req or {}
            self._task_resource_requests = _task_req or {}

    @property
    def id(self):
        if self._java_resource_profile is not None:
            return self._java_resource_profile.id()
        else:
            raise RuntimeError("SparkContext must be created to get the id, get the id "
                               "after adding the ResourceProfile to an RDD")

    @property
    def taskResources(self):
        if self._java_resource_profile is not None:
            taskRes = self._java_resource_profile.taskResourcesJMap()
            result = {}
            for k, v in taskRes.items():
                result[k] = TaskResourceRequest(v.resourceName(), v.amount())
            return result
        else:
            return self._task_resource_requests

    @property
    def executorResources(self):
        if self._java_resource_profile is not None:
            execRes = self._java_resource_profile.executorResourcesJMap()
            result = {}
            for k, v in execRes.items():
                result[k] = ExecutorResourceRequest(v.resourceName(), v.amount(),
                                                    v.discoveryScript(), v.vendor())
            return result
        else:
            return self._executor_resource_requests


class ResourceProfileBuilder(object):

    """
    Resource profile Builder to build a resource profile to associate with an RDD.
    A ResourceProfile allows the user to specify executor and task requirements for
    an RDD that will get applied during a stage. This allows the user to change the
    resource requirements between stages.

    .. versionadded:: 3.1.0

    Notes
    -----
    This API is evolving.
    """

    def __init__(self):
        from pyspark.context import SparkContext
        _jvm = SparkContext._jvm
        if _jvm is not None:
            self._jvm = _jvm
            self._java_resource_profile_builder = \
                _jvm.org.apache.spark.resource.ResourceProfileBuilder()
        else:
            self._jvm = None
            self._java_resource_profile_builder = None
            self._executor_resource_requests = {}
            self._task_resource_requests = {}

    def require(self, resourceRequest):
        if isinstance(resourceRequest, TaskResourceRequests):
            if self._java_resource_profile_builder is not None:
                if resourceRequest._java_task_resource_requests is not None:
                    self._java_resource_profile_builder.require(
                        resourceRequest._java_task_resource_requests)
                else:
                    taskReqs = TaskResourceRequests(self._jvm, resourceRequest.requests)
                    self._java_resource_profile_builder.require(
                        taskReqs._java_task_resource_requests)
            else:
                self._task_resource_requests.update(resourceRequest.requests)
        else:
            if self._java_resource_profile_builder is not None:
                if resourceRequest._java_executor_resource_requests is not None:
                    self._java_resource_profile_builder.require(
                        resourceRequest._java_executor_resource_requests)
                else:
                    execReqs = ExecutorResourceRequests(self._jvm, resourceRequest.requests)
                    self._java_resource_profile_builder.require(
                        execReqs._java_executor_resource_requests)
            else:
                self._executor_resource_requests.update(resourceRequest.requests)
        return self

    def clearExecutorResourceRequests(self):
        if self._java_resource_profile_builder is not None:
            self._java_resource_profile_builder.clearExecutorResourceRequests()
        else:
            self._executor_resource_requests = {}

    def clearTaskResourceRequests(self):
        if self._java_resource_profile_builder is not None:
            self._java_resource_profile_builder.clearTaskResourceRequests()
        else:
            self._task_resource_requests = {}

    @property
    def taskResources(self):
        if self._java_resource_profile_builder is not None:
            taskRes = self._java_resource_profile_builder.taskResourcesJMap()
            result = {}
            for k, v in taskRes.items():
                result[k] = TaskResourceRequest(v.resourceName(), v.amount())
            return result
        else:
            return self._task_resource_requests

    @property
    def executorResources(self):
        if self._java_resource_profile_builder is not None:
            result = {}
            execRes = self._java_resource_profile_builder.executorResourcesJMap()
            for k, v in execRes.items():
                result[k] = ExecutorResourceRequest(v.resourceName(), v.amount(),
                                                    v.discoveryScript(), v.vendor())
            return result
        else:
            return self._executor_resource_requests

    @property
    def build(self):
        if self._java_resource_profile_builder is not None:
            jresourceProfile = self._java_resource_profile_builder.build()
            return ResourceProfile(_java_resource_profile=jresourceProfile)
        else:
            return ResourceProfile(_exec_req=self._executor_resource_requests,
                                   _task_req=self._task_resource_requests)