from datasets.config import DEFAULT_MAX_BATCH_SIZE
from data_juicer.core.monitor import Monitor
[docs]class Adapter:
MAX_BATCH_SIZE = 10000
[docs] def __init__(self, cfg: dict):
self.cfg = cfg
self.idle_resources = Monitor.monitor_current_resources()
[docs] @staticmethod
def execute_and_probe(dataset, operators, sample_interval=0.5):
"""
Process the input dataset and probe related information for each OP in
the specified operator list.
For now, we support the following targets to probe:
"resource": resource utilization for each OP.
"speed": average processing speed for each OP.
The probe result is a list and each item in the list is the probe
result for each OP.
"""
if operators is None or len(operators) == 0:
return []
# resource utilization list
resource_util_list = []
# probe for each OP
for op in operators:
# set num_proc to 1 for each OP to focus on the influence of batch
# size only.
old_num_proc = op.num_proc
op.num_proc = 1
# number of test samples
sample_num = len(dataset)
# run single op and monitor the resource utilization
dataset, resource_util_per_op = Monitor.monitor_func(
op.run, args=(dataset, ), sample_interval=sample_interval)
# calculate speed
resource_util_per_op[
'speed'] = sample_num / resource_util_per_op['time']
resource_util_list.append(resource_util_per_op)
# restore to the original num_proc
op.num_proc = old_num_proc
return resource_util_list
[docs] @staticmethod
def take_batch(dataset, config):
"""
Split the dataset into batches based on configuration and load factor.
:param dataset: The dataset to be split
:param config: Configuration settings, including batch size
:return: An iterator of batches
"""
# get initial batch size
batch_size = config.get('batch_size', DEFAULT_MAX_BATCH_SIZE)
# should be in [1, 10000]
batch_size = min(max(batch_size, 1), Adapter.MAX_BATCH_SIZE)
# check if there are enough samples
num_samples = len(dataset)
if batch_size >= num_samples:
return dataset
else:
return dataset.take(batch_size)
[docs] def adapt_workloads(self, dataset, operators):
"""
Manage the scheduling and load balancing for the dataset processing.
:param dataset: The dataset that needs to be processed
:param operators: Operators in the data recipe
"""
# TODO: set batch size to 1 for all OPs for probing
load_analysis_res, probed_batch_size = self.probe_small_batch(
dataset, operators)
# calculate batch size for each OP according to the analysis results
bs_per_op = self.batch_size_strategy(load_analysis_res,
base_bs=probed_batch_size)
return bs_per_op
[docs] def probe_small_batch(self, dataset, operators):
"""
Perform small batch pre-execution to probe available resources,
current load and estimated OP speed, returning load factors and speed
ranks for each OP.
:param dataset: The dataset to pre-execute small batch on
:param operators: The OP list to be pre-execution and probe
:return: A list of probe results for each OP and the length of data
batch to probe.
"""
# take a small batch
data_batch = self.take_batch(dataset, self.cfg)
# process and monitor the resource utilization
resource_util_list = self.execute_and_probe(data_batch, operators)
# analyze resource utilization
analysis_res = Monitor.analyze_resource_util_list(resource_util_list)
return analysis_res, len(data_batch)
[docs] def batch_size_strategy(self, load_analysis_res, base_bs=1, util_th=0.9):
"""
Decide the batch size for each op according to their workload analysis
result and expected utilization threshold. We need to guarantee that
the resource utilization won't exceed the threshold. Now we only
consider the buckets effect, which means the max batch size is decided
by the max utilization of all types of resources except GPU util
(decided by num_proc).
"""
batch_size_per_op = []
# compute left utils according to the util_th
left_utils = {}
for key in self.idle_resources:
if 'util.' not in key or 'GPU' in key:
continue
left_utils[key] = max(0, util_th - self.idle_resources[key])
for item in load_analysis_res:
max_util = 1e-5
max_key = min(left_utils.items(), key=lambda it: it[1])[0]
analysis_res = item['resource_analysis']
for key in analysis_res:
if 'util.' not in key or 'GPU' in key:
continue
used_util = max(
0, analysis_res[key]['max'] - self.idle_resources[key])
if used_util > max_util:
max_util = used_util
max_key = key
load_factor = left_utils[max_key] / max_util
bs_this_op = min(max(int(base_bs * load_factor), 1),
self.MAX_BATCH_SIZE)
batch_size_per_op.append(bs_this_op)
return batch_size_per_op