[DNL] Experimenting on what a PTD transport buffer would look like #43
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
meta-cla
bot
added
the
CLA Signed
LucasLLC
commented
Sep 29, 2025
| return await self.store.get_meta(key, request) | ||
|
|
||
| @endpoint | ||
| async def handshake(self, file_store_name): |
LucasLLC
commented
Sep 29, 2025
| return | ||
| logger.info(f"Finalizing handshake from {file_store_name}") | ||
|
|
||
| file_store = torch.distributed.FileStore(file_store_name, 2) |
LucasLLC
commented
Sep 29, 2025
| # we allocate on the fly | ||
| tensor = await transport_buffer.read_into(tensor=None) | ||
|
|
||
| pg = self.pgs[transport_buffer.file_store_name] |
There was a problem hiding this comment.
from storage_volume
LucasLLC
commented
Sep 29, 2025
|
|
||
| if request.tensor_slice is None: | ||
| await transport_buffer.write_from(self.kv[key]) | ||
| await transport_buffer.write_from( |
LucasLLC
commented
Sep 29, 2025
| volume_coord = self.volume_id_to_coord[volume_id] | ||
| return self.storage_volumes.slice(**volume_coord) | ||
| storage_volume = self.storage_volumes.slice(**volume_coord) | ||
| storage_volume.volume_id = volume_id |
LucasLLC
commented
Sep 29, 2025
|
|
||
|
|
||
| import torch | ||
| from torchstore.utils import _gloo_factory |
LucasLLC
commented
Sep 29, 2025
|
|
||
| local_pgs = {} | ||
|
|
||
| class TorchDistributedBuffer(TransportBuffer): |
LucasLLC
commented
Sep 29, 2025
| # TODO: eventually this should be dependent on the connections available to a storage_volume | ||
|
|
||
| #TODO: | ||
| if True: |
LucasLLC
commented
Sep 29, 2025
|
|
||
| return global_tensor | ||
|
|
||
| def _gloo_factory( |
LucasLLC
commented
Sep 29, 2025
| else: | ||
| transport_buffer.allocate(request.tensor_val) | ||
|
|
||
| if isinstance(transport_buffer, TorchDistributedBuffer): |
LucasLLC
commented
Sep 29, 2025
| if transport_buffer.is_object: | ||
| return transport_buffer.objects | ||
|
|
||
| if isinstance(transport_buffer, TorchDistributedBuffer): |
|
closed in favor of #44 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Evaluating the lift for what it would take to implement a transport protocol around PTD.
The main drawback here is having to initiate a process group handshake on the first transport request across two actors. The handshake could be cache'd for future references, but this might not be super scalable.
Currently only basic tests pass and cache-ing logic is wrong for pg group creation
Ignore rdma changes, I'm not planning on landing this PR