Parallelize sequence-properties property computation across CPU cores#16
Parallelize sequence-properties property computation across CPU cores#16PaulNewling wants to merge 10 commits into
Conversation
…eption propagation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request parallelizes sequence property computations across CPU cores using a ProcessPoolExecutor in Python, while ensuring byte-identical outputs. It introduces a golden-master regression test suite, unit tests for parallel invariance, and updates the Tengo workflow to request 8 CPUs and 8GiB of memory. The review feedback highlights two important improvements for the parallel execution: using the spawn multiprocessing context instead of fork to prevent deadlocks when using polars in child processes (along with better executor lifecycle management on exceptions), and using os.sched_getaffinity(0) to resolve the worker count in containerized environments to avoid CPU thrashing.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| with ProcessPoolExecutor(max_workers=workers) as ex: | ||
| try: | ||
| return list(ex.map(fn, items, chunksize=chunksize)) | ||
| except BrokenProcessPool as exc: | ||
| raise RuntimeError( | ||
| f"compute worker pool broke while processing {len(items)} items " | ||
| f"with {workers} workers — a worker process was killed (most likely " | ||
| f"out of memory). Reduce the input size or raise the step's memory." | ||
| ) from exc |
There was a problem hiding this comment.
Using the default fork start method on Linux with multithreaded libraries like polars can lead to deadlocks in child processes. Additionally, if a worker process raises an exception, the default behavior of ProcessPoolExecutor is to wait for all pending tasks to complete before propagating the exception, which can cause significant hangs on large datasets.
We can resolve both issues by:
- Explicitly using the
spawnmultiprocessing context to safely start worker processes. - Managing the executor lifecycle to cancel pending futures and shut down immediately when an exception is encountered.
| with ProcessPoolExecutor(max_workers=workers) as ex: | |
| try: | |
| return list(ex.map(fn, items, chunksize=chunksize)) | |
| except BrokenProcessPool as exc: | |
| raise RuntimeError( | |
| f"compute worker pool broke while processing {len(items)} items " | |
| f"with {workers} workers — a worker process was killed (most likely " | |
| f"out of memory). Reduce the input size or raise the step's memory." | |
| ) from exc | |
| import multiprocessing | |
| ctx = multiprocessing.get_context("spawn") | |
| ex = ProcessPoolExecutor(max_workers=workers, mp_context=ctx) | |
| try: | |
| return list(ex.map(fn, items, chunksize=chunksize)) | |
| except BrokenProcessPool as exc: | |
| ex.shutdown(wait=False, cancel_futures=True) | |
| raise RuntimeError( | |
| f"compute worker pool broke while processing {len(items)} items " | |
| f"with {workers} workers — a worker process was killed (most likely " | |
| f"out of memory). Reduce the input size or raise the step's memory." | |
| ) from exc | |
| except Exception: | |
| ex.shutdown(wait=False, cancel_futures=True) | |
| raise | |
| finally: | |
| ex.shutdown(wait=True) |
| def resolve_workers(workers: int | None) -> int: | ||
| """How many worker processes to use. Explicit arg wins (used by tests and | ||
| by main.py once it reads the platform's CPU allocation). Falls back to the | ||
| PL_COMPUTE_WORKERS env var, then os.cpu_count(). The RESULT never depends on | ||
| this number — only the wall-clock does — so an over- or under-estimate is a | ||
| speed concern, never a correctness one. | ||
| """ | ||
| if workers is not None: | ||
| return max(1, int(workers)) | ||
| env = os.environ.get("PL_COMPUTE_WORKERS") | ||
| if env and env.isdigit() and int(env) > 0: | ||
| return int(env) | ||
| return max(1, os.cpu_count() or 1) |
There was a problem hiding this comment.
In containerized environments (such as Docker, Kubernetes, or cloud runners), os.cpu_count() returns the physical host's CPU count rather than the container's allocated CPU limit/quota. Spawning too many worker processes can lead to severe CPU thrashing and out-of-memory (OOM) issues.
Using os.sched_getaffinity(0) (where available) correctly respects the container's CPU limits.
| def resolve_workers(workers: int | None) -> int: | |
| """How many worker processes to use. Explicit arg wins (used by tests and | |
| by main.py once it reads the platform's CPU allocation). Falls back to the | |
| PL_COMPUTE_WORKERS env var, then os.cpu_count(). The RESULT never depends on | |
| this number — only the wall-clock does — so an over- or under-estimate is a | |
| speed concern, never a correctness one. | |
| """ | |
| if workers is not None: | |
| return max(1, int(workers)) | |
| env = os.environ.get("PL_COMPUTE_WORKERS") | |
| if env and env.isdigit() and int(env) > 0: | |
| return int(env) | |
| return max(1, os.cpu_count() or 1) | |
| def resolve_workers(workers: int | None) -> int: | |
| """How many worker processes to use. Explicit arg wins (used by tests and | |
| by main.py once it reads the platform's CPU allocation). Falls back to the | |
| PL_COMPUTE_WORKERS env var, then the container-aware CPU count (via affinity), | |
| and finally os.cpu_count(). | |
| """ | |
| if workers is not None: | |
| return max(1, int(workers)) | |
| env = os.environ.get("PL_COMPUTE_WORKERS") | |
| if env and env.isdigit() and int(env) > 0: | |
| return int(env) | |
| try: | |
| return max(1, len(os.sched_getaffinity(0))) | |
| except AttributeError: | |
| return max(1, os.cpu_count() or 1) |
1 participant
Summary
Make the
sequence-propertiesblock compute large datasets faster on the first run by parallelizing the Python property kernel across CPU cores — without changing any output byte, so the platform's content-addressed cache still dedups. Today the compute loops row-by-row on a single core; a big bulk-MiXCR / peptide input has to grind through serially.software/src/pipeline.py— per-row compute now runs through aProcessPoolExecutor(_pmap), with results reassembled in input order and a sequential fallback below 2000 rows. The existing sort + charge/pI quantization boundary is untouched, so output stays byte-identical regardless of worker count. A worker that dies uncleanly (e.g. OOM-killed) now surfaces as a clearRuntimeErrorwith item/worker counts instead of a crypticBrokenProcessPool.software/src/main.py— worker count is resolved from the environment (PL_COMPUTE_WORKERS→os.cpu_count()). No new CLI argument, so the exec step's command line — and therefore its content-addressed id — is unchanged.workflow/src/main.tpl.tengo— the Python compute step now requestscpu(8)/mem("8GiB"). These are scheduling meta-inputs (verified againstworkflow-tengorender.lib.tengo: meta-inputs are excluded from dedup), so the bump does not invalidate cached results and the CID stays stable run-to-run.Determinism rule throughout: parallelism is across rows (each row is an independent, pure computation; reassembled in input order; final sort restores byte order) — never inside a float reduction. Worker count is a speed knob only.
Safety net (built first, Phase A)
software/tests/regression/test_golden_master.py) — byte-exact (sha256) full-output snapshot for every mode / coverage tier / edge case (peptide incl. NA cases, antibody full+Fv with per-clone missing-region & empty-CDR3, CDR3-only, partial, TCR). Generated once from the pre-refactor code and never regenerated — direct proof the parallelization changed no bytes.test_scale_invariants.py) — row/key preservation + byte-stability at 5000 rows.test_parallel_invariance.py) —workers=1≡workers=4byte-identical on both the sequential and the real pool path (n=3000), plus exception-propagation.Streaming (memory) — deferred by decision
Measured peak memory against the new 8 GiB ceiling:
8 GiB comfortably covers datasets up to ~1M clonotypes; it fails only at multi-million-row extremes. A streaming I/O path (peak memory O(batch) instead of O(dataset)) is fully specified in the plan but intentionally not implemented — build it only if inputs reach several million rows. Full plan + decision:
docs/superpowers/plans/2026-06-10-sequence-properties-parallel-compute.md.Verification
uv run pytest(fromsoftware/): 221 passed (incl. slow pool-path and 5000-row byte-stability cases).pnpm run build:dev: green (7/7 turbo tasks).minorfor.software,.workflow, and the root block package.Follow-ups (non-blocking)
PL_COMPUTE_WORKERSis set but unparseable (currently falls back silently toos.cpu_count()).