o bi]!@sddlZddlZddlmZddlmZddlmZddlm Z m Z m Z ddl Z ddl mZddlmZe e jje jjfZde e eefd e eeffd d ZeGd d d ZeeGdddejZdS)N)deepcopy) dataclass) signature)DictListUnion)placement_group) DeveloperAPIbundlesreturncCs8i}|D]}|D]\}}||d|||<q q|S)zSum all resources in a list of resource bundles. Args: bundles: List of resource bundles. Returns: Dict containing all resources summed up. r)itemsget)r resourcesbundlekvrW/home/ubuntu/.local/lib/python3.10/site-packages/ray/air/execution/resources/request.py _sum_bundless rc@seZdZdZ d"deeeeee ffdefddZ e ddZ e e d e fd d Ze e d eeee ffd d Ze d eee ffddZe e d efddZddZddZd#ddZddZddZddZd efdd Zd!S)$ResourceRequestaRequest for resources. This class is used to define a resource request. A resource request comprises one or more bundles of resources and instructions on the scheduling behavior. The resource request can be submitted to a resource manager, which will schedule the resources. Depending on the resource backend, this may instruct Ray to scale up (autoscaling). Resource requests are compatible with the most fine-grained low-level resource backend, which are Ray placement groups. Args: bundles: A list of bundles which represent the resources requirements. E.g. ``[{"CPU": 1, "GPU": 1}]``. strategy: The scheduling strategy to acquire the bundles. - "PACK": Packs Bundles into as few nodes as possible. - "SPREAD": Places Bundles across distinct nodes as even as possible. - "STRICT_PACK": Packs Bundles into one node. The group is not allowed to span multiple nodes. - "STRICT_SPREAD": Packs Bundles across distinct nodes. *args: Passed to the call of ``placement_group()``, if applicable. **kwargs: Passed to the call of ``placement_group()``, if applicable. PACKr strategycOsx|stddd|D|_|jds$d|_|jd|js#tdnd|_||_||_||_d|_d|_| dS)Nz6Cannot initialize a ResourceRequest with zero bundles.cSsg|] }dd|DqS)cSs"i|] \}}|dkr|t|qS)r)float).0rrrrr Gs"z7ResourceRequest.__init__...)r )rrrrr Fsz,ResourceRequest.__init__..rTzOCannot initialize a ResourceRequest with an empty head and zero worker bundles.F) ValueError_bundles_head_bundle_is_emptypop _strategy_args_kwargs_hash_bound_bind)selfr rargskwargsrrr__init__;s(   zResourceRequest.__init__cC|jS)zReturns True if head bundle is empty while child bundles need resources. This is considered an internal API within Tune. )rr&rrrhead_bundle_is_emptybsz$ResourceRequest.head_bundle_is_emptyr cCs|jrdS|jdddS)z.Returns the number of cpus in the head bundle.grCPU)rrr r+rrr head_cpuskszResourceRequest.head_cpuscC t|jS)z'Returns a deep copy of resource bundles)rrr+rrrr qs zResourceRequest.bundlescCr/)z3Returns a dict containing the sums of all resources)rrr+rrrrequired_resourcesws z"ResourceRequest.required_resourcescCr*)zReturns the placement strategy)r r+rrrr|szResourceRequest.strategyc CsVtt}z|j|j|jg|jRi|j|_WdSty*}zt d|d}~ww)aBind the args and kwargs to the `placement_group()` signature. We bind the args and kwargs, so we can compare equality of two resource requests. The main reason for this is that the `placement_group()` API can evolve independently from the ResourceRequest API (e.g. adding new arguments). Then, `ResourceRequest(bundles, strategy, arg=arg)` should be the same as `ResourceRequest(bundles, strategy, arg)`. ztInvalid definition for resource request. Please check that you passed valid arguments to the ResourceRequest object.N) rrbindrr r!r"r$ Exception RuntimeError)r&sigexcrrrr%s" zResourceRequest._bindcCst|jji|jjSN)rr$r'r(r+rrrto_placement_groupsz"ResourceRequest.to_placement_groupothercCs"t|to|j|jko|j|jkSr6) isinstancerr$r,)r&r8rrr__eq__s   zResourceRequest.__eq__cCs2|jsttj|jj|jjddddd|_|jS)N)r'r(Tr) sort_keysindent ensure_ascii)r#hashjsondumpsr$r'r(r+rrr__hash__szResourceRequest.__hash__cCs&|j}|dd|dd|S)Nr#r$)__dict__copyrr&staterrr __getstate__s   zResourceRequest.__getstate__cCs$|j|d|_d|_|dSr6)rBupdater#r$r%rDrrr __setstate__s  zResourceRequest.__setstate__cCsd|jd|jdS)Nz)r$r,r+rrr__repr__s zResourceRequest.__repr__N)r)r8r)__name__ __module__ __qualname____doc__rrstrrintrr)propertyr,r r.r r0rr%r7r:rArFrHrIrrrrrs8 '   rc@sZeZdZUdZeed<deedeeefddZ dede e e fd e defd d Zd S) AcquiredResourcesaBase class for resources that have been acquired. Acquired resources can be associated to Ray objects, which can then be scheduled using these resources. Internally this can point e.g. to a placement group, a placement group bundle index, or just raw resources. The main API is the `annotate_remote_entities` method. This will associate remote Ray objects (tasks and actors) with the acquired resources by setting the Ray remote options to use the acquired resources. resource_requestentitiesr cCs|jj}t|t|jj}t||kr!tdt|d|dg}|jjr:||j|didd|dd}tt ||D]\}\}}||j|||dqA|S)aAReturn remote ray entities (tasks/actors) to use the acquired resources. The first entity will be associated with the first bundle, the second entity will be associated with the second bundle, etc. Args: entities: Remote Ray entities to annotate with the acquired resources. z%The number of callables to annotate (z1) cannot exceed the number of available bundles (z).r) bundle_indexN) rRr lenrOr,r3append_annotate_remote_entity enumeratezip)r&rSr num_bundles annotatedientityrrrrannotate_remote_entitiess&    z*AcquiredResources.annotate_remote_entitiesr^rrTcCstr6)NotImplementedError)r&r^rrTrrrrXsz)AcquiredResources._annotate_remote_entityN)rJrKrLrMr__annotations__rRemoteRayEntityrr_rrNrrOrXrrrrrQs"    ) rQ)abcr?rCr dataclassesrinspectrtypingrrrrayray.utilrray.util.annotationsr remote_functionRemoteFunctionactor ActorClassrbrNrrrABCrQrrrrs"     &#