#
# 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 typing import overload, Dict, Union, Optional
from py4j.java_gateway import JavaObject
from pyspark.resource.requests import (
TaskResourceRequest,
TaskResourceRequests,
ExecutorResourceRequests,
ExecutorResourceRequest,
)
[docs]class ResourceProfile:
"""
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.
Examples
--------
Create Executor resource requests.
>>> executor_requests = (
... ExecutorResourceRequests()
... .cores(2)
... .memory("6g")
... .memoryOverhead("1g")
... .pysparkMemory("2g")
... .offheapMemory("3g")
... .resource("gpu", 2, "testGpus", "nvidia.com")
... )
Create task resource requasts.
>>> task_requests = TaskResourceRequests().cpus(2).resource("gpu", 2)
Create a resource profile.
>>> builder = ResourceProfileBuilder()
>>> resource_profile = builder.require(executor_requests).require(task_requests).build
Create an RDD with the resource profile.
>>> rdd = sc.parallelize(range(10)).withResources(resource_profile)
>>> rdd.getResourceProfile()
<pyspark.resource.profile.ResourceProfile object ...>
>>> rdd.getResourceProfile().taskResources
{'cpus': <...TaskResourceRequest...>, 'gpu': <...TaskResourceRequest...>}
>>> rdd.getResourceProfile().executorResources
{'gpu': <...ExecutorResourceRequest...>,
'cores': <...ExecutorResourceRequest...>,
'offHeap': <...ExecutorResourceRequest...>,
'memoryOverhead': <...ExecutorResourceRequest...>,
'pyspark.memory': <...ExecutorResourceRequest...>,
'memory': <...ExecutorResourceRequest...>}
"""
@overload
def __init__(self, _java_resource_profile: JavaObject):
...
@overload
def __init__(
self,
_java_resource_profile: None = ...,
_exec_req: Optional[Dict[str, ExecutorResourceRequest]] = ...,
_task_req: Optional[Dict[str, TaskResourceRequest]] = ...,
):
...
def __init__(
self,
_java_resource_profile: Optional[JavaObject] = None,
_exec_req: Optional[Dict[str, ExecutorResourceRequest]] = None,
_task_req: Optional[Dict[str, TaskResourceRequest]] = 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) -> int:
"""
Returns
-------
int
A unique id of this :class:`ResourceProfile`
"""
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) -> Dict[str, TaskResourceRequest]:
"""
Returns
-------
dict
a dictionary of resources to :class:`TaskResourceRequest`
"""
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) -> Dict[str, ExecutorResourceRequest]:
"""
Returns
-------
dict
a dictionary of resources to :class:`ExecutorResourceRequest`
"""
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
[docs]class ResourceProfileBuilder:
"""
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
See Also
--------
:class:`pyspark.resource.ResourceProfile`
Notes
-----
This API is evolving.
"""
def __init__(self) -> None:
from pyspark.context import SparkContext
# TODO: ignore[attr-defined] will be removed, once SparkContext is inlined
_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: Dict[str, ExecutorResourceRequest] = {}
self._task_resource_requests: Dict[str, TaskResourceRequest] = {}
def require(
self, resourceRequest: Union[ExecutorResourceRequests, TaskResourceRequests]
) -> "ResourceProfileBuilder":
"""
Add executor resource requests
Parameters
----------
resourceRequest : :class:`ExecutorResourceRequests` or :class:`TaskResourceRequests`
The detailed executor resource requests, see :class:`ExecutorResourceRequests`
Returns
-------
dict
a dictionary of resources to :class:`ExecutorResourceRequest`
"""
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:
r = resourceRequest._java_executor_resource_requests
if r is not None:
self._java_resource_profile_builder.require(r)
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) -> None:
if self._java_resource_profile_builder is not None:
self._java_resource_profile_builder.clearExecutorResourceRequests()
else:
self._executor_resource_requests = {}
def clearTaskResourceRequests(self) -> None:
if self._java_resource_profile_builder is not None:
self._java_resource_profile_builder.clearTaskResourceRequests()
else:
self._task_resource_requests = {}
@property
def taskResources(self) -> Dict[str, TaskResourceRequest]:
"""
Returns
-------
dict
a dictionary of resources to :class:`TaskResourceRequest`
"""
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) -> Dict[str, ExecutorResourceRequest]:
"""
Returns
-------
dict
a dictionary of resources to :class:`ExecutorResourceRequest`
"""
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) -> ResourceProfile:
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
)
def _test() -> None:
import doctest
import sys
from pyspark import SparkContext
globs = globals().copy()
globs["sc"] = SparkContext("local[4]", "profile tests")
(failure_count, test_count) = doctest.testmod(
globs=globs, optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE
)
globs["sc"].stop()
if failure_count:
sys.exit(-1)
if __name__ == "__main__":
_test()