fix: fix TPM & RPM

Author: ChenZiHong-Gavin

graphgen/models/llm/limitter.py

@@ -1,17 +1,17 @@
+import asyncio
 import time
 from datetime import datetime, timedelta
-import asyncio
 
 from graphgen.utils import logger
 
 
 class RPM:
-
     def __init__(self, rpm: int = 1000):
         self.rpm = rpm
-        self.record = {'rpm_slot': self.get_minute_slot(), 'counter': 0}
+        self.record = {"rpm_slot": self.get_minute_slot(), "counter": 0}
 
-    def get_minute_slot(self):
+    @staticmethod
+    def get_minute_slot():
         current_time = time.time()
         dt_object = datetime.fromtimestamp(current_time)
         total_minutes_since_midnight = dt_object.hour * 60 + dt_object.minute
@@ -22,37 +22,35 @@ async def wait(self, silent=False):
         dt_object = datetime.fromtimestamp(current)
         minute_slot = self.get_minute_slot()
 
-        if self.record['rpm_slot'] == minute_slot:
+        if self.record["rpm_slot"] == minute_slot:
             # check RPM exceed
-            if self.record['counter'] >= self.rpm:
+            if self.record["counter"] >= self.rpm:
                 # wait until next minute
-                next_minute = dt_object.replace(
-                    second=0, microsecond=0) + timedelta(minutes=1)
+                next_minute = dt_object.replace(second=0, microsecond=0) + timedelta(
+                    minutes=1
+                )
                 _next = next_minute.timestamp()
                 sleep_time = abs(_next - current)
                 if not silent:
-                    logger.info('RPM sleep %s', sleep_time)
+                    logger.info("RPM sleep %s", sleep_time)
                 await asyncio.sleep(sleep_time)
 
-                self.record = {
-                    'rpm_slot': self.get_minute_slot(),
-                    'counter': 0
-                }
+                self.record = {"rpm_slot": self.get_minute_slot(), "counter": 0}
         else:
-            self.record = {'rpm_slot': self.get_minute_slot(), 'counter': 0}
-        self.record['counter'] += 1
+            self.record = {"rpm_slot": self.get_minute_slot(), "counter": 0}
+        self.record["counter"] += 1
 
         if not silent:
             logger.debug(self.record)
 
 
 class TPM:
-
-    def __init__(self, tpm: int = 20000):
+    def __init__(self, tpm: int = 100000):
         self.tpm = tpm
-        self.record = {'tpm_slot': self.get_minute_slot(), 'counter': 0}
+        self.record = {"tpm_slot": self.get_minute_slot(), "counter": 0}
 
-    def get_minute_slot(self):
+    @staticmethod
+    def get_minute_slot():
         current_time = time.time()
         dt_object = datetime.fromtimestamp(current_time)
         total_minutes_since_midnight = dt_object.hour * 60 + dt_object.minute
@@ -64,25 +62,24 @@ async def wait(self, token_count, silent=False):
         minute_slot = self.get_minute_slot()
 
         # get next slot, skip
-        if self.record['tpm_slot'] != minute_slot:
-            self.record = {'tpm_slot': minute_slot, 'counter': token_count}
+        if self.record["tpm_slot"] != minute_slot:
+            self.record = {"tpm_slot": minute_slot, "counter": token_count}
             return
 
         # check RPM exceed
-        self.record['counter'] += token_count
-        if self.record['counter'] > self.tpm:
+        self.record["counter"] += token_count
+        if self.record["counter"] > self.tpm:
             # wait until next minute
-            next_minute = dt_object.replace(
-                second=0, microsecond=0) + timedelta(minutes=1)
+            next_minute = dt_object.replace(second=0, microsecond=0) + timedelta(
+                minutes=1
+            )
             _next = next_minute.timestamp()
             sleep_time = abs(_next - current)
-            logger.info('TPM sleep %s', sleep_time)
+            logger.info("Current TPM: %s, limit: %s", self.record["counter"], self.tpm)
+            logger.warning("TPM limit exceeded, wait %s seconds", sleep_time)
             await asyncio.sleep(sleep_time)
 
-            self.record = {
-                'tpm_slot': self.get_minute_slot(),
-                'counter': token_count
-            }
+            self.record = {"tpm_slot": self.get_minute_slot(), "counter": token_count}
 
         if not silent:
             logger.debug(self.record)

graphgen/models/llm/openai_client.py

@@ -39,6 +39,8 @@ def __init__(
         seed: Optional[int] = None,
         topk_per_token: int = 5,  # number of topk tokens to generate for each token
         request_limit: bool = False,
+        rpm: RPM = None,
+        tpm: TPM = None,
         **kwargs: Any,
     ):
         super().__init__(**kwargs)
@@ -51,8 +53,8 @@ def __init__(
 
         self.token_usage: list = []
         self.request_limit = request_limit
-        self.rpm = RPM(rpm=1000)
-        self.tpm = TPM(tpm=50000)
+        self.rpm = rpm or RPM()
+        self.tpm = tpm or TPM()
 
         self.__post_init__()
 

webui/app.py

@@ -468,7 +468,7 @@ def sum_tokens(client):
                         label="TPM",
                         minimum=5000,
                         maximum=5000000,
-                        value=50000,
+                        value=100000,
                         step=1000,
                         interactive=True,
                         visible=True,