fix: add oci blackout when rate limited (#228)

* fix: add oci blackout when rate limited

* code review

Author: quitrk

skynet/env.py

@@ -111,7 +111,7 @@ def tobool(val: str | None):
 
 # jobs
 job_timeout = int(os.environ.get('JOB_TIMEOUT', 60 * 5))  # 5 minutes default
-max_concurrency = int(os.environ.get('MAX_CONCURRENCY', 10))
+max_concurrency = int(os.environ.get('MAX_CONCURRENCY', 5))
 
 # summaries
 summary_minimum_payload_length = int(os.environ.get('SUMMARY_MINIMUM_PAYLOAD_LENGTH', 100))
@@ -137,6 +137,7 @@ def tobool(val: str | None):
 oci_compartment_id = os.environ.get('OCI_COMPARTMENT_ID')
 oci_auth_type = os.environ.get('OCI_AUTH_TYPE', 'API_KEY')
 oci_config_profile = os.environ.get('OCI_CONFIG_PROFILE', 'DEFAULT')
+oci_blackout_fallback_duration = int(os.environ.get('OCI_BLACKOUT_FALLBACK_DURATION', 30))
 oci_available = oci_model_id and oci_service_endpoint and oci_compartment_id and oci_auth_type and oci_config_profile
 use_oci = oci_available and llama_path.startswith('oci://')
 

skynet/modules/stt/streaming_whisper/app.py

@@ -17,16 +17,18 @@ async def websocket_endpoint(websocket: WebSocket, meeting_id: str, auth_token:
         while connection.connected:
             try:
                 chunk = await websocket.receive_bytes()
-            except WebSocketDisconnect as dc:
+            except WebSocketDisconnect:
                 log.info(f'Meeting {connection.meeting_id} has ended')
                 await ws_connection_manager.disconnect(connection, already_closed=True)
-                break 
+                break
             except WebSocketException as wserr:
                 log.warning(f'Error on websocket {connection.meeting_id}. Error {wserr.__class__}: \n{wserr}')
                 await ws_connection_manager.disconnect(connection)
                 break
             except Exception as err:
-                log.warning(f'Expected bytes, received something else, disconnecting {connection.meeting_id}. Error {err.__class__}: \n{err}')
+                log.warning(
+                    f'Expected bytes, received something else, disconnecting {connection.meeting_id}. Error {err.__class__}: \n{err}'
+                )
                 await ws_connection_manager.disconnect(connection)
                 break
             if len(chunk) == 1 and ord(b'' + chunk) == 0:

skynet/modules/stt/streaming_whisper/connection_manager.py

@@ -41,7 +41,7 @@ async def connect(self, websocket: WebSocket, meeting_id: str, auth_token: str |
 
     async def process(self, connection: MeetingConnection, chunk: bytes, chunk_timestamp: int):
         log.debug(f'Processing chunk for meeting {connection.meeting_id}')
-        
+
         try:
             results = await connection.process(chunk, chunk_timestamp)
             await self.send(connection, results)
@@ -55,13 +55,15 @@ async def send(self, connection: MeetingConnection, results: list[utils.Transcri
                 try:
                     await connection.ws.send_json(result.model_dump())
                 except WebSocketDisconnect as e:
-                    log.warning(f'Meeting {connection.meeting_id}: the connection was closed before sending all results: {e}')
+                    log.warning(
+                        f'Meeting {connection.meeting_id}: the connection was closed before sending all results: {e}'
+                    )
                     await self.disconnect(connection, True)
                     break
                 except Exception as ex:
                     log.error(f'Meeting {connection.meeting_id}: exception while sending transcription results {ex}')
 
-    async def disconnect(self, connection: MeetingConnection, already_closed = False):
+    async def disconnect(self, connection: MeetingConnection, already_closed=False):
         try:
             self.connections.remove(connection)
         except ValueError:

skynet/modules/stt/vox/connection_manager.py

@@ -1,7 +1,10 @@
 from fastapi import WebSocketDisconnect
 
 from skynet.logs import get_logger
-from skynet.modules.stt.streaming_whisper.connection_manager import ConnectionManager as BaseConnectionManager, MeetingConnection
+from skynet.modules.stt.streaming_whisper.connection_manager import (
+    ConnectionManager as BaseConnectionManager,
+    MeetingConnection,
+)
 from skynet.modules.stt.streaming_whisper.utils.utils import TranscriptionResponse
 
 log = get_logger(__name__)
@@ -25,7 +28,9 @@ async def send(self, connection: MeetingConnection, results: list[TranscriptionR
                 )
                 log.debug(f'Participant {result.participant_id} result: {result.text}')
             except WebSocketDisconnect as e:
-                log.warning(f'Session {connection.meeting_id}: the connection was closed before sending all results: {e}')
+                log.warning(
+                    f'Session {connection.meeting_id}: the connection was closed before sending all results: {e}'
+                )
                 self.disconnect(connection, True)
             except Exception as ex:
                 log.error(f'Session {connection.meeting_id}: exception while sending transcription results {ex}')

skynet/modules/ttt/llm_selector.py

@@ -33,7 +33,7 @@ def override_job_processor(job_id: str, processor: Processors) -> None:
         overriden_processors[job_id] = processor
 
     @staticmethod
-    def get_job_processor(customer_id: str, job_id: Optional[str] = None) -> Processors:
+    def get_job_processor(customer_id: str, job_id: Optional[str] = None, oci_blackout: bool = False) -> Processors:
         if job_id and job_id in overriden_processors:
             return overriden_processors[job_id]
 
@@ -50,7 +50,7 @@ def get_job_processor(customer_id: str, job_id: Optional[str] = None) -> Process
         if api_type == CredentialsType.LOCAL.value:
             return Processors.LOCAL
 
-        if oci_available:
+        if oci_available and not oci_blackout:
             return Processors.OCI
 
         return Processors.LOCAL
@@ -62,8 +62,9 @@ def select(
         max_completion_tokens: Optional[int] = None,
         temperature: Optional[float] = 0,
         stream: Optional[bool] = False,
+        oci_blackout: bool = False,
     ) -> BaseChatModel:
-        processor = LLMSelector.get_job_processor(customer_id, job_id)
+        processor = LLMSelector.get_job_processor(customer_id, job_id, oci_blackout)
         options = get_credentials(customer_id)
 
         if processor == Processors.OPENAI:

skynet/modules/ttt/processor.py

@@ -1,4 +1,5 @@
 import json
+from datetime import datetime, timedelta, timezone
 from operator import itemgetter
 from typing import List, Optional
 
@@ -10,11 +11,13 @@
 from langchain_core.documents import Document
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.output_parsers import StrOutputParser
+
+from oci.exceptions import TransientServiceError
 from openai.types.chat import ChatCompletionMessageParam
 
 from skynet.constants import response_prefix
 
-from skynet.env import llama_n_ctx, modules, use_oci
+from skynet.env import llama_n_ctx, modules, oci_blackout_fallback_duration, use_oci
 from skynet.logs import get_logger
 from skynet.modules.monitoring import MAP_REDUCE_CHUNKING_COUNTER
 from skynet.modules.ttt.assistant.constants import assistant_rag_question_extractor
@@ -44,6 +47,31 @@
 
 log = get_logger(__name__)
 
+# Global OCI blackout state management
+_oci_blackout_until: Optional[datetime] = None
+
+
+def set_oci_blackout(duration_seconds: int) -> None:
+    """Set OCI blackout for the specified duration."""
+    global _oci_blackout_until
+    _oci_blackout_until = datetime.now(timezone.utc) + timedelta(seconds=duration_seconds)
+    log.warning(f"OCI blackout set until {_oci_blackout_until} ({duration_seconds} seconds)")
+
+
+def is_oci_blackout_active() -> bool:
+    """Check if OCI is currently in blackout period."""
+    global _oci_blackout_until
+    if _oci_blackout_until is None:
+        return False
+
+    now = datetime.now(timezone.utc)
+    if now >= _oci_blackout_until:
+        _oci_blackout_until = None  # Clear expired blackout
+        log.info("OCI blackout period expired, resuming normal processing")
+        return False
+
+    return True
+
 
 hint_type_to_prompt = {
     JobType.SUMMARY: {
@@ -201,7 +229,12 @@ async def process(job: Job) -> str:
     job_type = job.type
     customer_id = job.metadata.customer_id
 
-    llm = LLMSelector.select(customer_id, job_id=job.id, **{'max_completion_tokens': payload.max_completion_tokens})
+    llm = LLMSelector.select(
+        customer_id,
+        job_id=job.id,
+        oci_blackout=is_oci_blackout_active(),
+        **{'max_completion_tokens': payload.max_completion_tokens},
+    )
 
     try:
         if job_type == JobType.ASSIST:
@@ -212,6 +245,18 @@ async def process(job: Job) -> str:
             result = await process_text(llm, payload)
         else:
             raise ValueError(f'Invalid job type {job_type}')
+    except TransientServiceError as e:
+        log.warning(f"Job {job.id} hit TransientServiceError: {e}")
+
+        # Set blackout using fallback duration
+        blackout_duration = oci_blackout_fallback_duration
+        log.info(f"TransientServiceError detected, setting {blackout_duration}s blackout")
+        set_oci_blackout(blackout_duration)
+
+        # Switch current job to local processing
+        LLMSelector.override_job_processor(job.id, Processors.LOCAL)
+        return await process(job)
+
     except Exception as e:
         log.warning(f"Job {job.id} failed: {e}")
 

skynet/modules/ttt/processor_test.py

@@ -1,5 +1,7 @@
 import pytest
 
+from oci.exceptions import TransientServiceError
+
 from skynet.modules.ttt.summaries.v1.models import DocumentMetadata, DocumentPayload, Job, JobType, Processors
 
 
@@ -249,3 +251,52 @@ async def test_process_with_oci_error_fallback(self, process_fixture):
 
         assert LLMSelector.get_job_processor(job.metadata.customer_id, job.id) == Processors.LOCAL
         assert LLMSelector.select.call_count == 2
+
+    @pytest.mark.asyncio
+    async def test_process_with_transient_service_error_blackout(self, process_fixture):
+        '''Test that TransientServiceError triggers blackout and subsequent jobs use LOCAL processor.'''
+
+        from skynet.modules.ttt.llm_selector import LLMSelector
+        from skynet.modules.ttt.processor import process
+
+        # Create TransientServiceError with circuit breaker message
+        circuit_breaker_msg = (
+            'Circuit "test-id" OPEN until 2025-09-04 13:23:43.823175+00:00 (12 failures, 17 sec remaining)'
+        )
+        transient_error = TransientServiceError(status=429, code='429', headers={}, message=circuit_breaker_msg)
+
+        process_fixture.patch(
+            'skynet.modules.ttt.llm_selector.get_credentials',
+            return_value={'type': 'OCI'},
+        )
+        process_fixture.patch('skynet.modules.ttt.llm_selector.oci_available', True)
+        process_fixture.patch('skynet.modules.ttt.processor.use_oci', False)  # allow fallback
+        process_fixture.patch('skynet.modules.ttt.processor.summarize', side_effect=[transient_error, None, None])
+
+        job1 = Job(
+            payload=DocumentPayload(text="First job"),
+            metadata=DocumentMetadata(customer_id='test'),
+            type=JobType.SUMMARY,
+        )
+
+        job2 = Job(
+            payload=DocumentPayload(text="Second job"),
+            metadata=DocumentMetadata(customer_id='test'),
+            type=JobType.SUMMARY,
+        )
+
+        # First job should trigger TransientServiceError and set blackout
+        await process(job1)
+        assert LLMSelector.get_job_processor(job1.metadata.customer_id, job1.id) == Processors.LOCAL
+
+        # Second job should immediately go to LOCAL due to active blackout
+        await process(job2)
+        # Check that blackout causes LOCAL processor selection
+        from skynet.modules.ttt.processor import is_oci_blackout_active
+
+        blackout_active = is_oci_blackout_active()
+        assert blackout_active == True, "Blackout should be active after TransientServiceError"
+        assert (
+            LLMSelector.get_job_processor(job2.metadata.customer_id, job2.id, oci_blackout=blackout_active)
+            == Processors.LOCAL
+        )