【教程】python 使用 multiprocessing 加速 CPU 密集型计算任务的整体执行速度

[nilaoda]

前言

Python的全局解释器锁（GIL）会限制同一时间只能有一个线程执行Python字节码，这使得多线程无法充分利用多核CPU的优势。而多进程则可以利用多个CPU核心并行执行代码，从而提高计算效率。

需求

脚本想调用一个第三方解密模块用来解密很多小的文件，而这个解密模块解密效率不高，解密 1.31MB 的文件需要 500ms 左右。所以想利用多进程 IO 来加速解密。

代码逻辑

通过下载器下载，已经得到一个含有未加密和加密部分的 TS 文件（比如 input.ts ），该文件通常由若干小的 TS 组成。

---

优化的整体逻辑如下图所示：

每个进程只负责自己对应索引位置的数据，各个进程互不干扰。

---

第1步：从 TS 文件中识别出被加密的每一个小的 TS 切片所处的索引位置，返回类似于这样的结果：

(318251664, 318696660)
(318696660, 321374720)
(321374720, 323166360)
(323166360, 324743304)
(324743304, 328546732)
...

第2步：创建一个新的 TS 文件（比如 output.ts ），与 input.ts 大小一致：

# 获取TS体积
file_size = os.path.getsize(in_file)
# 创建一个新的文件，并维持原大小
with open(out_file, 'wb') as f:
    f.seek(file_size - 1)
    f.write(b'\0')

第3步：将 input.ts 中不加密的部分写入到 output.ts 的对应位置（即 seek(0) 后再写入）：

# 分块读取输入文件，并将前X字节写入输出文件
with open(in_file, "rb") as in_f, open(out_file, "r+b") as out_f:
    chunk_size = 2048
    bytes_read = 0
    out_f.seek(0)
    while bytes_read < to_write_length:
        chunk = in_f.read(chunk_size)
        if not chunk:
            break
        if bytes_read + len(chunk) > to_write_length:
            out_f.write(chunk[0:to_write_length-bytes_read])
        else:
            out_f.write(chunk)
        bytes_read += len(chunk)

第4步：开始多进程调用解密模块，并在 output.ts 对应的位置写入解密后的数据：

# 解密输入字节数组并返回解密后的字节数组
def dec_for_bytes(key, in_data: bytes) -> bytes:
    # 这里调用解密模块
    pass


def decrypt_segment(i, startindex, endindex, key, file_in_path, file_out_path, lock) -> bool:
    try:
        with open(file_in_path, 'rb') as input_file, open(file_out_path, 'r+b') as output_file:
            input_file.seek(startindex)
            enc_data = input_file.read(endindex - startindex)
            dec_data = dec_for_bytes(key, enc_data)
            with lock:
                output_file.seek(startindex)
                output_file.write(dec_data)
        return True
    except Exception as e:
        # print(e)
        return False


def decrypt_file(encrypted_ts_list, file_in_path, file_out_path, key, processes_count=4):
    results = []
    # 创建进程锁
    manager = multiprocessing.Manager()
    lock = manager.Lock()
    pool = multiprocessing.Pool(processes=processes_count)
    for i, r in enumerate(encrypted_ts_list):
        result = pool.apply_async(decrypt_segment, args=(i, r[0], r[1], key, file_in_path, file_out_path, lock))
        results.append(result)
    pool.close()
    pool.join()

    # 获取每个进程的结果
    results = [result.get() for result in results]

    # 检查每个进程的结果是否都为True
    if not all(results):
        raise Exception("At least one process failed!")

效果

单进程：

Input TS length: 1516044220
Parse TS file...
Encrypted segments found: 606
Create new TS file...
Processes count: 1
Decrypt segments...
606
Cost: 0:05:43.072584

---

多进程：

Input TS length: 1516044220
Parse TS file...
Encrypted segments found: 606
Create new TS file...
Processes count: 12
Decrypt segments...
606
Cost: 0:00:55.950846

提速效果明显。