如何公平拆分itertools.product输入以实现分块并行处理

• 有多组参数列表，通过itertools.product生成所有笛卡尔积组合并传入函数执行，示例代码如下：

import numpy as np
import itertools

# 实际列表元素无关，列表长度各异
listA = [1,2,3,4,5]
listB = [0.1,0.2,0.3]
listC = range(23)
# 实际共7个列表，部分长度为1故省略
listG = [5,1]

def theFunction(paramA, paramB, paramC, paramG):
    print(f"{paramA} {paramB} {paramC} {paramG}")

for pA, pB, pC, pG in itertools.product(listA, listB, listC, listG):
    theFunction(pA, pB, pC, pG)

• 组合总数随列表长度呈指数增长，需拆分为运行时长相近的任务块，通过公式numSplits = max(1, np.ceil(functionRuntime*numProducts/jobRuntime))计算拆分数量（允许±10%浮动）
• 限制条件：
  • 无法将itertools.product迭代器转为完整列表（内存不足）
  • 拆分需具备确定性：给定jobID可精准处理对应子集，无重复无遗漏
  • 避免任务块过短（减少初始化开销）
  • 当numSplits超过最长列表长度时，需避免任务时长差异过大的问题

核心思路 #

通过全局索引映射实现无全量生成的拆分：

1. 为每个参数组合分配唯一的全局整数索引
2. 按索引范围均匀拆分任务块，每个块处理连续的索引区间
3. 将索引反向映射回参数组合，直接遍历对应区间的组合，无需生成全量笛卡尔积

具体实现步骤 #

1. 预计算各列表的权重（索引映射基础） #

权重表示当前列表中单个元素对应的后续所有维度的组合数，例如：

• 对于列表[listA, listB, listC, listG]，listA的权重是len(listB)*len(listC)*len(listG)，listB的权重是len(listC)*len(listG)，以此类推，最后一个列表的权重为1。

2. 计算任务块的索引区间 #

• 总组合数total = np.prod([len(lst) for lst in all_lists])
• 每个任务块的基准大小block_size = total // numSplits
• 剩余组合数remainder = total % numSplits，前remainder个块的大小为block_size + 1，其余为block_size
• 给定job_id，计算起始索引start = job_id * block_size + min(job_id, remainder)，结束索引end = start + (block_size + 1 if job_id < remainder else block_size)

3. 索引转参数组合的函数 #

通过权重反向计算每个维度的参数：给定索引idx，依次对各列表的权重取商和余数，得到对应列表的元素索引，最终拼接成参数组合。

4. 任务执行逻辑 #

每个任务根据job_id获取索引区间，遍历区间内的每个索引，转成参数组合后执行目标函数。

完整代码实现 #

import numpy as np

# 定义所有参数列表（示例）
all_lists = [
    [1,2,3,4,5],       # listA
    [0.1,0.2,0.3],     # listB
    range(23),         # listC
    [5,1]              # listG
]

def theFunction(paramA, paramB, paramC, paramG):
    print(f"{paramA} {paramB} {paramC} {paramG}")

def compute_weights(lists):
    """计算每个列表的权重：当前元素对应的后续维度组合数"""
    weights = []
    current_weight = 1
    # 从后往前遍历计算权重
    for lst in reversed(lists):
        weights.append(current_weight)
        current_weight *= len(lst)
    # 反转回原顺序
    return weights[::-1]

def index_to_params(idx, lists, weights):
    """将全局索引转换为对应的参数组合"""
    params = []
    remaining = idx
    for lst, w in zip(lists, weights):
        elem_idx = remaining // w
        params.append(lst[elem_idx])
        remaining = remaining % w
    return params

def run_job(job_id, num_splits, lists, weights, total):
    """执行指定job_id的任务块"""
    block_size = total // num_splits
    remainder = total % num_splits
    
    # 计算当前任务的索引区间
    if job_id < remainder:
        start = job_id * (block_size + 1)
        end = start + (block_size + 1)
    else:
        start = remainder * (block_size + 1) + (job_id - remainder) * block_size
        end = start + block_size
    
    # 遍历区间内的所有索引，执行函数
    for idx in range(start, end):
        params = index_to_params(idx, lists, weights)
        theFunction(*params)

if __name__ == "__main__":
    # 输入参数
    function_runtime = 1  # 单函数执行时长（秒）
    job_runtime = 36000   # 单任务允许的最大运行时长（秒）
    
    # 计算总组合数和拆分数量
    total = np.prod([len(lst) for lst in all_lists])
    num_splits = max(1, int(np.ceil(function_runtime * total / job_runtime)))
    # 允许±10%的浮动，调整num_splits到最接近的整数（可选）
    num_splits = max(1, round(num_splits * np.clip(np.random.uniform(0.9, 1.1), 0.9, 1.1)))
    
    # 预计算权重
    weights = compute_weights(all_lists)
    
    # 示例：执行第0个任务
    run_job(0, num_splits, all_lists, weights, total)

方案优势 #

• 无全量生成：无需存储所有组合，内存占用极低
• 确定性：给定job_id可精准定位任务区间，无重复无遗漏
• 均匀拆分：任务块大小差异不超过1，保证运行时长相近
• 适配性强：无论numSplits是否超过最长列表长度，都能稳定工作
• 低初始化开销：每个任务只需预计算一次权重，无额外冗余操作

内容的提问来源于stack exchange，提问作者Maahk