`core.resharding.nvshmem_copy_service.planning.task_segmenter`#

Module Contents#

Classes#

TaskSegmenter

Splits large tasks (>256MB) into smaller segments to fit into the fixed-size communication slots.

Data#

`logger`
`REQUEST_ID_BASE`
`SEGMENT_ID_MULTIPLIER`
`MAX_REQUESTS`
`MAX_SEGMENTS_PER_REQUEST`

API#

core.resharding.nvshmem_copy_service.planning.task_segmenter.logger#: ‘getLogger(…)’

core.resharding.nvshmem_copy_service.planning.task_segmenter.REQUEST_ID_BASE#: 1000000000

core.resharding.nvshmem_copy_service.planning.task_segmenter.SEGMENT_ID_MULTIPLIER#: 1000

core.resharding.nvshmem_copy_service.planning.task_segmenter.MAX_REQUESTS#: 1000000

core.resharding.nvshmem_copy_service.planning.task_segmenter.MAX_SEGMENTS_PER_REQUEST#: 1000

class core.resharding.nvshmem_copy_service.planning.task_segmenter.TaskSegmenter#

Splits large tasks (>256MB) into smaller segments to fit into the fixed-size communication slots.

_encode_segment_id(task_id: int, segment_index: int) → int#

_calculate_num_segments(size: int) → int#

_validate_segmentation(task_id: int, size: int) → bool#

segment_send_request( req: core.resharding.nvshmem_copy_service.nvshmem_types.SendRequest, ) → List[core.resharding.nvshmem_copy_service.nvshmem_types.SendRequest]#: Splits a single send request into multiple segments if larger than MAX_SEGMENT_SIZE.

segment_receive_request( req: core.resharding.nvshmem_copy_service.nvshmem_types.ReceiveRequest, ) → List[core.resharding.nvshmem_copy_service.nvshmem_types.ReceiveRequest]#: Splits a single receive request into multiple segments if larger than MAX_SEGMENT_SIZE.