GA scoring lost its multiprocessing parallelism between v2 and latest

[lpatiny]

GA scoring lost its multiprocessing parallelism between v2 and latest

Runs on elucidation.cheminfo.org currently take ~14 min for a single elucidation
(gens_ga=10, offspring_ga=1024, pop_ga=512). Comparing the deployed image against
the older adrianmirza/elucidation:v2 tag, reward_function_ga used to be parallelised
across 8 processes and no longer is. Reporting it here because the change looks
unintentional — the GA parameters are byte-identical across the two builds, so nothing
else in the search accounts for a difference in cost.

I want to be upfront that I have not profiled the run, so I am not claiming a
measured factor. What follows is a code diff plus one runtime measurement, not an
attribution.

The change

adrianmirza/elucidation:v2 (image created 2025-08-18T18:24:33Z), /app/main.py:

from multiprocessing import Pool
...
def reward_function_ga(individuals, ga_models, target_1D_embs, atom_counts_orig):
    """
    Calculates the reward for a list of individuals entirely on the CPU,
    using multi-core parallelization for ALL heavy computations.
    """
    num_workers = 8
    chunks = np.array_split(individuals, num_workers)
    chunks = [chunk.tolist() for chunk in chunks if len(chunk) > 0]

    with Pool(processes=num_workers) as pool:
        embedding_worker = partial(gpu_encode_smiles_variable, ga_models=ga_models)
        list_of_emb_dicts = pool.map(embedding_worker, chunks)
        ...
        mf_loss_list = pool.map(mf_loss_calculator, individuals)
        is_radical_or_charged_list = pool.map(
            smiles_is_radical_or_is_charged_or_has_wrong_valence, individuals
        )

Current main.py (both adrianmirza/elucidation:latest and adrianmirza/celery:latest,
image created 2026-07-19T09:19:18Z — and main at e5355b6) does all three serially:

    cand_smiles_embs = gpu_encode_smiles_variable(individuals, ga_models)

    mf_loss = np.array(
        [calculate_mf_penalty(smi, atom_counts_orig) for smi in individuals]
    )
    ...
    is_radical_or_charged = np.array(
        [smiles_is_radical_or_is_charged_or_has_wrong_valence(smi) for smi in individuals]
    )

multiprocessing is no longer imported anywhere in main.py.

Verified by extracting the layers of all three tags from Docker Hub. adrianmirza/celery:latest
and adrianmirza/elucidation:latest share byte-identical layer digests and differ only in
Cmd, so this is the code the worker runs.

GA parameters are unchanged

param	v2	latest
seed	42	42
init_pop_ga	512	512
frac_graph_ga_mutate	0.3	0.3
gens_ga → generations	10	10
offspring_ga → offspring_size	1024	1024
pop_ga → population_size	512	512

So the search does the same amount of work in both versions.

Why this may matter more than the CPU graph suggests

Measured on the deployed worker, a single elucidation peaks around 460 % CPU. That is
torch's intra-op threading on the encoder matmuls, and it makes the box look busy. But
calculate_mf_penalty and smiles_is_radical_or_is_charged_or_has_wrong_valence are
pure-Python RDKit calls under the GIL — they thread not at all, and they now run strictly
serially over 1024 candidates per generation, ten generations per run.

Caveats — reasons not to just revert

I do not think a verbatim revert is obviously correct, for two reasons:

1. pool.map(partial(gpu_encode_smiles_variable, ga_models=ga_models), chunks) pickles
   the ga_models dict through the task queue to each of the 8 workers, and the with Pool(...) block builds and tears down the pool on every reward call. For a
   MolBind model dict that is not free, and it happens once per generation.
2. 8 processes each running a torch model that is itself intra-op threaded will
   oversubscribe the cores unless torch.set_num_threads(1) is set in the workers.
   Neither version sets any thread count (grep for set_num_threads, OMP_NUM_THREADS,
   MKL_NUM_THREADS finds nothing in either image).

So the shape I would expect to be fastest is: keep the encoder call serial and let torch
thread it, and parallelise only the two RDKit passes with a long-lived pool. But that is a
guess — you have the profiling context and I do not.

Minor, unrelated to the above: eval_batch rescans the whole cache each batch

While reading gafuncs.py for #1, CachedFunction.eval_batch rebuilds the entire
accumulated cache on every batch:

all_items = [
    {
        "smiles": k,
        "score": v,
        "molecular_formula": smiles_to_molecular_formula(k),
        "retrieved": k in self.initial_population,
    }
    for k, v in self.cache.items()
]

smiles_to_molecular_formula is an RDKit parse and runs for every cached SMILES on every
batch; k in self.initial_population is a linear scan of a 512-element list. Over a run
the cache reaches ~10 k entries, so this does roughly 60 k parses where 10 k would do.

To be clear about the size of this one: at ~100 µs per parse that is single-digit seconds,
not minutes. It is not a plausible explanation for the 14 min — it is just waste that is
easy to remove by computing the formula once when a SMILES enters the cache and making
initial_population a set.

What would settle it

If someone upstream can run one elucidation with cProfile (or just wall-clock timers
around gpu_encode_smiles_variable, the two RDKit passes, and eval_batch), that would
confirm or kill the hypothesis quickly. Happy to test a candidate image against our
deployment and report timings if that is useful.

Related: #1 (caching).