Skip to content

Implement direct server-to-server communication#331

Merged
borzunov merged 27 commits into
mainfrom
server-to-server
Jul 11, 2023
Merged

Implement direct server-to-server communication#331
borzunov merged 27 commits into
mainfrom
server-to-server

Conversation

@borzunov

@borzunov borzunov commented Jun 24, 2023

Copy link
Copy Markdown
Collaborator

Implement #226.

TODO:

  • Session pipes => session queues (manager.Queue())
  • Send block indices instead of uids to compactify next_servers

Comment thread src/petals/server/handler.py Outdated
module_backends: Dict[str, TransformerBackend],
*,
push_manager: SyncManager,
session_pipes: Dict[str, Tuple[PipeConnection, threading.Lock]],

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

rpc_push() may be received by a connection handler different from the one holding the inference session, so we use some multiprocess communication here.

Comment thread src/petals/server/handler.py Outdated
assert isinstance(block_uid, str) and isinstance(metadata, dict)
return block_uid, inputs, metadata

async def rpc_push(self, request: runtime_pb2.ExpertRequest, context: P2PContext) -> runtime_pb2.ExpertResponse:

@borzunov borzunov Jul 1, 2023

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TODO: This can be stream-to-unary handler, so that (a) the previous server doesn't have to make a new connection each time and (b) we don't have to parse metadata at this stage each time (now it's done to find session_id). I'm not sure if it affects performance much though, so I'd postpone that to a later PR.

span = None
try:
if not self._chosen_spans or not self._server_sessions or attempt_no >= 1:
# If there is a failed server session, this code closes it

@borzunov borzunov Jul 1, 2023

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This code was moved to a separate method InferenceSession._update_sequence() to simplify this method.

@borzunov
borzunov force-pushed the server-to-server branch from 73a0563 to 49bbce4 Compare July 1, 2023 23:50
@borzunov
borzunov force-pushed the server-to-server branch 2 times, most recently from 0e5ce37 to 2445aa5 Compare July 2, 2023 00:36
@borzunov
borzunov force-pushed the server-to-server branch from 2445aa5 to cef5662 Compare July 2, 2023 00:37
@borzunov
borzunov force-pushed the server-to-server branch from 7d73698 to 429146e Compare July 2, 2023 02:21
@borzunov
borzunov force-pushed the server-to-server branch from 429146e to 9a23c1e Compare July 2, 2023 02:24
@borzunov
borzunov force-pushed the server-to-server branch from bccdc21 to 845c172 Compare July 5, 2023 13:01
@borzunov
borzunov force-pushed the server-to-server branch from 845c172 to 0635968 Compare July 5, 2023 13:06
@borzunov
borzunov force-pushed the server-to-server branch from c5dfba2 to a637fd2 Compare July 9, 2023 22:58
@borzunov
borzunov force-pushed the server-to-server branch from a637fd2 to 7e204d0 Compare July 9, 2023 22:58
@borzunov
borzunov force-pushed the server-to-server branch from 9427cae to 5b3f180 Compare July 9, 2023 23:32
inputs = self.history # Pass full inputs including prefix
else:
inputs = inputs[:, -n_input_tokens:] # No need to pass prefix further

@borzunov borzunov Jul 10, 2023

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Refactored: This code was moved from InferenceSession.step() to _ServerInferenceSession.step(), since it's actually about one server only. The overall structure is more clear this way.

@borzunov
borzunov marked this pull request as ready for review July 10, 2023 00:16
Comment thread src/petals/client/inference_session.py Outdated
Comment thread src/petals/client/inference_session.py Outdated
Comment thread src/petals/client/inference_session.py Outdated
Comment thread src/petals/server/server.py
Comment thread src/petals/server/handler.py
@borzunov
borzunov merged commit 158013a into main Jul 11, 2023
@borzunov
borzunov deleted the server-to-server branch July 11, 2023 13:29
artek0chumak pushed a commit to artek0chumak/petals that referenced this pull request Jul 11, 2023
fthrvi pushed a commit to fthrvi/nakshatra that referenced this pull request Jun 1, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants