Parallelize sequence-properties property computation across CPU cores

.changeset/seqprops-parallel-compute.md

@@ -0,0 +1,7 @@
+---
+'@platforma-open/milaboratories.sequence-properties.software': minor
+'@platforma-open/milaboratories.sequence-properties.workflow': minor
+'@platforma-open/milaboratories.sequence-properties': minor
+---
+
+Parallelize the property computation across CPU cores and raise the compute step's CPU/memory request, so large datasets compute on the first run without the previous single-core bottleneck. Output is byte-identical to the prior version (worker count never affects results).

docs/superpowers/plans/2026-06-10-sequence-properties-parallel-compute.md.tasks.json

@@ -0,0 +1,46 @@
+{
+  "planPath": "docs/superpowers/plans/2026-06-10-sequence-properties-parallel-compute.md",
+  "tasks": [
+    {
+      "id": 1,
+      "subject": "Task 1: Freeze full-output golden master",
+      "status": "completed",
+      "description": "**Goal:** Committed byte-exact snapshot of full CLI output for every mode + coverage tier + edge case, generated from current pre-refactor code, asserted via sha256. Safety net for Phases B/C.\n\n**Files:** software/tests/regression/{__init__.py, golden_cases.py, regen_golden.py, test_golden_master.py}; software/tests/data/golden/<case>/{properties.tsv,aa_fraction.tsv,stats.json}\n\n**Verify:** uv run pytest tests/regression/test_golden_master.py -v\n\n```json:metadata\n{\"files\": [\"software/tests/regression/__init__.py\",\"software/tests/regression/golden_cases.py\",\"software/tests/regression/regen_golden.py\",\"software/tests/regression/test_golden_master.py\",\"software/tests/data/golden/\"], \"verifyCommand\": \"uv run pytest tests/regression/test_golden_master.py -v\", \"acceptanceCriteria\": [\"all golden cases byte-match via sha256\",\"cases cover all modes/tiers/edge-cases\",\"goldens generated from pre-refactor code and committed\"]}\n```"
+    },
+    {
+      "id": 2,
+      "subject": "Task 2: Structural invariants at scale",
+      "status": "completed",
+      "description": "**Goal:** Pin invariants the parallel/streaming refactor must preserve — one output row per input entity, same key set, AA-fraction 20 rows/peptide, two runs byte-identical — at 5000 rows.\n\n**Files:** software/tests/regression/test_scale_invariants.py\n\n**Verify:** uv run pytest tests/regression/test_scale_invariants.py -v\n\n```json:metadata\n{\"files\": [\"software/tests/regression/test_scale_invariants.py\"], \"verifyCommand\": \"uv run pytest tests/regression/test_scale_invariants.py -v\", \"acceptanceCriteria\": [\"row count + key set preserved\",\"aa_fraction 20 rows/peptide\",\"byte-stable across two runs\",\"seeded RNG, marked slow\"]}\n```"
+    },
+    {
+      "id": 3,
+      "subject": "Task 3: Parallelize per-row compute across a process pool",
+      "status": "completed",
+      "blockedBy": [1, 2],
+      "description": "**Goal:** Map existing pure per-row compute over a ProcessPoolExecutor, assemble in input order, keep sort + CID-quantization boundary, sequential fallback for small inputs/workers=1 — output byte-identical regardless of worker count.\n\n**Files:** Modify software/src/pipeline.py (resolve_workers, _pmap, _peptide_worker, _antibody_worker; thread workers through run/run_peptide/run_antibody_tcr — lines 71-91, 159-221, 414-439); Test software/tests/unit/test_parallel_invariance.py\n\n**Verify:** uv run pytest tests/unit/test_parallel_invariance.py tests/regression/ -v\n\n```json:metadata\n{\"files\": [\"software/src/pipeline.py\",\"software/tests/unit/test_parallel_invariance.py\"], \"verifyCommand\": \"uv run pytest tests/unit/test_parallel_invariance.py tests/regression/ -v\", \"acceptanceCriteria\": [\"workers=1 == workers=4 byte-identical\",\"golden master unchanged\",\"worker fns picklable/spawn-safe\",\"results in input order\"]}\n```"
+    },
+    {
+      "id": 4,
+      "subject": "Task 4: Resolve worker count from environment in CLI",
+      "status": "completed",
+      "blockedBy": [3],
+      "description": "**Goal:** CLI uses allocated cores at runtime without adding to argv — exec node command line stays constant so its CID does not change.\n\n**Files:** Modify software/src/main.py (import line 19; lines 44-46); Test software/tests/integration/test_cli.py\n\n**Verify:** uv run pytest tests/integration/test_cli.py -v\n\n```json:metadata\n{\"files\": [\"software/src/main.py\",\"software/tests/integration/test_cli.py\"], \"verifyCommand\": \"uv run pytest tests/integration/test_cli.py -v\", \"acceptanceCriteria\": [\"workers resolved from env, not argv\",\"argv unchanged (CID stable)\",\"PL_COMPUTE_WORKERS does not change output bytes\"]}\n```"
+    },
+    {
+      "id": 5,
+      "subject": "Task 5: Request more cores + memory on workflow exec step",
+      "status": "completed",
+      "blockedBy": [4],
+      "description": "**Goal:** Schedule the Python step with multiple cores + more memory; keep exec node identity stable.\n\n**Files:** Modify workflow/src/main.tpl.tengo (pyRun, lines 344-360): mem 4GiB->8GiB, cpu(1)->cpu(8); Create .changeset/seqprops-parallel-compute.md (.software patch + .workflow patch); Verify against core/platforma/sdk/workflow-tengo/src/exec/index.lib.tengo\n\n**Verify:** pnpm run build:dev (block root)\n\n```json:metadata\n{\"files\": [\"workflow/src/main.tpl.tengo\",\".changeset/seqprops-parallel-compute.md\"], \"verifyCommand\": \"pnpm run build:dev\", \"acceptanceCriteria\": [\"pyRun .cpu(8) + .mem(8GiB)\",\"no machine-dependent argv added\",\"cpu/mem-vs-CID verified against SDK\",\"build:dev succeeds\",\"changeset for workflow+software\"]}\n```"
+    },
+    {
+      "id": 6,
+      "subject": "Task 6 (OPTIONAL): Streaming compute path to bound memory",
+      "status": "pending",
+      "blockedBy": [5],
+      "description": "**OPTIONAL — decide at end.** Do not implement until Phase B merged + measured; the 8GiB bump may make this unnecessary.\n\n**Goal:** Process input in sorted batches, write output incrementally — peak memory O(batch). Output byte-identical (golden master is guard).\n\n**Files:** Modify software/src/io_layer.py, software/src/pipeline.py; Test software/tests/regression/test_streaming_equivalence.py\n\n**Verify (only if implemented):** uv run pytest tests/regression/test_streaming_equivalence.py tests/regression/test_golden_master.py -v\n\n```json:metadata\n{\"files\": [\"software/src/io_layer.py\",\"software/src/pipeline.py\",\"software/tests/regression/test_streaming_equivalence.py\"], \"verifyCommand\": \"uv run pytest tests/regression/test_streaming_equivalence.py tests/regression/test_golden_master.py -v\", \"acceptanceCriteria\": [\"decision gate recorded with RSS measurement\",\"streaming output byte-identical to in-memory\",\"golden master still passes\",\"env-gated, argv unchanged\"], \"optional\": true}\n```"
+    }
+  ],
+  "lastUpdated": "2026-06-10T00:00:05Z"
+}

software/src/main.py

@@ -16,7 +16,7 @@
 from pathlib import Path
 
 from io_layer import read_input_tsv, read_plan, write_output_tsv
-from pipeline import run
+from pipeline import resolve_workers, run
 
 
 def _configure_logging() -> None:
@@ -43,7 +43,9 @@ def main(argv: list[str] | None = None) -> int:
 
     reads = read_input_tsv(args.input)
     plan = read_plan(args.plan)
-    outputs = run(reads, plan)
+    workers = resolve_workers(None)
+    logging.getLogger(__name__).info("Using %d compute worker(s)", workers)
+    outputs = run(reads, plan, workers=workers)
     write_output_tsv(outputs["properties"], args.output, sort_keys=["entity_key"])
     write_output_tsv(outputs["aa_fraction"], args.aa_fraction, sort_keys=["entity_key", "aminoAcid"])
     Path(args.stats).write_text(json.dumps(outputs["stats"], sort_keys=True, separators=(",", ":")))

software/src/pipeline.py

@@ -15,7 +15,11 @@
 
 from __future__ import annotations
 
+import functools
 import logging
+import os
+from concurrent.futures import ProcessPoolExecutor
+from concurrent.futures.process import BrokenProcessPool
 from typing import Any
 
 import polars as pl
@@ -34,6 +38,54 @@
 PH = 7.0  # All charge values computed at pH 7 (spec default).
 
 
+# Parallelism. Below this row count the pool's process startup + pickle cost
+# outweighs the benefit, so we stay in-process. The threshold also keeps the
+# whole existing unit-test suite on the fast sequential path.
+# The threshold is compared against ROW count regardless of per-row cost — an
+# antibody/TCR clone does several times the work of a peptide (CDR3 × chains +
+# full-chain × chains + Fv), so 2000 is intentionally more conservative for
+# antibody mode than peptide mode.
+_PARALLEL_MIN_ROWS = 2000
+
+# Rows dispatched per pool task. Bounds per-task pickle/IPC overhead without
+# starving workers; never overridden today but kept as a tuning knob.
+_PARALLEL_CHUNKSIZE = 256
+
+
+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 _pmap(fn, items: list, workers: int, chunksize: int = _PARALLEL_CHUNKSIZE) -> list:
+    """Map fn over items, preserving input order. Sequential below the
+    threshold or when workers<=1; otherwise a process pool. ProcessPoolExecutor
+    .map() preserves input order, so results align with items by index — the
+    property the byte-stable output depends on.
+    """
+    if workers <= 1 or len(items) < _PARALLEL_MIN_ROWS:
+        return [fn(x) for x in items]
+    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
+
+
 # ---------------------------------------------------------------------------
 # CID quantization
 # ---------------------------------------------------------------------------
@@ -68,26 +120,22 @@ def _quantize_for_cid(df: pl.DataFrame) -> pl.DataFrame:
 # ---------------------------------------------------------------------------
 
 
-def run(reads: pl.DataFrame, plan: dict[str, Any]) -> dict[str, Any]:
-    """Dispatch by mode. Returns a dict with three entries:
-
-    - `properties` (DataFrame): one row per entity, columns per the plan.
-    - `aa_fraction` (DataFrame): long-format (entity_key, aminoAcid, value).
-      Empty body when mode is not peptide.
-    - `stats` (dict): dataset-level stats consumed by the workflow info layer
-      (e.g. R11c VHH detection — median CDR-H3 length per chain;
-      R9 — peptide count below the Instability Index length floor).
+def run(reads: pl.DataFrame, plan: dict[str, Any], workers: int | None = None) -> dict[str, Any]:
+    """Dispatch by mode. `workers` controls parallelism only — output is
+    identical for any value (see test_parallel_invariance). Returns a dict with
+    `properties`, `aa_fraction`, and `stats` entries (unchanged contract).
     """
+    n_workers = resolve_workers(workers)
     mode = plan["mode"]
     if mode == "peptide":
         log.info("Running peptide mode (%d entities)", reads.height)
-        return run_peptide(reads)
+        return run_peptide(reads, n_workers)
     log.info(
         "Running antibody/TCR mode (receptor=%s, %d clones)",
         plan.get("receptor", "IG"),
         reads.height,
     )
-    return run_antibody_tcr(reads, plan)
+    return run_antibody_tcr(reads, plan, n_workers)
 
 
 # ---------------------------------------------------------------------------
@@ -156,45 +204,48 @@ def _compute_peptide_row_from_ctx(ctx: SequenceContext) -> dict[str, float | Non
     }
 
 
-def run_peptide(reads: pl.DataFrame) -> dict[str, Any]:
-    """Compute peptide-mode outputs.
+def _peptide_worker(seq: str) -> tuple[dict[str, float | None], list[float | None] | None]:
+    """Picklable per-peptide unit: (properties row, 20 AA fractions in
+    STANDARD_AAS order | None). One SequenceContext per sequence, shared across
+    all 11 reads — same sharing the old inline loop relied on.
+    """
+    ctx = SequenceContext.from_seq(seq)
+    if ctx is None:
+        return (dict(_NA_PEPTIDE_ROW), None)
+    props = _compute_peptide_row_from_ctx(ctx)
+    fr = ctx.aa_fractions()
+    return (props, [fr[aa] for aa in STANDARD_AAS])
+
 
-    Builds one `SequenceContext` per sequence and reuses it for both the
-    properties row and the AA-fraction rows — so each sequence is `_prepare`d
-    once and BioPython objects are constructed once across all 11 reads.
-    Accumulates columnar arrays directly into a dict-of-lists (one allocation
-    per column, vs. one dict per row) and constructs the DataFrame from those.
+def run_peptide(reads: pl.DataFrame, workers: int = 1) -> dict[str, Any]:
+    """Compute peptide-mode outputs. Per-sequence work runs through `_pmap`
+    (sequential or pooled); results are reassembled in input order so the
+    serialized output is byte-identical regardless of worker count.
     """
     keys = reads["entity_key"].to_list()
     seqs = reads["sequence"].to_list()
     n = len(seqs)
 
     log.info("Computing peptide properties + AA fractions (%d sequences)", n)
-    prop_cols: dict[str, list[Any]] = {"entity_key": [], **{c: [] for c in PEPTIDE_PROPERTY_COLUMNS}}
+    results = _pmap(_peptide_worker, seqs, workers)
+
+    prop_cols: dict[str, list[Any]] = {"entity_key": list(keys), **{c: [] for c in PEPTIDE_PROPERTY_COLUMNS}}
     aa_entity: list[str] = []
     aa_amino: list[str] = []
     aa_value: list[float | None] = []
-    for k, s in zip(keys, seqs):
-        prop_cols["entity_key"].append(k)
-        ctx = SequenceContext.from_seq(s)
-        if ctx is None:
-            for c in PEPTIDE_PROPERTY_COLUMNS:
-                prop_cols[c].append(None)
-            # Emit one row per std AA with NA value, so the 2-axis PColumn
-            # keeps a uniform shape across entities.
+    for k, (props, fractions) in zip(keys, results):
+        for c in PEPTIDE_PROPERTY_COLUMNS:
+            prop_cols[c].append(props[c])
+        if fractions is None:
             for aa in STANDARD_AAS:
                 aa_entity.append(k)
                 aa_amino.append(aa)
                 aa_value.append(None)
         else:
-            row = _compute_peptide_row_from_ctx(ctx)
-            for c in PEPTIDE_PROPERTY_COLUMNS:
-                prop_cols[c].append(row[c])
-            fractions = ctx.aa_fractions()
-            for aa in STANDARD_AAS:
+            for aa, val in zip(STANDARD_AAS, fractions):
                 aa_entity.append(k)
                 aa_amino.append(aa)
-                aa_value.append(fractions[aa])
+                aa_value.append(val)
     properties = pl.DataFrame(
         prop_cols,
         schema={"entity_key": pl.Utf8, **{c: pl.Float64 for c in PEPTIDE_PROPERTY_COLUMNS}},
@@ -204,10 +255,6 @@ def run_peptide(reads: pl.DataFrame) -> dict[str, Any]:
         schema={"entity_key": pl.Utf8, "aminoAcid": pl.Utf8, "value": pl.Float64},
     )
 
-    # R9 — flag whether any *real* peptide falls below the Instability Index
-    # floor. `if s` filters None / "" so the banner does not fire on empty
-    # cells (which are no peptide, not a short peptide); `0 < effective_length`
-    # filters sequences that clean to empty (e.g. all-non-standard residues).
     has_below_floor = any(0 < effective_length(s) < INSTABILITY_MIN_LENGTH for s in seqs if s)
     stats = {
         "medianCdr3Length": {},
@@ -411,7 +458,15 @@ def _median_cdr3_length_by_chain(reads: pl.DataFrame, chains: list[str]) -> dict
     return out
 
 
-def run_antibody_tcr(reads: pl.DataFrame, plan: dict[str, Any]) -> dict[str, Any]:
+def _antibody_worker(record: dict, plan: dict[str, Any]) -> dict[str, Any]:
+    """Picklable per-clone unit. `plan` is bound per-call via functools.partial;
+    it is a plain dict and pickles cleanly. Delegates to the existing
+    _compute_row_for so the computation is unchanged.
+    """
+    return _compute_row_for(record, plan)
+
+
+def run_antibody_tcr(reads: pl.DataFrame, plan: dict[str, Any], workers: int = 1) -> dict[str, Any]:
     chains = plan.get("chains", [])
     full_chains = plan.get("fullChains", [])
     n = reads.height
@@ -421,10 +476,14 @@ def run_antibody_tcr(reads: pl.DataFrame, plan: dict[str, Any]) -> dict[str, Any
         log.info("Reconstructing full chains %s and computing full-chain properties", list(full_chains))
     if plan.get("hasFv"):
         log.info("Computing Fv properties (paired VH+VL)")
+
     out_cols = _planned_output_columns(plan)
+    records = list(reads.iter_rows(named=True))
+    worker = functools.partial(_antibody_worker, plan=plan)
+    rows = _pmap(worker, records, workers)
+
     columns: dict[str, list[Any]] = {"entity_key": [], **{c: [] for c in out_cols}}
-    for record in reads.iter_rows(named=True):
-        row = _compute_row_for(record, plan)
+    for row in rows:
         columns["entity_key"].append(row["entity_key"])
         for c in out_cols:
             columns[c].append(row[c])

software/tests/data/golden/antibody_cdr3_only/aa_fraction.tsv

@@ -0,0 +1 @@
+entity_key	aminoAcid	value

software/tests/data/golden/antibody_cdr3_only/input.tsv

@@ -0,0 +1,3 @@
+entity_key	A_CDR3	B_CDR3
+c1	CARDYW	CQQYNS
+c2	CARGFW	CQHFSS

software/tests/data/golden/antibody_cdr3_only/plan.json

@@ -0,0 +1 @@
+{"mode": "antibody_tcr_legacy_sc", "receptor": "IG", "chains": ["A", "B"], "fullChains": [], "hasFv": false}

software/tests/data/golden/antibody_cdr3_only/properties.tsv

@@ -0,0 +1,3 @@
+entity_key	charge_A_CDR3	chargeShift_A_CDR3	gravy_A_CDR3	charge_B_CDR3	chargeShift_B_CDR3	gravy_B_CDR3
+c1	-0.111	-0.245	-0.9833333333333334	-0.112	-0.236	-1.6833333333333333
+c2	0.895	-0.22	0.21666666666666665	0.111	-0.854	-0.5000000000000001

software/tests/data/golden/antibody_cdr3_only/stats.json

@@ -0,0 +1 @@
+{"medianCdr3Length":{"A":6.0,"B":6.0}}

software/tests/data/golden/antibody_full/aa_fraction.tsv

@@ -0,0 +1 @@
+entity_key	aminoAcid	value

software/tests/data/golden/antibody_full/input.tsv

@@ -0,0 +1,4 @@
+entity_key	A_FR1	A_CDR1	A_FR2	A_CDR2	A_FR3	A_CDR3	A_FR4	B_FR1	B_CDR1	B_FR2	B_CDR2	B_FR3	B_CDR3	B_FR4
+c1	EVQLVES	GFTFSSY	AMSWVRQ	ISGSGGS	TYYAESVKGRFTI	CARDYW	WGQGTLV	DIQMTQS	QSISSY	LNWYQQK	AASSLQS	GVPSRFSGSG	CQQYNS	FGQGTKV
+c2	EVQLVES	GFTFSSY	AMSWVRQ	ISGSGGS		CARGFW	WGQGTLV	DIQMTQS	QSISSY	LNWYQQK	AASSLQS	GVPSRFSGSG	CQHFSS	FGQGTKV
+c3	EVQLVES	GFTFSSY	AMSWVRQ	ISGSGGS	TYYAESVKGRFTI		WGQGTLV	DIQMTQS	QSISSY	LNWYQQK	AASSLQS	GVPSRFSGSG	CQQYNS	FGQGTKV

software/tests/data/golden/antibody_full/plan.json

@@ -0,0 +1 @@
+{"mode": "antibody_tcr_legacy_bulk", "receptor": "IG", "chains": ["A", "B"], "fullChains": ["A", "B"], "hasFv": true}

software/tests/data/golden/antibody_full/properties.tsv

@@ -0,0 +1,4 @@
+entity_key	charge_A_CDR3	chargeShift_A_CDR3	gravy_A_CDR3	charge_B_CDR3	chargeShift_B_CDR3	gravy_B_CDR3	charge_A_VDJRegion	pi_A_VDJRegion	gravy_A_VDJRegion	mw_A_VDJRegion	eox_A_VDJRegion	ered_A_VDJRegion	instability_A_VDJRegion	aliphatic_A_VDJRegion	aromaticity_A_VDJRegion	charge_B_VDJRegion	pi_B_VDJRegion	gravy_B_VDJRegion	mw_B_VDJRegion	eox_B_VDJRegion	ered_B_VDJRegion	instability_B_VDJRegion	aliphatic_B_VDJRegion	aromaticity_B_VDJRegion	charge_Fv	chargeShift_Fv	pi_Fv	eox_Fv	ered_Fv	mw_Fv
+c1	-0.111	-0.245	-0.9833333333333334	-0.112	-0.236	-1.6833333333333333	-0.837	6.007	-0.1111111111111111	6050.6567	22460.0	22460.0	38.7537037037037	61.29629629629629	0.18518518518518517	1.149	9.166	-0.686	5466.8949999999995	9970.0	9970.0	62.929999999999986	46.8	0.12	0.313	-1.836	7.634	32430.0	32430.0	11517.5517
+c2	0.895	-0.22	0.21666666666666665	0.111	-0.854	-0.5000000000000001										1.18	9.169	-0.544	5432.8804	8480.0	8480.0	62.36399999999999	46.8	0.12						
+c3				-0.112	-0.236	-1.6833333333333333										1.149	9.166	-0.686	5466.8949999999995	9970.0	9970.0	62.929999999999986	46.8	0.12						

software/tests/data/golden/antibody_full/stats.json

@@ -0,0 +1 @@
+{"medianCdr3Length":{"A":6.0,"B":6.0}}

software/tests/data/golden/antibody_partial/aa_fraction.tsv

@@ -0,0 +1 @@
+entity_key	aminoAcid	value

software/tests/data/golden/antibody_partial/input.tsv

@@ -0,0 +1,2 @@
+entity_key	A_FR1	A_CDR1	A_FR2	A_CDR2	A_FR3	A_CDR3	A_FR4	B_FR1	B_CDR1	B_FR2	B_CDR2	B_FR3	B_CDR3	B_FR4
+c1	EVQLVES	GFTFSSY	AMSWVRQ	ISGSGGS	TYYAESVKGRFTI	CARDYW	WGQGTLV	DIQMTQS	QSISSY	LNWYQQK	AASSLQS	GVPSRFSGSG	CQQYNS	FGQGTKV