From ea1db8bba42303fd5975f864b75fbda53fb25f2d Mon Sep 17 00:00:00 2001 From: Rhys Chappell Date: Tue, 26 May 2026 18:07:55 +0000 Subject: [PATCH] fix O(N^2) serialization in analysis step --- src/boltzgen/task/analyze/analyze.py | 27 +++++++++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/src/boltzgen/task/analyze/analyze.py b/src/boltzgen/task/analyze/analyze.py index 256696aa..88a7bb10 100755 --- a/src/boltzgen/task/analyze/analyze.py +++ b/src/boltzgen/task/analyze/analyze.py @@ -53,6 +53,23 @@ from boltzgen.data.write.mmcif import to_mmcif +_WORKER_ANALYZE = None + + +def _init_worker(analyze): + """Worker process with single pickled copy of the Analyze task""" + global _WORKER_ANALYZE + _WORKER_ANALYZE = analyze + torch.set_num_threads(1) + torch.set_num_interop_threads(1) + rdkit.Chem.SetDefaultPickleProperties(rdkit.Chem.PropertyPickleOptions.AllProps) + + +def _worker_compute_metrics(idx): + """Module-level. Only the index is pickled per submission""" + return _WORKER_ANALYZE.compute_metrics(idx) + + class Analyze(Task): """ The Analyze step of the BoltzGen pipeline. @@ -205,9 +222,15 @@ def run_parallel(self, num, num_processes): try: with ProcessPoolExecutor( - max_workers=num_processes, mp_context=ctx + max_workers=num_processes, + mp_context=ctx, + initializer=_init_worker, + initargs=(self,), ) as ex: - fut2idx = {ex.submit(self.compute_metrics, i): i for i in remaining} + # Submit only index + fut2idx = { + ex.submit(_worker_compute_metrics, i): i for i in remaining + } # Iterate over futures that actually *completed* (finished or raised) for f in as_completed(fut2idx):