python/pyspark/resource/profile.py - spark - Git at Google

 #
 # 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 threading import RLock
 from typing import overload, Dict, Union, Optional, TYPE_CHECKING

 from pyspark.resource.requests import (
     TaskResourceRequest,
     TaskResourceRequests,
     ExecutorResourceRequests,
     ExecutorResourceRequest,
 )

 if TYPE_CHECKING:
     from py4j.java_gateway import JavaObject


 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

     .. versionchanged:: 4.0.0
         Supports Spark Connect.

     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,
     ):
         # profile id
         self._id: Optional[int] = None
         # lock to protect _id
         self._lock = RLock()

         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`
         """
         with self._lock:
             if self._id is None:
                 if self._java_resource_profile is not None:
                     self._id = self._java_resource_profile.id()
                 else:
                     from pyspark.sql import is_remote

                     if is_remote():
                         from pyspark.sql.connect.resource.profile import ResourceProfile

                         # Utilize the connect ResourceProfile to create Spark ResourceProfile
                         # on the server and get the profile ID.
                         rp = ResourceProfile(
                             self._executor_resource_requests, self._task_resource_requests
                         )
                         self._id = rp.id
                     else:
                         raise RuntimeError("SparkContext must be created to get the profile id.")
             return self._id

     @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


 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.sql import is_remote

         _jvm = None
         if not is_remote():
             from pyspark.core.context import SparkContext

             _jvm = SparkContext._jvm

         if _jvm is not None:
             self._jvm = _jvm
             self._java_resource_profile_builder = getattr(
                 _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()
	#
	# 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 threading import RLock
	from typing import overload, Dict, Union, Optional, TYPE_CHECKING

	from pyspark.resource.requests import (
	TaskResourceRequest,
	TaskResourceRequests,
	ExecutorResourceRequests,
	ExecutorResourceRequest,
	)

	if TYPE_CHECKING:
	from py4j.java_gateway import JavaObject


	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

	.. versionchanged:: 4.0.0
	Supports Spark Connect.

	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,
	):
	# profile id
	self._id: Optional[int] = None
	# lock to protect _id
	self._lock = RLock()

	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`
	"""
	with self._lock:
	if self._id is None:
	if self._java_resource_profile is not None:
	self._id = self._java_resource_profile.id()
	else:
	from pyspark.sql import is_remote

	if is_remote():
	from pyspark.sql.connect.resource.profile import ResourceProfile

	# Utilize the connect ResourceProfile to create Spark ResourceProfile
	# on the server and get the profile ID.
	rp = ResourceProfile(
	self._executor_resource_requests, self._task_resource_requests
	)
	self._id = rp.id
	else:
	raise RuntimeError("SparkContext must be created to get the profile id.")
	return self._id

	@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


	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.sql import is_remote

	_jvm = None
	if not is_remote():
	from pyspark.core.context import SparkContext

	_jvm = SparkContext._jvm

	if _jvm is not None:
	self._jvm = _jvm
	self._java_resource_profile_builder = getattr(
	_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()