Repository navigation
Parallel embedding by default with cores-scaled workers; time-capped ETKDG with one random-coordinates retry - #184
Merged
Merged
Conversation
embed_params now sets EmbedParameters.timeout (EMBED_TIMEOUT_S, guarded for older RDKit) so a geometrically impossible species cannot stall the whole run. embed_with_retry wraps EmbedMultipleConfs and retries once with useRandomCoords when the first attempt embeds nothing; _embed_single now goes through it.
RDKitIsomer.embed_conformer and RDKitSdfIsomer.run now call embed_with_retry instead of EmbedMultipleConfs directly, so the timeout and random-coords retry apply everywhere Auto3D embeds a species, not only in the parallel worker.
Conformer embedding ran serially unless a caller found and set use_parallel_embedding, and the worker count behind it was a fixed 4 -- which left 124 of 128 cores idle on the bench box (P-C3). The switch is now on by default and the count is no longer a class default at all: parallel_workers=None resolves to min(cores, species, 32) in resolve_embedding_workers, called at dispatch, where both the machine and the number of species this run enumerated are known. The serial path is unchanged and still runs for inputs below parallel_embedding_threshold and for use_parallel_embedding=False. parallel_workers joins SENTINEL_FIELDS so None passes the bounds loop while a concrete count below 1 is still refused, and the log line names the resolved worker count rather than the unresolved request.
With parallel embedding on by default, the serial path needed a way back from the command line: the flag maps to use_parallel_embedding the same way --tf32/--no-tf32 maps to allow_tf32, defaulting to None so a value in a -c config file survives a flag the user did not pass. The shipped parameters.yaml now shows the new defaults, with parallel_workers spelled None -- this file's convention for unset, which load_yaml_config converts -- rather than a worker count the code no longer picks. cli.rst gains the option row and usage.rst a paragraph on the three embedding options and the per-species time cap.
Two consequences of the default flip, both found in review. smiles2mols runs in the caller's own process, and the spawn pool re-imports the calling script in every worker -- so a default-on parallel path broke every documented smiles2mols example, none of which has an `if __name__ == "__main__":` guard, as soon as the batch reached parallel_embedding_threshold. It now embeds in parallel only when the caller sets use_parallel_embedding explicitly, read from model_fields_set before replace() rebuilds the config and marks every field as set. main() and `auto3d run` still honor the default; they have always required the guard. BrokenProcessPool keeps its type, since an OOM kill must surface the way it always has, but now carries a note naming the guard and the serial fallback -- previously the only clue was a child-process traceback about freeze_support. Each worker also hands mpi_np threads to EmbedMultipleConfs, so one worker per core put cores x threads runnable threads on the box: a 4x oversubscription on a small machine. resolve_embedding_workers takes threads_per_worker and divides the core count by it before the cap, which the engine supplies from its own np. Also: the constructor defaults on RDKitIsomer and the factory stay off while the product default is on, so both docstrings now say which is which and the two tests that pin them say "constructor_default_off"; the threshold and the species term are documented as counting species after stereoisomer enumeration, not input molecules; the 60 s ETKDG cap and its one retry are documented as applying to every embedding call, serial or parallel, SDF engine included; and check_field_bounds no longer names the CLIConfig class that no longer exists.
Document the parallel embedding default, the per-worker-thread division behind resolve_embedding_workers, and smiles2mols' opt-in exception and broken-pool guidance. Record the ETKDG timeout and random-coords retry, now applied to every embedding call site, and the unbounded-embedding bug it fixes.
Fast suite wall clock dropped from 279 s to 116 s (pytest-reported; 131 s real) by moving the 16 slowest test items (14 functions, two parametrized) to @pytest.mark.slow. CI's `slow` job already runs `-m slow` on every PR, so their coverage is unchanged.
Parallel embedding is on by default on this branch, so an ordinary run now starts a ProcessPoolExecutor -- and it was the one pool in the package that armed neither PR_SET_PDEATHSIG nor the parent-sentinel watchdog. Under spawn each child holds a dup of the call queue's write end as well as its read end, so the parent's death never produces EOF, and a parent-only signal (the orchestrator's own p1.terminate(), an OOM kill, a plain kill) runs none of the `with` block's shutdown sentinels either: an idle worker blocked on get() forever and a busy one kept burning a core on ETKDG. Measured twice before the fix, both workers reparented to init and still sleeping 15 s later. A terminal Ctrl-C was never the gap -- that is a process-group signal the children share. The same initializer sets CoordsAsDouble, because the worker is where the pickling happens. RDKit pickles conformer coordinates as float32 by default, so the result travelling back from a worker lost its low bits: at SDF write precision that flipped one coordinate of one species in a 28-species comparison, which made a documented performance switch change the bytes of the output. Set in the initializer rather than at module import: the pickle properties are process-global, every spawned worker runs the initializer, and importing the module must not reconfigure RDKit for a caller that never starts a pool -- the same stance the module already takes on set_start_method. _embed_single no longer logs its own skip reasons. A pool child has no QueueHandler (the run-log wiring is for the optimizer workers), so those warnings fell through to logging.lastResort: raw on stderr, carrying no level or logger prefix, ungoverned by --quiet, and absent from the run log -- which recorded the parent's "produced no conformers" line, naming the species but not the reason, for a console total of two lines per bad species where the serial path shows one. The reason now rides a SpeciesSkipped back to the parent, which emits exactly one formatted warning and skips the generic empty-result branch for that species. The serial path (RDKitIsomer.embed_conformer) is unchanged: its caller is two frames away, so a sentinel return still reaches the right logger.
…word Two defects the whole-branch review found in the embedding work. embed_with_retry fired its random-coordinates retry whenever the first attempt returned nothing, including when it returned nothing because it had burned the whole 60 s cap -- so the per-species worst case was two caps, 120 s, against the ~67 s one impossible stereoisomer measured serially on the 2026-09-21 bench set, which is the number the cap was written to bound. The retry now runs on the remainder of the budget: attempt 1 is timed, the retry gets ceil(timeout_s - elapsed) seconds, and under a second left it is skipped with a DEBUG line. Both attempts together stay inside timeout_s, to within the second the rounding can add, and a species that fails fast still gets its recovery. smiles2mols inferred its parallel-embedding opt-in from "use_parallel_embedding" in args.model_fields_set, which asks a config object a question it cannot answer: replace() rebuilds the model through its constructor and every YAML-built config names its keys, so both mark every field as explicitly set and both silently took the spawn path -- including the parameters.yaml this repo ships. The opt-in is now a keyword-only argument, smiles2mols(smiles, args, parallel_embedding=True), and args.use_parallel_embedding is ignored by this entry point (one DEBUG line when it is True, so a debug log says why the run was serial). That field keeps governing main() and auto3d run. The regression test for the broken-pool note drives a real unguarded script through a subprocess: the note is a property of module re-import under spawn, and the existing test breaks the pool by killing a worker, which is the OOM shape rather than the commoner one.
Nine text items from the whole-branch review; no behavior change. The cap's cost is now stated where a user would look. usage.rst says that a species near the 60 s cap keeps the conformers it managed to embed, so its conformer count depends on machine load and on mpi_np and is not reproducible run to run -- the one case CONFORMER_RANDOM_SEED does not pin -- and it names EMBED_TIMEOUT_S for anyone who would rather have the reproducibility than the bound. The constant's own comment drops an aggregate extrapolation nothing measured and keeps the measurement: ~67 s per impossible species, serially, on the 2026-09-21 bench set. PARALLEL_EMBED_MAX_WORKERS gains the term nobody had written down, ~85 MiB RSS per warm worker on the 2026-10-02 box, which makes the 32-worker cap roughly a 2.5 GiB ceiling as well. cli.rst stops pointing the CLI user at parallel_workers as though it were a flag: it is a config-file / Python-API option and the row says so. The --parallel-embedding help string, and that row, now say the worker count is per RDKit thread, which is what the division by mpi_np means. "After stereoisomer enumeration" was one step short of what the gate counts -- the set it reads has already been through enantiomer removal -- so parallel_embedding_threshold's docstring is now the canonical statement carrying the full phrase, and the four places that repeat it (cli.rst, usage.rst, parameters.yaml, the changelog) repeat all of it. resolve_embedding_workers says that its one-worker floor can still overshoot a box with fewer cores than threads_per_worker, since closing that would mean overriding the caller's own mpi_np. config.py's prose still described CLIConfig as a live second declaration of the fields, including a claim that a test pins its Literals against ENGINE_CHOICES; that class and that comparison are both gone. Every reference now either names the class as deleted or names Auto3DOptions and both entry points, and the four stale mentions of __post_init__ name the validator that replaced it. The changelog's Fixed bullet for the embedding timeout no longer leans on a Changed bullet further down to define "the cap" and "the retry"; both are named inline. The Changed bullet picks up the remaining-budget retry, the smiles2mols keyword, and the three user-visible results of the pool initializer work -- workers die with their parent, a refused species reports its reason once in the run log, and serial and parallel embedding produce byte-identical output.
The note still told every caller to pass use_parallel_embedding=False, which is inert for smiles2mols now that its opt-in is the parallel_embedding keyword. Name the right switch for each entry point; the guard advice is unchanged.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Workstream 3 of the review remediation plan: embedding throughput. Parallel conformer embedding becomes the default with workers scaled to the machine, and every ETKDG call is time-capped and retried once with random initial coordinates.
Embedding time cap and retry (Task 17)
embed_with_retry, which sets RDKit'sEmbedParameters.timeouttoEMBED_TIMEOUT_S(60 s) and, when the first attempt embeds nothing, retries once with random initial coordinates and the same seed. A stereoisomer that ETKDG cannot embed used to run unbounded; on the 2026-09-21 bench set one such species cost about 67 s, serially, before failing.Parallel embedding by default (Task 18)
use_parallel_embeddingdefaults to on (spec decision D1).parallel_workers=Nonenow means "resolve at dispatch": min(cores, species, 32), divided by the per-worker RDKit thread count so workers times threads never exceed the cores. An explicit integer is honored as given. The serial path remains for--no-parallel-embedding/use_parallel_embedding=Falseand for inputs belowparallel_embedding_threshold(unchanged, 10).smiles2molskeeps its documented single-process contract: it embeds in parallel only when called with the new keywordparallel_embedding=True. Theuse_parallel_embeddingoption governsmain()only, because parallel embedding spawns worker processes that re-import the calling script, and the shippedsmiles2molsexamples are plain scripts without a main guard. An inferred opt-in was tried first and found to be defeated byreplace()and YAML-built configs, so the keyword replaced it. A broken worker pool now carries a note saying exactly that.auto3d rungains--parallel-embedding/--no-parallel-embedding, mapped like--tf32. YAML acceptsparallel_workers: None;parameters.yamlshows the new defaults. The usage guide and the CLI reference document the option and the time cap.Test tiers (rider Task 37)
slowmarker. CI already runs the two tiers separately on every PR, so their coverage is unchanged; the local default run is about four times shorter.Verification
ruff checkandruff format --checkclean.smiles2molsopt-in capture moved afterreplace()).main()) and runs in CI's slow job on every push.smiles2molswith twelve SMILES was run against the pre-fix and post-fix commits: a broken worker pool before, a serial run after.Notes for reviewers
main()and the CLI embed in parallel by default for inputs of ten or more species after enumeration.smiles2molsdoes not, unless asked.parallel_workers=Noneis resolved per run as min(cores divided by RDKit threads per worker, species, 32). An explicit integer is honored as given.os.cpu_count()ignores cgroup CPU quotas and affinity; a container with a small quota on a large host resolves more workers than it has CPU. Deferred to the performance-levers workstream.[Unreleased]for 3.2.0.