o
    Y۷i                     @   s   d dl Z d dlZd dlZd dlZd dlZd dlmZ d dlmZ d dl	m
Z
 G dd dejZe Zdd Zd	d
 Zdd Zdd Zdd Zdd Zdd Ze
jrhe dZee j e  ge_eZeZdS eZeZdS )    N)wraps)chain)configc                   @   s   e Zd ZdZdd ZdS )TLSzGUse a subclass to properly initialize the TLS variables in all threads.c                 C   s   d| _ d| _d S )NFr   )tracingindent)self r	   H/home/ubuntu/vllm_env/lib/python3.10/site-packages/numba/core/tracing.py__init__   s   
zTLS.__init__N)__name__
__module____qualname____doc__r   r	   r	   r	   r
   r   
   s    r   c                 C   s   t | dd}t | dd}t | dd}d}|r|j}n$t|jr-|jd dkr-|d jj}nt|jr>|jd dkr>|d j}|rhg }|rR|dkrR|| |d	 |r^|| |d	 || d
|}|dfS )z7Return function meta-data in a tuple.

    (name, type)r   Nr   __self__r   r   cls__main__. )getattrr   lenargs	__class__appendjoin)funcspecr   modulenamer   cnameqnamer	   r	   r
   find_function_info   s*   






r!   c                 C   s4   d}t | }t||kr|d | d |d  S |S )Ni@  z...)reprr   )valueMAX_SIZEsr	   r	   r
   chop1   s
   r'   c           
         s   t   |jrt t|jt|j d  |j  |  tt|jd t| |  fdd|jD }dd |t|d  D } fddt  D }d	dd t
|||D }dtjd | d	|d
g}dtjd | g}	||	fS )Nc                    s   g | ]}d |  |f qS z%s=%r)pop.0avaluesr	   r
   
<listcomp>@   s    z!create_events.<locals>.<listcomp>c                 S   s   g | ]}t |qS r	   )strr*   r	   r	   r
   r/   A       c                    s   g | ]
}d | | f qS r(   r	   )r+   kr-   r	   r
   r/   B   s    , c                 S   s   g | ]}|r|qS r	   r	   )r+   fr	   r	   r
   r/   C   r1   z>>  ()z<< )dictdefaultszipr   r   updatelistsortedkeysr   r   tlsr   )
fnamer   r   kwds
positional	anonymouskeywordsparamsenterleaver	   r-   r
   create_events9   s   "
"rH   c                  O   s   | dd}dd }t| r| d pd}|rtt
|s$t|ttfv r(||S t|tkrQd	\}}	}
|jr;||j}|jrC||j}	|jrK||j}
t||	|
S |S )
zFunction decorator to trace a function's entry and exit.

    *args: categories in which to trace this function. Example usage:

    @trace
    def function(...):...

    @trace('mycategory')
    def function(...):...


    	recursiveFc                    s   d t d fdd}d }dd }t tkr&t } dj nt tkr6t } d nt tkr>tt	
 |t |S )Ntracec                     s   tjr	tjr | i |S t | \}}zdt_t|| |\}}zd| t j	d7  _	zzdt_ | i |}W dt_ndt_w W nH   t
 \}}}	|d |j}
|
dkrl||
 |d ||j |jr|d |d	d
d |jD  |d  |d  |d ur|d |t| W t j	d8  _	d| nt j	d8  _	d| w W dt_|S dt_w )NTr      Fz => exception thrown
	raise r   r   r6   r3   c                 s   s    | ]}t |V  qd S N)r'   )r+   vr	   r	   r
   	<genexpr>z   s    z>dotrace.<locals>.decorator.<locals>.wrapper.<locals>.<genexpr>r7   z()z -> )isEnabledForloggingINFOr?   r   r!   rH   infor   r   sysexc_infor   r   r   r   r'   )r   rA   r@   ftyperF   rG   resulttyper$   	tracebackmnamer   loggerr   r	   r
   wrapper]   sP   






z+dotrace.<locals>.decorator.<locals>.wrapperc                 S      | S rL   r	   xr	   r	   r
   <lambda>   s    z,dotrace.<locals>.decorator.<locals>.<lambda>T)rP   	getLoggerrW   classmethod__get____func__staticmethodpropertyNotImplementedErrorinspectgetfullargspecr   )r   r\   rV   rewrapr	   rZ   r
   	decoratorY   s   
.
zdotrace.<locals>.decoratorr   Nc                 S   s   t | p	t | S rL   )rh   
isfunctionismethodr^   r	   r	   r
   r`      s   
 
zdotrace.<locals>.<lambda>)NNN)getr   rg   rh   ismodule
getmembersrl   setattrisclassdotracecallablerW   rb   re   rf   fgetfsetfdel)r   rA   rI   rk   arg0nr4   cpgetpsetpdelr	   r	   r
   rs   J   s"   D



rs   c                  O   s@   dd }t | r| d pd}t|st|ttfv r||S |S )z)Just a no-op in case tracing is disabled.c                 S   r]   rL   r	   )r   r	   r	   r
   rk         znotrace.<locals>.decoratorr   N)r   rt   rW   rb   re   )r   rA   rk   rx   r	   r	   r
   notrace   s
   r   c                 C   s.   dt jd | g} td}|d|  d S )Nz== r5   rJ   r   )r?   r   rP   ra   rR   r   )msgr[   r	   r	   r
   doevent   s   
r   c                 C   s   d S rL   r	   )r   r	   r	   r
   noevent   r~   r   rJ   )rP   rS   types	threadingrh   	functoolsr   	itertoolsr   
numba.corer   localr   r?   r!   r'   rH   rs   r   r   r   TRACEra   r[   setLevelrQ   StreamHandlerhandlersrJ   eventr	   r	   r	   r
   <module>   s2    s
