python - 多处理/psycopg2 TypeError : can't pickle _thread. RLock 对象

Question

我按照以下代码在 postgres 数据库上实现并行选择查询:

https://tech.geoblink.com/2017/07/06/parallelizing-queries-in-postgresql-with-python/

我的基本问题是我有大约 6k 个查询需要执行，我正在尝试优化这些选择查询的执行。最初它是一个包含所有 6k 谓词 ID 的 where id in (...) 查询，但我遇到了问题，查询在它运行的机器上耗尽了 > 4GB 的 RAM，所以我决定将其拆分为 6k 个单独的查询，这些查询在同步时保持稳定的内存使用。然而，明智地运行时间需要更长的时间，这对我的用例来说不是问题。尽管如此，我还是尽量减少时间。

我的代码是这样的:

class PostgresConnector(object):
    def __init__(self, db_url):
        self.db_url = db_url
        self.engine = self.init_connection()
        self.pool = self.init_pool()

    def init_pool(self):
        CPUS = multiprocessing.cpu_count()
        return multiprocessing.Pool(CPUS)

    def init_connection(self):
        LOGGER.info('Creating Postgres engine')
        return create_engine(self.db_url)

    def run_parallel_queries(self, queries):
        results = []
        try:
            for i in self.pool.imap_unordered(self.execute_parallel_query, queries):
                results.append(i)
        except Exception as exception:
            LOGGER.error('Error whilst executing %s queries in parallel: %s', len(queries), exception)
            raise
        finally:
            self.pool.close()
            self.pool.join()

        LOGGER.info('Parallel query ran producing %s sets of results of type: %s', len(results), type(results))

        return list(chain.from_iterable(results))

    def execute_parallel_query(self, query):
        con = psycopg2.connect(self.db_url)
        cur = con.cursor()
        cur.execute(query)
        records = cur.fetchall()
        con.close()

        return list(records)

但是无论何时运行，我都会收到以下错误:

TypeError: can't pickle _thread.RLock objects

我读过很多关于使用多处理和可腌制对象的类似问题，但我终究无法弄清楚我做错了什么。

池通常是每个进程一个(我认为这是最佳实践)，但每个连接器类的实例共享，因此它不会为每次使用 parallel_query 方法创建一个池。

类似问题的最佳答案:

Accessing a MySQL connection pool from Python multiprocessing

除了使用 MySql 而不是 Postgres 之外，显示了与我自己的几乎相同的实现。

我做错了什么吗？

谢谢!

编辑:

我找到了这个答案:

Python Postgres psycopg2 ThreadedConnectionPool exhausted

这非常详细，看起来我误解了 multiprocessing.Pool 与连接池(例如 ThreadedConnectionPool)给我的含义。然而，在第一个链接中，它没有提到需要任何连接池等。这个解决方案看起来不错，但对于我认为是一个相当简单的问题，似乎有很多代码？

编辑 2:

所以上面的链接解决了另一个问题，我很可能会遇到这个问题，所以我很高兴我发现了这个问题，但它并没有解决最初无法使用 imap_unordered 的问题到酸洗错误。非常令人沮丧。

最后，我认为可能值得注意的是，它在 Heroku 中运行，在 worker dyno 上，使用 Redis rq 进行调度、后台任务等，并使用 Postgres 的托管实例作为数据库。

Answer 1

简单来说，postgres连接和sqlalchemy连接池是线程安全的，但是它们不是fork安全的。

如果你想使用多进程，你应该在fork之后的每个子进程中初始化引擎。

如果你想共享引擎，你应该使用多线程。

引用Thread and process safety in psycopg2 documentation :

libpq connections shouldn’t be used by a forked processes, so when using a module such as multiprocessing or a forking web deploy method such as FastCGI make sure to create the connections after the fork.

如果您正在使用 multiprocessing.Pool，则有一个关键字参数初始化器，可用于在每个子进程上运行一次代码。试试这个:

class PostgresConnector(object):
    def __init__(self, db_url):
        self.db_url = db_url
        self.pool = self.init_pool()

    def init_pool(self):
        CPUS = multiprocessing.cpu_count()
        return multiprocessing.Pool(CPUS, initializer=self.init_connection(self.db_url))

    @classmethod
    def init_connection(cls, db_url):
        def _init_connection():
            LOGGER.info('Creating Postgres engine')
            cls.engine = create_engine(db_url)
        return _init_connection

    def run_parallel_queries(self, queries):
        results = []
        try:
            for i in self.pool.imap_unordered(self.execute_parallel_query, queries):
                results.append(i)
        except Exception as exception:
            LOGGER.error('Error whilst executing %s queries in parallel: %s', len(queries), exception)
            raise
        finally:
            pass
            #self.pool.close()
            #self.pool.join()

        LOGGER.info('Parallel query ran producing %s sets of results of type: %s', len(results), type(results))

        return list(chain.from_iterable(results))

    def execute_parallel_query(self, query):
        with self.engine.connect() as conn:
            with conn.begin():
                result = conn.execute(query)
                return result.fetchall()

    def __getstate__(self):
        # this is a hack, if you want to remove this method, you should
        # remove self.pool and just pass pool explicitly
        self_dict = self.__dict__.copy()
        del self_dict['pool']
        return self_dict

---

现在，解决 XY 问题。

Initially it was a single query with the where id in (...) contained all 6k predicate IDs but I ran into issues with the query using up > 4GB of RAM on the machine it ran on, so I decided to split it out into 6k individual queries which when synchronously keeps a steady memory usage.

您可能想要做的是以下选项之一:

1. 编写一个生成所有 6000 个 ID 的子查询，并在您的原始批量查询中使用该子查询。
2. 同上，但是把子查询写成CTE
3. 如果您的 ID 列表来自外部来源(即不是来自数据库)，那么您可以创建一个包含 6000 个 ID 的临时表，然后对临时表运行原始批量查询

但是，如果您坚持通过 python 运行 6000 个 ID，那么最快的查询很可能既不会一次性完成所有 6000 个 ID(这会耗尽内存)，也不会运行 6000 个单独的查询。相反，您可能想尝试对查询进行分 block 。例如一次发送 500 个 ID。您将不得不试验 block 大小，以确定您一次可以发送的最大 ID 数，同时仍然在您的内存预算范围内。