ԋg(ddlmZddlZddlZddlZddlmZddlmZddl m Z ddl m Z m Z ddlmZej eZGdd ZGd d Zy) ) annotationsN) defaultdict)parse_timedelta)Future)Deadlinewait_for) get_clientcFeZdZdZdZd dZd dZd dZd dZddZ d d Z y)QueueExtensionzAn extension for the scheduler to manage queues This adds the following routes to the scheduler * queue_create * queue_release * queue_put * queue_get * queue_size c||_t|_t|_t t |_|jjj|j|j|j|jd|jjj|j|jdy)N) queue_create queue_put queue_get queue_qsize)queue-future-release queue_release) schedulerdictqueuesclient_refcountrintfuture_refcounthandlersupdatecreateputgetqsizestream_handlersfuture_releaserelease)selfrs 2lib/python3.12/site-packages/distributed/queues.py__init__zQueueExtension.__init__s"f #v*3/ && $ !XX!XX#zz    &&--%)%8%84<< X Nctjd|||jvr3tj||j|<d|j |<y|j |xxdz cc<y)Nz Queue name: )maxsize)loggerdebugrasyncioQueuerr"nameclientr's r#rzQueueExtension.create0s^ |D6*+ t{{ " ' g >DKK )*D  &   &! + &r%cb||jvry|j|xxdzcc<|j|dk(rr|j|=|j|j}|j|=|Dcgc]}|ddk(s |d}}|r!|jj |d|zyyycc}w)Nr(rtypervaluequeue-%skeysr/)rr_queuerclient_releases_keys)r"r.r/futuresdr5s r#r!zQueueExtension.release8s t{{ "  T"a'"    % *$$T*kk$'..G D!(/I11V93HAgJDI33jSWFW3X +Js 0 B,>B,cKtj|}|||jjvrQt j dd{|j rtd|d||jjvrQd|d}|j||fxxdz cc<|jj|gd|znd |d}t|j|j||j d{y77w) Ng{Gz?zTask z unknown to scheduler.r)r1r2r(r3r4msgpacktimeout)rafterrtasksr+sleepexpired TimeoutErrorrclient_desires_keysrrr remaining)r"r.keydatar/r=deadlinerecords r#rzQueueExtension.putEs>>'* ?T^^111mmD)))##&se3I'JKKT^^111 '5F  s +q 0 + NN . .SE*tBS . T'$7Ft{{4(,,V4h>P>PQQQ* Rs+AD D 7DA;D>D?DDc|jj|g||j||fxxdzcc<|j||fdk(r1|jj|gd|z|j||f=yy)Nr4r(rr3)rrCrr7)r"r.rEr/s r#r zQueueExtension.future_releaseTs **f*E T3Y'1,'   c *a / NN / /cU:PTCT / U$$T3Y/ 0r%c"Kfd}|rj|}g}|durK|js|jd{}|j||js:nH| d} t | t |D]+} |jd{}|j|-|D cgc] } ||  }} |St j|j|d{}||}|S77jcc} w7w)Nc$|ddk(r|j}|d}jjj|}| |jnd}||d<|dk(r2|j j |d<|j j|d<|S) zAdd task status if knownr1rr2loststateerred exception traceback)copyrr?rrMexception_blamerOrP)rHrEtsrMr"s r#processz#QueueExtension.get..process\sf~)Wo^^))--c2$&N"'wG#*,*<*<*F*FF;'*,*<*<*F*FF;'Mr%TzNDask queues don't support simultaneous use of integer batch sizes and timeoutsr<)remptyrappendNotImplementedErrorranger) r"r.r/r=batchrTqoutrHmsg_os ` r#rzQueueExtension.get[s   D!AC}'')#$557]FJJv&'')&;.c22uA#$557]FJJv&&(++s!71:sC+J#DKK$5$9$9$;WMMFV_FM#+++NsHADD%D*1DDD7D/D5D 6DDDc<|j|jSN)rr)r"r.r/s r#rzQueueExtension.qsizes{{4 &&((r%NNr)NN)NNNNN)NNN)NNNF) __name__ __module__ __qualname____doc__r$rr!rr rrr%r#r r s-  &, Y R0&P)r%r cxeZdZdZddZdZedZdZdZ dZ dd Z dd Z dd Z d Zdd ZdZdZdZy)r,a^Distributed Queue This allows multiple clients to share futures or small bits of data between each other with a multi-producer/multi-consumer queue. All metadata is sequentialized through the scheduler. Elements of the Queue must be either Futures or msgpack-encodable data (ints, strings, lists, dicts). All data is sent through the scheduler so it is wise not to send large objects. To share large objects scatter the data and share the future instead. .. warning:: This object is experimental Parameters ---------- name: string (optional) Name used by other clients and the scheduler to identify the queue. If not given, a random name will be generated. client: Client (optional) Client used for communication with the scheduler. If not given, the default global client will be used. maxsize: int (optional) Number of items allowed in the queue. If 0 (the default), the queue size is unbounded. Examples -------- >>> from dask.distributed import Client, Queue # doctest: +SKIP >>> client = Client() # doctest: +SKIP >>> queue = Queue('x') # doctest: +SKIP >>> future = client.submit(f, x) # doctest: +SKIP >>> queue.put(future) # doctest: +SKIP See Also -------- Variable: shared variable between clients Nc||_|xs!dtjjz|_||_|j y)Nzqueue-)_clientuuiduuid4hexr.r' _maybe_startr-s r#r$zQueue.__init__s9 7Htzz|'7'77   r%c|jre|jjr)tj|j |_y|jj |jyyr`)r/ asynchronousr+ ensure_future_start_startedsyncr"s r#rmzQueue._maybe_startsI ;;{{'' ' 5 5dkkm D    - r%c|js t|_|jS|jS#t$rY|jSwxYwr`)rir ValueErrorrts r#r/z Queue.clientsJ|| )| ||t|| || s5 A  A cJ|jstt|dy)Nz object not properly initialized. This can happen if the object is being deserialized outside of the context of a Client or Worker.)r/ RuntimeErrorr1rts r#_verify_runningzQueue._verify_runnings.{{:,'' r%cK|jjj|j|jd{|S7w)N)r.r')r/rr r.r'rts r#rqz Queue._starts<kk##00dii0VVV  Ws?A AA cjtdrjjSfd}|jS)NrrcKSwr`rfrtsr#r]zQueue.__await__.._s  s)rmhasattrrr __await__)r"r]s` r#r~zQueue.__await__sA  4 $==**, , 3==? "r%c2Kt|trE|jjj |j ||j d{y|jjj |||j d{y7@7w)N)rEr=r.)rFr=r.) isinstancerr/rrrEr.)r"r2r=s r#_putz Queue._puts eV $++''11IIwTYY2  ++''11G$))2     s$ABB:B BBBc |jt|}|jj|j|fd|i|S)a9Put data into the queue Parameters ---------- timeout : number or string or timedelta, optional Time in seconds to wait before timing out. Instead of number of seconds, it is also possible to specify a timedelta in string format, e.g. "200ms". r=)ryrr/rsr)r"r2r=kwargss r#rz Queue.putsA !'*t{{ 5L'LVLLr%c |jt|}|jj|jf||d|S)a4Get data from the queue Parameters ---------- timeout : number or string or timedelta, optional Time in seconds to wait before timing out. Instead of number of seconds, it is also possible to specify a timedelta in string format, e.g. "200ms". batch : boolean, int (optional) If True then return all elements currently waiting in the queue. If an integer than return that many elements from the queue If False (default) then return one item at a time )r=rY)ryrr/rs_get)r"r=rYrs r#rz Queue.getsA !'*t{{ R7%R6RRr%c p|j|jj|jfi|S)z'Current number of elements in the queue)ryr/rs_qsize)r"rs r#rz Queue.qsizes. t{{ 6v66r%cKjjj|j|d{}fd}|dur ||}|St t ||}|S7.w)N)r=r.rYc|ddk(rvt|dj|d}|ddk(r"|jj|d|djj d j |dd |S|d}|S) Nr1rr2rM)rMrNrOrPr)opr.rE)rr/_state set_error_send_to_schedulerr.)r9r2r"s r#rTzQueue._get..processsyH$qz4;;ajIW:(LL**1[>1[>J ..1499QwZX L' Lr%F)r/rrr.listmap)r"r=rYresprTresults` r#rz Queue._get sq[[**44$))55    E>T]F #gt,-F - s6A*A(/A*cK|jjj|jd{}|S7w)N)r.)r/rrr.)r"rs r#rz Queue._qsize%s5{{,,88dii8HH Is 4?=?c|j|jjdk(r)|jjd|jdyy)Nrunningr)rr.)ryr/statusrr.rts r#closez Queue.close)s@  ;;   * KK * */499+U V +r%cJt||jd|jffSr`)r1r.r'rts r# __reduce__zQueue.__reduce__.s DzDIItT\\:::r%rar`)NF)rbrcrdrer$rmpropertyr/ryrqr~rrrrrrrrrfr%r#r,r,s`&P . # MS$7 2W ;r%r,) __future__rr+loggingrj collectionsr dask.utilsrdistributed.clientrdistributed.utilsrrdistributed.workerr getLoggerrbr)r r,rfr%r#rsK" #&%0)   8 $s)s)lh;h;r%