chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

fix black in pre-commit (#1940)

Browse Source
This commit is contained in:
Chayenne
2024-11-08 07:42:47 +08:00
committed by GitHub
parent dca87ec348
commit c77c1e05ba
29 changed files with 641 additions and 508 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/io_struct.py
+4 -2
View File
@@ -86,8 +86,10 @@ class GenerateReqInput:
self.parallel_sample_num = self.sampling_params.get("n", 1)
else: # isinstance(self.sampling_params, list):
self.parallel_sample_num = self.sampling_params[0].get("n", 1)
assert all(self.parallel_sample_num == sampling_params.get("n", 1) for sampling_params in self.sampling_params), (
"The parallel_sample_num should be the same for all samples in sample params.")
assert all(
self.parallel_sample_num == sampling_params.get("n", 1)
for sampling_params in self.sampling_params
), "The parallel_sample_num should be the same for all samples in sample params."
if self.parallel_sample_num > 1 and self.is_single:
self.is_single = False
python/sglang/srt/managers/schedule_batch.py
+2 -2
View File
@@ -911,8 +911,7 @@ class ScheduleBatch:
keep_indices = [
i
for i in range(len(self.reqs))
if not self.reqs[i].finished()
and self.reqs[i] is not being_chunked_req
if not self.reqs[i].finished() and self.reqs[i] is not being_chunked_req
]
if keep_indices is None or len(keep_indices) == 0:
@@ -1043,6 +1042,7 @@ class ScheduleBatch:
for req in self.reqs:
req.started_time = time.time()
@dataclasses.dataclass
class ModelWorkerBatch:
# The batch id
python/sglang/srt/managers/scheduler.py
+25 -22
View File
@@ -224,8 +224,8 @@ class Scheduler:
self.forward_ct = 0
self.forward_ct_decode = 0
self.num_generated_tokens = 0
self.last_stats_tic = time.time() # time of last stats for every iter
self.last_log_tic = time.time() # time of last log for print decode log
self.last_stats_tic = time.time() # time of last stats for every iter
self.last_log_tic = time.time() # time of last log for print decode log
self.stream_interval = server_args.stream_interval
# Init chunked prefill
@@ -566,9 +566,7 @@ class Scheduler:
and not self.last_batch.is_empty()
):
if self.being_chunked_req:
self.last_batch.filter_batch(
being_chunked_req=self.being_chunked_req
)
self.last_batch.filter_batch(being_chunked_req=self.being_chunked_req)
self.tree_cache.cache_unfinished_req(self.being_chunked_req)
# Inflight request keeps its rid but will get a new req_pool_idx.
self.req_to_token_pool.free(self.being_chunked_req.req_pool_idx)
@@ -628,9 +626,7 @@ class Scheduler:
has_inflight = self.being_chunked_req is not None
if has_inflight:
self.being_chunked_req.init_next_round_input()
self.being_chunked_req = adder.add_inflight_req(
self.being_chunked_req
)
self.being_chunked_req = adder.add_inflight_req(self.being_chunked_req)
if self.lora_paths:
lora_set = (
@@ -813,7 +809,8 @@ class Scheduler:
embeddings = self.tp_worker.forward_batch_embedding(model_worker_batch)
ret = embeddings, model_worker_batch.bid
return ret
def get_stats(self,batch: ScheduleBatch):
def get_stats(self, batch: ScheduleBatch):
# TODO: get stats for chunked prefill
now = time.time()
@@ -829,8 +826,8 @@ class Scheduler:
# set stats from prefill
if self.stats is not None:
# new_seq=self.stats.new_seq
cache_hit_rate=self.stats.cache_hit_rate
token_usage=self.stats.token_usage
cache_hit_rate = self.stats.cache_hit_rate
token_usage = self.stats.token_usage
# Iteration stats
num_prompt_tokens_iter = 0
num_generation_tokens_iter = 0
@@ -851,15 +848,19 @@ class Scheduler:
# _, next_token_ids, _ = result
if batch is not None:
num_generation_tokens_iter = len(batch.output_ids)
gen_throughput = round(num_generation_tokens_iter / (now - self.last_stats_tic), 2)
gen_throughput = round(
num_generation_tokens_iter / (now - self.last_stats_tic), 2
)
for i, req in enumerate(batch.reqs):
# NOTE: Batch forward mode is extend befor start decode,
if batch.forward_mode.is_extend():
num_prompt_tokens_iter=len(batch.input_ids)+sum(batch.prefix_lens)
num_prompt_tokens_iter = len(batch.input_ids) + sum(
batch.prefix_lens
)
time_to_first_tokens_iter.append(now - req.started_time)
else:
time_per_output_tokens_iter.append(now-self.last_stats_tic)
time_per_output_tokens_iter.append(now - self.last_stats_tic)
if req.finished():
time_e2e_requests.append(now - req.created_time)
@@ -867,9 +868,10 @@ class Scheduler:
num_prompt_tokens_requests.append(len(req.origin_input_ids))
num_generation_tokens_requests.append(len(req.output_ids))
finished_reason_requests.append(
req.finished_reason.to_json()
if req.finished_reason is not None
else None)
req.finished_reason.to_json()
if req.finished_reason is not None
else None
)
return Stats(
new_seq=new_seq,
@@ -893,7 +895,7 @@ class Scheduler:
max_running_requests=self.max_running_requests,
)
def log_stats(self,stats:Stats):
def log_stats(self, stats: Stats):
self.metrics_collector.log_stats(stats)
def process_batch_result(self, batch: ScheduleBatch, result):
@@ -1003,9 +1005,7 @@ class Scheduler:
if req.is_retracted:
continue
if self.server_args.enable_overlap_schedule and (
req.finished()
):
if self.server_args.enable_overlap_schedule and (req.finished()):
self.token_to_kv_pool.free(batch.out_cache_loc[i : i + 1])
continue
@@ -1031,7 +1031,10 @@ class Scheduler:
self.token_to_kv_pool.free_group_end()
self.forward_ct_decode = (self.forward_ct_decode + 1) % (1 << 30)
if self.tp_rank == 0 and self.forward_ct_decode % self.server_args.decode_log_interval == 0:
if (
self.tp_rank == 0
and self.forward_ct_decode % self.server_args.decode_log_interval == 0
):
self.print_decode_stats()
def add_logprob_return_values(
python/sglang/srt/managers/tokenizer_manager.py
+11 -7
View File
@@ -215,7 +215,7 @@ class TokenizerManager:
logprob_start_len,
top_logprobs_num,
obj.stream,
obj.lora_path
obj.lora_path,
)
elif isinstance(obj, EmbeddingReqInput):
tokenized_obj = TokenizedEmbeddingReqInput(
@@ -290,7 +290,9 @@ class TokenizerManager:
# Tokenize all requests
objs = [obj[i] for i in range(batch_size)]
tokenized_objs = await asyncio.gather(*(self._tokenize_one_request(obj) for obj in objs))
tokenized_objs = await asyncio.gather(
*(self._tokenize_one_request(obj) for obj in objs)
)
# Cache the common prefix for parallel sampling
for i in range(batch_size):
@@ -322,7 +324,9 @@ class TokenizerManager:
rid_to_index = {rid: i for i, rid in enumerate(rids)}
task_map = {asyncio.create_task(gen.__anext__()): gen for gen in generators}
while task_map:
done, _ = await asyncio.wait(task_map.keys(), return_when=asyncio.FIRST_COMPLETED)
done, _ = await asyncio.wait(
task_map.keys(), return_when=asyncio.FIRST_COMPLETED
)
for task in done:
gen = task_map.pop(task)
@@ -367,7 +371,7 @@ class TokenizerManager:
if self.server_args.dp_size == 1:
res = await self.mem_pool_size
return res.size
else: # self.server_args.dp_size > 1
else: # self.server_args.dp_size > 1
self.mem_pool_size_tmp = []
res = await self.mem_pool_size
ret = [r.size for r in res]
@@ -399,7 +403,7 @@ class TokenizerManager:
self.server_args.load_format = obj.load_format
self.model_path = obj.model_path
return result.success, result.message
else: # self.server_args.dp_size > 1
else: # self.server_args.dp_size > 1
self.model_update_tmp = []
result = await self.model_update_result
@@ -470,7 +474,7 @@ class TokenizerManager:
if isinstance(recv_obj, UpdateWeightReqOutput):
if self.server_args.dp_size == 1:
self.model_update_result.set_result(recv_obj)
else: # self.server_args.dp_size > 1
else: # self.server_args.dp_size > 1
self.model_update_tmp.append(recv_obj)
# set future if the all results are recevied
if len(self.model_update_tmp) == self.server_args.dp_size:
@@ -479,7 +483,7 @@ class TokenizerManager:
elif isinstance(recv_obj, GetMemPoolSizeReqOutput):
if self.server_args.dp_size == 1:
self.mem_pool_size.set_result(recv_obj)
else: # self.sever_args.dp_size > 1
else: # self.sever_args.dp_size > 1
self.mem_pool_size_tmp.append(recv_obj)
# set future if the all results are received
if len(self.mem_pool_size_tmp) == self.server_args.dp_size: