# # 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. """ @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: 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]: 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]: 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 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: Optional[Dict[str, ExecutorResourceRequest]] = {} self._task_resource_requests: Optional[Dict[str, TaskResourceRequest]] = {} def require( self, resourceRequest: Union[ExecutorResourceRequest, TaskResourceRequests] ) -> "ResourceProfileBuilder": 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( # type: ignore[union-attr] resourceRequest.requests ) else: if self._java_resource_profile_builder is not None: r = resourceRequest._java_executor_resource_requests # type: ignore[attr-defined] if r is not None: self._java_resource_profile_builder.require(r) else: execReqs = ExecutorResourceRequests( self._jvm, resourceRequest.requests # type: ignore[attr-defined] ) self._java_resource_profile_builder.require( execReqs._java_executor_resource_requests ) else: self._executor_resource_requests.update( # type: ignore[union-attr] resourceRequest.requests # type: ignore[attr-defined] ) 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) -> Optional[Dict[str, 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) -> Optional[Dict[str, 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 )