
    .`i                        d dl Z d dlmZ d dlmZmZmZmZmZ d dl	m
Z
 d dlmZ d dlmZ d dlmZ d dlmZmZ d d	lmZ  ee          Zd
edefdZ e            Ze                    d ee          gej        j        deiej        j        deiej        j        deiej         j        deii          defd            Z!e                    d          defd            Z"defdZ#dS )    N)
HTTPStatus)	APIRouterDependsFastAPIHTTPExceptionRequest)JSONResponse)EngineClient)ErrorResponse)validate_json_request)get_scaling_elastic_epset_scaling_elastic_ep)init_loggerrequestreturnc                 $    | j         j        j        S N)appstateengine_client)r   s    /home/jaya/work/projects/VOICE-AGENT/VIET/agent-env/lib/python3.11/site-packages/vllm/entrypoints/serve/elastic_ep/api_router.pyr   r      s    ;**    z/scale_elastic_epmodel)dependencies	responsesraw_requestc                 f  K   	 |                                   d {V }n)# t           j        $ r}t          dd          |d }~ww xY w|                    d          }|                    dd          }|t          dd          t	          |t
                    r|dk    rt          dd	          t	          |t
                    r|dk    rt          dd
          t          d           t          |           }	 |                    ||           d {V  t          dd| di          t          d           S # t          $ r}t          dd| d          |d }~wt          $ r2}t                              d|           t          dd          |d }~ww xY w# t          d           w xY w)Ni  zInvalid JSON format)status_codedetailnew_data_parallel_sizedrain_timeoutx   z"new_data_parallel_size is requiredr   z1new_data_parallel_size must be a positive integerz(drain_timeout must be a positive integerTmessagez
Scaled to z data parallel enginesFi  z0Scale failed due to request drain timeout after z secondszScale failed: %si  zScale failed)jsonJSONDecodeErrorr   get
isinstanceintr   r   scale_elastic_epr	   TimeoutError	Exceptionloggererror)r   bodyer    r!   clients         r   r)   r)       sh     R %%'''''''' R R R4IJJJPQQR "XX&>??HH_c22M%$H
 
 
 	
 ,c22 
6LPQ6Q6QF
 
 
 	

 mS)) 
]a-?-?$N
 
 
 	

 4   ;''F&%%&<mLLLLLLLLLV(>VVV
 
 	u%%%%    -"- - -
 
 
 		
  K K K'+++NCCCJK 	u%%%%sD    AA  A=0D= =
FEF*-FFF F0z/is_scaling_elastic_epc                 @   K   t          dt                      i          S )Nis_scaling_elastic_ep)r	   r   )r   s    r   r2   r2   Z   s!      02H2J2JKLLLr   r   c                 :    |                      t                     d S r   )include_routerrouter)r   s    r   attach_routerr6   _   s    vr   )$r$   httpr   fastapir   r   r   r   r   fastapi.responsesr	   vllm.engine.protocolr
   'vllm.entrypoints.openai.engine.protocolr   vllm.entrypoints.openai.utilsr   ,vllm.entrypoints.serve.elastic_ep.middlewarer   r   vllm.loggerr   __name__r,   r   r5   postOKvaluedictBAD_REQUESTREQUEST_TIMEOUTINTERNAL_SERVER_ERRORr)   r2   r6    r   r   <module>rH      s  
        G G G G G G G G G G G G G G * * * * * * - - - - - -      @ ? ? ? ? ?        $ # # # # #	X		+7 +| + + + + 
 '/001gt_$w&>"(7M*B(.-0H	  	 	-& -& -& -&	 	-&` %&&MW M M M '&Mw      r   