TensorDictBase.isend#

TensorDictBase.isend(dst: int | None = None, *, group: 'torch.distributed.ProcessGroup' | None = None, group_dst: int | None = None, init_tag: int = 0, pseudo_rand: bool = False, return_early: bool = False, consolidated: bool = False) → int | List['Work'] | None#

Sends the content of the tensordict asynchronously.

Parameters:

dst (int, optional) – the global rank of the destination worker where the content should be sent. Mutually exclusive with group_dst; exactly one of the two must be provided.

Keyword Arguments:

group (torch.distributed.ProcessGroup, optional) – if set, the specified process group will be used for communication. Otherwise, the default process group will be used. Defaults to None.
group_dst (int, optional) – the rank of the destination worker relative to group. Requires group to be passed and is mutually exclusive with dst. When set, the p2p calls are issued directly on the group’s backend, so group may be a standalone ProcessGroup built against a store (never registered through init_process_group() or new_group()). Defaults to None.
init_tag (int) – the initial tag to be used to mark the tensors. Note that this will be incremented by as much as the number of tensors contained in the TensorDict.
pseudo_rand (bool) – if True, the sequence of tags will be pseudo- random, allowing to send multiple data from different nodes without overlap. Notice that the generation of these pseudo-random numbers is expensive (1e-5 sec/number), meaning that it could slow down the runtime of your algorithm. Defaults to False.
return_early (bool, optional) – if True, a list of futures will be returned instead of the tag of the last tensor sent. Defaults to False.
consolidated (bool) – if True, sends the consolidated storage as a single tensor (1 message), as send() does. The tensordict is consolidated first if needed, which copies every leaf on each call: a sender that keeps a consolidated copy and refreshes it with update_() avoids that copy. The receiver must use irecv(consolidated=True) or recv(consolidated=True) and must already hold a consolidated tensordict with matching schema. With return_early=True, the returned list holds the single future; otherwise None is returned once the message is sent. init_tag and pseudo_rand are ignored. Defaults to False.

Example

>>> import torch
>>> from tensordict import TensorDict
>>> from torch import multiprocessing as mp
>>> def client():
...     torch.distributed.init_process_group(
...         "gloo",
...         rank=1,
...         world_size=2,
...         init_method=f"tcp://localhost:10003",
...     )
...
...     td = TensorDict(
...         {
...             ("a", "b"): torch.randn(2),
...             "c": torch.randn(2, 3),
...             "_": torch.ones(2, 1, 5),
...         },
...         [2],
...     )
...     td.isend(0)
...
>>>
>>> def server(queue, return_premature=True):
...     torch.distributed.init_process_group(
...         "gloo",
...         rank=0,
...         world_size=2,
...         init_method=f"tcp://localhost:10003",
...     )
...     td = TensorDict(
...         {
...             ("a", "b"): torch.zeros(2),
...             "c": torch.zeros(2, 3),
...             "_": torch.zeros(2, 1, 5),
...         },
...         [2],
...     )
...     out = td.irecv(1, return_premature=return_premature)
...     if return_premature:
...         for fut in out:
...             fut.wait()
...     assert (td != 0).all()
...     queue.put("yuppie")
...
>>>
>>> if __name__ == "__main__":
...     queue = mp.Queue(1)
...     main_worker = mp.Process(
...         target=server,
...         args=(queue, )
...         )
...     secondary_worker = mp.Process(target=client)
...
...     main_worker.start()
...     secondary_worker.start()
...     out = queue.get(timeout=10)
...     assert out == "yuppie"
...     main_worker.join()
...     secondary_worker.join()

Every tensordict class and every tensorclass has this method. For the other methods, see Methods by task.

TensorDictBase.isend#

Docs

Tutorials

Resources