tni%vdZddlZddlmZmZmZddlmZmZddl m Z GddZ Gdd Z dS) z,Module for running endpoints asynchronously.N)AnyDictOptional) FINAL_STATES is_completed) ClientSessionceZdZdZdedededefdZdded eee ffd Z d efd Z d Z dde d e fdZd e fdZd efdZdS)Jobz5Class representing a job for an asynchronous endpoint endpoint_idjob_idsessionheadersctddlm}||_||_||_||_||_d|_d|_dS)a Initialize a Job instance. Args: endpoint_id: The identifier for the endpoint. job_id: The identifier for the job. session: The aiohttp ClientSession. headers: Headers to use for requests. r)endpoint_url_baseN)runpodrr r r r job_status job_output)selfr r r rrs c/var/www/html/gpu-tools/venv/lib/python3.11/site-packages/runpod/endpoint/asyncio/asyncio_runner.py__init__z Job.__init__sY      '  !2 statussourcereturnc>K|jd|jd|d|j}|j||jd{V}|d{V}t|dr(|d|_|dd|_ |S)zReturns the raw json of the status, raises an exception if invalid. Args: source: The URL source path of the job status. /rNroutput) rr r r getrjsonrrr)rr status_url job_states r _fetch_jobzJob._fetch_job&s% Q Q(8 Q Q6 Q QDK Q Q ,**:t|*LLLLLLLL #..********  (+ , , <'1DO'mmHd;;DOrcfK|j|jS|d{V}|dS)zYGets jobs' status Returns: COMPLETED, FAILED or IN_PROGRESS Nr)rr#)rr"s rrz Job.status8sD ? &? "//++++++++ ""rcKt|d{VsCtjdd{Vt|d{VAdSdS)N)rrasynciosleep)rs r_wait_for_completionzJob._wait_for_completionDsT[[]]22222233 #-"" " " " " " " "T[[]]22222233 # # # # #rrtimeoutc.K|j|jS tj||d{Vn'#tj$r}t d|d}~wwxYw|d{V}|ddS)zWaits for serverless API job to complete or fail Returns: Output of job Raises: KeyError if job Failed NzJob timed out.r)rr'wait_forr) TimeoutErrorr#r)rr*excjob_datas rrz Job.outputHs ? &? " :"4#<#<#>#>HH H H H H H H H H# : : :/00c 9 :********||Hd+++s-AA$AA$cNK tjdd{V|dd{V}|dtvs't |dgdkr%|dgD] }|dWVn|dtvrdS) z>Returns a generator that yields the output of the job request.Tr&Nstream)rrrr)r'r(r#rlenr)rstream_partialchunks rr1z Job.stream[s -"" " " " " " " "#'??(?#C#CCCCCCCNx( <<~))(B77881<<+//"==**E/)))))*)\99 rcK|jd|jd|j}|j||j4d{V}|d{Vcdddd{VS#1d{VswxYwYdS)zUCancels current job Returns: Output of cancel operation rz/cancel/rN)rr r r postrr )r cancel_urlresps rcancelz Job.cancelis  .XX1AXX4;XX <$$Z$FF % % % % % % %$$$$$$$ % % % % % % % % % % % % % % % % % % % % % % % % % % % % % %sA00 A:=A:N)r)r)__name__ __module__ __qualname____doc__strrdictrrrr#rr)intrr1r9rrr r s??C}W[.s$sCx.$ #c # # # ####,,C,,,,,& c    %d%%%%%%rr c^eZdZdZ d dededeefdZdede fd Z defd Z defd Z dS) EndpointzClass for running endpointNr r api_keycddlm}m}||_||_||_|d|jd|_|p||_|jt ddd|jd |_dS) z Initialize an async Endpoint instance. Args: endpoint_id: The identifier for the endpoint. session: The aiohttp ClientSession. api_key: Optional API key for this endpoint instance. r)rDrrz/runNz(API key must be provided or set globallyzapplication/jsonzBearer )z Content-Type Authorization)rrDrr r endpoint_url RuntimeErrorr)rr r rDglobal_api_keyrs rrzEndpoint.__init__ws        ' !20II43CIII0. < IJJ J/5t|55   rendpoint_inputrcnK|j|j|jd|i4d{V }|d{V}dddd{Vn#1d{VswxYwY|j}|d|d<t |j|d|j|S)z Runs endpoint with specified input. Args: endpoint_input: any dictionary with input Returns: Newly created job input)rr Nidz X-Request-ID)r r6rGrr copyr r )rrJr8 json_resp job_headerss rrunz Endpoint.runsG<$$  t|7N:S%   * * * * * * * * "iikk))))))I * * * * * * * * * * * * * * * * * * * * * * * * * * * l'')) &/o N#4#Yt_dlKPPPsA  A*-A*cK|jd|jd}|j||j4d{V}|d{Vcdddd{VS#1d{VswxYwYdS)z\ Checks health of endpoint Returns: Health of endpoint rz/healthrN)rr r rrr )r health_urlr8s rhealthzEndpoint.healths.JJ1AJJJ <##J #EE % % % % % % %$$$$$$ % % % % % % % % % % % % % % % % % % % % % % % % % % % % % %A)) A36A3cK|jd|jd}|j||j4d{V}|d{Vcdddd{VS#1d{VswxYwYdS)zU Purges queue of endpoint Returns: Purge status rz /purge-queuerN)rr r r6rr )r purge_urlr8s r purge_queuezEndpoint.purge_queues-NN0@NNN <$$Y $EE % % % % % % %$$$$$$ % % % % % % % % % % % % % % % % % % % % % % % % % % % % % %rU)N) r:r;r<r=r>rrrr?r rQrTrXrArrrCrCts$$+/  C - "3-    >QQQQQQ( %d % % % % %4 % % % % % %rrC) r=r'typingrrrrunpod.endpoint.helpersrrrunpod.http_clientrr rCrArrr\s22&&&&&&&&&&>>>>>>>>,,,,,,e%e%e%e%e%e%e%e%PL%L%L%L%L%L%L%L%L%L%r