adds stream manager

validators/__init__.py

@@ -1,3 +1,4 @@
 from .base import QueryValidatorParams, ValidatorAPI, MockValidator
 from .sn1_validator_wrapper import S1ValidatorAPI
 from .streamer import AsyncResponseDataStreamer
+from .stream_manager import StreamManager

validators/sn1_validator_wrapper.py

@@ -1,3 +1,4 @@
+import asyncio
 import random
 import bittensor as bt
 from prompting.validator import Validator
@@ -7,7 +8,7 @@ from .base import QueryValidatorParams, ValidatorAPI
 from aiohttp.web_response import Response, StreamResponse
 from .streamer import AsyncResponseDataStreamer
 from .validator_utils import get_top_incentive_uids
-
+from .stream_manager import StreamManager
 
 class S1ValidatorAPI(ValidatorAPI):
     def __init__(self):
@@ -47,13 +48,18 @@ 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()))                        
+        # 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
-        streamer = AsyncResponseDataStreamer(async_iterator=random_stream, selected_uid=random_uid)        
-        response = await streamer.stream(params.request)                        
-        return response
+        stream_manager = StreamManager()
+        await stream_manager.process_streams(params.request, streams_responses, uids)
+        
+        # response = await streamer.stream(params.request)   
+        # return response
+        
+        return None
+
 
     async def query_validator(self, params: QueryValidatorParams) -> Response:
         return await self.get_stream_response(params)

validators/stream_manager.py

@@ -0,0 +1,32 @@
+import asyncio
+from .streamer import AsyncResponseDataStreamer
+from typing import List, AsyncIterator
+from aiohttp.web import Request
+
+class StreamManager:
+    
+    def __init__(self):
+        ...
+    
+    
+    
+    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.")
+        
+                        
+
+
+

validators/streamer.py

@@ -1,10 +1,11 @@
 import json
 import time
+import asyncio
 import traceback
 import bittensor as bt
 from pydantic import BaseModel
 from datetime import datetime
-from typing import AsyncIterator, Optional, List
+from typing import AsyncIterator, Optional, List, Union
 from aiohttp import web, web_response
 from prompting.protocol import StreamPromptingSynapse
 
@@ -35,7 +36,7 @@ class StreamError(BaseModel):
 
 
 class AsyncResponseDataStreamer:
-    def __init__(self, async_iterator: AsyncIterator, selected_uid:int, delay: float = 0.1):
+    def __init__(self, async_iterator: AsyncIterator, selected_uid:int, lock: asyncio.Lock, delay: float = 0.1):
         self.async_iterator = async_iterator        
         self.delay = delay
         self.selected_uid = selected_uid
@@ -43,16 +44,51 @@ class AsyncResponseDataStreamer:
         self.accumulated_chunks_timings: List[float] = []
         self.finish_reason: str = None
         self.sequence_number: int = 0
+        self.lock = lock
+        self.lock_acquired = False
+        
+    
+    def ensure_response_is_created(self, initiated_response: web.StreamResponse) -> web.StreamResponse:
+        # Creates response if it was not created
+        if initiated_response == None:
+            initiated_response = web_response.StreamResponse(status=200, reason="OK")
+            initiated_response.headers["Content-Type"] = "application/json"
+            return initiated_response
+        
+        return initiated_response
+    
+        
+    async def write_to_stream(self, request: web.Request, initiated_response: web.StreamResponse, stream_chunk: StreamChunk, lock: asyncio.Lock) -> web.StreamResponse:
+        # Try to acquire the lock and sets the lock_acquired flag. Only the stream that acquires the lock should write to the response            
+        if lock.locked() == False:            
+            self.lock_acquired = await lock.acquire()
+        
+        if initiated_response == None and self.lock_acquired:
+            initiated_response = self.ensure_response_is_created(initiated_response)
+            # Prepare and send the headers
+            await initiated_response.prepare(request)
+            
+        if self.lock_acquired:
+            await initiated_response.write(stream_chunk.encode('utf-8'))
+        else:
+            bt.logging.info(f"Stream of uid {stream_chunk.selected_uid} was not the first to return, skipping...")
+        
+        return initiated_response
+        
+    
 
     async def stream(self, request: web.Request) -> web_response.StreamResponse:    
-        response = web_response.StreamResponse(status=200, reason="OK")
-        response.headers["Content-Type"] = "application/json"
-        await response.prepare(request)  # Prepare and send the headers
-                        
+        # response = web_response.StreamResponse(status=200, reason="OK")
+        # response.headers["Content-Type"] = "application/json"
+        # await response.prepare(request)  # Prepare and send the headers
+
         try:
             start_time = time.time()
+            client_response: web.Response = None
+            final_response: Union[StreamChunk, StreamError]
+            
             async for chunk in self.async_iterator:
-                 if isinstance(chunk, list):
+                 if isinstance(chunk, str):
                     # Chunks are currently returned in string arrays, so we need to concatenate them
                     concatenated_chunks = "".join(chunk)
                     self.accumulated_chunks.append(concatenated_chunks)                    
@@ -60,8 +96,9 @@ class AsyncResponseDataStreamer:
                     # Gets new response state
                     self.sequence_number += 1
                     new_response_state = self._create_chunk_response(concatenated_chunks)                    
-                    # Writes the new response state to the response                                    
-                    await response.write(new_response_state.encode('utf-8'))
+                    # Writes the new response state to the response
+                    client_response = await self.write_to_stream(request, client_response, new_response_state, self.lock)
+                    #await response.write(new_response_state.encode('utf-8'))
                                 
             if chunk is not None and isinstance(chunk, StreamPromptingSynapse):
                 if len(self.accumulated_chunks) == 0:                    
@@ -72,19 +109,33 @@ class AsyncResponseDataStreamer:
                 self.sequence_number += 1
                 # Assuming the last chunk holds the last value yielded which should be a synapse with the completion filled
                 synapse = chunk 
-                final_state = self._create_chunk_response(synapse.completion)
-                await response.write(final_state.encode('utf-8'))
+                final_response = self._create_chunk_response(synapse.completion)
+                
+                if synapse.completion:                
+                    client_response = await self.write_to_stream(request, client_response, final_response, self.lock)
+            else:
+                raise ValueError("Stream did not return a valid synapse.")
                                         
         except Exception as e:
             bt.logging.error(
-                f"Encountered an error in {self.__class__.__name__}:get_stream_response:\n{traceback.format_exc()}"
+                f"Encountered an error while processing stream for uid {self.selected_uid} get_stream_response:\n{traceback.format_exc()}"
             )
-            response.set_status(500, reason="Internal error")
             error_response = self._create_error_response(str(e))
-            response.write(error_response.encode('utf-8'))
+            final_response = error_response
+            
+            # Only the stream that acquires the lock should write the error response
+            if self.lock.locked():
+                self.ensure_response_is_created(client_response)
+                client_response.set_status(500, reason="Internal error")
+                client_response.write(error_response.encode('utf-8'))            
         finally:
-            await response.write_eof()  # Ensure to close the response properly
-            return response
+            # Only the stream that acquires the lock should close the response
+            if self.lock.locked():                
+                self.ensure_response_is_created(client_response)
+                # Ensure to close the response properly
+                await client_response.write_eof()
+                
+            return final_response
 
     def _create_chunk_response(self, chunk: str) -> StreamChunk:
         """