updates streamer and stream manager

validators/sn1_validator_wrapper.py

@@ -48,17 +48,12 @@ class S1ValidatorAPI(ValidatorAPI):
             deserialize=False,
             streaming=True,
         )
-        # uid_stream_dict = dict(zip(uids, streams_responses))
-        # random_uid, random_stream = random.choice(list(uid_stream_dict.items()))                        
         
         # Creates a streamer from the selected stream
         stream_manager = StreamManager()
-        await stream_manager.process_streams(params.request, streams_responses, uids)
+        selected_stream = await stream_manager.process_streams(params.request, streams_responses, uids)
         
-        # response = await streamer.stream(params.request)   
-        # return response
-        
-        return None
+        return selected_stream
 
 
     async def query_validator(self, params: QueryValidatorParams) -> Response:

validators/stream_manager.py

@@ -1,32 +1,28 @@
 import asyncio
+import bittensor as bt
 from .streamer import AsyncResponseDataStreamer
+from .database import LogDatabase
 from typing import List, AsyncIterator
 from aiohttp.web import Request
 
+
 class StreamManager:
-    
-    def __init__(self):
-        ...
-    
-    
+    def __init__(self, log_database_path: str = "requests_db.jsonl"):
+        self.log_database = LogDatabase(log_database_path)
     
     async def process_streams(self, request:Request, streams_responses: List[AsyncIterator], stream_uids: List[int]):
-        
-        # create local lock for returning responses to the front-end
-        # creates n number of streamers
-        # organizes responses
-        # logs responses locally
-        # returns selected response to the front-end
-        
         lock = asyncio.Lock()
         
         streamers = [AsyncResponseDataStreamer(async_iterator=stream, selected_uid=stream_uid, lock=lock) for stream, stream_uid in zip(streams_responses, stream_uids)]            
         completed_streams = await asyncio.gather(*[streamer.stream(request) for streamer in streamers])
         
         lock.release()
-        print(f"Stream {stream_uids} completed the operation.")
+        bt.logging.info(f"Streams from uids: {stream_uids} processing completed.")
         
-                        
-
-
 
+        await self.log_database.add_streams_to_db(completed_streams)
+        # Gets the first stream that acquired the lock, meaning the first stream that was able to return a non-empty chunk
+        _, selected_stream = next(((streamer, completed_stream) for streamer, completed_stream in zip(streamers, completed_streams) if streamer.lock_acquired), None)
+        
+        return selected_stream 
+        

validators/streamer.py

@@ -34,6 +34,7 @@ class StreamError(BaseModel):
         data = json.dumps(self.dict(), indent=4)
         return data.encode(encoding)
 
+ProcessedStreamResponse = Union[StreamChunk, StreamError]
 
 class AsyncResponseDataStreamer:
     def __init__(self, async_iterator: AsyncIterator, selected_uid:int, lock: asyncio.Lock, delay: float = 0.1):
@@ -76,8 +77,7 @@ class AsyncResponseDataStreamer:
         return initiated_response
         
     
-
-    async def stream(self, request: web.Request) -> web_response.StreamResponse:    
+    async def stream(self, request: web.Request) -> ProcessedStreamResponse:
         # response = web_response.StreamResponse(status=200, reason="OK")
         # response.headers["Content-Type"] = "application/json"
         # await response.prepare(request)  # Prepare and send the headers
@@ -85,7 +85,7 @@ class AsyncResponseDataStreamer:
         try:
             start_time = time.time()
             client_response: web.Response = None
-            final_response: Union[StreamChunk, StreamError]
+            final_response: ProcessedStreamResponse
             
             async for chunk in self.async_iterator:
                  if isinstance(chunk, str):
@@ -124,13 +124,13 @@ class AsyncResponseDataStreamer:
             final_response = error_response
             
             # Only the stream that acquires the lock should write the error response
-            if self.lock.locked():
+            if self.lock_acquired:
                 self.ensure_response_is_created(client_response)
                 client_response.set_status(500, reason="Internal error")
                 client_response.write(error_response.encode('utf-8'))            
         finally:
             # Only the stream that acquires the lock should close the response
-            if self.lock.locked():                
+            if self.lock_acquired:                
                 self.ensure_response_is_created(client_response)
                 # Ensure to close the response properly
                 await client_response.write_eof()