Merge pull request #1 from macrocosm-os/stream

forward.py

@@ -0,0 +1,244 @@
+import time
+import sys
+import asyncio
+import numpy as np
+import bittensor as bt
+import traceback
+from typing import List, Dict, Awaitable
+from prompting.agent import HumanAgent
+from prompting.dendrite import DendriteResponseEvent
+from prompting.conversation import create_task
+from prompting.protocol import StreamPromptingSynapse
+from prompting.rewards import RewardResult
+from prompting.utils.uids import get_random_uids
+from prompting.utils.logging import log_event
+from prompting.utils.misc import async_log, serialize_exception_to_string
+from dataclasses import dataclass
+
+@async_log
+async def generate_reference(agent):    
+    loop = asyncio.get_running_loop()
+    result = await loop.run_in_executor(None, agent.task.generate_reference, agent.llm_pipeline)
+    return result    
+
+@async_log
+async def execute_dendrite_call(dendrite_call):
+    responses = await dendrite_call
+    return responses
+
+@dataclass
+class StreamResult:
+    synapse: StreamPromptingSynapse = None
+    exception: BaseException = None
+    uid: int = None
+
+
+async def process_response(uid: int, async_generator: Awaitable):
+    """Process a single response asynchronously."""
+    try:
+        chunk = None  # Initialize chunk with a default value
+        async for chunk in async_generator:  # most important loop, as this is where we acquire the final synapse.
+            bt.logging.debug(f"\nchunk for uid {uid}: {chunk}")
+
+        if chunk is not None:
+            synapse = chunk  # last object yielded is the synapse itself with completion filled
+
+            # Assuming chunk holds the last value yielded which should be a synapse
+            if isinstance(synapse, StreamPromptingSynapse):
+                return synapse
+
+        bt.logging.debug(
+            f"Synapse is not StreamPromptingSynapse. Miner uid {uid} completion set to '' "
+        )
+    except Exception as e:
+        # bt.logging.error(f"Error in generating reference or handling responses: {e}", exc_info=True)
+        traceback_details = traceback.format_exc()
+        bt.logging.error(
+            f"Error in generating reference or handling responses for uid {uid}: {e}\n{traceback_details}"
+        )
+
+        failed_synapse = StreamPromptingSynapse(
+            roles=["user"], messages=["failure"], completion=""
+        )
+
+        return failed_synapse
+
+
+@async_log
+async def handle_response(responses: Dict[int, Awaitable]) -> List[StreamResult]:
+    """The handle_response function is responsible for creating asyncio tasks around acquiring streamed miner chunks
+    and processing them asynchronously. It then pairs the results with their original UIDs and returns a list of StreamResults.
+
+    Args:
+        responses (Dict[int, Awaitable]): Responses contains awaitables that are used to acquire streamed miner chunks.
+
+    Raises:
+        ValueError
+
+    Returns:
+        List[StreamResult]: DataClass containing the synapse, exception, and uid
+    """
+    tasks_with_uid = [
+        (uid, responses[uid]) for uid, _ in responses.items()
+    ]  # Pair UIDs with their tasks
+
+    # Start tasks, preserving order and their associated UIDs
+    tasks = [process_response(uid, resp) for uid, resp in tasks_with_uid]
+
+    results = await asyncio.gather(*tasks, return_exceptions=True)
+
+    mapped_results = []
+    # Pair each result with its original uid
+    for (uid, _), result in zip(tasks_with_uid, results):
+        # If the result is a StreamPromptingSynapse, the response was successful and the stream result is added without exceptions
+        if isinstance(result, StreamPromptingSynapse):
+            mapped_results.append(StreamResult(synapse=result, uid=uid))
+
+        # If the result is an exception, the response was unsuccessful and the stream result is added with the exception and an empty synapse
+        elif isinstance(result, BaseException):
+            failed_synapse = StreamPromptingSynapse(
+                roles=["user"], messages=["failure"], completion=""
+            )
+            mapped_results.append(
+                StreamResult(synapse=failed_synapse, exception=result, uid=uid)
+            )
+
+        # If the result is neither an error or a StreamSynapse, log the error and raise a ValueError
+        else:
+            bt.logging.error(f"Unexpected result type for UID {uid}: {result}")
+            raise ValueError(f"Unexpected result type for UID {uid}: {result}")
+
+    return mapped_results
+
+
+@async_log
+async def generate_reference(agent: HumanAgent):
+    loop = asyncio.get_running_loop()
+    result = await loop.run_in_executor(
+        None, agent.task.generate_reference, agent.llm_pipeline
+    )
+    return result
+
+
+def log_stream_results(stream_results: List[StreamResult]):
+    failed_responses = [
+        response for response in stream_results if response.exception is not None
+    ]
+    empty_responses = [
+        response
+        for response in stream_results
+        if response.exception is None and response.synapse.completion == ""
+    ]
+    non_empty_responses = [
+        response
+        for response in stream_results
+        if response.exception is None and response.synapse.completion != ""
+    ]
+
+    bt.logging.info(f"Total of non_empty responses: ({len(non_empty_responses)})")
+    bt.logging.info(f"Total of empty responses: ({len(empty_responses)})")
+    bt.logging.info(
+        f"Total of failed responses: ({len(failed_responses)}):\n {failed_responses}"
+    )
+
+    for failed_response in failed_responses:
+        formatted_exception = serialize_exception_to_string(failed_response.exception)
+        bt.logging.error(
+            f"Failed response for uid {failed_response.uid}: {formatted_exception}"
+        )
+
+
+async def run_step(
+    self, agent: HumanAgent, k: int, timeout: float, exclude: list = None
+):
+    """Executes a single step of the agent, which consists of:
+    - Getting a list of uids to query
+    - Querying the network
+    - Rewarding the network
+    - Updating the scores
+    - Logging the event
+
+    Args:
+        agent (HumanAgent): The agent to run the step for.
+        k (int): The number of uids to query.
+        timeout (float): The timeout for the queries.
+        exclude (list, optional): The list of uids to exclude from the query. Defaults to [].
+    """
+
+    bt.logging.debug("run_step", agent.task.name)
+
+    # Record event start time.
+    start_time = time.time()
+    # Get the list of uids to query for this step.
+    uids = get_random_uids(self, k=k, exclude=exclude or []).to(self.device)
+    uids_cpu = uids.cpu().tolist()
+
+    axons = [self.metagraph.axons[uid] for uid in uids]
+
+    # Directly call dendrite and process responses in parallel
+    streams_responses = await self.dendrite(
+        axons=axons,
+        synapse=StreamPromptingSynapse(roles=["user"], messages=[agent.challenge]),
+        timeout=timeout,
+        deserialize=False,
+        streaming=True,
+    )
+
+    # Prepare the task for handling stream responses
+    handle_stream_responses_task = asyncio.create_task(
+        handle_response(responses=dict(zip(uids_cpu, streams_responses)))
+    )
+
+    if not agent.task.static_reference:
+        reference_generation_task = generate_reference(agent)
+        _, stream_results = await asyncio.gather(
+            reference_generation_task, handle_stream_responses_task
+        )
+    else:
+        stream_results = await handle_stream_responses_task
+
+    log_stream_results(stream_results)
+
+    all_synapses_results = [stream_result.synapse for stream_result in stream_results]   
+            
+    # Encapsulate the responses in a response event (dataclass)
+    response_event = DendriteResponseEvent(
+        responses=all_synapses_results, uids=uids, timeout=timeout
+    )
+
+    bt.logging.info(f"Created DendriteResponseEvent:\n {response_event}")
+    # Reward the responses and get the reward result (dataclass)
+    # This contains a list of RewardEvents but can be exported as a dict (column-wise) for logging etc
+    reward_result = RewardResult(
+        self.reward_pipeline,
+        agent=agent,
+        response_event=response_event,
+        device=self.device,
+    )
+    bt.logging.info(f"Created RewardResult:\n {reward_result}")
+
+    # The original idea was that the agent is 'satisfied' when it gets a good enough response (e.g. reward critera is met, such as ROUGE>threshold)
+    agent.update_progress(
+        top_reward=reward_result.rewards.max(),
+        top_response=response_event.completions[reward_result.rewards.argmax()],
+    )
+
+    self.update_scores(reward_result.rewards, uids)
+
+    stream_results_uids = [stream_result.uid for stream_result in stream_results]
+    stream_results_exceptions = [
+        serialize_exception_to_string(stream_result.exception)
+        for stream_result in stream_results
+    ]
+    # Log the step event.
+    event = {
+        "block": self.block,
+        "step_time": time.time() - start_time,
+        "stream_results_uids": stream_results_uids,
+        "stream_results_exceptions": stream_results_exceptions,
+        **agent.__state_dict__(full=self.config.neuron.log_full),
+        **reward_result.__state_dict__(full=self.config.neuron.log_full),
+        **response_event.__state_dict__(),
+    }
+
+    return event

server.py

@@ -3,6 +3,7 @@
 
 import os
 import re
+import time
 import asyncio
 import json
 import traceback
@@ -24,6 +25,9 @@ from aiohttp.web_response import Response
 curl -X POST http://0.0.0.0:10000/chat/ -H "api_key: hello" -d '{"k": 5, "timeout": 3, "roles": ["user"], "messages": ["hello world"]}'
 
 curl -X POST http://0.0.0.0:10000/chat/ -H "api_key: hey-michal" -d '{"k": 5, "timeout": 3, "roles": ["user"], "messages": ["on what exact date did the 21st century begin?"]}'
+
+# stream
+curl --no-buffer -X POST http://129.146.127.82:10000/echo/ -H "api_key: hey-michal" -d '{"k": 3, "timeout": 0.2, "roles": ["user"], "messages": ["i need to tell you something important but first"]}'
 ```
 
 TROUBLESHOOT
@@ -31,11 +35,17 @@ check if port is open
 ```
 sudo ufw allow 10000/tcp
 sudo ufw allow 10000/tcp
-``` 
+```
 # run
 ```
 EXPECTED_ACCESS_KEY="hey-michal" pm2 start app.py --interpreter python3 --name app -- --neuron.model_id mock --wallet.name sn1 --wallet.hotkey v1 --netuid 1 --neuron.tasks math --neuron.task_p 1 --neuron.device cpu
 ```
+
+basic testing
+```
+EXPECTED_ACCESS_KEY="hey-michal" python app.py --neuron.model_id mock --wallet.name sn1 --wallet.hotkey v1 --netuid 1 --neuron.tasks math --neuron.task_p 1 --neuron.device cpu 
+```
+add --mock to test the echo stream
 """
 
 EXPECTED_ACCESS_KEY = os.environ.get('EXPECTED_ACCESS_KEY')
@@ -210,6 +220,62 @@ async def chat(request: web.Request) -> Response:
 
 
 
+async def echo_stream(request):
+    
+    bt.logging.info(f'echo_stream()')
+    # Check access key
+    access_key = request.headers.get("api_key")
+    if EXPECTED_ACCESS_KEY is not None and access_key != EXPECTED_ACCESS_KEY:
+        bt.logging.error(f'Invalid access key: {access_key}')
+        return Response(status=401, reason="Invalid access key")
+
+    try:
+        request_data = await request.json()
+    except ValueError:
+        bt.logging.error(f'Invalid request data: {request_data}')
+        return Response(status=400)
+
+    bt.logging.info(f'Request data: {request_data}')
+    k = request_data.get('k', 1)
+    exclude = request_data.get('exclude', [])
+    timeout = request_data.get('timeout', 0.2)    
+    message = '\n\n'.join(request_data['messages'])
+
+    # Create a StreamResponse
+    response = web.StreamResponse(status=200, reason='OK', headers={'Content-Type': 'text/plain'})
+    await response.prepare(request)
+
+    completion = ''
+    # Echo the message k times with a timeout between each chunk
+    for _ in range(k):
+        for word in message.split():
+            chunk = f'{word} '
+            await response.write(chunk.encode('utf-8'))
+            completion += chunk 
+            time.sleep(timeout)
+            bt.logging.info(f"Echoed: {chunk}")
+
+    completion = completion.strip()
+
+    # Prepare final JSON chunk
+    json_chunk = json.dumps({
+        "uids": [0],
+        "completion": completion,
+        "completions": [completion.strip()],
+        "timings": [0],
+        "status_messages": ['Went well!'],
+        "status_codes": [200],
+        "completion_is_valid": [True],
+        "task_name": 'echo',
+        "ensemble_result": {}
+    })
+    
+    # Send the final JSON as part of the stream
+    await response.write(f"\n\nJSON_RESPONSE_BEGIN:\n{json_chunk}".encode('utf-8'))
+
+    # Finalize the response
+    await response.write_eof()
+    return response
 
 class ValidatorApplication(web.Application):
     def __init__(self, *a, **kw):
@@ -218,7 +284,10 @@ class ValidatorApplication(web.Application):
 
 
 validator_app = ValidatorApplication()
-validator_app.add_routes([web.post('/chat/', chat)])
+validator_app.add_routes([
+    web.post('/chat/', chat),
+    web.post('/echo/', echo_stream)
+    ])
 
 bt.logging.info("Starting validator application.")
 bt.logging.info(validator_app)