为何多进程URL特征提取脚本运行一段时间后大幅变慢？

使用featurizer.py从URL提取特征并写入num_features.csv，通过data_collection.py的多进程循环调用处理10万条URL。初期运行正常，但写入2000-3000条数据后，速度降至原来的1/10。怀疑和CSV文件大小有关，但不确定是否影响多进程功能。

featurizer.py 最小复现代码 #

class UrlFeaturizer(object):
    def __init__(self, url):
        self.url = url
        try:
            self.response = requests.get(
                prepend_protocols(self.url), headers=headers, timeout=5
            )
        except Exception:
            self.response = None
        try:
            self.whois = whois.query(self.url).__dict__
        except Exception:
            self.whois = None
        try:
            self.soup_c = BeautifulSoup(
                self.response.content,
                features="lxml",
                from_encoding=self.response.encoding,
            )
        except Exception:
            self.soup_c = None

    def lookup_whois(self) -> int:
        return int(False) if self.whois else int(True)

    def lookup_domain_age(self) -> int:
        if self.whois and self.whois["creation_date"]:
            return (date.today() - self.whois["creation_date"].date()).days
        return

    def verify_ssl(self) -> bool:
        try:
            ssl_cert = ssl.get_server_certificate((self.url, 443), timeout=10)
            return int(True) if ssl_cert else int(False)
        except Exception:
            return

    def check_security(self) -> bool:
        try:
            requests.head(f"https://{self.url}", timeout=10)
            return int(True)
        except Exception:
            return int(False)

    def has_com_tld(self):
        return int(True) if extract_tld(self.url) == "com" else int(False)

    def run(self, dataset=None):
        data = {
            "url": self.url,
            "uses_whois_privacy": self.lookup_whois(),
            "domain_age": self.lookup_domain_age(),
            "has_ssl": self.verify_ssl(),
            "is_secure": self.check_security(),
            "has_com_tld": self.has_com_tld(),
            "label": Path(dataset).stem,
        }
        return data.keys(), data

data_collection.py 代码 #

import csv
import multiprocessing as mp
import multiprocessing.managers
import os
import pathlib

import pandas as pd
from tqdm.auto import tqdm

from discovery.featurizer import UrlFeaturizer
from utils import test_internet_connection
from threading import Event


keys = UrlFeaturizer("1.1.1.1").run("")[0]


def worker(url: str, dataset: str, q: multiprocessing.managers.AutoProxy) -> type(None):
    """
    Pass a URL through UrlNumFeaturizer for feature extraction.

    :param url: The URL passed for feature extraction.
    :param dataset: A dataset containing URLs with the same label.
    :param q: The Queue.
    :return: None
    """
    try:
        res = UrlFeaturizer(url).run(dataset)[1]
        q.put(res)
    except AttributeError:
        pass
    return


def listener(q: multiprocessing.managers.AutoProxy) -> type(None):
    """
    Listen for messages on the Queue to determine whether to write to file.

    :param q: The Queue.
    :return: None
    """
    with open("num_features.csv", "a") as f:
        while 1:
            m = q.get()
            if m == "kill":
                break
            csv_out = csv.DictWriter(f, keys)
            csv_out.writerow(m)
            f.flush()
    return


def start(
    pool: multiprocessing.pool.Pool,
    q: multiprocessing.managers.AutoProxy,
    jobs: list[multiprocessing.pool.ApplyResult],
) -> type(None):
    """
    Start collecting data from the first URL in the list.

    :param pool: The process pool.
    :param q: The Queue.
    :param jobs: The list that will contain multiprocessing results.
    :return: None
    """
    completed = 0
    completion = Event()

    def my_callback(result):
        nonlocal completed

        pbar.update()
        completed += 1

        if completed == len(args):
            completion.set()
            pool.close()
        elif completed % 100 == 0:
            if not test_internet_connection():
                os.system("say Internet connection lost.")
                completion.set()
                pool.terminate()
    datasets = ["benign_domains.csv", "dmca_domains.csv",]
    args = []
    for dataset in datasets:
        urls = pd.read_csv(dataset, header=None).iloc[:, 0].to_list()
        for url in urls:
            args.append((url, dataset, q))
    with tqdm(range(len(args))) as pbar:
        for arg in args:
            job = pool.apply_async(worker, arg, callback=my_callback)
            jobs.append(job)
        completion.wait()
        pool.join()
    q.put("kill")
    sort_csv()
    return


def resume(
    pool: multiprocessing.pool.Pool,
    q: multiprocessing.managers.AutoProxy,
    jobs: list[multiprocessing.pool.ApplyResult],
) -> type(None):
    """
    Resume collecting data from the last iterated URL if the program was
    interrupted.

    :param pool: The process pool.
    :param q: The Queue.
    :param jobs: The list that will contain multiprocessing results.
    :return: None
    """
    completed = 0
    completion = Event()

    def my_callback(result):
        nonlocal completed

        pbar.update()
        completed += 1

        if completed == len(args):
            completion.set()
            pool.close()
        elif completed % 100 == 0:
            if not test_internet_connection():
                os.system("say Internet connection lost.")
                completion.set()
                pool.terminate()
    processed_urls = (
        pd.read_csv(
            "num_features.csv",
            usecols=[0],
            header=None,
        )
        .iloc[:, 0]
        .to_list()
    )
    datasets = ["benign_domains.csv", "dmca_domains.csv",]
    args = []
    for dataset in datasets:
        unprocessed_urls = pd.read_csv(dataset, header=None).iloc[:, 0].to_list()
        urls = [item for item in unprocessed_urls if item not in processed_urls]
        for url in urls:
            args.append((url, dataset, q))
    with tqdm(range(len(args))) as pbar:
        for arg in args:
            job = pool.apply_async(worker, arg, callback=my_callback)
            jobs.append(job)
        completion.wait()
        pool.join()
    q.put("kill")
    sort_csv()
    return


def write_header_to_csv(log_file: pathlib.PosixPath) -> type(None):
    """
    Write a header to the CSV file if it's not already there.

    :param log_file: The CSV file that the collected data is written to.
    :return: None
    """
    if log_file.is_file():
        pass
    else:
        with open("num_features.csv", "a") as f:
            csv_out = csv.DictWriter(f, keys)
            csv_out.writeheader()
            f.flush()
    return


def sort_csv() -> type(None):
    """
    Group the values in the CSV by the 'label' column.

    :return: None
    """
    df = pd.read_csv("num_features.csv")
    df = df.sort_values(by=["label"])
    df.to_csv("num_features.csv", index=None)
    return


def collect_feature_data() -> type(None):
    """
    Write extracted URL features to file.

    :return: None
    """
    log_file = pathlib.Path("num_features.csv")
    write_header_to_csv(log_file)
    manager = mp.Manager()
    q = manager.Queue()
    pool = mp.Pool(mp.cpu_count())
    watcher = pool.apply_async(listener, (q,))
    jobs = []
    if log_file.is_file() and os.path.getsize("num_features.csv") > 730:
        resume(pool, q, jobs)
    else:
        start(pool, q, jobs)
    return


if __name__ == "__main__":
    collect_feature_data()

核心原因排查 #

1. CSV写入冗余开销：listener每次写入都重新创建csv.DictWriter对象，且每次写入后强制flush()触发磁盘同步，文件越大同步开销越高。
2. 频繁排序阻塞：start和resume完成批次任务后都会调用sort_csv()，需要加载整个CSV到内存排序后重写，文件增大时内存和IO开销剧增。
3. 已处理URL查询低效：resume中用列表存储已处理URL，成员查询是O(n)复杂度，数据量越大耗时越长。
4. 跨进程队列开销：mp.Manager().Queue()的跨进程通信开销随队列任务积累逐渐增加。

优化措施 #

1. 优化CSV写入逻辑 #

修改listener，仅初始化一次csv.DictWriter，取消实时flush()，改为任务结束时统一刷新：

def listener(q: multiprocessing.managers.AutoProxy) -> type(None):
    with open("num_features.csv", "a") as f:
        csv_out = csv.DictWriter(f, keys)
        while 1:
            m = q.get()
            if m == "kill":
                break
            csv_out.writerow(m)
        f.flush()
    return

2. 调整排序时机 #

删除start和resume中的sort_csv()调用，仅在所有任务完成后执行一次排序：

def collect_feature_data() -> type(None):
    log_file = pathlib.Path("num_features.csv")
    write_header_to_csv(log_file)
    manager = mp.Manager()
    q = manager.Queue()
    pool = mp.Pool(mp.cpu_count())
    watcher = pool.apply_async(listener, (q,))
    jobs = []
    if log_file.is_file() and os.path.getsize("num_features.csv") > 730:
        resume(pool, q, jobs)
    else:
        start(pool, q, jobs)
    sort_csv()
    return

3. 加速已处理URL查询 #

将resume中的processed_urls转为集合，成员查询复杂度降至O(1)：

processed_urls = set(
    pd.read_csv(
        "num_features.csv",
        usecols=[0],
        header=None,
    )
    .iloc[:, 0]
    .to_list()
)

4. 优化跨进程队列 #

使用mp.Queue()替代mp.Manager().Queue()，减少跨进程通信的额外开销（注意进程池内传递mp.Queue的正确方式）。

5. 串行网络请求异步化 #

将UrlFeaturizer.__init__中的串行网络请求（requests.get、whois.query等）改为异步实现（如用aiohttp），减少单进程的等待时间，提升整体吞吐量。

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