`core.resharding.copy_services.base`#

Module Contents#

Classes#

Abstract interface for submitting and executing batched P2P copy operations.

API#

class core.resharding.copy_services.base.CopyService#

Bases: abc.ABC

Abstract interface for submitting and executing batched P2P copy operations.

abstractmethod submit_send(src_tensor: torch.Tensor, dest_rank: int)#: Register a tensor send from the current rank to dest_rank.

abstractmethod submit_recv(dest_tensor: torch.Tensor, src_rank: int)#: Register a tensor receive into dest_tensor from src_rank.

abstractmethod run()#: Execute all previously submitted send/recv operations as a single batch.