wehub-resource-sync
113 lines
4.4 KiB
Python
113 lines
4.4 KiB
Python
"""Utility functions for DistributedItemSampler."""
|
|
|
|
|
|
def count_split(total, num_workers, worker_id, batch_size=1):
|
|
"""Calculate the number of assigned items after splitting them by batch
|
|
size evenly. It will return the number for this worker and also a sum of
|
|
previous workers.
|
|
"""
|
|
quotient, remainder = divmod(total, num_workers * batch_size)
|
|
if batch_size == 1:
|
|
assigned = quotient + (worker_id < remainder)
|
|
else:
|
|
batch_count, last_batch = divmod(remainder, batch_size)
|
|
assigned = quotient * batch_size + (
|
|
batch_size
|
|
if worker_id < batch_count
|
|
else (last_batch if worker_id == batch_count else 0)
|
|
)
|
|
prefix_sum = quotient * worker_id * batch_size + min(
|
|
worker_id * batch_size, remainder
|
|
)
|
|
return (assigned, prefix_sum)
|
|
|
|
|
|
def calculate_range(
|
|
distributed,
|
|
total,
|
|
num_replicas,
|
|
rank,
|
|
num_workers,
|
|
worker_id,
|
|
batch_size,
|
|
drop_last,
|
|
drop_uneven_inputs,
|
|
):
|
|
"""Calculates the range of items to be assigned to the current worker.
|
|
|
|
This function evenly distributes `total` items among multiple workers,
|
|
batching them using `batch_size`. Each replica has `num_workers` workers.
|
|
The batches generated by workers within the same replica are combined into
|
|
the replica`s output. The `drop_last` parameter determines whether
|
|
incomplete batches should be dropped. If `drop_last` is True, incomplete
|
|
batches are discarded. The `drop_uneven_inputs` parameter determines if the
|
|
number of batches assigned to each replica should be the same. If
|
|
`drop_uneven_inputs` is True, excessive batches for some replicas will be
|
|
dropped.
|
|
|
|
Args:
|
|
distributed (bool): Whether it's in distributed mode.
|
|
total (int): The total number of items.
|
|
num_replicas (int): The total number of replicas.
|
|
rank (int): The rank of the current replica.
|
|
num_workers (int): The number of workers per replica.
|
|
worker_id (int): The ID of the current worker.
|
|
batch_size (int): The desired batch size.
|
|
drop_last (bool): Whether to drop incomplete batches.
|
|
drop_uneven_inputs (bool): Whether to drop excessive batches for some
|
|
replicas.
|
|
|
|
Returns:
|
|
tuple: A tuple containing three numbers:
|
|
- start_offset (int): The starting offset of the range assigned to
|
|
the current worker.
|
|
- assigned_count (int): The length of the range assigned to the
|
|
current worker.
|
|
- output_count (int): The number of items that the current worker
|
|
will produce after dropping.
|
|
"""
|
|
# Check if it's distributed mode.
|
|
if not distributed:
|
|
if not drop_last:
|
|
return (0, total, total)
|
|
else:
|
|
return (0, total, total // batch_size * batch_size)
|
|
# First, equally distribute items into all replicas.
|
|
assigned_count, start_offset = count_split(
|
|
total, num_replicas, rank, batch_size
|
|
)
|
|
# Calculate the number of outputs when drop_uneven_inputs is True.
|
|
# `assigned_count` is the number of items distributed to the current
|
|
# process. `output_count` is the number of items should be output
|
|
# by this process after dropping.
|
|
if not drop_uneven_inputs:
|
|
if not drop_last:
|
|
output_count = assigned_count
|
|
else:
|
|
output_count = assigned_count // batch_size * batch_size
|
|
else:
|
|
if not drop_last:
|
|
min_item_count, _ = count_split(
|
|
total, num_replicas, num_replicas - 1, batch_size
|
|
)
|
|
min_batch_count = (min_item_count + batch_size - 1) // batch_size
|
|
output_count = min(min_batch_count * batch_size, assigned_count)
|
|
else:
|
|
output_count = total // (batch_size * num_replicas) * batch_size
|
|
# If there are multiple workers, equally distribute the batches to
|
|
# all workers.
|
|
if num_workers > 1:
|
|
# Equally distribute the dropped number too.
|
|
dropped_items, prev_dropped_items = count_split(
|
|
assigned_count - output_count, num_workers, worker_id
|
|
)
|
|
output_count, prev_output_count = count_split(
|
|
output_count,
|
|
num_workers,
|
|
worker_id,
|
|
batch_size,
|
|
)
|
|
assigned_count = output_count + dropped_items
|
|
start_offset += prev_output_count + prev_dropped_items
|
|
return (start_offset, assigned_count, output_count)
|