o
    پi                     @   s   d dl Z d dlZd dlZd dlmZmZ d dlmZ ee	Z
G dd deZ	ddee dee defd	d
Z	ddedededeee  fddZdS )    N)ListOptional)	BaseModelc                   @   s   e Zd ZU dZeed< eed< eed< eed< eed< eed< eed< eed	< eed
< eed< eed< dZe	e ed< dZ
e	e ed< dZe	e ed< dZe	ee  ed< edefddZ	ddededefddZdS )BenchmarkResultzIPydantic model for benchmark results table data, for a single isl and osl
model_pathrun_name
batch_size	input_len
output_lenlatencyinput_throughputoutput_throughputoverall_throughput	last_ttftlast_gen_throughputN
acc_lengthprofile_link_extendprofile_link_decodeserver_argsreturnc                   C   s   dS )Nzi
Note: To view the traces through perfetto-ui, please:
    1. open with Google Chrome
    2. allow popup
 r   r   r   S/home/ubuntu/.local/lib/python3.10/site-packages/sglang/test/nightly_bench_utils.pyhelp_str$   s   zBenchmarkResult.help_str base_url
relay_basec                    s0  d}|d }d}| j dkrt| j dnd}d| j| j  d }d| j|  d | }	d| j d | }
d	tf fd
dd}| jsF| jrl| j| jg}tdd |D r^t	
dd| fdd|D }d|}d| j d| j d| jdd| jdd| jdd| d|dd|	dd|
dd| dS )z6Convert this benchmark result to a markdown table row.      gffffff?r   zn/ai  g    .Ai  
trace_filec                    sN   ddl m} tj| }  d| }r# d||dd }|S |}|S )Nr   )quote/z?src=r   )safe)urllib.parser   ospathrelpath)r   r   rel_pathraw_file_link
relay_link)r   r   	trace_dirr   r   'get_perfetto_relay_link_from_trace_file9   s   zPBenchmarkResult.to_markdown_row.<locals>.get_perfetto_relay_link_from_trace_filezNA | NAc                 s   s    | ]}|d u V  qd S Nr   .0r   r   r   r   	<genexpr>J   s    z2BenchmarkResult.to_markdown_row.<locals>.<genexpr>zSome trace files are Noneztrace_files=c                    s$   g | ]}|rd  | dndqS )z[trace]()zN/Ar   r,   )r*   r   r   
<listcomp>L   s    z3BenchmarkResult.to_markdown_row.<locals>.<listcomp>z | z| z.2fz |
)r   roundr   r   r   strr   r   anyloggererrorjoinr	   r   )selfr)   r   r   hourly_cost_per_gpuhourly_cost
input_utilaccept_lengthitl
input_costoutput_costprofile_linktrace_filestrace_files_relay_linksr   )r   r*   r   r)   r   to_markdown_row,   s$   

	XzBenchmarkResult.to_markdown_row)r   r   )__name__
__module____qualname____doc__r2   __annotations__intfloatr   r   r   r   r   r   staticmethodr   rB   r   r   r   r   r      s6   
 r   resultsvariantr   c           	      C   s   |d j }|d jr|d jdkr|d|d j d7 }tdd}|r,|d| d7 }|r6|d| d7 }d	| d
}|d7 }|d7 }|D ]}tddd}tddd}||| ||7 }qF|S )zSGenerate a markdown report from a list of BenchmarkResult object from a single run.r   defaultz (r/   
GPU_CONFIGr   z []z### 
z| batch size | input len | latency (s) | input throughput (tok/s)  | output throughput (tok/s) | acc length | ITL (ms) | input cost ($/1M) | output cost ($/1M) | profile (extend) | profile (decode)|
z| ---------- | --------- | ----------- | ------------------------- | ------------------------- | ---------- | -------- | ----------------- | ------------------ | ---------------- | --------------- |
TRACE_BASE_URLr    PERFETTO_RELAY_URL)r   r   r#   getenvrstriprB   )	r)   rK   rL   model_header
gpu_configsummaryresultr   r   r   r   r   generate_markdown_report[   s*   
rY   pydantic_result_filenamer   r   c                 C   sT  g }| D ]}d}d}|j rag }g }	t|j D ])}
|
ds#|
dr@d|
 v s/d|
 v r5||
 qd|
 v r@|	|
 q|  |	  |rUtj|j |d }|	ratj|j |	d }t	||j
|j|j|j|j|j|j|j|j|j|j|||d}||  qt|d	d
d}tj||ddd W d   dS 1 sw   Y  dS )z5Save benchmark results as JSON using Pydantic models.Nz.trace.json.gzz.trace.jsonextendprefilldecoder   )r   r   r   r	   r
   r   r   r   r   r   r   r   r   r   r   wzutf-8)encodingr   F)indentensure_ascii)r?   r#   listdirendswithlowerappendsortr$   r6   r   r   r   r	   r
   r   r   r   r   r   r   r   
model_dumpopenjsondump)rK   rZ   r   r   json_resultsresr   r   extend_filesdecode_filesfilebenchmark_resultfr   r   r   save_results_as_pydantic_models}   sR   
"rr   r+   )ri   loggingr#   typingr   r   pydanticr   	getLoggerrC   r4   r   r2   rY   rr   r   r   r   r   <module>   s2    
	K
&
