-
Notifications
You must be signed in to change notification settings - Fork 144
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge branch 'main' into feat/evict-req-on-client-disconnect-streamin…
…g-case
- Loading branch information
Showing
33 changed files
with
948 additions
and
130 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,3 @@ | ||
uvloop | ||
tenacity | ||
jsonargparse |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,3 +1,3 @@ | ||
fastapi >=0.100 | ||
httpx | ||
uvicorn >=0.29.0 | ||
uvicorn[standard] >=0.29.0 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,3 @@ | ||
from .base import Callback, CallbackRunner, EventTypes, NoopCallback | ||
|
||
__all__ = ["Callback", "CallbackRunner", "EventTypes", "NoopCallback"] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,80 @@ | ||
import dataclasses | ||
import logging | ||
from abc import ABC | ||
from typing import List, Union | ||
|
||
logger = logging.getLogger(__name__) | ||
|
||
|
||
@dataclasses.dataclass | ||
class EventTypes: | ||
BEFORE_SETUP = "on_before_setup" | ||
AFTER_SETUP = "on_after_setup" | ||
BEFORE_DECODE_REQUEST = "on_before_decode_request" | ||
AFTER_DECODE_REQUEST = "on_after_decode_request" | ||
BEFORE_ENCODE_RESPONSE = "on_before_encode_response" | ||
AFTER_ENCODE_RESPONSE = "on_after_encode_response" | ||
BEFORE_PREDICT = "on_before_predict" | ||
AFTER_PREDICT = "on_after_predict" | ||
ON_SERVER_START = "on_server_start" | ||
ON_SERVER_END = "on_server_end" | ||
|
||
|
||
class Callback(ABC): | ||
def on_before_setup(self, *args, **kwargs): | ||
"""Called before setup is started.""" | ||
|
||
def on_after_setup(self, *args, **kwargs): | ||
"""Called after setup is completed.""" | ||
|
||
def on_before_decode_request(self, *args, **kwargs): | ||
"""Called before request decoding is started.""" | ||
|
||
def on_after_decode_request(self, *args, **kwargs): | ||
"""Called after request decoding is completed.""" | ||
|
||
def on_before_encode_response(self, *args, **kwargs): | ||
"""Called before response encoding is started.""" | ||
|
||
def on_after_encode_response(self, *args, **kwargs): | ||
"""Called after response encoding is completed.""" | ||
|
||
def on_before_predict(self, *args, **kwargs): | ||
"""Called before prediction is started.""" | ||
|
||
def on_after_predict(self, *args, **kwargs): | ||
"""Called after prediction is completed.""" | ||
|
||
def on_server_start(self, *args, **kwargs): | ||
"""Called before server starts.""" | ||
|
||
def on_server_end(self, *args, **kwargs): | ||
"""Called when server terminates.""" | ||
|
||
|
||
class CallbackRunner: | ||
def __init__(self, callbacks: Union[Callback, List[Callback]] = None): | ||
self._callbacks = [] | ||
if callbacks: | ||
self._add_callbacks(callbacks) | ||
|
||
def _add_callbacks(self, callbacks: Union[Callback, List[Callback]]): | ||
if not isinstance(callbacks, list): | ||
callbacks = [callbacks] | ||
for callback in callbacks: | ||
if not isinstance(callback, Callback): | ||
raise ValueError(f"Invalid callback type: {callback}") | ||
self._callbacks.extend(callbacks) | ||
|
||
def trigger_event(self, event_name, *args, **kwargs): | ||
"""Triggers an event, invoking all registered callbacks for that event.""" | ||
for callback in self._callbacks: | ||
try: | ||
getattr(callback, event_name)(*args, **kwargs) | ||
except Exception: | ||
# Handle exceptions to prevent one callback from disrupting others | ||
logger.exception(f"Error in callback '{callback}' during event '{event_name}'") | ||
|
||
|
||
class NoopCallback(Callback): | ||
"""This callback does nothing.""" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,3 @@ | ||
from litserve.callbacks.defaults.metric_callback import PredictionTimeLogger | ||
|
||
__all__ = ["PredictionTimeLogger"] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,21 @@ | ||
import time | ||
import typing | ||
from logging import getLogger | ||
|
||
from ..base import Callback | ||
|
||
if typing.TYPE_CHECKING: | ||
from litserve import LitAPI | ||
|
||
logger = getLogger(__name__) | ||
|
||
|
||
class PredictionTimeLogger(Callback): | ||
def on_before_predict(self, lit_api: "LitAPI"): | ||
t0 = time.perf_counter() | ||
self._start_time = t0 | ||
|
||
def on_after_predict(self, lit_api: "LitAPI"): | ||
t1 = time.perf_counter() | ||
elapsed = t1 - self._start_time | ||
print(f"Prediction took {elapsed:.2f} seconds", flush=True) |
Oops, something went wrong.