steffenc commited on
Commit
7ea7f29
·
1 Parent(s): b338bed

force streaming calls to all endpoints

Browse files
Files changed (4) hide show
  1. server.py +16 -14
  2. utils.py +25 -25
  3. validators/base.py +3 -3
  4. validators/sn1_validator_wrapper.py +64 -20
server.py CHANGED
@@ -2,21 +2,16 @@ import asyncio
2
  import utils
3
  import bittensor as bt
4
  from aiohttp import web
5
- from aiohttp.web_response import Response
6
  from validators import S1ValidatorAPI, QueryValidatorParams, ValidatorAPI
7
  from middlewares import api_key_middleware, json_parsing_middleware
8
 
9
  """
10
- # test
11
  ```
12
- curl -X POST http://0.0.0.0:10000/chat/ -H "api_key: hello" -d '{"k": 5, "timeout": 3, "roles": ["user"], "messages": ["hello world"]}'
13
 
14
- 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?"]}'
15
-
16
- curl -X POST http://0.0.0.0:10000/chat/ -H "api_key: hey-michal" -d '{"k": 5, "timeout": 15, "roles": ["user"], "messages": ["who you really are?"]}'
17
-
18
- # stream
19
- 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"]}'
20
  ```
21
 
22
  TROUBLESHOOT
@@ -25,7 +20,15 @@ check if port is open
25
  sudo ufw allow 10000/tcp
26
  sudo ufw allow 10000/tcp
27
  ```
28
- # run
 
 
 
 
 
 
 
 
29
  ```
30
  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
31
  ```
@@ -38,7 +41,7 @@ add --mock to test the echo stream
38
  """
39
 
40
 
41
- async def chat(request: web.Request) -> Response:
42
  """
43
  Chat endpoint for the validator.
44
  """
@@ -51,9 +54,8 @@ async def chat(request: web.Request) -> Response:
51
  return response
52
 
53
 
54
- async def echo_stream(request, request_data):
55
- request_data = request["data"]
56
- return await utils.echo_stream(request_data)
57
 
58
 
59
  class ValidatorApplication(web.Application):
 
2
  import utils
3
  import bittensor as bt
4
  from aiohttp import web
 
5
  from validators import S1ValidatorAPI, QueryValidatorParams, ValidatorAPI
6
  from middlewares import api_key_middleware, json_parsing_middleware
7
 
8
  """
9
+ # Test chat endpoint with curl
10
  ```
11
+ curl --no-buffer -X POST http://0.0.0.0:10000/chat/ -H "api_key: hey-michal" -d '{"k": 5, "timeout": 15, "roles": ["user"], "messages": ["on what exact date did the 21st century begin??"]}'
12
 
13
+ # echo stream test endpoint
14
+ curl --no-buffer -X POST http://0.0.0.0: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"]}'
 
 
 
 
15
  ```
16
 
17
  TROUBLESHOOT
 
20
  sudo ufw allow 10000/tcp
21
  sudo ufw allow 10000/tcp
22
  ```
23
+
24
+ ---
25
+
26
+ # Run Chattensor
27
+
28
+ ## With vanilla python:
29
+ python server.py --neuron.model_id mock --wallet.name sn1 --wallet.hotkey v1 --netuid 1 --neuron.tasks math --neuron.task_p 1 --neuron.device cpu --subtensor.network local
30
+
31
+ ## With PM2:
32
  ```
33
  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
34
  ```
 
41
  """
42
 
43
 
44
+ async def chat(request: web.Request) -> web.StreamResponse:
45
  """
46
  Chat endpoint for the validator.
47
  """
 
54
  return response
55
 
56
 
57
+ async def echo_stream(request: web.Request) -> web.StreamResponse:
58
+ return await utils.echo_stream(request)
 
59
 
60
 
61
  class ValidatorApplication(web.Application):
utils.py CHANGED
@@ -1,8 +1,10 @@
1
  import re
2
- import bittensor as bt
3
  import time
4
  import json
 
 
5
  from aiohttp import web
 
6
  from collections import Counter
7
  from prompting.rewards import DateRewardModel, FloatDiffModel
8
 
@@ -134,47 +136,45 @@ def guess_task_name(challenge: str):
134
  return "qa"
135
 
136
 
137
- async def echo_stream(request_data: dict):
138
- k = request_data.get("k", 1)
139
- exclude = request_data.get("exclude", [])
140
- timeout = request_data.get("timeout", 0.2)
141
  message = "\n\n".join(request_data["messages"])
142
 
143
  # Create a StreamResponse
144
  response = web.StreamResponse(
145
- status=200, reason="OK", headers={"Content-Type": "text/plain"}
146
  )
147
- await response.prepare()
148
 
149
  completion = ""
 
 
 
150
  # Echo the message k times with a timeout between each chunk
151
  for _ in range(k):
152
  for word in message.split():
153
- chunk = f"{word} "
154
  await response.write(chunk.encode("utf-8"))
155
  completion += chunk
156
- time.sleep(timeout)
157
  bt.logging.info(f"Echoed: {chunk}")
158
-
 
 
 
159
  completion = completion.strip()
160
 
161
- # Prepare final JSON chunk
162
- json_chunk = json.dumps(
163
- {
164
- "uids": [0],
165
- "completion": completion,
166
- "completions": [completion.strip()],
167
- "timings": [0],
168
- "status_messages": ["Went well!"],
169
- "status_codes": [200],
170
- "completion_is_valid": [True],
171
- "task_name": "echo",
172
- "ensemble_result": {},
173
- }
174
- )
175
 
176
  # Send the final JSON as part of the stream
177
- await response.write(f"\n\nJSON_RESPONSE_BEGIN:\n{json_chunk}".encode("utf-8"))
178
 
179
  # Finalize the response
180
  await response.write_eof()
 
1
  import re
 
2
  import time
3
  import json
4
+ import asyncio
5
+ import bittensor as bt
6
  from aiohttp import web
7
+ from responses import TextStreamResponse
8
  from collections import Counter
9
  from prompting.rewards import DateRewardModel, FloatDiffModel
10
 
 
136
  return "qa"
137
 
138
 
139
+ async def echo_stream(request: web.Request) -> web.StreamResponse:
140
+ request_data = request["data"]
141
+ k = request_data.get("k", 1)
 
142
  message = "\n\n".join(request_data["messages"])
143
 
144
  # Create a StreamResponse
145
  response = web.StreamResponse(
146
+ status=200, reason="OK", headers={"Content-Type": "application/json"}
147
  )
148
+ await response.prepare(request)
149
 
150
  completion = ""
151
+ chunks = []
152
+ chunks_timings = []
153
+ start_time = time.time()
154
  # Echo the message k times with a timeout between each chunk
155
  for _ in range(k):
156
  for word in message.split():
157
+ chunk = f"{word} "
158
  await response.write(chunk.encode("utf-8"))
159
  completion += chunk
160
+ await asyncio.sleep(.3)
161
  bt.logging.info(f"Echoed: {chunk}")
162
+
163
+ chunks.append(chunk)
164
+ chunks_timings.append(time.time() - start_time)
165
+
166
  completion = completion.strip()
167
 
168
+ # Prepare final JSON chunk
169
+ response_data = TextStreamResponse(
170
+ streamed_chunks=chunks,
171
+ streamed_chunks_timings=chunks_timings,
172
+ completion=completion,
173
+ timing = time.time()- start_time
174
+ ).to_dict()
 
 
 
 
 
 
 
175
 
176
  # Send the final JSON as part of the stream
177
+ await response.write(json.dumps(response_data).encode("utf-8"))
178
 
179
  # Finalize the response
180
  await response.write_eof()
validators/base.py CHANGED
@@ -1,7 +1,7 @@
1
  from abc import ABC, abstractmethod
2
  from typing import List
3
  from dataclasses import dataclass
4
- from aiohttp.web import Response, Request
5
 
6
 
7
  @dataclass
@@ -31,10 +31,10 @@ class QueryValidatorParams:
31
 
32
  class ValidatorAPI(ABC):
33
  @abstractmethod
34
- async def query_validator(self, params: QueryValidatorParams) -> Response:
35
  pass
36
 
37
 
38
  class MockValidator(ValidatorAPI):
39
- async def query_validator(self, params: QueryValidatorParams) -> Response:
40
  ...
 
1
  from abc import ABC, abstractmethod
2
  from typing import List
3
  from dataclasses import dataclass
4
+ from aiohttp.web import Response, Request, StreamResponse
5
 
6
 
7
  @dataclass
 
31
 
32
  class ValidatorAPI(ABC):
33
  @abstractmethod
34
+ async def query_validator(self, params: QueryValidatorParams) -> StreamResponse:
35
  pass
36
 
37
 
38
  class MockValidator(ValidatorAPI):
39
+ async def query_validator(self, params: QueryValidatorParams) -> StreamResponse:
40
  ...
validators/sn1_validator_wrapper.py CHANGED
@@ -2,7 +2,7 @@ import json
2
  import utils
3
  import torch
4
  import traceback
5
- import asyncio
6
  import random
7
  import bittensor as bt
8
  from typing import Awaitable
@@ -13,6 +13,16 @@ from prompting.dendrite import DendriteResponseEvent
13
  from .base import QueryValidatorParams, ValidatorAPI
14
  from aiohttp.web_response import Response, StreamResponse
15
  from deprecated import deprecated
 
 
 
 
 
 
 
 
 
 
16
 
17
 
18
  class S1ValidatorAPI(ValidatorAPI):
@@ -77,13 +87,39 @@ class S1ValidatorAPI(ValidatorAPI):
77
 
78
  async def process_response(
79
  self, response: StreamResponse, async_generator: Awaitable
80
- ):
81
  """Process a single response asynchronously."""
82
- chunk = None # Initialize chunk with a default value
83
- async for chunk in async_generator:
84
- if chunk is not None and hasattr(chunk, 'completion'):
85
- # Directly write the string encoded as UTF-8 bytes
86
- await response.write(chunk.completion.encode('utf-8'))
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
87
 
88
 
89
  async def get_stream_response(self, params: QueryValidatorParams) -> StreamResponse:
@@ -94,7 +130,7 @@ class S1ValidatorAPI(ValidatorAPI):
94
 
95
  try:
96
  # Guess the task name of current request
97
- task_name = utils.guess_task_name(params.messages[-1])
98
 
99
  # Get the list of uids to query for this step.
100
  uids = get_random_uids(
@@ -103,7 +139,8 @@ class S1ValidatorAPI(ValidatorAPI):
103
  axons = [self.validator.metagraph.axons[uid] for uid in uids]
104
 
105
  # Make calls to the network with the prompt.
106
- bt.logging.info(f"Calling dendrite")
 
107
 
108
  streams_responses = await self.validator.dendrite(
109
  axons=axons,
@@ -115,8 +152,22 @@ class S1ValidatorAPI(ValidatorAPI):
115
  streaming=True,
116
  )
117
 
118
- random_stream = random.choice(streams_responses)
119
- await self.process_response(response, random_stream)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
120
  except Exception as e:
121
  bt.logging.error(
122
  f"Encountered an error in {self.__class__.__name__}:get_stream_response:\n{traceback.format_exc()}"
@@ -128,12 +179,5 @@ class S1ValidatorAPI(ValidatorAPI):
128
 
129
  return response
130
 
131
- async def query_validator(self, params: QueryValidatorParams) -> Response:
132
- # TODO: SET STREAM AS DEFAULT
133
- stream = params.request.get("stream", True)
134
-
135
- if stream:
136
- return await self.get_stream_response(params)
137
- else:
138
- # DEPRECATED
139
- return await self.get_response(params)
 
2
  import utils
3
  import torch
4
  import traceback
5
+ import time
6
  import random
7
  import bittensor as bt
8
  from typing import Awaitable
 
13
  from .base import QueryValidatorParams, ValidatorAPI
14
  from aiohttp.web_response import Response, StreamResponse
15
  from deprecated import deprecated
16
+ from dataclasses import dataclass
17
+ from typing import List
18
+ from responses import TextStreamResponse
19
+
20
+
21
+ @dataclass
22
+ class ProcessedStreamResponse:
23
+ streamed_chunks: List[str]
24
+ streamed_chunks_timings: List[float]
25
+ synapse: StreamPromptingSynapse
26
 
27
 
28
  class S1ValidatorAPI(ValidatorAPI):
 
87
 
88
  async def process_response(
89
  self, response: StreamResponse, async_generator: Awaitable
90
+ ) -> ProcessedStreamResponse:
91
  """Process a single response asynchronously."""
92
+ # Initialize chunk with a default value
93
+ chunk = None
94
+ # Initialize chunk array to accumulate streamed chunks
95
+ chunks = []
96
+ chunks_timings = []
97
+
98
+ start_time = time.time()
99
+ last_sent_index = 0
100
+ async for chunk in async_generator:
101
+ if isinstance(chunk, list):
102
+ # Chunks are currently returned in string arrays, so we need to concatenate them
103
+ concatenated_chunks = "".join(chunk)
104
+ new_data = concatenated_chunks[last_sent_index:]
105
+
106
+ if new_data:
107
+ await response.write(new_data.encode('utf-8'))
108
+ bt.logging.info(f"Received new chunk from miner: {chunk}")
109
+ last_sent_index += len(new_data)
110
+ chunks.extend(chunk)
111
+ chunks_timings.append(time.time() - start_time)
112
+
113
+ if chunk is not None and isinstance(chunk, StreamPromptingSynapse):
114
+ # Assuming the last chunk holds the last value yielded which should be a synapse with the completion filled
115
+ return ProcessedStreamResponse(
116
+ synapse=chunk,
117
+ streamed_chunks=chunks,
118
+ streamed_chunks_timings=chunks_timings
119
+ )
120
+ else:
121
+ raise ValueError("The last chunkis not a StreamPrompting synapse")
122
+
123
 
124
 
125
  async def get_stream_response(self, params: QueryValidatorParams) -> StreamResponse:
 
130
 
131
  try:
132
  # Guess the task name of current request
133
+ # task_name = utils.guess_task_name(params.messages[-1])
134
 
135
  # Get the list of uids to query for this step.
136
  uids = get_random_uids(
 
139
  axons = [self.validator.metagraph.axons[uid] for uid in uids]
140
 
141
  # Make calls to the network with the prompt.
142
+ bt.logging.info(f"Calling dendrite")
143
+ start_time = time.time()
144
 
145
  streams_responses = await self.validator.dendrite(
146
  axons=axons,
 
152
  streaming=True,
153
  )
154
 
155
+ uid_stream_dict = dict(zip(uids, streams_responses))
156
+
157
+ random_uid, random_stream = random.choice(list(uid_stream_dict.items()))
158
+ processed_response = await self.process_response(response, random_stream)
159
+
160
+ # Prepare final JSON chunk
161
+ response_data = json.dumps(TextStreamResponse(
162
+ streamed_chunks=processed_response.streamed_chunks,
163
+ streamed_chunks_timings=processed_response.streamed_chunks_timings,
164
+ uid = random_uid,
165
+ completion=processed_response.synapse.completion,
166
+ timing = time.time()- start_time
167
+ ).to_dict())
168
+
169
+ # Send the final JSON as part of the stream
170
+ await response.write(json.dumps(response_data).encode("utf-8"))
171
  except Exception as e:
172
  bt.logging.error(
173
  f"Encountered an error in {self.__class__.__name__}:get_stream_response:\n{traceback.format_exc()}"
 
179
 
180
  return response
181
 
182
+ async def query_validator(self, params: QueryValidatorParams) -> Response:
183
+ return await self.get_stream_response(params)