core.resharding.nvshmem_copy_service.planning.workload_packer

Module Contents

Classes

WorkloadPacker

Packs individual SendRequests into WorkloadGroups (batches) destined for the same PE, respecting size limits.

API

class core.resharding.nvshmem_copy_service.planning.workload_packer.WorkloadPacker

Packs individual SendRequests into WorkloadGroups (batches) destined for the same PE, respecting size limits.

pack_workloads(

send_requests: List[core.resharding.nvshmem_copy_service.nvshmem_types.SendRequest],

n_pes: int,

) → Dict[int, List[core.resharding.nvshmem_copy_service.nvshmem_types.WorkloadGroup]]

Groups requests by destination PE and packs them into batches. Returns a map: dest_pe -> list of batches

_pack_single_destination(

tasks: List[core.resharding.nvshmem_copy_service.nvshmem_types.SendRequest],

dest_pe: int,

) → List[core.resharding.nvshmem_copy_service.nvshmem_types.WorkloadGroup]