From 7b2901be9ef77aefabd2eb6d247798f1e62ea30b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 1 Aug 2023 12:20:25 -0700 Subject: [PATCH] custom timeout decorator --- litellm/__init__.py | 1 + litellm/__pycache__/__init__.cpython-311.pyc | Bin 721 -> 767 bytes litellm/__pycache__/main.cpython-311.pyc | Bin 11556 -> 11550 bytes litellm/__pycache__/timeout.cpython-311.pyc | Bin 0 -> 4905 bytes litellm/main.py | 7 +- litellm/tests/test_exceptions.py | 17 ++-- litellm/tests/test_timeout.py | 26 ++++++ litellm/timeout.py | 80 +++++++++++++++++++ requirements.txt | 1 - setup.py | 5 +- 10 files changed, 121 insertions(+), 16 deletions(-) create mode 100644 litellm/__pycache__/timeout.cpython-311.pyc create mode 100644 litellm/tests/test_timeout.py create mode 100644 litellm/timeout.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 7ed52d7cd..933a6ca8a 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -25,6 +25,7 @@ open_ai_embedding_models = [ 'text-embedding-ada-002' ] +from .timeout import timeout from .utils import client, logging, exception_type # Import all the symbols from main.py from .main import * # Import all the symbols from main.py diff --git a/litellm/__pycache__/__init__.cpython-311.pyc b/litellm/__pycache__/__init__.cpython-311.pyc index 3e9ac33f0fa3155ad41786391a11d4482cb3d8ac..6dd968755a8e42d6602ea098bc7df7701dae8373 100644 GIT binary patch delta 164 zcmcb}`k$3|IWI340}$MfIGIv3kyny2W}>>4I!6=-BSQ*D3ReziE>{#+E_W1nE>9E> zSd2S|H-sUDCzwH#cVbNm7vnAVlFZ!H{L+$%@3k2DCQC7{ktkvYns|$~v?Md9xQGSF uE@A}{Y?E&?MsSG&nT$YOykN35Q?j7!1qK;pbYt=YCaKBWnZ#Ixfr{#+E_W0+BSQ*D3RezK2tx{Y zFoPz~#Ez25?2MX>yps(X*9hNYEiK8+DK26L8d1aoB3LJXWsKkw0WukZxVUd}I8!pC R@#G6k5|f`YiLnR)1p(#dc|UK05Ucf1qC*%d*d}EnJ;$C+o>}s-9}6>uIKOjk8v!o*B_Z{9G2+ z=hK&&wxRI=52{ayGZf``_$Je|+jxfu%`^+Ono+*1&+2(Gfm%r(`M=f^`H66R3vV!b z-OpovAAH4TO3H7DNBGz$@tY$U0p8(3-^6E4KOcuL*`C-mpXh7#-JGpzdIElBbrygZ zM~Y!D=vr;z^xOd-xx`X&H1}f7o^868V;g3p)so}zmrrLCQ1rhBH*yEq4BYYeMrToT zd{7=uviJ&;x*yFPX$r#~n1TCI(hSa^U`>|b%H;7t32C_i-%h?9s1f-r6q9FI5&Fq} z%z?YfQfQ9cjK~^1NN!CZAe|+GWrQvzRw&3KFEP>J&q}9bnerq(ZNiEC zOL5giQ~gF*8C=gV2NP>L^I|w{&ekiNLQ0E^!H>o7e5-0~%A{N&s1iJlfQhS4K7;g| z~n~hdaEEj%XBs~TDiF$5_MGgL*UkoVOt57M;UZ%eRFMOSZg9sH_NG)bp=Lz2gkDz8hs#IK)zSEpb@Bi#*&2 zt((TJ{L{x6`w5DbJQONZcmr^#;?~x`*Ph@Dzf^cQ^j7Ab#M@K-DL7TFh0@s%dUyQy S1iW9J_QjC-SylQentuTh-d$n< delta 1409 zcmb7DO=uid9G|zl$?i7$k?npaCh6?9yUFaXNzGPi(loZ<7t$))73o@8#(8hkF|#wv zn@Q@$h!77wwe^X3(BeVbf}UK#Ab1eGhzAq!Aaf{sP!Pe3;6Via|KDh-QYpTf`SG9s z_t(4Hd9RbYo=omj;5>io3p!uFk}8o;uFPIii9!`RYK`}Eom@ZP$;*FY6*`5G!qP7m z#m%|GD@0jSXq3j}6Q=}3CA#k#QIz-L3l6wZOwj5j=-4*tOt1`S#b|;iKS*5*fhK&x z`Cq=$G*06LVEpX}%?OJe&?f(KoxY>%I89U1i_rbsMA&2|L9^RRXNr3ZDT8KasEy=E z^&=qL8u62$?Ob?rlW}fR_Tvh$wTPeSIox3F;i=K@5`lfV&1c1W_L0$9(BO08n1r9!DK)WV-RkOm&E<2bM!3|FO^F)k9n)Yg zFQE}e^y9kgny#nY9wfFT(uGR-I0_#}m_}FvFq^OxS~!S!LOfmIcTIonb2JXp7pdaP z-3@-a&ESnH65M zxo{l)vDF%ig1I53AshhkW1P7I$99=)wx8u}V3~&Qv93NayPmVoY<^h$KA{${N6
    *W>)0=+|bRT1Tb@phf?+ zEnGGM+sSQ&O#n9vpaw@SB6Rwe|Cd3gkm2{WON3O!#@rv|oY>Hof^spg3|)@xvUxjy~GtHAchEJvqZ>UhlFJEcvOq`5o(mP`U zsnbSn+O<*CP1Pb$RcZ5qs1T^4Qa}3FuGF9Y*~k*sNJy3Xv;8y8BBA~2IWwM#Cy*7r zo;h>pKF;f&*Z7agWSl_z+Xw$t{~IIZ-}vG;q21XiL+1`rNtURL%IZv(fhVW4MJ~%R zIL7OIQOF8zU(h4PXg2EhBf40OWn*qXs>h36*)Hg_T34P`MJ)j^wjyMcEcp&m<10k% z(qi{HLLS2HwQRS_evf1li)Z7POtqwIQqDANrC=0{YtooqC~9WO9+OlpXHvyBsg#G? zx@qJj#ZV>1?d23*SFY-sH5va8Vx7(~PBdS#OH{L<+f%S$GoUFLTE1lD?1E{)u&_Xt zd7O`2o}rqe;-5Oof23%!;_gmQf&9 zxd_YR60{#69Z8XV+LCQbR4G`RWY1`lHLv9g`GTfO+KrqxkHn-x9$vH@a11ZVii@XU zN!Nf;Tcfa1fr-%WHpi23kETiIENP)&&q!BI0H~)KY8qw{4dd+fpv|1-Jw}5XJYIu2h`HRy>zHe#NIxgKh|ylFFzxnk!8~CA6m2TUas{|tiKNcuhhf_%o8UG2_>h=ZGsOV3ADSJ zn3Iq%n&v#lZ=mA{maS0RiMnJ8rZYfGh8$!TxgF({yyoyaY;z)zejp%FRF0Fl?4=er zBsek4Hs@trm5I|GFi*dN06p4?c!->cm%*GqR|P(UL1d)ExOnHQQH-V#721=9ng7Za z3{Dc94&8}6qHVP;9{N9YEFm_`tuJ4`TqZ4^Ngi&IF9R>U{?reFtl)gDnBpgD|ZKN%lNV?0%BiT~CZO5+k+5 zNV8|Ka>+t>{WLXx<^=il#N<1BxKI>B|NR8+ja2}?!^A=%cs7Vr;7EKK9I4EpBUOb* zOgPf~0t&C6aPN~l+^=}AdVpvFc>OjT#3YCoU2YWMwu=H8cA8>UlHuqscsSyMQm|z( zaXNydkxgTBgzJi4(&!!-?=YDNymp#P(sL@`XpX&!(S)`f2ZYuw0Bu)n@sZ@AXCr$X zBa_Yaz}nQ^@me~wz7GIi$G_d;*wLNO34qTLz?lwi57q=Yz=7elL|_ICnt>0t9y$jT zaOf7)!Q_{FGUUOoGaQ#;DaJpFc_1Trsu;(C0|;F8xXy%&uAn1%eiREL(&Db|m4d0x zsscd2c^2!g*wvxcp%$Bp?P`&2z;lkod)td%jFb2de^WFDvTXx`O)&prCrOOnwd-Bu zjjnOfM(nH29;ud)nmhU`-87A2bprq!LBTu`GthIxIK&HAK-!OU2ym@+2!QMRW0KI{ zBHMu}h9r7xeA=ZFav9%T$;jTJG=sS%Z5;j_FdtdOiwAvc&T zCaWmLv=l0TI2gd&{cyPr_4oEQ!|Uq>m3hR5v-EZ~L)g>GL=Nf&mAK{2W&DIt6|D$g z6#^uQ66Oj~mf;h|mSxWkyA5qY)~-Vebb~&{)RbWz-qQ5E6O-kfu2_~VTc{+-GcY;{ z{l(tS3^eH@JquIT*8reViOFhoHTqN>cp?t0^EGjxF79oJdu!s}XX&9v`d~dh(MV5J z_|;eie(u3wYQjo%`t?d(2F!Oznask|Yp+*gST#__bnfua+ehI(J6CL$Wm^vd6-$$_ zA;Y_xa=I|nLJ$KDQYVdtxNMoG(V23EoNiiLhM`AcpM|m4Gh>H?y^E=k8e!x-t~v<- zY6CH`8ha`ZKM{xP;%Gx0t%;-Uv3*a(ee3qUoA-WJ7f&_BQ#Jqc%sYtifG4oXKfz1b z9yc=qh##mWEc6ulQQs0%X4VnUD{ZTySNl6Ltp^s_xv_J_EXFRkg{i`)dZ6lvXQ?q% zuB}*emZ$niPBC(tPEWvguvaKsAYP)59o}8#T5LOFV6a>43#mG&b1uQ0g#auT&2sO$ zdaqDVA8n+MVuXTd2%s|6Or$=(aO*<#((0uO{I;IMPwV2zhIq2(U+y^oL+A_BnJBuQ zp%gV)cH**JbpKC)enOUiTvBvz#`PP_iB2GL!l@SdGT;j2Qp@acqOz=-IX{NP;hQ&Cq7;IN zS5ig`2L5?qyP3<0xu21ysatM3b7Ec@34fzd%;`?Tgl}J^Fsb3kFU11J5!`QXFVMSE zVx;#tz*6xG@hwkD_Xl%2hideh5^AD9^?Zc5@yG8H~#)bMrz@!N%q&mSCb6Y z!dH{*u7$5Ak!s-!h>Y&}Yk!M7&h)kjfM(y&=8HSx0i8)OpyOA;+R5iQ`uS_7T#^3= DmrbUR literal 0 HcmV?d00001 diff --git a/litellm/main.py b/litellm/main.py index ec2de634f..0cc0ac5e4 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6,8 +6,7 @@ import traceback import dotenv import traceback import litellm -from litellm import client, logging, exception_type -from litellm import success_callback, failure_callback +from litellm import client, logging, exception_type, timeout, success_callback, failure_callback import random ####### ENVIRONMENT VARIABLES ################### dotenv.load_dotenv() # Loading env variables using dotenv @@ -59,7 +58,7 @@ def get_optional_params( ####### COMPLETION ENDPOINTS ################ ############################################# @client -@func_set_timeout(180, allowOverride=True) ## https://pypi.org/project/func-timeout/ - timeouts, in case calls hang (e.g. Azure) +@timeout(60) ## set timeouts, in case calls hang (e.g. Azure) def completion( model, messages, # required params # Optional OpenAI params: see https://platform.openai.com/docs/api-reference/chat/create @@ -67,7 +66,7 @@ def completion( temperature=1, top_p=1, n=1, stream=False, stop=None, max_tokens=float('inf'), presence_penalty=0, frequency_penalty=0, logit_bias={}, user="", # Optional liteLLM function params - *, forceTimeout=60, azure=False, logger_fn=None, verbose=False + *, force_timeout=60, azure=False, logger_fn=None, verbose=False ): try: # check if user passed in any of the OpenAI optional params diff --git a/litellm/tests/test_exceptions.py b/litellm/tests/test_exceptions.py index 38be0e2c1..526e85289 100644 --- a/litellm/tests/test_exceptions.py +++ b/litellm/tests/test_exceptions.py @@ -1,11 +1,3 @@ -from openai.error import AuthenticationError, InvalidRequestError, RateLimitError, OpenAIError -import os -import sys -import traceback -sys.path.insert(0, os.path.abspath('../..')) # Adds the parent directory to the system path -import litellm -from litellm import embedding, completion -from concurrent.futures import ThreadPoolExecutor #### What this tests #### # This tests exception mapping -> trigger an exception from an llm provider -> assert if output is of the expected type @@ -16,6 +8,15 @@ from concurrent.futures import ThreadPoolExecutor # Approach: Run each model through the test -> assert if the correct error (always the same one) is triggered +from openai.error import AuthenticationError, InvalidRequestError, RateLimitError, OpenAIError +import os +import sys +import traceback +sys.path.insert(0, os.path.abspath('../..')) # Adds the parent directory to the system path +import litellm +from litellm import embedding, completion +from concurrent.futures import ThreadPoolExecutor + models = ["gpt-3.5-turbo", "chatgpt-test", "claude-instant-1", "command-nightly", "replicate/llama-2-70b-chat:2c1608e18606fad2812020dc541930f2d0495ce32eee50074220b87300bc16e1"] # Test 1: Rate Limit Errors diff --git a/litellm/tests/test_timeout.py b/litellm/tests/test_timeout.py new file mode 100644 index 000000000..fd36effc3 --- /dev/null +++ b/litellm/tests/test_timeout.py @@ -0,0 +1,26 @@ +#### What this tests #### +# This tests the timeout decorator + +import sys, os +import traceback +sys.path.insert(0, os.path.abspath('../..')) # Adds the parent directory to the system path +import time +from litellm import timeout + +@timeout(10) +def stop_after_10_s(force_timeout=60): + print("Stopping after 10 seconds") + time.sleep(10) + return + + +start_time = time.time() + +try: + stop_after_10_s(force_timeout=1) +except: + pass + +end_time = time.time() + +print(f"total time: {end_time-start_time}") \ No newline at end of file diff --git a/litellm/timeout.py b/litellm/timeout.py new file mode 100644 index 000000000..8cbe650ed --- /dev/null +++ b/litellm/timeout.py @@ -0,0 +1,80 @@ +""" +Module containing "timeout" decorator for sync and async callables. +""" + +import asyncio + +from concurrent import futures +from inspect import iscoroutinefunction +from functools import wraps +from threading import Thread +from openai.error import Timeout + + +def timeout( + timeout_duration: float = None, exception_to_raise = Timeout +): + """ + Wraps a function to raise the specified exception if execution time + is greater than the specified timeout. + + Works with both synchronous and asynchronous callables, but with synchronous ones will introduce + some overhead due to the backend use of threads and asyncio. + + :param float timeout_duration: Timeout duration in seconds. If none callable won't time out. + :param OpenAIError exception_to_raise: Exception to raise when the callable times out. + Defaults to TimeoutError. + :return: The decorated function. + :rtype: callable + """ + + def decorator(func): + @wraps(func) + def wrapper(*args, **kwargs): + async def async_func(): + return func(*args, **kwargs) + + thread = _LoopWrapper() + thread.start() + future = asyncio.run_coroutine_threadsafe(async_func(), thread.loop) + try: + local_timeout_duration = timeout_duration + if "force_timeout" in kwargs: + local_timeout_duration = kwargs["force_timeout"] + result = future.result(timeout=local_timeout_duration) + except futures.TimeoutError: + thread.stop_loop() + raise exception_to_raise() + thread.stop_loop() + return result + + @wraps(func) + async def async_wrapper(*args, **kwargs): + try: + value = await asyncio.wait_for( + func(*args, **kwargs), timeout=timeout_duration + ) + return value + except asyncio.TimeoutError: + raise exception_to_raise() + + if iscoroutinefunction(func): + return async_wrapper + return wrapper + + return decorator + + +class _LoopWrapper(Thread): + def __init__(self): + super().__init__(daemon=True) + self.loop = asyncio.new_event_loop() + + def run(self) -> None: + self.loop.run_forever() + self.loop.call_soon_threadsafe(self.loop.close) + + def stop_loop(self): + for task in asyncio.all_tasks(self.loop): + task.cancel() + self.loop.call_soon_threadsafe(self.loop.stop) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 37bc975e4..313eb0f5a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,5 @@ openai cohere -func_timeout anthropic replicate pytest diff --git a/setup.py b/setup.py index 230bd8596..c84031889 100644 --- a/setup.py +++ b/setup.py @@ -2,7 +2,7 @@ from setuptools import setup, find_packages setup( name='litellm', - version='0.1.206', + version='0.1.207', description='Library to easily interface with LLM API providers', author='BerriAI', packages=[ @@ -11,11 +11,10 @@ setup( install_requires=[ 'openai', 'cohere', - 'func_timeout', 'pytest', 'anthropic', 'replicate', 'python-dotenv', - 'openai[datalib]' + 'openai[datalib]', ], )