o
    )i                     @   s   d dl mZ d dlmZmZ d dl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 d dlmZmZ d d	lmZ eg df ZG d
d deZG dd deeZG dd deeZdS )    )Future)CallableUnionN)
VllmConfig)ExecutorBase)ExecutorWithExternalLauncher)UniProcExecutor)KVCacheConfigKVCacheSpec)ModelRunnerOutputc                   @   s   e Zd ZdZededed  fddZdee	 ddfdd	Z
d
efddZdee fddZdeeeef  fddZdeeee f fddZedefddZddefddZdS )Executorz
    Abstract class for v1 executors, mainly define some methods for v1.
    For methods shared by v0 and v1, define them in ExecutorBasevllm_configreturnc                 C   s   | j }|j}t|trt|tstd| d|}|S |dkr*ddlm} |}|S |dkr8ddl	m
} |}|S |dkr@t}|S |d	krHt}|S td
| )NzEdistributed_executor_backend must be a subclass of ExecutorBase. Got .Zrayr   )RayDistributedExecutormp)MultiprocExecutoruniZexternal_launcherz&Unknown distributed executor backend: )parallel_configdistributed_executor_backend
isinstancetype
issubclassr   	TypeErrorZ)vllm.v1.executor.ray_distributed_executorr   Z#vllm.v1.executor.multiproc_executorr   r   r   
ValueError)r   r   r   Zexecutor_classr   r    r   e/home/app/PaddleOCR-VL-test/.venv_paddleocr/lib/python3.10/site-packages/vllm/v1/executor/abstract.py	get_class   s:   


zExecutor.get_classkv_cache_configsNc                 C   s   | j d|fd |  d dS )zp
        Initialize the KV caches and begin the model execution loop of the
        underlying workers.
        initialize_from_configargsZcompile_or_warm_up_modelNZcollective_rpc)selfr   r   r   r   r   :   s   zExecutor.initialize_from_configcallbackc                 C      dS )zk
        Register a function to be called if the executor enters a permanent
        failed state.
        Nr   )r#   r$   r   r   r   register_failure_callbackD   s   z"Executor.register_failure_callbackc                 C      |  d}|S )Ndetermine_available_memoryr"   r#   outputr   r   r   r(   K      
z#Executor.determine_available_memoryc                 C   r'   )NZget_kv_cache_specr"   r)   r   r   r   get_kv_cache_specsO   r+   zExecutor.get_kv_cache_specsc                 C   s   | j d|fd}|d S )Nexecute_modelr    r   r"   )r#   Zscheduler_outputr*   r   r   r   r-   S   s   zExecutor.execute_modelc                 C   r%   )N   r   )r#   r   r   r   max_concurrent_batches[   s   zExecutor.max_concurrent_batchesTis_startc                 C   s   | j d|fd d S )Nprofiler    r"   )r#   r0   r   r   r   r1   _   s   zExecutor.profile)T)__name__
__module____qualname____doc__staticmethodr   r   r   listr	   r   FailureCallbackr&   intr(   dictstrr
   r,   r   r   r   r-   propertyr/   boolr1   r   r   r   r   r      s$    


r   c                   @   s   e Zd ZdS )r   N)r2   r3   r4   r   r   r   r   r   c   s    r   c                       s&   e Zd Zdee f fddZ  ZS )r   r   c                    sP   t   }ddlm} | j}tj|gdtjd}tj	||tj
jd | gS )Nr   )get_world_groupcpu)ZdeviceZdtype)groupop)superr(   Zvllm.distributed.parallel_stater>   	cpu_grouptorchZtensorZint64distZ
all_reduceZReduceOpZMINitem)r#   Zmemoryr>   rC   Zmemory_tensor	__class__r   r   r(   i   s   

z7ExecutorWithExternalLauncher.determine_available_memory)r2   r3   r4   r7   r9   r(   __classcell__r   r   rG   r   r   g   s    r   )concurrent.futuresr   typingr   r   rD   Ztorch.distributeddistributedrE   Zvllm.configr   Zvllm.executor.executor_baser   Zvllm.executor.uniproc_executorr   ZExecutorWithExternalLauncherV0r   ZUniProcExecutorV0Zvllm.v1.kv_cache_interfacer	   r
   Zvllm.v1.outputsr   r8   r   r   r   r   r   <module>   s   M