## 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.#fromtypingimportoverload,Dict,Union,Optionalfrompy4j.java_gatewayimportJavaObjectfrompyspark.resource.requestsimport(TaskResourceRequest,TaskResourceRequests,ExecutorResourceRequests,ExecutorResourceRequest,)
[docs]classResourceProfile:""" 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...>} """@overloaddef__init__(self,_java_resource_profile:JavaObject):...@overloaddef__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_profileisnotNone:self._java_resource_profile=_java_resource_profileelse:self._java_resource_profile=Noneself._executor_resource_requests=_exec_reqor{}self._task_resource_requests=_task_reqor{}@propertydefid(self)->int:""" Returns ------- int A unique id of this :class:`ResourceProfile` """ifself._java_resource_profileisnotNone:returnself._java_resource_profile.id()else:raiseRuntimeError("SparkContext must be created to get the id, get the id ""after adding the ResourceProfile to an RDD")@propertydeftaskResources(self)->Dict[str,TaskResourceRequest]:""" Returns ------- dict a dictionary of resources to :class:`TaskResourceRequest` """ifself._java_resource_profileisnotNone:taskRes=self._java_resource_profile.taskResourcesJMap()result={}fork,vintaskRes.items():result[k]=TaskResourceRequest(v.resourceName(),v.amount())returnresultelse:returnself._task_resource_requests@propertydefexecutorResources(self)->Dict[str,ExecutorResourceRequest]:""" Returns ------- dict a dictionary of resources to :class:`ExecutorResourceRequest` """ifself._java_resource_profileisnotNone:execRes=self._java_resource_profile.executorResourcesJMap()result={}fork,vinexecRes.items():result[k]=ExecutorResourceRequest(v.resourceName(),v.amount(),v.discoveryScript(),v.vendor())returnresultelse:returnself._executor_resource_requests
[docs]classResourceProfileBuilder:""" 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:frompyspark.contextimportSparkContext# TODO: ignore[attr-defined] will be removed, once SparkContext is inlined_jvm=SparkContext._jvmif_jvmisnotNone:self._jvm=_jvmself._java_resource_profile_builder=(_jvm.org.apache.spark.resource.ResourceProfileBuilder())else:self._jvm=Noneself._java_resource_profile_builder=Noneself._executor_resource_requests:Dict[str,ExecutorResourceRequest]={}self._task_resource_requests:Dict[str,TaskResourceRequest]={}defrequire(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` """ifisinstance(resourceRequest,TaskResourceRequests):ifself._java_resource_profile_builderisnotNone:ifresourceRequest._java_task_resource_requestsisnotNone: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:ifself._java_resource_profile_builderisnotNone:r=resourceRequest._java_executor_resource_requestsifrisnotNone: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)returnselfdefclearExecutorResourceRequests(self)->None:ifself._java_resource_profile_builderisnotNone:self._java_resource_profile_builder.clearExecutorResourceRequests()else:self._executor_resource_requests={}defclearTaskResourceRequests(self)->None:ifself._java_resource_profile_builderisnotNone:self._java_resource_profile_builder.clearTaskResourceRequests()else:self._task_resource_requests={}@propertydeftaskResources(self)->Dict[str,TaskResourceRequest]:""" Returns ------- dict a dictionary of resources to :class:`TaskResourceRequest` """ifself._java_resource_profile_builderisnotNone:taskRes=self._java_resource_profile_builder.taskResourcesJMap()result={}fork,vintaskRes.items():result[k]=TaskResourceRequest(v.resourceName(),v.amount())returnresultelse:returnself._task_resource_requests@propertydefexecutorResources(self)->Dict[str,ExecutorResourceRequest]:""" Returns ------- dict a dictionary of resources to :class:`ExecutorResourceRequest` """ifself._java_resource_profile_builderisnotNone:result={}execRes=self._java_resource_profile_builder.executorResourcesJMap()fork,vinexecRes.items():result[k]=ExecutorResourceRequest(v.resourceName(),v.amount(),v.discoveryScript(),v.vendor())returnresultelse:returnself._executor_resource_requests@propertydefbuild(self)->ResourceProfile:ifself._java_resource_profile_builderisnotNone:jresourceProfile=self._java_resource_profile_builder.build()returnResourceProfile(_java_resource_profile=jresourceProfile)else:returnResourceProfile(_exec_req=self._executor_resource_requests,_task_req=self._task_resource_requests)