fix end of request behaviour and add back broadcasting tokens to other nodes

72c3fdab464db138d930219ca7f3e463e28f8233 · 2024-11-21 22:24:13 +0400 · Alex Cheema

Files touched

M exo/orchestration/standard_node.py

Diff

commit 72c3fdab464db138d930219ca7f3e463e28f8233
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Thu Nov 21 22:24:13 2024 +0400

    fix end of request behaviour and add back broadcasting tokens to other nodes
---
 exo/orchestration/standard_node.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/exo/orchestration/standard_node.py b/exo/orchestration/standard_node.py
index 0d0b4187..1094a2a9 100644
--- a/exo/orchestration/standard_node.py
+++ b/exo/orchestration/standard_node.py
@@ -115,10 +115,11 @@ class StandardNode(Node):
       token = await self.inference_engine.sample(result)
       await self.inference_engine.ensure_shard(shard)
       self.buffered_token_output[request_id][0].append(token.item())
-      self.trigger_on_token_callbacks(request_id, self.buffered_token_output[request_id][0], is_finished)
       if DEBUG >= 2: print(f"[{request_id}] result size: {result.size}, is finished: {is_finished}, buffered tokens: {len(self.buffered_token_output[request_id][0])}")
       is_finished = token.item() == self.inference_engine.tokenizer.eos_token_id
       forward = token.reshape(1, -1)
+      self.trigger_on_token_callbacks(request_id, self.buffered_token_output[request_id][0], is_finished)
+      asyncio.create_task(self.broadcast_result(request_id, self.buffered_token_output[request_id][0], is_finished))
     else:
       forward = result

← c7735523 target_shard to next_shard · back to Exo · remvoing console log 7e6c69fd →