Skip to content

Commit

Permalink
alternative suggested by lite
Browse files Browse the repository at this point in the history
  • Loading branch information
masahi committed Mar 15, 2024
1 parent 3ddfbd8 commit ad482bf
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 9 deletions.
5 changes: 4 additions & 1 deletion serve/mlc_serve/engine/async_connector.py
Original file line number Diff line number Diff line change
Expand Up @@ -150,7 +150,10 @@ async def _add_request(self, request: Request) -> ResultQueue:
queue = asyncio.Queue()
self.result_queues[request.request_id] = queue

await asyncio.to_thread(self.engine.add, [request])
try:
await asyncio.to_thread(self.engine.add, [request])
except TextGenerationError as e:
raise asyncio.CancelledError(e)

return queue

Expand Down
20 changes: 12 additions & 8 deletions serve/mlc_serve/engine/staging_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,9 @@
ScopedInferenceEngine,
SequenceOutput,
)
from .error import TextGenerationError
from .engine_common import get_new_request_state, prepare_output
from .model_module import ModelModule, TokenizerModule
from ..model.base import get_model_artifact_config
from .staging_engine_worker import (
AddRequestsCommand,
CancelRequestCommand,
Expand Down Expand Up @@ -119,13 +119,17 @@ def add(self, requests: list[Request]):
assert isinstance(req.stopping_criteria.stop_sequences, list)

# If the request violates the tokenization, this returns None, so skip.
state = get_new_request_state(
req,
self.conversation_template,
self.tokenizer,
self.model_artifact_config.vocab_size,
)
new_request_states.append(state)
try:
state = get_new_request_state(
req,
self.conversation_template,
self.tokenizer,
self.model_artifact_config.vocab_size,
)
new_request_states.append(state)
except Exception as e:
LOG.warn("Failed to add a request", request_id=req.request_id)
raise TextGenerationError(str(e))

self.command_queue.put(AddRequestsCommand(request_states=new_request_states))

Expand Down

0 comments on commit ad482bf

Please sign in to comment.