o $iP$@s ddlZddlmZmZmZmZmZmZddlm Z ddl m Z m Z ddl mZmZeeZedZedZeeee e gee feee e e gee fedee ffZeGd d d ZeGd d d eZeGd ddeZdeeefdefddZdS)N)AnyCallableIterableOptionalTypeVarUnion) TaskContext)BlockUserDefinedFunction) DeveloperAPI PublicAPITU.c@s eZdZdS)ComputeStrategyN)__name__ __module__ __qualname__rrW/home/ubuntu/veenaModal/venv/lib/python3.10/site-packages/ray/data/_internal/compute.pyrsrc@sFeZdZdZ d deefddZdedefdd Z de fd d Z dS) TaskPoolStrategya;Specify the task-based compute strategy for a Dataset transform. TaskPoolStrategy executes dataset transformations using Ray tasks that are scheduled through a pool. Provide ``size`` to cap the number of concurrent tasks; leave it unset to allow Ray Data to scale the task count automatically. NsizecCs$|dur |dkr td|||_dS)zConstruct TaskPoolStrategy for a Dataset transform. Args: size: Specify the maximum size of the task pool. Nz`size` must be >= 1) ValueErrorr)selfrrrr__init__)s  zTaskPoolStrategy.__init__otherreturncCs(t|tr |j|jkp|dko|jduS)Ntasks) isinstancerrrrrrr__eq__7szTaskPoolStrategy.__eq__cCsd|jdS)NzTaskPoolStrategy(size=))rrrrr__repr__<szTaskPoolStrategy.__repr__N) rrr__doc__rintrrboolr strr#rrrrrs  rc @sveZdZdZddddddddeedeedeedeed eed ef d d Zd edefddZ de fddZ dS)ActorPoolStrategyae Specify the actor-based compute strategy for a Dataset transform. ActorPoolStrategy specifies that an autoscaling pool of actors should be used for a given Dataset transform. This is useful for stateful setup of callable classes. For a fixed-sized pool of size ``n``, use ``ActorPoolStrategy(size=n)``. To autoscale from ``m`` to ``n`` actors, use ``ActorPoolStrategy(min_size=m, max_size=n)``. To autoscale from ``m`` to ``n`` actors, with an initial size of ``initial``, use ``ActorPoolStrategy(min_size=m, max_size=n, initial_size=initial)``. To increase opportunities for pipelining task dependency prefetching with computation and avoiding actor startup delays, set max_tasks_in_flight_per_actor to 2 or greater; to try to decrease the delay due to queueing of tasks on the worker actors, set max_tasks_in_flight_per_actor to 1. The `enable_true_multi_threading` argument primarily exists to prevent GPU OOM issues with multi-threaded actors. The life cycle of an actor task involves 3 main steps: 1. Batching Inputs 2. Running actor UDF 3. Batching Outputs The `enable_true_multi_threading` flag affects step 2. If set to `True`, then the UDF can be run concurrently. By default, it is set to `False`, so at most 1 actor UDF is running at a time per actor. The `max_concurrency` flag on `ray.remote` affects steps 1 and 3. Below is a matrix summary: - [`enable_true_multi_threading=False or True`, `max_concurrency=1`] = 1 actor task running per actor. So at most 1 of steps 1, 2, or 3 is running at any point in time. - [`enable_true_multi_threading=False`, `max_concurrency>1`] = multiple tasks running per actor (respecting GIL) but UDF runs 1 at a time. This is useful for doing CPU and GPU work, where you want to use a large batch size but want to hide the overhead of *batching* the inputs. In this case, CPU *batching* is done concurrently, while GPU *inference* is done 1 at a time. Concretely, steps 1 and 3 can have multiple threads, while step 2 is done serially. - [`enable_true_multi_threading=True`, `max_concurrency>1`] = multiple tasks running per actor. Unlike bullet #3 ^, the UDF runs concurrently (respecting GIL). No restrictions on steps 1, 2, or 3 NOTE: `enable_true_multi_threading` does not apply to async actors NF)rmin_sizemax_size initial_sizemax_tasks_in_flight_per_actorenable_true_multi_threadingrr*r+r,r-r.cCs<|dur#|dkr td||dus|dus|durtd|}|}|}|dur0|dkr0td||durD|dur:d}||krDtd|||durQ|dkrQtd||pTd|_|p[td|_|dur||jkrrtd |d |jd |jtdkr||jkrtd |d |jd |p|j|_||_d |_d|_||_dS)aConstruct ActorPoolStrategy for a Dataset transform. Args: size: Specify a fixed size actor pool of this size. It is an error to specify both `size` and `min_size` or `max_size`. min_size: The minimum size of the actor pool. max_size: The maximum size of the actor pool. initial_size: The initial number of actors to start with. If not specified, defaults to min_size. Must be between min_size and max_size. max_tasks_in_flight_per_actor: The maximum number of tasks to concurrently send to a single actor worker. Increasing this will increase opportunities for pipelining task dependency prefetching with computation and avoiding actor startup delays, but will also increase queueing delay. enable_true_multi_threading: If enable_true_multi_threading=True, no more than 1 actor task runs per actor. Otherwise, respects the `max_concurrency` argument. Nrzsize must be >= 1zMmin_size, max_size, and initial_size cannot be set at the same time as `size`zmin_size must be >= 1zmin_size must be <= max_sizez1max_tasks_in_flight_per_actor must be >= 1, got: infzinitial_size (z) must be >= min_size (r!z) must be <= max_size (rg?) rr*floatr+r,r- num_workersready_to_total_workers_ratior.)rrr*r+r,r-r.rrrrmsN       zActorPoolStrategy.__init__rrcCsFt|to"|j|jko"|j|jko"|j|jko"|j|jko"|j|jkSr$)rr)r*r+r,r.r-rrrrr s     zActorPoolStrategy.__eq__cCs>d|jd|jd|jd|jd|jd|jd|jdS) NzActorPoolStrategy(min_size=z , max_size=z, initial_size=z , max_tasks_in_flight_per_actor=z )num_workers=z, enable_true_multi_threading=z, ready_to_total_workers_ratio=r!)r*r+r,r-r1r.r2r"rrrr#s zActorPoolStrategy.__repr__) rrrr%rr&r'rrr r(r#rrrrr)@s.. I r) compute_specrcCsTt|ttfstd|d|r|dkrtS|dkrtSt|tr&|Std)NzXIn Ray 2.5, the compute spec must be either TaskPoolStrategy or ActorPoolStrategy, was: .ractorsz;compute must be one of [`tasks`, `actors`, ComputeStrategy])rrr)rr)r3rrr get_computes  r6)loggingtypingrrrrrr'ray.data._internal.execution.interfacesrray.data.blockr r ray.util.annotationsr r getLoggerrloggerr rBlockTransformrrr)r(r6rrrrs,