Patch: deberta success rate throttle

cleanlab_studio/errors.py

@@ -86,6 +86,12 @@ def __init__(self, message: str, status_code: int) -> None:
         self.status_code = status_code
 
 
+class TlmPartialSuccess(APIError):
+    """TLM request partially succeeded. Still returns result to user."""
+
+    pass
+
+
 class UnsupportedVersionError(HandledError):
     def __init__(self) -> None:
         super().__init__(

cleanlab_studio/internal/api/api.py

@@ -10,6 +10,7 @@
     InvalidProjectConfiguration,
     RateLimitError,
     TlmBadRequest,
+    TlmPartialSuccess,
     TlmServerError,
 )
 from cleanlab_studio.internal.tlm.concurrency import TlmRateHandler
@@ -651,6 +652,9 @@ async def tlm_prompt(
             await handle_tlm_client_error_from_resp(res, batch_index)
             await handle_tlm_api_error_from_resp(res, batch_index)
 
+            if not res_json.get("deberta_success", True):
+                raise TlmPartialSuccess("Partial failure on deberta call -- slowdown request rate.")
+
     finally:
         if local_scoped_client:
             await client_session.close()

cleanlab_studio/internal/tlm/concurrency.py

@@ -2,7 +2,7 @@
 from types import TracebackType
 from typing import Optional, Type
 
-from cleanlab_studio.errors import RateLimitError, TlmServerError
+from cleanlab_studio.errors import RateLimitError, TlmPartialSuccess, TlmServerError
 
 
 class TlmRateHandler:
@@ -50,6 +50,8 @@ async def __aexit__(
         If request failed due to 503, decrease congestion window.
         Else if request failed for other reason, don't change congestion window, just exit.
         """
+        swallow_exception: bool = False
+
         if exc_type is None:
             await self._increase_congestion_window()
 
@@ -60,10 +62,14 @@ async def __aexit__(
         ):
             await self._decrease_congestion_window()
 
+        elif isinstance(exc, TlmPartialSuccess):
+            await self._decrease_congestion_window()
+            swallow_exception = True
+
         # release acquired send semaphore from aenter
         self._send_semaphore.release()
 
-        return False
+        return swallow_exception
 
     async def _increase_congestion_window(
         self,

cleanlab_studio/version.py

@@ -1,7 +1,7 @@
 # Note to developers:
 # Consider if backend's MIN_CLI_VERSION needs updating when pushing any changes to this file.
 
-__version__ = "2.0.5"
+__version__ = "2.0.6"
 
 SCHEMA_VERSION = "0.2.0"
 MIN_SCHEMA_VERSION = "0.1.0"