From 3739e1b4d2fab224a7f958bc205ed70a7de47b33 Mon Sep 17 00:00:00 2001 From: GeorgWa Date: Mon, 31 Aug 2026 00:14:52 +0200 Subject: [PATCH 1/5] Share the modification registry with multiprocessing workers Workers are started with the "spawn" start method and re-import alphabase, so they only ever saw the modifications in modification.tsv. Every runtime change to the registry -- custom modifications, modloss filtering, lower-case AAs, a custom TSV -- was silently absent in workers, raising KeyError on mod lookups. MOD_DF is the single source of truth and update_all_by_MOD_DF() already rebuilds everything derived from it, so a copy of MOD_DF is a complete snapshot of the registry. get_modification_state()/set_modification_state() move that snapshot between processes, and spawn_pool() installs it in each worker as it starts. Because the snapshot is the whole table rather than a record of which mutators ran, it stays correct as new mutators are added. This replaces the fallback in SpecLibBase.calc_precursor_isotope_intensity, which silently dropped to a single process whenever custom modifications were present -- and which never covered the other three mutators anyway, since has_custom_mods() only tests for the User-added classification. has_custom_mods itself stays exported. Co-Authored-By: Claude Opus 5 (1M context) --- alphabase/constants/modification.py | 40 +++++++- alphabase/peptide/precursor.py | 6 +- alphabase/psm_reader/sage_reader.py | 4 +- alphabase/spectral_library/base.py | 8 -- alphabase/spectral_library/decoy.py | 4 +- alphabase/utils.py | 36 +++++++ tests/unit/peptide/test_precursor_mp.py | 125 ++++++++++++++++++++++++ 7 files changed, 207 insertions(+), 16 deletions(-) create mode 100644 tests/unit/peptide/test_precursor_mp.py diff --git a/alphabase/constants/modification.py b/alphabase/constants/modification.py index 63b14e08..2d0bf718 100644 --- a/alphabase/constants/modification.py +++ b/alphabase/constants/modification.py @@ -1,3 +1,11 @@ +"""Modification registry. + +`MOD_DF` is the single source of truth: every other module-level lookup here is +derived from it and rebuilt in place by :func:`update_all_by_MOD_DF`. Anything +that mutates the registry must go through that function, which is what makes +:func:`get_modification_state` a complete snapshot of the registry. +""" + import os from typing import List, Union @@ -463,7 +471,7 @@ def _check_mass_sanity( composition_mass = calc_mass_from_formula(composition) if not np.allclose(composition_mass, MOD_MASS[mod_name], atol=1e-5): raise ValueError( - f"Modification mass of {mod_name} is inconsistent with the composition formula: {composition}, df version {MOD_DF.loc[mod_name,['composition']]}" + f"Modification mass of {mod_name} is inconsistent with the composition formula: {composition}, df version {MOD_DF.loc[mod_name, ['composition']]}" f" calculated_mass={composition_mass}, mod_mass={MOD_MASS[mod_name]}" ) @@ -533,6 +541,36 @@ def has_custom_mods(): return len(MOD_DF[MOD_DF["classification"] == _MOD_CLASSIFICATION_USER_ADDED]) > 0 +def get_modification_state() -> pd.DataFrame: + """Snapshot the modification registry so it can be handed to another process. + + Returns + ------- + pd.DataFrame + A copy of :data:`MOD_DF`. It carries every runtime change to the + registry, not just user-added modifications. + """ + return MOD_DF.copy() + + +def set_modification_state(mod_df: pd.DataFrame) -> None: + """Install a registry snapshot and rebuild every derived lookup. + + Processes started with the "spawn" start method re-import alphabase and so + see only the modifications in `modification.tsv`. Without this, every + runtime change -- custom modifications, modloss filtering, lower-case AAs, + a custom TSV -- is silently absent in workers. + + Parameters + ---------- + mod_df : pd.DataFrame + Snapshot as returned by :func:`get_modification_state`. + """ + global MOD_DF + MOD_DF = mod_df + update_all_by_MOD_DF() + + def add_new_modifications(new_mods: Union[list, dict]): """Add new modifications into :data:`MOD_DF`. diff --git a/alphabase/peptide/precursor.py b/alphabase/peptide/precursor.py index 82dad36b..b76e313e 100644 --- a/alphabase/peptide/precursor.py +++ b/alphabase/peptide/precursor.py @@ -1,4 +1,3 @@ -import multiprocessing as mp import typing from functools import partial @@ -13,6 +12,7 @@ from alphabase.constants.modification import MOD_Composition, ModificationKeys from alphabase.numba_wrapper import numba_njit from alphabase.peptide.mass_calc import calc_peptide_masses_for_same_len_seqs +from alphabase.utils import spawn_pool def refine_precursor_df( @@ -537,7 +537,7 @@ def calc_precursor_isotope_info_mp( ) df_list = [] df_group = precursor_df.groupby("nAA") - with mp.get_context("spawn").Pool(processes) as p: + with spawn_pool(processes) as p: processing = p.imap( partial( calc_precursor_isotope_info, @@ -681,7 +681,7 @@ def calc_precursor_isotope_intensity_mp( df_list = [] df_group = precursor_df.groupby("nAA") - with mp.get_context("spawn").Pool(mp_process_num) as p: + with spawn_pool(mp_process_num) as p: processing = p.imap( partial( calc_precursor_isotope_intensity, diff --git a/alphabase/psm_reader/sage_reader.py b/alphabase/psm_reader/sage_reader.py index 7ed527a5..c8d75236 100644 --- a/alphabase/psm_reader/sage_reader.py +++ b/alphabase/psm_reader/sage_reader.py @@ -1,7 +1,6 @@ """SageReader for reading Sage output files.""" import logging -import multiprocessing as mp import re from abc import ABC from functools import partial @@ -17,6 +16,7 @@ PSMReaderBase, psm_reader_provider, ) +from alphabase.utils import spawn_pool class SageModificationTranslator: @@ -492,7 +492,7 @@ def _apply_translate_modifications_mp( Whether to show a progress bar. Defaults to True """ - with mp.get_context("spawn").Pool(mp_process_num) as p: + with spawn_pool(mp_process_num) as p: processing = p.imap( partial( _apply_translate_modifications, diff --git a/alphabase/spectral_library/base.py b/alphabase/spectral_library/base.py index 6b781f4d..366b70a3 100644 --- a/alphabase/spectral_library/base.py +++ b/alphabase/spectral_library/base.py @@ -7,7 +7,6 @@ import numpy as np import pandas as pd -from alphabase.constants.modification import has_custom_mods from alphabase.io.hdf import HDF_File from alphabase.peptide.fragment import ( calc_fragment_count, @@ -429,13 +428,6 @@ def calc_precursor_isotope_intensity( mp_process_num > 1 and len(self.precursor_df) > mp_batch_size ) - if do_multiprocessing and has_custom_mods(): - logging.warning( - "Multiprocessing not compatible with custom modifications yet, falling back to single process." - ) - do_multiprocessing = False - # TODO enable multiprocessing also in this case - if do_multiprocessing: (self._precursor_df) = calc_precursor_isotope_intensity_mp( self.precursor_df, diff --git a/alphabase/spectral_library/decoy.py b/alphabase/spectral_library/decoy.py index 2389ad23..e91d0504 100644 --- a/alphabase/spectral_library/decoy.py +++ b/alphabase/spectral_library/decoy.py @@ -1,10 +1,10 @@ import copy -import multiprocessing as mp from typing import Any import pandas as pd from alphabase.spectral_library.base import SpecLibBase +from alphabase.utils import spawn_pool def _batchify_series(series, mp_batch_size): @@ -195,7 +195,7 @@ def decoy_sequence( ) series_list = [] - with mp.get_context("spawn").Pool(mp_process_num) as p: + with spawn_pool(mp_process_num) as p: processing = p.imap(self.generator, sequence_batches) for df in processing: series_list.append(df) diff --git a/alphabase/utils.py b/alphabase/utils.py index 9b4ea67d..abed2fd6 100644 --- a/alphabase/utils.py +++ b/alphabase/utils.py @@ -1,10 +1,16 @@ import io import itertools +import multiprocessing as mp import warnings import pandas as pd import tqdm +from alphabase.constants.modification import ( + get_modification_state, + set_modification_state, +) + class AlphabaseDeprecationWarning(DeprecationWarning): pass @@ -70,3 +76,33 @@ def _get_delimiter(file_path: str) -> str: return "," else: return "\t" + + +def spawn_pool(processes: int, *, context=None, **kwargs): + """Create a worker pool that shares this process's modification registry. + + Workers are started with the "spawn" start method and re-import alphabase + from scratch, so they would otherwise lose every runtime change to the + modification registry. The snapshot is taken once here and installed in each + worker as it starts. + + Parameters + ---------- + processes : int + Number of worker processes. + + context : multiprocessing context, optional + Substitute context, e.g. `torch.multiprocessing.get_context("spawn")`, + whose reducers are needed to share model tensors with workers. + Defaults to the standard "spawn" context. + + **kwargs + Forwarded to `Pool`. + """ + ctx = context if context is not None else mp.get_context("spawn") + return ctx.Pool( + processes, + initializer=set_modification_state, + initargs=(get_modification_state(),), + **kwargs, + ) diff --git a/tests/unit/peptide/test_precursor_mp.py b/tests/unit/peptide/test_precursor_mp.py new file mode 100644 index 00000000..5133f707 --- /dev/null +++ b/tests/unit/peptide/test_precursor_mp.py @@ -0,0 +1,125 @@ +"""Multiprocessing workers must see the same modification registry as the parent. + +Workers are started with the "spawn" start method and re-import alphabase, so +they only see `modification.tsv` unless the registry is handed to them +explicitly. These tests cover each way the registry can change at runtime. +""" + +import os + +import pandas as pd +import pytest + +from alphabase.constants._const import CONST_FILE_FOLDER +from alphabase.constants.modification import ( + add_modifications_for_lower_case_AA, + add_new_modifications, + get_modification_state, + keep_modloss_by_importance, + load_mod_df, + set_modification_state, +) +from alphabase.peptide.precursor import ( + calc_precursor_isotope_intensity, + calc_precursor_isotope_intensity_mp, + update_precursor_mz, +) +from alphabase.utils import spawn_pool + +CUSTOM_MOD = "TestCustomMod@K" +CUSTOM_MOD_COMPOSITION = "H(4)O(2)" + + +@pytest.fixture +def restore_registry(): + """Undo the global registry changes each test makes.""" + snapshot = get_modification_state() + yield + set_modification_state(snapshot) + + +def _worker_registry(_): + """Return the worker process's own view of the registry.""" + return get_modification_state() + + +def _add_custom_mod(): + add_new_modifications({CUSTOM_MOD: {"composition": CUSTOM_MOD_COMPOSITION}}) + + +def _filter_modloss(): + # `load_mod_df` applies level 1 at import, keeping 2 modloss values; level 0 + # keeps 867, so a worker that missed this change is detectably different. + keep_modloss_by_importance(0.0) + + +def _add_lower_case_aa(): + add_modifications_for_lower_case_AA() + + +def _load_custom_tsv(): + """Load a TSV that differs from the default one shipped with alphabase.""" + default_tsv = os.path.join(CONST_FILE_FOLDER, "modification.tsv") + mod_df = pd.read_table(default_tsv, keep_default_na=False) + extra = mod_df.iloc[[0]].copy() + extra["mod_name"] = "TestTsvMod@K" + extra["composition"] = CUSTOM_MOD_COMPOSITION + custom_tsv = os.path.join( + os.path.dirname(default_tsv), "_test_modification_tmp.tsv" + ) + pd.concat([mod_df, extra], ignore_index=True).to_csv( + custom_tsv, sep="\t", index=False + ) + try: + load_mod_df(custom_tsv) + finally: + os.remove(custom_tsv) + + +@pytest.mark.parametrize( + "mutate", + [_add_custom_mod, _filter_modloss, _add_lower_case_aa, _load_custom_tsv], + ids=["custom_mods", "modloss_filtering", "lower_case_AA", "custom_tsv"], +) +def test_worker_registry_matches_parent(mutate, restore_registry): + # Given a registry changed at runtime + mutate() + expected = get_modification_state() + + # When a worker process reports its own registry + with spawn_pool(2) as pool: + registries = pool.map(_worker_registry, [None, None]) + + # Then it is identical to the parent's + for registry in registries: + pd.testing.assert_frame_equal(registry, expected) + + +def _precursor_df(n_precursors=40): + df = pd.DataFrame( + { + "sequence": ["PEPTIDEK"] * n_precursors, + "mods": [CUSTOM_MOD] * n_precursors, + "mod_sites": ["8"] * n_precursors, + "charge": [2] * n_precursors, + } + ) + df["nAA"] = df["sequence"].str.len() + return update_precursor_mz(df) + + +def test_isotope_intensity_mp_matches_single_process(restore_registry): + # Given a library whose precursors carry a custom modification + _add_custom_mod() + isotope_cols = [f"i_{i}" for i in range(6)] + + # When the isotope intensities are calculated with and without multiprocessing + single = calc_precursor_isotope_intensity(_precursor_df(), max_isotope=6) + multi = calc_precursor_isotope_intensity_mp( + _precursor_df(), max_isotope=6, mp_batch_size=10, mp_process_num=2 + ) + + # Then the results agree + pd.testing.assert_frame_equal( + single.sort_index()[isotope_cols], multi.sort_index()[isotope_cols] + ) From 1178032aa8f7deb001b2a8ac7ac0a229b06dd15a Mon Sep 17 00:00:00 2001 From: GeorgWa Date: Mon, 31 Aug 2026 00:54:51 +0200 Subject: [PATCH 2/5] Take the worker pool out of the multiprocessing interface All four pools had the same shape -- batchify, imap, optional progress bar, concat -- reimplemented each time, with three separate batching helpers between them. parallel_imap() and parallel_apply() carry that shape once, so callers no longer construct a pool, choose a context, or think about what state a worker needs; spawn_pool is now private. Public signatures are unchanged. The helpers this removes were all private: _batchify_df (defined twice, with different signatures), _batchify_series and _count_batchify_df. Fixes two progress bar bugs found on the way: SpecLibBase.calc_precursor_isotope_info() passed process_bar= to a parameter named progress_bar=, so the multiprocessing branch raised TypeError and was dead. It also dropped mp_batch_size, which is now forwarded. calc_precursor_isotope_intensity_mp() documented progress_bar as a flag but peptdeep passes a callback. A callable was truthy, so it fell into the tqdm branch and the callback was silently discarded. Progress now accepts True for a tqdm bar, a callable to supply your own, or anything falsy for none, which also restores the callback contract calc_precursor_isotope_info_mp already had. Co-Authored-By: Claude Opus 5 (1M context) --- alphabase/peptide/precursor.py | 85 ++++++------------- alphabase/psm_reader/sage_reader.py | 53 +++--------- alphabase/spectral_library/base.py | 3 +- alphabase/spectral_library/decoy.py | 23 ++--- alphabase/utils.py | 108 +++++++++++++++++++++--- tests/unit/peptide/test_precursor_mp.py | 47 +++++++++-- 6 files changed, 187 insertions(+), 132 deletions(-) diff --git a/alphabase/peptide/precursor.py b/alphabase/peptide/precursor.py index b76e313e..a8ece2c3 100644 --- a/alphabase/peptide/precursor.py +++ b/alphabase/peptide/precursor.py @@ -3,7 +3,6 @@ import numpy as np import pandas as pd -from tqdm import tqdm from xxhash import xxh64_intdigest from alphabase.constants.aa import AA_Composition @@ -12,7 +11,7 @@ from alphabase.constants.modification import MOD_Composition, ModificationKeys from alphabase.numba_wrapper import numba_njit from alphabase.peptide.mass_calc import calc_peptide_masses_for_same_len_seqs -from alphabase.utils import spawn_pool +from alphabase.utils import parallel_apply def refine_precursor_df( @@ -475,24 +474,6 @@ def calc_precursor_isotope_info( return precursor_df -def _batchify_df(df_group, mp_batch_size): - """Internal funciton for multiprocessing""" - for _, df in df_group: - for i in range(0, len(df), mp_batch_size): - yield df.iloc[i : i + mp_batch_size, :] - - -def _count_batchify_df(df_group, mp_batch_size): - """Internal funciton for multiprocessing""" - count = 0 - for _, df in df_group: - for _ in range(0, len(df), mp_batch_size): - count += 1 - return count - - -# `progress_bar` should be replaced by more advanced tqdm wrappers created by Sander -# I will leave it to alphabase.utils def calc_precursor_isotope_info_mp( precursor_df: pd.DataFrame, processes: int = 8, @@ -535,23 +516,17 @@ def calc_precursor_isotope_info_mp( precursor_df=precursor_df, min_right_most_intensity=min_right_most_intensity, ) - df_list = [] - df_group = precursor_df.groupby("nAA") - with spawn_pool(processes) as p: - processing = p.imap( - partial( - calc_precursor_isotope_info, - min_right_most_intensity=min_right_most_intensity, - ), - _batchify_df(df_group, mp_batch_size), - ) - if progress_bar: - processing = progress_bar( - processing, _count_batchify_df(df_group, mp_batch_size) - ) - for df in processing: - df_list.append(df) - return pd.concat(df_list) + return parallel_apply( + partial( + calc_precursor_isotope_info, + min_right_most_intensity=min_right_most_intensity, + ), + precursor_df, + processes=processes, + batch_size=mp_batch_size, + group_by="nAA", + progress=progress_bar, + ) def calc_precursor_isotope_intensity( @@ -678,28 +653,20 @@ def calc_precursor_isotope_intensity_mp( normalize=normalize, ) - df_list = [] - df_group = precursor_df.groupby("nAA") - - with spawn_pool(mp_process_num) as p: - processing = p.imap( - partial( - calc_precursor_isotope_intensity, - max_isotope=max_isotope, - min_right_most_intensity=min_right_most_intensity, - normalize=normalize, - ), - _batchify_df(df_group, mp_batch_size), - ) - - if progress_bar: - df_list = list( - tqdm(processing, total=_count_batchify_df(df_group, mp_batch_size)) - ) - else: - df_list = list(processing) - - return pd.concat(df_list, ignore_index=True) + return parallel_apply( + partial( + calc_precursor_isotope_intensity, + max_isotope=max_isotope, + min_right_most_intensity=min_right_most_intensity, + normalize=normalize, + ), + precursor_df, + processes=mp_process_num, + batch_size=mp_batch_size, + group_by="nAA", + progress=progress_bar, + ignore_index=True, + ) calc_precursor_isotope = calc_precursor_isotope_intensity diff --git a/alphabase/psm_reader/sage_reader.py b/alphabase/psm_reader/sage_reader.py index c8d75236..04c55843 100644 --- a/alphabase/psm_reader/sage_reader.py +++ b/alphabase/psm_reader/sage_reader.py @@ -4,11 +4,10 @@ import re from abc import ABC from functools import partial -from typing import Generator, List, Optional, Tuple +from typing import List, Optional, Tuple import numpy as np import pandas as pd -from tqdm import tqdm from alphabase.constants.modification import MOD_DF, ModificationKeys from alphabase.psm_reader.keys import PsmDfCols @@ -16,7 +15,7 @@ PSMReaderBase, psm_reader_provider, ) -from alphabase.utils import spawn_pool +from alphabase.utils import parallel_apply class SageModificationTranslator: @@ -443,27 +442,6 @@ def _apply_translate_modifications( return psm_df -def _batchify_df(df: pd.DataFrame, mp_batch_size: int) -> Generator: - """Internal funciton for applying translation modifications in parallel. - - Parameters - ---------- - df : pd.DataFrame - The PSM dataframe. - - mp_batch_size : int - The batch size for parallel processing. - - Returns - ------- - typing.Generator - A generator for the batchified dataframe. - - """ - for i in range(0, len(df), mp_batch_size): - yield df.iloc[i : i + mp_batch_size, :] - - def _apply_translate_modifications_mp( psm_df: pd.DataFrame, mod_translation_df: pd.DataFrame, @@ -492,22 +470,17 @@ def _apply_translate_modifications_mp( Whether to show a progress bar. Defaults to True """ - with spawn_pool(mp_process_num) as p: - processing = p.imap( - partial( - _apply_translate_modifications, - mod_translation_df=mod_translation_df, - ), - _batchify_df(psm_df, mp_batch_size), - ) - if progress_bar: - df_list = list( - tqdm(processing, total=int(np.ceil(len(psm_df) / mp_batch_size))) - ) - else: - df_list = list(processing) - - return pd.concat(df_list, ignore_index=True) + return parallel_apply( + partial( + _apply_translate_modifications, + mod_translation_df=mod_translation_df, + ), + psm_df, + processes=mp_process_num, + batch_size=mp_batch_size, + progress=progress_bar, + ignore_index=True, + ) def _get_annotated_mod_df() -> pd.DataFrame: diff --git a/alphabase/spectral_library/base.py b/alphabase/spectral_library/base.py index 366b70a3..19fee371 100644 --- a/alphabase/spectral_library/base.py +++ b/alphabase/spectral_library/base.py @@ -477,7 +477,8 @@ def calc_precursor_isotope_info( (self._precursor_df) = calc_precursor_isotope_info_mp( self.precursor_df, processes=mp_process_num, - process_bar=mp_process_bar, + mp_batch_size=mp_batch_size, + progress_bar=mp_process_bar, ) else: (self._precursor_df) = calc_precursor_isotope_info(self.precursor_df) diff --git a/alphabase/spectral_library/decoy.py b/alphabase/spectral_library/decoy.py index e91d0504..d4000735 100644 --- a/alphabase/spectral_library/decoy.py +++ b/alphabase/spectral_library/decoy.py @@ -4,13 +4,7 @@ import pandas as pd from alphabase.spectral_library.base import SpecLibBase -from alphabase.utils import spawn_pool - - -def _batchify_series(series, mp_batch_size): - """Internal funciton for multiprocessing""" - for i in range(0, len(series), mp_batch_size): - yield series.iloc[i : i + mp_batch_size] +from alphabase.utils import parallel_apply class BaseDecoyGenerator: @@ -190,16 +184,13 @@ def decoy_sequence( self._remove_target_seqs() return - sequence_batches = list( - _batchify_series(self._precursor_df["sequence"], mp_batch_size) + self._precursor_df["sequence"] = parallel_apply( + self.generator, + self._precursor_df["sequence"], + processes=mp_process_num, + batch_size=mp_batch_size, + progress=False, ) - - series_list = [] - with spawn_pool(mp_process_num) as p: - processing = p.imap(self.generator, sequence_batches) - for df in processing: - series_list.append(df) - self._precursor_df["sequence"] = pd.concat(series_list) self._remove_target_seqs() def _remove_target_seqs(self): diff --git a/alphabase/utils.py b/alphabase/utils.py index abed2fd6..d18fac00 100644 --- a/alphabase/utils.py +++ b/alphabase/utils.py @@ -78,31 +78,119 @@ def _get_delimiter(file_path: str) -> str: return "\t" -def spawn_pool(processes: int, *, context=None, **kwargs): +def _spawn_pool(processes: int, *, context=None): """Create a worker pool that shares this process's modification registry. Workers are started with the "spawn" start method and re-import alphabase from scratch, so they would otherwise lose every runtime change to the modification registry. The snapshot is taken once here and installed in each worker as it starts. + """ + ctx = context if context is not None else mp.get_context("spawn") + return ctx.Pool( + processes, + initializer=set_modification_state, + initargs=(get_modification_state(),), + ) + + +def _batchify(obj, batch_size: int, group_by=None): + """Yield row batches of a DataFrame or Series, optionally within groups.""" + groups = (group for _, group in obj.groupby(group_by)) if group_by else iter((obj,)) + for group in groups: + for i in range(0, len(group), batch_size): + yield group.iloc[i : i + batch_size] + + +def _batch_count(obj, batch_size: int, group_by=None) -> int: + sizes = obj.groupby(group_by).size().values if group_by else [len(obj)] + return sum((size + batch_size - 1) // batch_size for size in sizes) + + +def _with_progress(iterator, total, progress): + """Wrap `iterator` in a progress bar. + + `progress` is True for a tqdm bar, a callable `progress(iterator, total)` to + supply your own, or anything falsy for no bar. + """ + if progress is True: + return tqdm.tqdm(iterator, total=total) + if callable(progress): + return progress(iterator, total) + return iterator + + +def parallel_imap( + func, + iterable, + *, + processes: int, + total: int = None, + unordered: bool = False, + progress=True, + context=None, +): + """Map `func` over `iterable` in workers that share the modification registry. + + Results are yielded as they arrive, so callers holding large objects only + keep one batch in memory at a time. Parameters ---------- processes : int Number of worker processes. + total : int, optional + Number of items, for the progress bar. + + unordered : bool, optional + Yield results as they finish rather than in input order. + + progress : bool or callable, optional + See :func:`_with_progress`. + context : multiprocessing context, optional Substitute context, e.g. `torch.multiprocessing.get_context("spawn")`, whose reducers are needed to share model tensors with workers. - Defaults to the standard "spawn" context. + """ + with _spawn_pool(processes, context=context) as pool: + mapper = pool.imap_unordered if unordered else pool.imap + yield from _with_progress(mapper(func, iterable), total, progress) + + +def parallel_apply( + func, + obj, + *, + processes: int, + batch_size: int, + group_by=None, + progress=True, + context=None, + ignore_index: bool = False, +): + """Apply `func` to row batches of `obj` in parallel and concatenate the result. + + Parameters + ---------- + obj : pd.DataFrame or pd.Series + Object to split into batches. + + group_by : optional + Column to group by before batching, so each batch is within one group. - **kwargs - Forwarded to `Pool`. + See :func:`parallel_imap` for the remaining parameters. """ - ctx = context if context is not None else mp.get_context("spawn") - return ctx.Pool( - processes, - initializer=set_modification_state, - initargs=(get_modification_state(),), - **kwargs, + return pd.concat( + list( + parallel_imap( + func, + _batchify(obj, batch_size, group_by), + processes=processes, + total=_batch_count(obj, batch_size, group_by), + progress=progress, + context=context, + ) + ), + ignore_index=ignore_index, ) diff --git a/tests/unit/peptide/test_precursor_mp.py b/tests/unit/peptide/test_precursor_mp.py index 5133f707..faaf33aa 100644 --- a/tests/unit/peptide/test_precursor_mp.py +++ b/tests/unit/peptide/test_precursor_mp.py @@ -24,7 +24,8 @@ calc_precursor_isotope_intensity_mp, update_precursor_mz, ) -from alphabase.utils import spawn_pool +from alphabase.spectral_library.base import SpecLibBase +from alphabase.utils import parallel_imap CUSTOM_MOD = "TestCustomMod@K" CUSTOM_MOD_COMPOSITION = "H(4)O(2)" @@ -87,20 +88,21 @@ def test_worker_registry_matches_parent(mutate, restore_registry): expected = get_modification_state() # When a worker process reports its own registry - with spawn_pool(2) as pool: - registries = pool.map(_worker_registry, [None, None]) + registries = list( + parallel_imap(_worker_registry, [None, None], processes=2, progress=False) + ) # Then it is identical to the parent's for registry in registries: pd.testing.assert_frame_equal(registry, expected) -def _precursor_df(n_precursors=40): +def _precursor_df(n_precursors=40, mod=CUSTOM_MOD): df = pd.DataFrame( { "sequence": ["PEPTIDEK"] * n_precursors, - "mods": [CUSTOM_MOD] * n_precursors, - "mod_sites": ["8"] * n_precursors, + "mods": [mod] * n_precursors, + "mod_sites": ["8" if mod else ""] * n_precursors, "charge": [2] * n_precursors, } ) @@ -123,3 +125,36 @@ def test_isotope_intensity_mp_matches_single_process(restore_registry): pd.testing.assert_frame_equal( single.sort_index()[isotope_cols], multi.sort_index()[isotope_cols] ) + + +def test_caller_supplied_progress_bar_is_used(): + # Given a caller supplying its own progress bar rather than the default + seen_totals = [] + + def progress(iterator, total): + seen_totals.append(total) + return iterator + + # When isotope intensities are calculated with multiprocessing + calc_precursor_isotope_intensity_mp( + _precursor_df(40, mod=""), + max_isotope=6, + mp_batch_size=10, + mp_process_num=2, + progress_bar=progress, + ) + + # Then it is actually driven, rather than being treated as a plain flag + assert seen_totals == [4] + + +def test_speclib_isotope_info_runs_with_multiprocessing(): + # Given a library large enough to take the multiprocessing branch + lib = SpecLibBase() + lib._precursor_df = _precursor_df(20_000, mod="") + + # When isotope info is calculated + lib.calc_precursor_isotope_info(mp_process_num=2, mp_batch_size=1000) + + # Then it completes; it used to raise TypeError on an unknown keyword + assert "isotope_apex_offset" in lib.precursor_df.columns From 39e37de0d3c1adb8143a0f5c0c8d22919f7bb9b8 Mon Sep 17 00:00:00 2001 From: GeorgWa Date: Mon, 31 Aug 2026 10:57:44 +0200 Subject: [PATCH 3/5] Write the comments in Simplified Technical English Co-Authored-By: Claude Opus 5 (1M context) --- alphabase/constants/modification.py | 24 +++++++------- alphabase/utils.py | 42 ++++++++++++------------- tests/unit/peptide/test_precursor_mp.py | 42 ++++++++++++------------- 3 files changed, 53 insertions(+), 55 deletions(-) diff --git a/alphabase/constants/modification.py b/alphabase/constants/modification.py index 2d0bf718..c06fee4a 100644 --- a/alphabase/constants/modification.py +++ b/alphabase/constants/modification.py @@ -1,9 +1,8 @@ """Modification registry. -`MOD_DF` is the single source of truth: every other module-level lookup here is -derived from it and rebuilt in place by :func:`update_all_by_MOD_DF`. Anything -that mutates the registry must go through that function, which is what makes -:func:`get_modification_state` a complete snapshot of the registry. +`update_all_by_MOD_DF` calculates all other module-level lookups from `MOD_DF`. +Thus a copy of `MOD_DF` is a complete snapshot of the registry. Change the +registry only with `update_all_by_MOD_DF`, to keep this true. """ import os @@ -542,29 +541,28 @@ def has_custom_mods(): def get_modification_state() -> pd.DataFrame: - """Snapshot the modification registry so it can be handed to another process. + """Copy the modification registry, to send it to a different process. Returns ------- pd.DataFrame - A copy of :data:`MOD_DF`. It carries every runtime change to the - registry, not just user-added modifications. + A copy of :data:`MOD_DF`. It contains all changes made at run time, not + only the user-added modifications. """ return MOD_DF.copy() def set_modification_state(mod_df: pd.DataFrame) -> None: - """Install a registry snapshot and rebuild every derived lookup. + """Install a registry copy and calculate the derived lookups again. - Processes started with the "spawn" start method re-import alphabase and so - see only the modifications in `modification.tsv`. Without this, every - runtime change -- custom modifications, modloss filtering, lower-case AAs, - a custom TSV -- is silently absent in workers. + A process that starts with the "spawn" method imports alphabase again. Thus + it knows only the modifications in `modification.tsv`. Use this function to + give the process all changes that were made at run time. Parameters ---------- mod_df : pd.DataFrame - Snapshot as returned by :func:`get_modification_state`. + A copy from :func:`get_modification_state`. """ global MOD_DF MOD_DF = mod_df diff --git a/alphabase/utils.py b/alphabase/utils.py index d18fac00..1d655d8c 100644 --- a/alphabase/utils.py +++ b/alphabase/utils.py @@ -79,12 +79,11 @@ def _get_delimiter(file_path: str) -> str: def _spawn_pool(processes: int, *, context=None): - """Create a worker pool that shares this process's modification registry. + """Create a worker pool that shares the modification registry. - Workers are started with the "spawn" start method and re-import alphabase - from scratch, so they would otherwise lose every runtime change to the - modification registry. The snapshot is taken once here and installed in each - worker as it starts. + A worker starts with the "spawn" method and imports alphabase again. Thus it + loses all changes that were made to the registry at run time. This function + copies the registry one time, then installs it in each worker at start. """ ctx = context if context is not None else mp.get_context("spawn") return ctx.Pool( @@ -95,7 +94,7 @@ def _spawn_pool(processes: int, *, context=None): def _batchify(obj, batch_size: int, group_by=None): - """Yield row batches of a DataFrame or Series, optionally within groups.""" + """Divide a DataFrame or Series into batches of rows.""" groups = (group for _, group in obj.groupby(group_by)) if group_by else iter((obj,)) for group in groups: for i in range(0, len(group), batch_size): @@ -108,10 +107,11 @@ def _batch_count(obj, batch_size: int, group_by=None) -> int: def _with_progress(iterator, total, progress): - """Wrap `iterator` in a progress bar. + """Add a progress bar to `iterator`. - `progress` is True for a tqdm bar, a callable `progress(iterator, total)` to - supply your own, or anything falsy for no bar. + Set `progress` to True for a tqdm bar. Set it to a callable + `progress(iterator, total)` to supply a different bar. Set it to a false + value for no bar. """ if progress is True: return tqdm.tqdm(iterator, total=total) @@ -130,28 +130,28 @@ def parallel_imap( progress=True, context=None, ): - """Map `func` over `iterable` in workers that share the modification registry. + """Apply `func` to each item of `iterable` in workers. - Results are yielded as they arrive, so callers holding large objects only - keep one batch in memory at a time. + The workers share the modification registry. This function gives each result + when it is ready. Thus the caller keeps only one batch in memory. Parameters ---------- processes : int - Number of worker processes. + The number of worker processes. total : int, optional - Number of items, for the progress bar. + The number of items, for the progress bar. unordered : bool, optional - Yield results as they finish rather than in input order. + Give each result when it is ready, not in the sequence of the input. progress : bool or callable, optional See :func:`_with_progress`. context : multiprocessing context, optional - Substitute context, e.g. `torch.multiprocessing.get_context("spawn")`, - whose reducers are needed to share model tensors with workers. + A different context, for example `torch.multiprocessing.get_context`. + Its reducers are necessary to share model tensors with the workers. """ with _spawn_pool(processes, context=context) as pool: mapper = pool.imap_unordered if unordered else pool.imap @@ -169,17 +169,17 @@ def parallel_apply( context=None, ignore_index: bool = False, ): - """Apply `func` to row batches of `obj` in parallel and concatenate the result. + """Apply `func` to batches of rows of `obj`, then join the results. Parameters ---------- obj : pd.DataFrame or pd.Series - Object to split into batches. + The object to divide into batches. group_by : optional - Column to group by before batching, so each batch is within one group. + A column to group by. Each batch then stays in one group. - See :func:`parallel_imap` for the remaining parameters. + See :func:`parallel_imap` for the other parameters. """ return pd.concat( list( diff --git a/tests/unit/peptide/test_precursor_mp.py b/tests/unit/peptide/test_precursor_mp.py index faaf33aa..103f0142 100644 --- a/tests/unit/peptide/test_precursor_mp.py +++ b/tests/unit/peptide/test_precursor_mp.py @@ -1,8 +1,8 @@ -"""Multiprocessing workers must see the same modification registry as the parent. +"""A worker must have the same modification registry as the parent process. -Workers are started with the "spawn" start method and re-import alphabase, so -they only see `modification.tsv` unless the registry is handed to them -explicitly. These tests cover each way the registry can change at runtime. +A worker starts with the "spawn" method and imports alphabase again. Thus it +knows only `modification.tsv`, until the parent sends it the registry. These +tests examine each change that is possible at run time. """ import os @@ -33,14 +33,14 @@ @pytest.fixture def restore_registry(): - """Undo the global registry changes each test makes.""" + """Undo the changes that a test makes to the global registry.""" snapshot = get_modification_state() yield set_modification_state(snapshot) def _worker_registry(_): - """Return the worker process's own view of the registry.""" + """Give the registry of the worker process.""" return get_modification_state() @@ -49,8 +49,8 @@ def _add_custom_mod(): def _filter_modloss(): - # `load_mod_df` applies level 1 at import, keeping 2 modloss values; level 0 - # keeps 867, so a worker that missed this change is detectably different. + # At import, `load_mod_df` uses level 1 and keeps 2 modloss values. Level 0 + # keeps 867. Thus a worker without this change is different. keep_modloss_by_importance(0.0) @@ -59,7 +59,7 @@ def _add_lower_case_aa(): def _load_custom_tsv(): - """Load a TSV that differs from the default one shipped with alphabase.""" + """Load a TSV that is different from the default TSV of alphabase.""" default_tsv = os.path.join(CONST_FILE_FOLDER, "modification.tsv") mod_df = pd.read_table(default_tsv, keep_default_na=False) extra = mod_df.iloc[[0]].copy() @@ -83,16 +83,16 @@ def _load_custom_tsv(): ids=["custom_mods", "modloss_filtering", "lower_case_AA", "custom_tsv"], ) def test_worker_registry_matches_parent(mutate, restore_registry): - # Given a registry changed at runtime + # Given a registry that changed at run time mutate() expected = get_modification_state() - # When a worker process reports its own registry + # When a worker process gives its own registry registries = list( parallel_imap(_worker_registry, [None, None], processes=2, progress=False) ) - # Then it is identical to the parent's + # Then it is the same as the registry of the parent for registry in registries: pd.testing.assert_frame_equal(registry, expected) @@ -111,31 +111,31 @@ def _precursor_df(n_precursors=40, mod=CUSTOM_MOD): def test_isotope_intensity_mp_matches_single_process(restore_registry): - # Given a library whose precursors carry a custom modification + # Given a library with a custom modification on its precursors _add_custom_mod() isotope_cols = [f"i_{i}" for i in range(6)] - # When the isotope intensities are calculated with and without multiprocessing + # When the isotope intensities are calculated with and without workers single = calc_precursor_isotope_intensity(_precursor_df(), max_isotope=6) multi = calc_precursor_isotope_intensity_mp( _precursor_df(), max_isotope=6, mp_batch_size=10, mp_process_num=2 ) - # Then the results agree + # Then the two results are the same pd.testing.assert_frame_equal( single.sort_index()[isotope_cols], multi.sort_index()[isotope_cols] ) def test_caller_supplied_progress_bar_is_used(): - # Given a caller supplying its own progress bar rather than the default + # Given a caller that supplies its own progress bar seen_totals = [] def progress(iterator, total): seen_totals.append(total) return iterator - # When isotope intensities are calculated with multiprocessing + # When the isotope intensities are calculated with workers calc_precursor_isotope_intensity_mp( _precursor_df(40, mod=""), max_isotope=6, @@ -144,17 +144,17 @@ def progress(iterator, total): progress_bar=progress, ) - # Then it is actually driven, rather than being treated as a plain flag + # Then the code uses the bar, and does not use it as a flag assert seen_totals == [4] def test_speclib_isotope_info_runs_with_multiprocessing(): - # Given a library large enough to take the multiprocessing branch + # Given a library that is large enough to use the workers lib = SpecLibBase() lib._precursor_df = _precursor_df(20_000, mod="") - # When isotope info is calculated + # When the isotope info is calculated lib.calc_precursor_isotope_info(mp_process_num=2, mp_batch_size=1000) - # Then it completes; it used to raise TypeError on an unknown keyword + # Then it completes. Before, an unknown keyword caused a TypeError. assert "isotope_apex_offset" in lib.precursor_df.columns From 4d4a71aaeb7c7751688fbe30cc4004507756ce5f Mon Sep 17 00:00:00 2001 From: GeorgWa Date: Mon, 31 Aug 2026 11:07:37 +0200 Subject: [PATCH 4/5] Skip the numba tests when numba is absent The "not full" install has no numba. Three tests call isotope functions that need it. Co-Authored-By: Claude Opus 5 (1M context) --- tests/unit/peptide/test_precursor_mp.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unit/peptide/test_precursor_mp.py b/tests/unit/peptide/test_precursor_mp.py index 103f0142..edc1fa98 100644 --- a/tests/unit/peptide/test_precursor_mp.py +++ b/tests/unit/peptide/test_precursor_mp.py @@ -110,6 +110,7 @@ def _precursor_df(n_precursors=40, mod=CUSTOM_MOD): return update_precursor_mz(df) +@pytest.mark.requires_numba def test_isotope_intensity_mp_matches_single_process(restore_registry): # Given a library with a custom modification on its precursors _add_custom_mod() @@ -127,6 +128,7 @@ def test_isotope_intensity_mp_matches_single_process(restore_registry): ) +@pytest.mark.requires_numba def test_caller_supplied_progress_bar_is_used(): # Given a caller that supplies its own progress bar seen_totals = [] @@ -148,6 +150,7 @@ def progress(iterator, total): assert seen_totals == [4] +@pytest.mark.requires_numba def test_speclib_isotope_info_runs_with_multiprocessing(): # Given a library that is large enough to use the workers lib = SpecLibBase() From d7ca78d52ffa907a38e96e7656cf0cb507aef2c6 Mon Sep 17 00:00:00 2001 From: GeorgWa Date: Tue, 15 Sep 2026 20:22:27 +0200 Subject: [PATCH 5/5] Trim the parallel helpers to what the callers use Six functions become four. `_batchify` returns a list, so `len()` gives the batch count and `_batch_count` no longer duplicates the grouping arithmetic. `parallel_imap` had one production caller, `parallel_apply`, which now uses the pool directly. `context` and `unordered` had no caller at all. Tell the spawn story once, in `set_modification_state`, and lead both registry functions with what they are for. The registry test now also compares a lookup that `update_all_by_MOD_DF` derives, so it fails if a worker installs `MOD_DF` without the rebuild. Verified by removing the rebuild: all four cases fail. Co-Authored-By: Claude Fable 5.1 --- alphabase/constants/modification.py | 16 ++-- alphabase/utils.py | 110 ++++++------------------ tests/unit/peptide/test_precursor_mp.py | 33 +++---- 3 files changed, 48 insertions(+), 111 deletions(-) diff --git a/alphabase/constants/modification.py b/alphabase/constants/modification.py index 7c3cf86a..f38c83dc 100644 --- a/alphabase/constants/modification.py +++ b/alphabase/constants/modification.py @@ -541,28 +541,28 @@ def has_custom_mods(): def get_modification_state() -> pd.DataFrame: - """Copy the modification registry, to send it to a different process. + """Snapshot the modification registry, for a multiprocessing worker or a test. Returns ------- pd.DataFrame - A copy of :data:`MOD_DF`. It contains all changes made at run time, not - only the user-added modifications. + A copy of :data:`MOD_DF`, including all run-time changes. """ return MOD_DF.copy() def set_modification_state(mod_df: pd.DataFrame) -> None: - """Install a registry copy and calculate the derived lookups again. + """Install a registry snapshot and rebuild the derived lookups. - A process that starts with the "spawn" method imports alphabase again. Thus - it knows only the modifications in `modification.tsv`. Use this function to - give the process all changes that were made at run time. + A multiprocessing worker started with "spawn" imports alphabase afresh and + so knows only `modification.tsv`. Installing the parent's snapshot gives it + every run-time change: custom modifications, modloss filtering, lower-case + amino acids, a custom TSV. Parameters ---------- mod_df : pd.DataFrame - A copy from :func:`get_modification_state`. + A snapshot from :func:`get_modification_state`. """ global MOD_DF MOD_DF = mod_df diff --git a/alphabase/utils.py b/alphabase/utils.py index ffa99c12..07fa1a77 100644 --- a/alphabase/utils.py +++ b/alphabase/utils.py @@ -126,84 +126,32 @@ def _sanitize_missing_values(df: pd.DataFrame) -> pd.DataFrame: return df -def _spawn_pool(processes: int, *, context=None): - """Create a worker pool that shares the modification registry. +def _spawn_pool(processes: int): + """Create a pool whose workers start with the parent's modification registry. - A worker starts with the "spawn" method and imports alphabase again. Thus it - loses all changes that were made to the registry at run time. This function - copies the registry one time, then installs it in each worker at start. + See :func:`set_modification_state` for why a spawned worker needs it. """ - ctx = context if context is not None else mp.get_context("spawn") - return ctx.Pool( - processes, - initializer=set_modification_state, - initargs=(get_modification_state(),), + registry = get_modification_state() + return mp.get_context("spawn").Pool( + processes, initializer=set_modification_state, initargs=(registry,) ) -def _batchify(obj, batch_size: int, group_by=None): - """Divide a DataFrame or Series into batches of rows.""" - groups = (group for _, group in obj.groupby(group_by)) if group_by else iter((obj,)) - for group in groups: - for i in range(0, len(group), batch_size): - yield group.iloc[i : i + batch_size] - - -def _batch_count(obj, batch_size: int, group_by=None) -> int: - sizes = obj.groupby(group_by).size().values if group_by else [len(obj)] - return sum((size + batch_size - 1) // batch_size for size in sizes) +def _batchify(obj, batch_size: int, group_by=None) -> list: + """Give the row batches of a DataFrame or Series, each within one group.""" + groups = [group for _, group in obj.groupby(group_by)] if group_by else [obj] + return [ + group.iloc[i : i + batch_size] + for group in groups + for i in range(0, len(group), batch_size) + ] def _with_progress(iterator, total, progress): - """Add a progress bar to `iterator`. - - Set `progress` to True for a tqdm bar. Set it to a callable - `progress(iterator, total)` to supply a different bar. Set it to a false - value for no bar. - """ + """True gives a tqdm bar, a callable `progress(iterator, total)` its own, falsy none.""" if progress is True: return tqdm.tqdm(iterator, total=total) - if callable(progress): - return progress(iterator, total) - return iterator - - -def parallel_imap( - func, - iterable, - *, - processes: int, - total: int = None, - unordered: bool = False, - progress=True, - context=None, -): - """Apply `func` to each item of `iterable` in workers. - - The workers share the modification registry. This function gives each result - when it is ready. Thus the caller keeps only one batch in memory. - - Parameters - ---------- - processes : int - The number of worker processes. - - total : int, optional - The number of items, for the progress bar. - - unordered : bool, optional - Give each result when it is ready, not in the sequence of the input. - - progress : bool or callable, optional - See :func:`_with_progress`. - - context : multiprocessing context, optional - A different context, for example `torch.multiprocessing.get_context`. - Its reducers are necessary to share model tensors with the workers. - """ - with _spawn_pool(processes, context=context) as pool: - mapper = pool.imap_unordered if unordered else pool.imap - yield from _with_progress(mapper(func, iterable), total, progress) + return progress(iterator, total) if callable(progress) else iterator def parallel_apply( @@ -214,31 +162,25 @@ def parallel_apply( batch_size: int, group_by=None, progress=True, - context=None, ignore_index: bool = False, ): - """Apply `func` to batches of rows of `obj`, then join the results. + """Apply `func` to row batches of `obj` in spawned workers, then join the results. Parameters ---------- obj : pd.DataFrame or pd.Series The object to divide into batches. + processes : int + The number of worker processes. + group_by : optional A column to group by. Each batch then stays in one group. - See :func:`parallel_imap` for the other parameters. + progress : bool or callable, optional + See :func:`_with_progress`. """ - return pd.concat( - list( - parallel_imap( - func, - _batchify(obj, batch_size, group_by), - processes=processes, - total=_batch_count(obj, batch_size, group_by), - progress=progress, - context=context, - ) - ), - ignore_index=ignore_index, - ) + batches = _batchify(obj, batch_size, group_by) + with _spawn_pool(processes) as pool: + results = _with_progress(pool.imap(func, batches), len(batches), progress) + return pd.concat(list(results), ignore_index=ignore_index) diff --git a/tests/unit/peptide/test_precursor_mp.py b/tests/unit/peptide/test_precursor_mp.py index edc1fa98..609a989b 100644 --- a/tests/unit/peptide/test_precursor_mp.py +++ b/tests/unit/peptide/test_precursor_mp.py @@ -1,9 +1,4 @@ -"""A worker must have the same modification registry as the parent process. - -A worker starts with the "spawn" method and imports alphabase again. Thus it -knows only `modification.tsv`, until the parent sends it the registry. These -tests examine each change that is possible at run time. -""" +"""A spawned worker must end up with the parent's modification registry.""" import os @@ -12,6 +7,7 @@ from alphabase.constants._const import CONST_FILE_FOLDER from alphabase.constants.modification import ( + MOD_MASS, add_modifications_for_lower_case_AA, add_new_modifications, get_modification_state, @@ -25,7 +21,7 @@ update_precursor_mz, ) from alphabase.spectral_library.base import SpecLibBase -from alphabase.utils import parallel_imap +from alphabase.utils import _spawn_pool CUSTOM_MOD = "TestCustomMod@K" CUSTOM_MOD_COMPOSITION = "H(4)O(2)" @@ -40,8 +36,8 @@ def restore_registry(): def _worker_registry(_): - """Give the registry of the worker process.""" - return get_modification_state() + """Give the worker's registry, and one lookup that is derived from it.""" + return get_modification_state(), dict(MOD_MASS) def _add_custom_mod(): @@ -49,8 +45,7 @@ def _add_custom_mod(): def _filter_modloss(): - # At import, `load_mod_df` uses level 1 and keeps 2 modloss values. Level 0 - # keeps 867. Thus a worker without this change is different. + # level 0 differs from the level `load_mod_df` uses at import keep_modloss_by_importance(0.0) @@ -85,16 +80,16 @@ def _load_custom_tsv(): def test_worker_registry_matches_parent(mutate, restore_registry): # Given a registry that changed at run time mutate() - expected = get_modification_state() + expected_df, expected_mass = get_modification_state(), dict(MOD_MASS) - # When a worker process gives its own registry - registries = list( - parallel_imap(_worker_registry, [None, None], processes=2, progress=False) - ) + # When each worker reports its registry and a lookup derived from it + with _spawn_pool(2) as pool: + reports = pool.map(_worker_registry, [None, None]) - # Then it is the same as the registry of the parent - for registry in registries: - pd.testing.assert_frame_equal(registry, expected) + # Then both match the parent, so the worker rebuilt its derived lookups too + for mod_df, mod_mass in reports: + pd.testing.assert_frame_equal(mod_df, expected_df) + assert mod_mass == expected_mass def _precursor_df(n_precursors=40, mod=CUSTOM_MOD):