fix: refactor

Author: exiaohu

volcenginesdkarkruntime/_utils/_model_breaker.py

@@ -1,29 +1,56 @@
-from datetime import datetime, timedelta
-import random
+import asyncio
+import math
+import time
 
 
-class ModelBreaker:
+class ModelBreaker(object):
     def __init__(self):
         # 初始化 allow_time 为当前时间
-        self.allow_time = datetime.now()
+        self._allow_time = time.perf_counter()
 
-    def allow(self):
-        # 检查当前时间是否在 allow_time 之后
-        return datetime.now() > self.allow_time
+        self._waiters = 0
+        self._permits = 0
+        self._base = 0
 
-    def reset(self, duration):
-        # 将 allow_time 重置为当前时间加上指定的持续时间
-        self.allow_time = datetime.now() + timedelta(seconds=duration.total_seconds())
+    def _allow(self) -> bool:
+        # 检查当前时间是否在 allow_time 之后
+        return time.perf_counter() > self._allow_time
 
-    def get_allowed_duration(self):
+    def _get_allowed_duration(self) -> float:
         # 计算当前时间与 allow_time 之间的持续时间
-        allow_duration = self.allow_time - datetime.now()
+        allow_duration = self._allow_time - time.perf_counter()
+
+        # 如果持续时间为负，返回零
+        if allow_duration < 0:
+            return 0
+        return allow_duration
+
+    def _acquire(self) -> int:
+        self._waiters += 1
+        return self._waiters
+
+    def _release(self) -> None:
+        self._waiters -= 1
+        self._permits += 1
+
+    def _jitter(self, i: int) -> float:
+        if i <= self._base:
+            return 0
+        return math.log2(i - self._base)
+
+    def reset(self, duration: float) -> None:
+        # 将 allow_time 重置为当前时间加上指定的持续时间
+        self._allow_time = time.perf_counter() + duration
+        self._base = self._permits
 
-        # 添加一个随机抖动，该抖动的范围为 0 到 10 秒，概率密度函数为 f(x) = x/50 (0 <= x <= 10)
-        # 约有 1/100 的请求会在结束熔断的第一秒内发起重试
-        jitter = timedelta(seconds=random.triangular(0, 10, 10))
+    def wait(self) -> None:
+        i = self._acquire()
+        while not self._allow():
+            time.sleep(self._get_allowed_duration() + self._jitter(i))
+        self._release()
 
-        # 如果持续时间为负，则返回一个零时长的 timedelta 对象
-        if allow_duration.total_seconds() < 0:
-            return timedelta(0) + jitter
-        return allow_duration + jitter
+    async def asyncwait(self) -> None:
+        i = self._acquire()
+        while not self._allow():
+            await asyncio.sleep(self._get_allowed_duration() + self._jitter(i))
+        self._release()

volcenginesdkarkruntime/resources/batch_chat/completions.py

@@ -146,10 +146,7 @@ def create(
         last_time = self._get_request_last_time(timeout)
         model_breaker = self._client.get_model_breaker(model)
         while True:
-            while not model_breaker.allow():
-                if datetime.now() + timedelta(seconds=model_breaker.get_allowed_duration().total_seconds()) > last_time:
-                    raise ArkAPITimeoutError()
-                time.sleep(model_breaker.get_allowed_duration().total_seconds())
+            model_breaker.wait()
             if datetime.now() > last_time:
                 raise ArkAPITimeoutError()
             try:
@@ -193,7 +190,7 @@ def create(
             except ArkAPIStatusError as err:
                 retry_after = _get_retry_after(err.response)
                 if retry_after is not None:
-                    model_breaker.reset(timedelta(seconds=retry_after))
+                    model_breaker.reset(retry_after)
                 if _should_retry(err.response):
                     continue
                 else:
@@ -288,10 +285,7 @@ async def create(
         last_time = self._get_request_last_time(timeout)
         model_breaker = await self._client.get_model_breaker(model)
         while True:
-            while not model_breaker.allow():
-                if datetime.now() + timedelta(seconds=model_breaker.get_allowed_duration().total_seconds()) > last_time:
-                    raise ArkAPITimeoutError()
-                await asyncio.sleep(model_breaker.get_allowed_duration().total_seconds())
+            await model_breaker.asyncwait()
             if datetime.now() > last_time:
                 raise ArkAPITimeoutError()
             try:
@@ -335,7 +329,7 @@ async def create(
             except ArkAPIStatusError as err:
                 retry_after = _get_retry_after(err.response)
                 if retry_after is not None:
-                    model_breaker.reset(timedelta(seconds=retry_after))
+                    model_breaker.reset(retry_after)
                 if _should_retry(err.response):
                     continue
                 else: