`core.resharding.copy_services.base`#

Module Contents#

Classes#

`SendOp`	Single send operation pending in a CopyService queue.
`RecvOp`	Single receive operation pending in a CopyService queue.
`CopyService`	Abstract interface for submitting and executing batched P2P copy operations.

Functions#

match_local_ops_by_task_id

Pair same-rank send/recv ops by task_id, raising on any mismatch.

API#

class core.resharding.copy_services.base.SendOp#

Single send operation pending in a CopyService queue.

task_id: int | None#: None

tensor: torch.Tensor#: None

dest_rank: int#: None

class core.resharding.copy_services.base.RecvOp#

Single receive operation pending in a CopyService queue.

task_id: int | None#: None

tensor: torch.Tensor#: None

src_rank: int#: None

class core.resharding.copy_services.base.CopyService(group=None)#

Bases: abc.ABC

Abstract interface for submitting and executing batched P2P copy operations.

All backends accept an optional task_id on submit calls. The task_id is a globally unique identifier shared between the matching send and recv for the same transfer. It is required for local (same-rank) copy matching and for the NVSHMEM backend’s scheduling. Backends that do not need it for remote transfers simply ignore it.

Initialization

requires_process_group_barrier#: True

supports_idle_ranks#: True

supports_multiple_runs_per_plan#: False

abstractmethod submit_send( src_tensor: torch.Tensor, dest_rank: int, task_id: Optional[int] = None, )#: Register a tensor send from the current rank to dest_rank.

abstractmethod submit_recv( dest_tensor: torch.Tensor, src_rank: int, task_id: Optional[int] = None, )#: Register a tensor receive into dest_tensor from src_rank.

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

close() → None#: Release backend-owned resources. The default implementation is a no-op.

set_model_roles(*, is_source: bool, is_destination: bool) → None#

Provide this rank’s model participation to topology-aware backends.

Point-to-point backends infer everything they need from submitted operations and keep this default no-op. Collective cross-group backends may override it to construct their source/destination meshes.

set_plan(plan: object, *, transform: object | None = None) → None#

Associate subsequent submissions with an immutable reshard plan.

Backends may use this identity to cache collective setup that is valid for the lifetime of the plan and transform. Point-to-point backends do not require per-plan setup and keep this default no-op.

execute_plan( plan: object, src_tensors: Mapping[str, torch.Tensor], dst_tensors: Mapping[str, torch.Tensor], *, transform: object | None = None, ) → bool#

Execute a plan natively, returning whether it was handled.

The default keeps the established submit/run path. Backends whose native primitive needs whole tensors and mesh metadata can override this hook without changing the public ReFIT API or the slice-copy contract used by other transports.

core.resharding.copy_services.base.match_local_ops_by_task_id( local_sends: list, local_recvs: list, backend_name: str, rank: int, ) → list[tuple]#