minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

feat(grpc): add generation request semantics (#32588)

Browse Source 
Signed-off-by: Connor Carpenter <connorc@nvidia.com>
Co-authored-by: ishandhanani <82981111+ishandhanani@users.noreply.github.com>
Co-authored-by: Alex Nails <alex.nails@radixark.ai>
This commit is contained in:
Connor Carpenter
2026-08-04 18:53:45 -07:00
committed by GitHub
co-authored by ishandhanani Alex Nails
parent 29831d58ef
commit a0b04dbe4c
10 changed files with 982 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/entrypoints/grpc_bridge.py
+25 -5
View File
@@ -293,14 +293,23 @@ class RuntimeHandle:
async def _run_generate(self, obj, chunk_callback, stream: bool, request):
ready_event = None
gen = None
try:
ready_event = self._install_on_ready(chunk_callback) if stream else None
ready_event = self._install_on_ready(chunk_callback)
gen = self.tokenizer_manager.generate_request(obj, request=request)
if stream:
completed_choices = set()
expected_choices = obj.batch_size * obj.parallel_sample_num
async for chunk in gen:
finished = (
choice_finished = (
chunk.get("meta_info", {}).get("finish_reason") is not None
)
if choice_finished:
choice_id = chunk.get(
"index", chunk.get("meta_info", {}).get("id")
)
completed_choices.add(choice_id)
finished = len(completed_choices) >= expected_choices
keep_going = await self._send_with_backpressure(
chunk_callback,
ready_event,
@@ -314,15 +323,26 @@ class RuntimeHandle:
self._safe_callback(chunk_callback, {}, finished=True)
else:
result = await gen.__anext__()
self._safe_callback(chunk_callback, result, finished=True)
chunks = result if isinstance(result, list) else [result]
for index, chunk in enumerate(chunks):
keep_going = await self._send_with_backpressure(
chunk_callback,
ready_event,
chunk,
finished=index == len(chunks) - 1,
timeout_abort_rid=obj.rid,
)
if not keep_going:
return
except StopAsyncIteration:
self._safe_callback(chunk_callback, {}, finished=True)
except Exception as e:
logger.error("gRPC generate error for rid=%s: %s", obj.rid, e)
self._send_native_error(chunk_callback, str(e))
finally:
if stream:
self._uninstall_on_ready(chunk_callback)
if gen is not None:
await gen.aclose()
self._uninstall_on_ready(chunk_callback)
async def _run_embed(self, obj, chunk_callback, request):
try:
python/sglang/srt/entrypoints/http_server.py
+3 -2
View File
@@ -895,7 +895,7 @@ async def generate_request(obj: GenerateReqInput, request: Request):
"error": {
"message": str(e),
"type": "invalid_request_error",
"code": 400,
"code": getattr(e, "status_code", 400),
"retryable": False,
}
}
@@ -2048,7 +2048,8 @@ async def vertex_generate(
def _create_error_response(e):
return ORJSONResponse(
{"error": {"message": str(e)}}, status_code=HTTPStatus.BAD_REQUEST
{"error": {"message": str(e)}},
status_code=getattr(e, "status_code", HTTPStatus.BAD_REQUEST),
)
python/sglang/srt/managers/io_struct.py
+26 -14
View File
@@ -158,8 +158,9 @@ MultimodalDataInputFormat = Union[
@dataclass
class GenerateReqInput:
# Request ID(s). If omitted, generated during normalization. For batch
# requests, a string is expanded to per-item IDs using it as a prefix.
# Logical request ID(s). If omitted, generated during normalization. For
# batch requests, a string is expanded to one ID per original batch item.
# Parallel-sampling child IDs are internal to TokenizerManager.
rid: Optional[Union[str, List[str]]] = field(default=None, kw_only=True)
# Stable identity shared by requests in the same session. Unlike
# session_params, this does not alter or reconstruct the prompt.
@@ -276,6 +277,9 @@ class GenerateReqInput:
background: bool = False
# Require reasoning for the request (hybrid reasoning model only)
require_reasoning: bool = False
# Per-request thinking budget. Requires strict thinking so the runtime can
# enforce the limit rather than silently treating it as metadata.
max_thinking_tokens: Optional[int] = None
# Priority for the request
priority: Optional[int] = None
@@ -319,12 +323,17 @@ class GenerateReqInput:
# Batch-level: List[List[int]] (one per request). After __getitem__: List[int].
multi_item_delimiter_indices: Optional[Union[List[List[int]], List[int]]] = None
def regenerate_rid(self):
def regenerate_rid(self, prefix: Optional[str] = None):
"""Generate a new request ID and return it."""
def new_rid() -> str:
suffix = uuid.uuid4().hex
return f"{prefix}_{suffix}" if prefix is not None else suffix
if isinstance(self.rid, list):
self.rid = [uuid.uuid4().hex for _ in range(len(self.rid))]
self.rid = [new_rid() for _ in range(len(self.rid))]
else:
self.rid = uuid.uuid4().hex
self.rid = new_rid()
return self.rid
def _validate_rid_uniqueness(self):
@@ -480,7 +489,7 @@ class GenerateReqInput:
# Expand input based on type
self._expand_inputs(num)
self._normalize_rid(num)
self._normalize_rid()
self._normalize_lora_paths(num)
self._normalize_image_data(num)
self._normalize_video_data(num)
@@ -590,16 +599,16 @@ class GenerateReqInput:
else: # Already a list
self.sampling_params = self.sampling_params * self.parallel_sample_num
def _normalize_rid(self, num):
"""Normalize request IDs for batch processing."""
def _normalize_rid(self):
"""Normalize one logical request ID per original batch item."""
if self.rid is None:
self.rid = [uuid.uuid4().hex for _ in range(num)]
self.rid = [uuid.uuid4().hex for _ in range(self.batch_size)]
elif isinstance(self.rid, str):
new_rids = [f"{self.rid}_{i}" for i in range(num)]
self.rid = new_rids
if self.batch_size == 1:
self.rid = [self.rid]
else:
self.rid = [f"{self.rid}_{i}" for i in range(self.batch_size)]
elif isinstance(self.rid, list):
# Note: the length of rid shall be the same as the batch_size,
# as the rid would be expanded for parallel sampling in tokenizer_manager
if len(self.rid) != self.batch_size:
raise ValueError(
"The specified rids length mismatch with the batch_size for batch processing."
@@ -751,8 +760,9 @@ class GenerateReqInput:
cache = self.__dict__.setdefault("_sub_obj_cache", {})
if i in cache:
return cache[i]
logical_index = i % self.batch_size
sub = GenerateReqInput(
rid=self.rid[i],
rid=self.rid[logical_index],
session_id=self.session_id,
text=self.text[i] if self.text is not None else None,
input_ids=self.input_ids[i] if self.input_ids is not None else None,
@@ -813,6 +823,8 @@ class GenerateReqInput:
disagg_prefill_dp_rank=self.disagg_prefill_dp_rank,
conversation_id=self.conversation_id,
http_worker_ipc=self.http_worker_ipc,
require_reasoning=self.require_reasoning,
max_thinking_tokens=self.max_thinking_tokens,
priority=self.priority,
extra_key=self.extra_key[i] if self.extra_key is not None else None,
no_logs=self.no_logs,
python/sglang/srt/managers/tokenizer_manager.py
+266 -37
View File
@@ -195,6 +195,10 @@ _INCREMENTAL_STREAMING_META_INFO_KEYS = (
)
class RequestAbortedError(ValueError):
status_code = 499
@dataclasses.dataclass
class ReqState:
"""Store the state a request."""
@@ -206,6 +210,9 @@ class ReqState:
# For performance metrics
time_stats: APIServerReqTimeStats
abort_requested: bool = False
lifecycle_id: object = dataclasses.field(default_factory=object)
dispatched: bool = False
last_completion_tokens: int = 1
ttft_observed: bool = False
@@ -545,6 +552,10 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
def init_running_status(self):
# Request states
self.rid_to_state: Dict[str, ReqState] = {}
# Parallel sampling keeps one caller-visible logical RID per original
# prompt while the scheduler operates on separate prefix/sample RIDs.
self.logical_rid_to_child_rids: Dict[str, set[str]] = {}
self.child_rid_to_logical_rid: Dict[str, str] = {}
self.event_loop = None
self.asyncio_tasks = set()
@@ -740,6 +751,15 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
# Normalize the request
obj.normalize_batch_and_arguments()
self._set_default_priority(obj)
if (
isinstance(obj, GenerateReqInput)
and obj.max_thinking_tokens is not None
and not self.server_args.enable_strict_thinking
):
raise ValueError(
"max_thinking_tokens requires the server to be launched with "
"--enable-strict-thinking"
)
if isinstance(obj, GenerateReqInput) and obj.routed_dp_rank is not None:
dp_size = self.elastic_worker_count
@@ -752,7 +772,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
f"routed_dp_rank={obj.routed_dp_rank} out of range [0, {dp_size})"
)
self._init_req_state(obj, request)
request_lifecycles = self._init_req_state(obj, request)
try:
if self.server_args.language_only:
self._handle_epd_disaggregation_encode_request(obj)
@@ -762,13 +782,16 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
async with self.is_pause_cond:
await self.is_pause_cond.wait_for(lambda: not self.is_pause)
self._raise_if_logical_request_aborted(obj)
async with self.model_update_lock.reader_lock:
await self._validate_and_resolve_lora(obj)
self._raise_if_logical_request_aborted(obj)
# Tokenize the request and send it to the scheduler
if obj.is_single:
tokenized_obj = await self._tokenize_one_request(obj)
self._raise_if_logical_rid_aborted(obj.rid)
state = self.rid_to_state[obj.rid]
if obj.return_prompt_token_ids:
state.prompt_token_ids = list(tokenized_obj.input_ids)
@@ -778,7 +801,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
else:
async for response in self._handle_batch_request(obj, request):
yield response
except Exception:
except BaseException:
# _init_req_state created a rid_to_state entry per (sub-)request up
# front. The normal remover is the scheduler-response path
# (_handle_batch_output), so a failure *before* a request reaches the
@@ -786,7 +809,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
# request -- would otherwise leak those entries forever. Drop any that
# are still pending; entries already removed on the normal completion
# path are left untouched (pop is a no-op).
self._discard_pending_req_states(obj)
self._discard_pending_req_states(obj, request_lifecycles)
raise
def _detect_input_format(
@@ -1308,6 +1331,11 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
sampling_kwargs = {**self.preferred_sampling_params, **obj.sampling_params}
else:
sampling_kwargs = obj.sampling_params
if isinstance(obj, GenerateReqInput) and obj.max_thinking_tokens is not None:
sampling_kwargs = dict(sampling_kwargs)
custom_params = dict(sampling_kwargs.get("custom_params") or {})
custom_params["thinking_budget"] = obj.max_thinking_tokens
sampling_kwargs["custom_params"] = custom_params
sampling_params = self.sampling_params_class(**sampling_kwargs)
sampling_params.normalize(self.tokenizer)
sampling_params.verify(self.model_config.vocab_size)
@@ -1518,6 +1546,9 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
time_stats = tokenized_obj.time_stats
tokenized_obj.wrap_pickle_fields()
self._dispatch_to_scheduler(tokenized_obj)
state = self.rid_to_state.get(tokenized_obj.rid)
if state is not None:
state.dispatched = True
tokenized_obj.time_stats = time_stats
tokenized_obj.time_stats.set_api_server_dispatch_finish_time()
@@ -1539,6 +1570,10 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
batch_req = BatchTokenizedEmbeddingReqInput(batch=tokenized_objs)
self._dispatch_to_scheduler(batch_req)
for tokenized_obj in tokenized_objs:
state = self.rid_to_state.get(tokenized_obj.rid)
if state is not None:
state.dispatched = True
for tokenized_obj, time_stat in zip(tokenized_objs, time_stats):
tokenized_obj.time_stats = time_stat
set_time_batch(tokenized_objs, "set_api_server_dispatch_finish_time")
@@ -1610,7 +1645,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
# Delete the key to prevent resending abort request to the scheduler and
# to ensure aborted request state is cleaned up.
if state.obj.rid in self.rid_to_state:
del self.rid_to_state[state.obj.rid]
self._remove_req_state(state.obj.rid)
# Mark ongoing LoRA request as finished.
if self.enable_lora and state.obj.lora_path:
@@ -1744,6 +1779,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
if getattr(obj, "parallel_sample_num", 1) == 1:
if self._should_use_batch_tokenization(batch_size, obj):
tokenized_objs = await self._batch_tokenize_and_process(batch_size, obj)
self._raise_if_logical_request_aborted(obj)
self._send_batch_request(tokenized_objs)
# Set up generators for each request in the batch
@@ -1766,6 +1802,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
for i in range(batch_size):
tmp_obj = obj[i]
tokenized_obj = await self._tokenize_one_request(tmp_obj)
self._raise_if_logical_rid_aborted(tmp_obj.rid)
state = self.rid_to_state[tmp_obj.rid]
if tmp_obj.return_prompt_token_ids:
state.prompt_token_ids = list(tokenized_obj.input_ids)
@@ -1786,9 +1823,12 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
tokenized_objs = await asyncio.gather(
*(self._tokenize_one_request(obj) for obj in objs)
)
self._raise_if_logical_request_aborted(obj)
# Cache the common prefix for parallel sampling
for i in range(batch_size):
logical_rid = objs[i].rid
self._raise_if_logical_rid_aborted(logical_rid)
tmp_obj = copy.copy(objs[i])
tokenized_obj = copy.copy(tokenized_objs[i])
# Ensure independent mm_items so wrap_shm_features won't mutate the original
@@ -1797,17 +1837,20 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
tokenized_obj.mm_inputs.mm_items = [
copy.copy(item) for item in tokenized_obj.mm_inputs.mm_items
]
tokenized_obj.rid = tmp_obj.regenerate_rid()
tokenized_obj.rid = tmp_obj.regenerate_rid(prefix=logical_rid)
tokenized_obj.sampling_params = copy.copy(tokenized_obj.sampling_params)
tokenized_obj.sampling_params.max_new_tokens = 0
tokenized_obj.stream = False
self._init_req_state(tmp_obj)
self._init_child_req_state(logical_rid, tmp_obj)
self._send_one_request(tokenized_obj)
await self._wait_one_response(tmp_obj, request).__anext__()
self._raise_if_logical_rid_aborted(logical_rid)
# Expand requests, assign new rids for them, and send them
for i in range(batch_size):
logical_rid = objs[i].rid
for _ in range(obj.parallel_sample_num):
self._raise_if_logical_rid_aborted(logical_rid)
tmp_obj = copy.copy(objs[i])
tokenized_obj = copy.copy(tokenized_objs[i])
# Ensure independent mm_items so wrap_shm_features won't mutate the original
@@ -1816,8 +1859,8 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
tokenized_obj.mm_inputs.mm_items = [
copy.copy(item) for item in tokenized_obj.mm_inputs.mm_items
]
tokenized_obj.rid = tmp_obj.regenerate_rid()
self._init_req_state(tmp_obj)
tokenized_obj.rid = tmp_obj.regenerate_rid(prefix=logical_rid)
self._init_child_req_state(logical_rid, tmp_obj)
state = self.rid_to_state[tmp_obj.rid]
tokenized_obj.time_stats = state.time_stats
if tmp_obj.return_prompt_token_ids:
@@ -1826,17 +1869,38 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
generators.append(self._wait_one_response(tmp_obj, request))
rids.append(tmp_obj.rid)
self.rid_to_state[objs[i].rid].time_stats.set_finished_time()
del self.rid_to_state[objs[i].rid]
parent_state = self.rid_to_state.get(logical_rid)
if parent_state is not None:
parent_state.time_stats.set_finished_time()
self._remove_req_state(logical_rid)
# Wait for all requests
is_stream = hasattr(obj, "stream") and obj.stream
if not is_stream:
outputs = await asyncio.gather(*(gen.__anext__() for gen in generators))
outputs = await self._collect_batch_responses(generators)
yield outputs
else:
rid_to_index = {rid: i for i, rid in enumerate(rids)}
task_map = {asyncio.create_task(gen.__anext__()): gen for gen in generators}
async for response in self._stream_batch_responses(generators, rids):
yield response
async def _collect_batch_responses(self, generators):
tasks = [asyncio.create_task(gen.__anext__()) for gen in generators]
try:
return await asyncio.gather(*tasks)
finally:
for task in tasks:
if not task.done():
task.cancel()
await asyncio.gather(*tasks, return_exceptions=True)
await asyncio.gather(
*(gen.aclose() for gen in generators),
return_exceptions=True,
)
async def _stream_batch_responses(self, generators, rids):
rid_to_index = {rid: i for i, rid in enumerate(rids)}
task_map = {asyncio.create_task(gen.__anext__()): gen for gen in generators}
try:
while task_map:
done, _ = await asyncio.wait(
task_map.keys(), return_when=asyncio.FIRST_COMPLETED
@@ -1852,20 +1916,55 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
task_map[new_task] = gen
except StopAsyncIteration:
pass
finally:
pending_tasks = list(task_map)
for task in pending_tasks:
task.cancel()
if pending_tasks:
await asyncio.gather(*pending_tasks, return_exceptions=True)
await asyncio.gather(
*(gen.aclose() for gen in generators),
return_exceptions=True,
)
def abort_request(self, rid: str = "", abort_all: bool = False):
# Empty rid would startswith-match every request on the scheduler.
if not abort_all and not rid:
logger.warning("Ignore abort_request with empty rid and abort_all=False")
return
if (
not abort_all
and self.server_args.tokenizer_worker_num == 1
and rid not in self.rid_to_state
):
if abort_all:
for state_rid, state in self.rid_to_state.items():
if state_rid not in self.child_rid_to_logical_rid:
state.abort_requested = True
target_rids = (rid,)
elif rid in self.child_rid_to_logical_rid:
# Preserve direct child aborts for internal callers.
target_rids = (rid,)
elif rid in self.rid_to_state:
state = self.rid_to_state[rid]
state.abort_requested = True
parallel_sample_num = getattr(state.obj, "parallel_sample_num", None)
if parallel_sample_num is None:
sampling_params = getattr(state.obj, "sampling_params", None)
parallel_sample_num = (
sampling_params.get("n", 1)
if isinstance(sampling_params, dict)
else 1
)
if parallel_sample_num > 1:
# Snapshot because scheduler abort echoes remove child ownership.
target_rids = tuple(sorted(self.logical_rid_to_child_rids.get(rid, ())))
else:
target_rids = (rid,)
elif child_rids := self.logical_rid_to_child_rids.get(rid):
target_rids = tuple(sorted(child_rids))
elif self.server_args.tokenizer_worker_num == 1:
return
req = AbortReq(rid=rid, abort_all=abort_all)
self._dispatch_to_scheduler(req)
else:
target_rids = (rid,)
for target_rid in target_rids:
self._dispatch_to_scheduler(AbortReq(rid=target_rid, abort_all=abort_all))
if self.enable_metrics:
# TODO: also use custom_labels from the request
self.metrics_collector.observe_one_aborted_request(
@@ -2352,7 +2451,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
)
)
del self.rid_to_state[rid]
self._remove_req_state(rid)
# Mark ongoing LoRA request as finished.
if self.enable_lora and state.obj.lora_path:
@@ -3088,7 +3187,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
"output_ids": output_ids,
"meta_info": meta_info,
}
del self.rid_to_state[recv_obj.rid]
self._remove_req_state(recv_obj.rid)
state.out_list.append(out)
state.event.set()
@@ -3244,11 +3343,80 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
obj.lora_id[i] if isinstance(obj.lora_id, list) else obj.lora_id
)
@staticmethod
def _logical_rids(obj) -> List[str]:
if not hasattr(obj, "is_single") or obj.is_single:
return [obj.rid]
return list(obj.rid)
def _register_child_rid(self, logical_rid: str, child_rid: str) -> None:
if child_rid == logical_rid:
raise ValueError(
"Parallel-sampling child RID must differ from its logical RID"
)
owner = self.child_rid_to_logical_rid.get(child_rid)
if owner is not None and owner != logical_rid:
raise ValueError(
f"Request ID {child_rid} is already owned by logical request {owner}"
)
self.child_rid_to_logical_rid[child_rid] = logical_rid
self.logical_rid_to_child_rids.setdefault(logical_rid, set()).add(child_rid)
def _init_child_req_state(
self,
logical_rid: str,
obj: Union[GenerateReqInput, EmbeddingReqInput],
request: Optional[fastapi.Request] = None,
) -> None:
self._raise_if_logical_rid_aborted(logical_rid)
logical_state = self.rid_to_state[logical_rid]
self._init_req_state(
obj,
request,
lifecycle_id=logical_state.lifecycle_id,
)
try:
self._register_child_rid(logical_rid, obj.rid)
except BaseException:
self._remove_req_state(obj.rid)
raise
def _remove_req_state(
self,
rid: str,
lifecycle_id: Optional[object] = None,
) -> Optional[ReqState]:
"""Remove a request state and its parallel-sampling ownership."""
state = self.rid_to_state.get(rid)
if state is None or (
lifecycle_id is not None and state.lifecycle_id is not lifecycle_id
):
return None
self.rid_to_state.pop(rid)
logical_rid = self.child_rid_to_logical_rid.pop(rid, None)
if logical_rid is not None:
children = self.logical_rid_to_child_rids.get(logical_rid)
if children is not None:
children.discard(rid)
if not children:
self.logical_rid_to_child_rids.pop(logical_rid, None)
return state
def _raise_if_logical_rid_aborted(self, logical_rid: str) -> None:
state = self.rid_to_state.get(logical_rid)
if state is None or state.abort_requested:
raise RequestAbortedError(f"Request {logical_rid} was aborted")
def _raise_if_logical_request_aborted(self, obj) -> None:
for logical_rid in self._logical_rids(obj):
self._raise_if_logical_rid_aborted(logical_rid)
def _init_req_state(
self,
obj: Union[GenerateReqInput, EmbeddingReqInput],
request: Optional[fastapi.Request] = None,
):
lifecycle_id: Optional[object] = None,
) -> Dict[str, object]:
created_time = obj.received_time
external_trace_header = None
@@ -3279,29 +3447,90 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
for i in range(len(obj.rid))
]
for rid, sub_obj, bootstrap_room in items:
if rid in self.rid_to_state:
rids = [rid for rid, _, _ in items]
seen_rids = set()
for rid in rids:
if rid in seen_rids:
raise ValueError(f"Duplicate request ID detected: {rid}")
seen_rids.add(rid)
if (
rid in self.rid_to_state
or rid in self.logical_rid_to_child_rids
or rid in self.child_rid_to_logical_rid
):
raise ValueError(f"Duplicate request ID detected: {rid}")
# Mutate only after every RID passes duplicate validation so a rejected
# batch cannot leave a partial rid_to_state insertion behind.
lifecycle_ids = {}
for rid, sub_obj, bootstrap_room in items:
time_stats = APIServerReqTimeStats(disagg_mode=self.disaggregation_mode)
state = ReqState([], False, asyncio.Event(), sub_obj, time_stats)
state = ReqState(
[],
False,
asyncio.Event(),
sub_obj,
time_stats,
lifecycle_id=lifecycle_id if lifecycle_id is not None else object(),
)
self.rid_to_state[rid] = state
lifecycle_ids[rid] = state.lifecycle_id
if self.enable_trace:
time_stats.init_trace_ctx(rid, bootstrap_room, external_trace_header)
time_stats.set_created_time(created_time)
return lifecycle_ids
def _discard_pending_req_states(self, obj):
"""Drop rid_to_state entries created by _init_req_state for *obj*.
def _discard_pending_req_states(
self,
obj,
lifecycle_ids: Optional[Dict[str, object]] = None,
):
"""Drop all logical and child state owned by *obj*.
Safe to call after a partial/failed dispatch: only entries still present
are removed, and the scheduler-response path looks up state with
``.get(...)`` so a later output for a discarded rid is ignored, not fatal.
Safe to call after a partial/failed dispatch: only requests known to have
reached the scheduler are aborted, all owned state is removed, and a later
output for a discarded RID is ignored by the scheduler-response path.
"""
if not hasattr(obj, "is_single") or obj.is_single:
rids = [obj.rid]
else:
rids = obj.rid
for rid in rids:
self.rid_to_state.pop(rid, None)
if lifecycle_ids is None:
lifecycle_ids = {
logical_rid: state.lifecycle_id
for logical_rid in self._logical_rids(obj)
if (state := self.rid_to_state.get(logical_rid)) is not None
}
for logical_rid in self._logical_rids(obj):
lifecycle_id = lifecycle_ids.get(logical_rid)
if lifecycle_id is None:
continue
child_rids = tuple(
child_rid
for child_rid in self.logical_rid_to_child_rids.get(logical_rid, ())
if (
(state := self.rid_to_state.get(child_rid)) is not None
and state.lifecycle_id is lifecycle_id
)
)
logical_state = self.rid_to_state.get(logical_rid)
owns_logical_state = (
logical_state is not None and logical_state.lifecycle_id is lifecycle_id
)
target_rids = tuple(
rid for rid in child_rids if self.rid_to_state[rid].dispatched
)
if not child_rids and owns_logical_state and logical_state.dispatched:
target_rids = (logical_rid,)
for target_rid in target_rids:
try:
self._dispatch_to_scheduler(
AbortReq(rid=target_rid, abort_all=False)
)
except Exception:
logger.exception(
"Failed to abort request rid=%s",
target_rid,
)
for child_rid in child_rids:
self._remove_req_state(child_rid, lifecycle_id)
self._remove_req_state(logical_rid, lifecycle_id)
def _should_dispatch_to_encoder(
self, obj: Union[GenerateReqInput, EmbeddingReqInput]