Skip to content

Commit

Permalink
add validation for stream=False and yield usage
Browse files Browse the repository at this point in the history
  • Loading branch information
aniketmaurya authored Jan 8, 2025
1 parent e9b1ee1 commit d07c076
Showing 1 changed file with 23 additions and 0 deletions.
23 changes: 23 additions & 0 deletions src/litserve/loops/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -264,6 +264,29 @@ def pre_setup(self, lit_api: LitAPI, spec: Optional[LitSpec]):
return

original = lit_api.unbatch.__code__ is LitAPI.unbatch.__code__
if not lit_api.stream and any([
inspect.isgeneratorfunction(lit_api.predict),
inspect.isgeneratorfunction(lit_api.encode_response),
(original or inspect.isgeneratorfunction(lit_api.unbatch)),
]):
raise ValueError(
"""When `stream=False`, `lit_api.predict`, `lit_api.encode_response` and `lit_api.unbatch` must not be
generator functions.
Correct example:
def predict(self, inputs):
...
return {"output": output}
Incorrect example:
def predict(self, inputs):
...
for i in range(max_token_length):
yield prediction
"""
)
if (
lit_api.stream
and lit_api.max_batch_size > 1
Expand Down

0 comments on commit d07c076

Please sign in to comment.