Merge pull request #22 from macrocosm-os/features/distributed-streaming

common/utils.py

@@ -153,7 +153,6 @@ async def echo_stream(request: web.Request) -> web.StreamResponse:
     k = request_data.get("k", 1)
     message = "\n\n".join(request_data["messages"])
 
-
     echo_iterator = EchoAsyncIterator(message, k, delay=0.3)
     streamer = AsyncResponseDataStreamer(echo_iterator, selected_uid=0, delay=0.3)
 

requirements.txt

@@ -2,3 +2,4 @@ git+https://github.com/macrocosm-os/prompting.git
 aiohttp
 deprecated
 aiohttp_apispec>=2.2.3
+aiofiles

server.py

@@ -2,7 +2,13 @@ import asyncio
 
 import bittensor as bt
 from aiohttp import web
-from aiohttp_apispec import docs, request_schema, response_schema, setup_aiohttp_apispec, validation_middleware
+from aiohttp_apispec import (
+    docs,
+    request_schema,
+    response_schema,
+    setup_aiohttp_apispec,
+    validation_middleware,
+)
 
 from common import utils
 from common.middlewares import api_key_middleware, json_parsing_middleware
@@ -10,11 +16,7 @@ from common.schemas import QueryChatSchema, StreamChunkSchema, StreamErrorSchema
 from validators import QueryValidatorParams, S1ValidatorAPI, ValidatorAPI
 
 
-@docs(
-    tags=["Prompting API"],
-    summary="Chat",
-    description="Chat endpoint."
-)
+@docs(tags=["Prompting API"], summary="Chat", description="Chat endpoint.")
 @request_schema(QueryChatSchema)
 @response_schema(StreamChunkSchema, 200)
 @response_schema(StreamErrorSchema, 400)
@@ -32,7 +34,7 @@ async def chat(request: web.Request) -> web.StreamResponse:
 @docs(
     tags=["Prompting API"],
     summary="Echo test",
-    description="Echo endpoint for testing purposes."
+    description="Echo endpoint for testing purposes.",
 )
 @request_schema(QueryChatSchema)
 @response_schema(StreamChunkSchema, 200)
@@ -45,7 +47,9 @@ class ValidatorApplication(web.Application):
     def __init__(self, validator_instance=None, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-        self["validator"] = validator_instance if validator_instance else S1ValidatorAPI()
+        self["validator"] = (
+            validator_instance if validator_instance else S1ValidatorAPI()
+        )
 
         # Add middlewares to application
         self.add_routes(

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/database.py

@@ -0,0 +1,85 @@
+import os
+import json
+import aiofiles
+import bittensor as bt
+from .streamer import ProcessedStreamResponse
+
+
+class LogDatabase:
+    """
+    A class to manage a log database stored as a JSONL (JSON Lines) file.
+
+    Attributes:
+        log_database_path (str): The path to the log database file.
+
+    Methods:
+        ensure_db_exists(file_path):
+            Ensures that the log database file exists. If it doesn't, an empty file is created.
+        
+        add_streams_to_db(stream_responses: ProcessedStreamResponse):
+            Asynchronously adds stream responses to the log database.
+        
+        append_dicts_to_file(file_path, dictionaries):
+            Asynchronously appends a list of dictionaries to the specified file.
+    """
+    
+    def __init__(self, log_database_path: str):
+        """
+        Initializes the LogDatabase with the given log database file path.
+
+        Args:
+            log_database_path (str): The path to the log database file.
+        """
+        self.log_database_path = log_database_path
+        self.ensure_db_exists(log_database_path)
+
+    def ensure_db_exists(self, file_path):
+        """
+        Ensures that the log database file exists. If it doesn't, creates an empty JSONL file.
+
+        Args:
+            file_path (str): The path to the log database file.
+        """
+        if not os.path.exists(file_path):
+            # Create an empty JSONL file
+            with open(file_path, "w") as file:
+                pass
+            # TODO: change log to debug
+            bt.logging.info(f"File '{file_path}' created.")
+        else:
+            bt.logging.info(f"File '{file_path}' already exists.")
+
+    async def add_streams_to_db(self, stream_responses: ProcessedStreamResponse):
+        """
+        Asynchronously adds stream responses to the log database.
+
+        Args:
+            stream_responses (ProcessedStreamResponse): A list of processed stream responses to add to the log database.
+
+        Raises:
+            Exception: If an error occurs while adding streams to the database.
+        """
+        bt.logging.info(f"Writing streams to the database...")
+        try:
+            stream_responses_dict = [
+                dict(stream_response) for stream_response in stream_responses
+            ]
+            await self.append_dicts_to_file(
+                self.log_database_path, stream_responses_dict
+            )
+            bt.logging.success("Streams added to the database.")
+        except Exception as e:
+            bt.logging.error(f"Error while adding streams to the database: {e}")
+            raise e
+
+    async def append_dicts_to_file(self, file_path, dictionaries):
+        """
+        Asynchronously appends a list of dictionaries to the specified file.
+
+        Args:
+            file_path (str): The path to the file where dictionaries will be appended.
+            dictionaries (list): A list of dictionaries to append to the file.
+        """
+        async with aiofiles.open(file_path, mode="a") as file:
+            for dictionary in dictionaries:
+                await file.write(json.dumps(dictionary) + "\n")

validators/sn1_validator_wrapper.py

@@ -1,3 +1,4 @@
+import asyncio
 import random
 import bittensor as bt
 from prompting.validator import Validator
@@ -7,13 +8,14 @@ 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):
         self.validator = Validator()
 
-    def sample_uids(self, params: QueryValidatorParams):    
+    def sample_uids(self, params: QueryValidatorParams):
         if params.sampling_mode == "random":
             uids = get_random_uids(
                 self.validator, k=params.k_miners, exclude=params.exclude or []
@@ -22,9 +24,11 @@ class S1ValidatorAPI(ValidatorAPI):
         if params.sampling_mode == "top_incentive":
             metagraph = self.validator.metagraph
             vpermit_tao_limit = self.validator.config.neuron.vpermit_tao_limit
-            
-            top_uids = get_top_incentive_uids(metagraph, k=params.k_miners, vpermit_tao_limit=vpermit_tao_limit)
-            
+
+            top_uids = get_top_incentive_uids(
+                metagraph, k=params.k_miners, vpermit_tao_limit=vpermit_tao_limit
+            )
+
             return top_uids
 
     async def get_stream_response(self, params: QueryValidatorParams) -> StreamResponse:
@@ -32,11 +36,13 @@ class S1ValidatorAPI(ValidatorAPI):
         # task_name = utils.guess_task_name(params.messages[-1])
 
         # Get the list of uids to query for this step.
-        uids =  self.sample_uids(params)
+        uids = self.sample_uids(params)
         axons = [self.validator.metagraph.axons[uid] for uid in uids]
 
         # Make calls to the network with the prompt.
-        bt.logging.info(f"Sampling dendrite by {params.sampling_mode} with roles {params.roles} and messages {params.messages}")
+        bt.logging.info(
+            f"Sampling dendrite by {params.sampling_mode} with roles {params.roles} and messages {params.messages}"
+        )
 
         streams_responses = await self.validator.dendrite(
             axons=axons,
@@ -47,13 +53,14 @@ 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
-        streamer = AsyncResponseDataStreamer(async_iterator=random_stream, selected_uid=random_uid)        
-        response = await streamer.stream(params.request)                        
-        return response
+        stream_manager = StreamManager()
+        selected_stream = await stream_manager.process_streams(
+            params.request, streams_responses, uids
+        )
+
+        return selected_stream
 
     async def query_validator(self, params: QueryValidatorParams) -> Response:
         return await self.get_stream_response(params)

validators/stream_manager.py

@@ -0,0 +1,73 @@
+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:
+    """
+    A class to manage the processing of multiple asynchronous data streams and log their responses.
+
+    Attributes:
+        log_database (LogDatabase): The log database to store stream responses.
+
+    Methods:
+        process_streams(request, streams_responses, stream_uids):
+            Processes multiple asynchronous streams, logs their responses, and returns the selected stream response.
+    """
+    
+    def __init__(self, log_database_path: str = "requests_db.jsonl"):
+        """
+        Initializes the StreamManager with the given log database file path.
+
+        Args:
+            log_database_path (str): The path to the log database file, defaults to "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],
+    ):
+        """
+        Processes multiple asynchronous streams, logs their responses, and returns the selected stream response (stream from first non-empty chunk).
+
+        Args:
+            request (Request): The web request object.
+            streams_responses (List[AsyncIterator]): A list of asynchronous iterators representing the streams.
+            stream_uids (List[int]): A list of unique IDs for the streams.
+
+        Returns:
+            ProcessedStreamResponse: The response from the selected stream.
+        """
+        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()
+        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(
+            (
+                completed_stream
+                for streamer, completed_stream in zip(streamers, completed_streams)
+                if streamer.lock_acquired
+            ),
+            None,
+        )
+
+        return selected_stream

validators/streamer.py

@@ -1,15 +1,28 @@
 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
 
 
 class StreamChunk(BaseModel):
+    """
+    A model representing a chunk of streaming data.
+
+    Attributes:
+        delta (str): The change in the stream.
+        finish_reason (Optional[str]): The reason for finishing the stream.
+        accumulated_chunks (List[str]): List of accumulated chunks.
+        accumulated_chunks_timings (List[float]): Timings for the accumulated chunks.
+        timestamp (str): The timestamp of the chunk.
+        sequence_number (int): The sequence number of the chunk.
+        selected_uid (int): The selected user ID.
+    """
     delta: str
     finish_reason: Optional[str]
     accumulated_chunks: List[str]
@@ -17,74 +30,207 @@ class StreamChunk(BaseModel):
     timestamp: str
     sequence_number: int
     selected_uid: int
-    
+
     def encode(self, encoding: str) -> bytes:
+        """
+        Encodes the StreamChunk instance to a JSON-formatted bytes object.
+
+        Args:
+            encoding (str): The encoding to use.
+
+        Returns:
+            bytes: The encoded JSON data.
+        """
         data = json.dumps(self.dict(), indent=4)
         return data.encode(encoding)
 
 
 class StreamError(BaseModel):
+    """
+    A model representing an error in the streaming data.
+
+    Attributes:
+        error (str): The error message.
+        timestamp (str): The timestamp of the error.
+        sequence_number (int): The sequence number at the time of error.
+        finish_reason (str): The reason for finishing the stream, defaults to "error".
+    """
     error: str
     timestamp: str
     sequence_number: int
-    finish_reason: str = 'error'
+    finish_reason: str = "error"
 
     def encode(self, encoding: str) -> bytes:
         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, delay: float = 0.1):
-        self.async_iterator = async_iterator        
+    """
+    A class to manage asynchronous streaming of response data.
+
+    Attributes:
+        async_iterator (AsyncIterator): An asynchronous iterator for streaming data.
+        selected_uid (int): The selected user ID.
+        lock (asyncio.Lock): An asyncio lock to ensure exclusive access.
+        delay (float): Delay between processing chunks, defaults to 0.1 seconds.
+        accumulated_chunks (List[str]): List of accumulated chunks.
+        accumulated_chunks_timings (List[float]): Timings for the accumulated chunks.
+        finish_reason (str): The reason for finishing the stream.
+        sequence_number (int): The sequence number of the stream.
+        lock_acquired (bool): Flag indicating if the lock was acquired.
+    """
+    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
         self.accumulated_chunks: List[str] = []
         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:
+        """
+        Ensures that a StreamResponse is created if it does not already exist.
+
+        Args:
+            initiated_response (web.StreamResponse): The initiated response.
+
+        Returns:
+            web.StreamResponse: The ensured response.
+        """
+        # 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:
+        """
+        Writes a stream chunk to the response if the lock is acquired.
+
+        Args:
+            request (web.Request): The web request object.
+            initiated_response (web.StreamResponse): The initiated response.
+            stream_chunk (StreamChunk): The chunk of stream data to write.
+            lock (asyncio.Lock): The lock to ensure exclusive access.
 
-    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
-                        
+        Returns:
+            web.StreamResponse: The response with the written chunk.
+        """
+        # 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.debug(
+                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) -> ProcessedStreamResponse:
+        """
+        Streams data from the async iterator and writes it to the response.
+
+        Args:
+            request (web.Request): The web request object.
+
+        Returns:
+            ProcessedStreamResponse: The final processed stream response.
+
+        Raises:
+            ValueError: If the stream does not return a valid synapse.
+        """
         try:
             start_time = time.time()
+            client_response: web.Response = None
+            final_response: ProcessedStreamResponse
+
             async for chunk in self.async_iterator:
-                 if isinstance(chunk, list):
-                    # Chunks are currently returned in string arrays, so we need to concatenate them
-                    concatenated_chunks = "".join(chunk)
-                    self.accumulated_chunks.append(concatenated_chunks)                    
+                if isinstance(chunk, str):
+                    # If chunk is empty, skip
+                    if not chunk:
+                        continue
+                    
+                    self.accumulated_chunks.append(chunk)
                     self.accumulated_chunks_timings.append(time.time() - start_time)
                     # 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'))
-                                
+                    new_response_state = self._create_chunk_response(
+                        chunk
+                    )
+                    # Writes the new response state to the response                    
+                    client_response = await self.write_to_stream(
+                        request, client_response, new_response_state, self.lock
+                    )
+
             if chunk is not None and isinstance(chunk, StreamPromptingSynapse):
-                if len(self.accumulated_chunks) == 0:                    
-                    self.accumulated_chunks.append(chunk.completion)                    
+                if len(self.accumulated_chunks) == 0:
+                    self.accumulated_chunks.append(chunk.completion)
                     self.accumulated_chunks_timings.append(time.time() - start_time)
-                
+
                 self.finish_reason = "completed"
                 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'))
-                                        
+                synapse = chunk
+                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_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:
-            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_acquired:
+                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:
         """
@@ -100,7 +246,7 @@ class AsyncResponseDataStreamer:
             accumulated_chunks_timings=self.accumulated_chunks_timings,
             timestamp=self._current_timestamp(),
             sequence_number=self.sequence_number,
-            selected_uid=self.selected_uid
+            selected_uid=self.selected_uid,
         )
 
     def _create_error_response(self, error_message: str) -> StreamError:
@@ -113,7 +259,7 @@ class AsyncResponseDataStreamer:
         return StreamError(
             error=error_message,
             timestamp=self._current_timestamp(),
-            sequence_number=self.sequence_number
+            sequence_number=self.sequence_number,
         )
 
     def _current_timestamp(self) -> str:
@@ -122,4 +268,4 @@ class AsyncResponseDataStreamer:
 
         :return: Current timestamp as a string.
         """
-        return datetime.utcnow().isoformat()
+        return datetime.utcnow().isoformat()

validators/validator_utils.py

@@ -3,21 +3,33 @@ from prompting.utils.uids import check_uid_availability
 
 
 def get_top_incentive_uids(metagraph, k: int, vpermit_tao_limit: int) -> List[int]:
-    miners_uids = list(map(int, filter(lambda uid: check_uid_availability(metagraph, uid, vpermit_tao_limit), metagraph.uids)))
-    
+    miners_uids = list(
+        map(
+            int,
+            filter(
+                lambda uid: check_uid_availability(metagraph, uid, vpermit_tao_limit),
+                metagraph.uids,
+            ),
+        )
+    )
+
     # Builds a dictionary of uids and their corresponding incentives
     all_miners_incentives = {
         "miners_uids": miners_uids,
-        "incentives": list(map(lambda uid: metagraph.I[uid], miners_uids))
+        "incentives": list(map(lambda uid: metagraph.I[uid], miners_uids)),
     }
-    
+
     # Zip the uids and their corresponding incentives into a list of tuples
-    uid_incentive_pairs = list(zip(all_miners_incentives['miners_uids'], all_miners_incentives['incentives']))
+    uid_incentive_pairs = list(
+        zip(all_miners_incentives["miners_uids"], all_miners_incentives["incentives"])
+    )
 
     # Sort the list of tuples by the incentive value in descending order
-    uid_incentive_pairs_sorted = sorted(uid_incentive_pairs, key=lambda x: x[1], reverse=True)
+    uid_incentive_pairs_sorted = sorted(
+        uid_incentive_pairs, key=lambda x: x[1], reverse=True
+    )
 
     # Extract the top 10 uids
     top_k_uids = [uid for uid, incentive in uid_incentive_pairs_sorted[:k]]
-    
-    return top_k_uids     
+
+    return top_k_uids