Parallelize and right-size the sequence-properties compute step

.changeset/deterministic-threads.md

@@ -0,0 +1,25 @@
+---
+'@platforma-open/milaboratories.sequence-properties.software': patch
+'@platforma-open/milaboratories.sequence-properties.workflow': patch
+---
+
+Parallelize the compute step and size its resources to the input, keeping output byte-identical.
+
+Memory now scales with the input: `memFormula(clamp(max(lineCount * perRow, size * 64), 2GiB, 64GiB))` replaces the flat `mem("12GiB")`. The two terms cover two cost sources — a row-scaled term for per-clone structures (count matrix, output columns, aa_fraction; measured ~3 KiB/row peptide, ~4.3 KiB/row antibody), and a residue-scaled term (`size` ≈ total residues) for the O(residues) cleaning transients. The size term keeps long-sequence/low-row inputs (amplicon variants) off the 2 GiB floor, where a `lineCount`-only formula would OOM them. Large datasets get the RAM they need; small ones stop over-reserving. A `fallback: "12GiB"` preserves the old behavior on backends that cannot evaluate formulas.
+
+CPU scales with the input too: both modes request up to 4 cores (`cpuFormula`), matching their measured thread speedups (peptide ~1.7x, antibody ~1.5x). `POLARS_MAX_THREADS` — which sizes both the polars pool and the numpy pI-bisection workers — takes the cores the backend actually grants, via the `{system.cpu}` command expression, so a sub-cap allocation never oversubscribes its quota. A static `env` covers backends without command expressions.
+
+Determinism holds regardless of core count. Polars parallelizes only order-stable work (CSV read, chain reconstruction, aa_fraction reshape, the unique-key sort, CSV write), and `main.py` pins the BLAS/OpenMP intra-op threads to 1 so the one order-sensitive step — the `counts @ weights` reduction — never splits across threads. The emitted bytes, and the resource CID, stay identical.
+
+Compute-engine optimizations, all output-preserving:
+
+- `_charge_raw` computes `10**ph` once and reuses it via scalar factors instead of a per-amino-acid `10**(ph±pk)` — ~2.4x faster on the charge/pI path (FP drift ~1e-15, absorbed by the 3-dp quantization).
+- Peptide mode cleans its column once, sharing the intermediate between the count substrate and the instability index instead of cleaning twice.
+- The count and instability scatters use `np.bincount` instead of the slower unbuffered `np.add.at` — bit-identical, ~3x faster.
+- The per-chain median-CDR3-length stat uses a vectorized polars `count_matches` instead of one Python `effective_length` call per row.
+- Peptide `aa_fraction` uses a polars unpivot instead of a 20×N Python-list explosion, cutting peak memory on the peptide path.
+- The pI bisection (`isoelectric_point` / `fv_isoelectric_point`) runs row-parallel across threads — contiguous, data-independent blocks, each depending only on its own rows, so the result is bit-identical regardless of worker count.
+
+Together these cut peptide compute ~2.3x and make antibody scale from ~1.1x to ~1.5x on 4 cores (10.0s → 4.6s at 1M rows).
+
+Output stays within the quantized-equal contract, verified byte-for-byte between 1 and 4 threads by an extended determinism test and against the pre-change implementation (peptide and antibody, including edge cases).

pnpm-workspace.yaml

@@ -11,12 +11,12 @@ catalog:
   "@milaboratories/ts-builder": 1.6.0
   "@milaboratories/ts-configs": 1.3.0
   "@platforma-sdk/workflow-tengo": 6.6.5
-  "@platforma-sdk/block-tools": 2.11.6
-  "@platforma-sdk/model": 1.79.20
-  "@platforma-sdk/ui-vue": 1.79.20
-  "@platforma-sdk/test": 1.79.23
+  "@platforma-sdk/block-tools": 2.11.9
+  "@platforma-sdk/model": 1.79.24
+  "@platforma-sdk/ui-vue": 1.79.24
+  "@platforma-sdk/test": 1.79.26
   "@milaboratories/helpers": 1.14.2
-  "@platforma-sdk/tengo-builder": 4.0.11
+  "@platforma-sdk/tengo-builder": 4.0.14
   "@platforma-sdk/package-builder": 3.14.0
   "@platforma-sdk/blocks-deps-updater": 2.2.0
   "@platforma-sdk/eslint-config": 1.2.0

software/src/main.py

@@ -11,11 +11,34 @@
 
 import os
 
-# Pin the polars thread pool to 1 BEFORE polars (or anything importing it,
-# including `pipeline`) is loaded — POLARS_MAX_THREADS is read once at import.
-# The block reserves cpu(1), and single-threaded execution keeps output
-# byte-stable (no thread-count-dependent reduction order). setdefault so an
-# explicit override from the environment still wins.
+# Thread configuration MUST be set before numpy / polars are imported — both
+# read their thread-pool size once, at import time.
+#
+# BLAS / OpenMP intra-op threads -> 1 (forced). The vectorized engine's only
+# order-sensitive step is the `counts @ weights` matvec (vectorized.py): if a
+# single dot-product reduction is split across threads, the float summation
+# order — and therefore the emitted TSV bytes and the resource CID — becomes
+# thread-count dependent. Pinning every BLAS backend to one intra-op thread
+# keeps that reduction bit-identical. This is the load-bearing half of the
+# determinism contract, so it is forced, not setdefault.
+for _thread_var in (
+    "OMP_NUM_THREADS",
+    "OPENBLAS_NUM_THREADS",
+    "MKL_NUM_THREADS",
+    "NUMEXPR_NUM_THREADS",
+    "VECLIB_MAXIMUM_THREADS",
+):
+    os.environ[_thread_var] = "1"
+
+# Polars parallelism is safe to scale across cores: every polars op the pipeline
+# uses (CSV read, concat_str chain reconstruction, the aa_fraction explode, the
+# unique-key output sort, CSV write) is deterministic regardless of thread count
+# — none is a cross-row float reduction, and the sorts are total orders on unique
+# keys. This value also sizes the numpy pI-bisection worker pool (see
+# vectorized._n_workers). The workflow sets POLARS_MAX_THREADS to the cores the
+# backend actually grants (the {system.cpu} command expression = the resolved
+# cpuFormula value); setdefault to 1 keeps local / test runs single-threaded
+# unless they opt in.
 os.environ.setdefault("POLARS_MAX_THREADS", "1")
 
 import argparse
@@ -42,6 +65,10 @@ def _configure_logging() -> None:
 
 def main(argv: list[str] | None = None) -> int:
     _configure_logging()
+    logging.getLogger(__name__).info(
+        "thread config: POLARS_MAX_THREADS=%s, BLAS intra-op threads pinned to 1",
+        os.environ["POLARS_MAX_THREADS"],
+    )
     parser = argparse.ArgumentParser(prog="compute-properties")
     parser.add_argument("--input", required=True, help="path to input entity TSV")
     parser.add_argument("--plan", required=True, help="path to plan JSON")

software/src/pipeline.py

@@ -33,6 +33,12 @@
 
 PH = 7.0  # All charge values computed at pH 7 (spec default).
 
+# Regex character class matching a standard residue (either case), derived from
+# the single source of truth STANDARD_AAS so it cannot drift from
+# properties.effective_length / the count substrate. STANDARD_AAS is all letters,
+# so it needs no regex escaping.
+_STANDARD_AA_CLASS = "[" + STANDARD_AAS + STANDARD_AAS.lower() + "]"
+
 
 # ---------------------------------------------------------------------------
 # CID quantization
@@ -186,7 +192,13 @@ def run_peptide(reads: pl.DataFrame) -> dict[str, Any]:
     n = len(seqs)
 
     log.info("Computing peptide properties + AA fractions (%d sequences)", n)
-    sub = vec.build_counts(seqs)
+    # Clean the column ONCE and share the _Cleaned between the count substrate and
+    # the instability index — the same pattern the antibody full-chain path uses.
+    # The old code called build_counts(seqs) and instability_index(seqs)
+    # separately, each running a full _clean_vectorized pass (the dominant per-row
+    # cost), so the column was cleaned twice.
+    cleaned = vec._clean_vectorized(seqs)
+    sub = vec.counts_from_cleaned(cleaned)
 
     eox, ered = vec.extinction(sub)
     prop_arrays: dict[str, np.ndarray] = {
@@ -197,7 +209,7 @@ def run_peptide(reads: pl.DataFrame) -> dict[str, Any]:
         "pi_peptide": vec.isoelectric_point(sub, IPC2_PEPTIDE, include_cys=True),
         "eox_peptide": eox,
         "ered_peptide": ered,
-        "instability_peptide": vec.instability_index(seqs),
+        "instability_peptide": vec.instability_from_cleaned(cleaned),
         "aliphatic_peptide": vec.aliphatic_index(sub),
         "aromaticity_peptide": vec.aromaticity(sub),
     }
@@ -210,19 +222,27 @@ def run_peptide(reads: pl.DataFrame) -> dict[str, Any]:
 
     # AA-fraction long frame from the (N, 20) mole-fraction matrix: 20 rows per
     # entity (one per STANDARD_AAS), value NaN -> null for invalid entities so
-    # the 2-axis PColumn keeps a uniform shape across entities. The long-format
-    # entity axis needs Python-level repetition, so materialize keys here.
+    # the 2-axis PColumn keeps a uniform shape across entities.
+    #
+    # Built polars-native: assemble a wide frame (entity_key + one Float64 column
+    # per standard AA, in STANDARD_AAS order), then unpivot to the long
+    # (entity_key, aminoAcid, value) contract. This replaces the old
+    # `[k for k in keys for _ in STANDARD_AAS]` construction, which materialized
+    # two 20*N-element Python lists — the dominant peptide-mode memory transient —
+    # and lets polars parallelize the reshape. The unpivot reorders the rows, but
+    # write_output_tsv sorts by (entity_key, aminoAcid), so the emitted bytes are
+    # unchanged. np.ascontiguousarray keeps each column slice contiguous, so the
+    # pl.Series build never hits a strided-array path.
     fractions = vec.aa_fractions(sub)  # (N, 20), NaN for invalid rows
-    keys = key_series.to_list()
-    aa_entity = [k for k in keys for _ in STANDARD_AAS]
-    aa_amino = list(STANDARD_AAS) * n
-    aa_value = fractions.reshape(-1)  # row-major: entity 0's 20 AAs, then entity 1's, ...
-    aa_fraction = pl.DataFrame(
-        {
-            "entity_key": pl.Series(values=aa_entity, dtype=pl.Utf8),
-            "aminoAcid": pl.Series(values=aa_amino, dtype=pl.Utf8),
-            "value": _na_to_null(aa_value),
-        }
+    wide = pl.DataFrame(
+        {"entity_key": key_series}
+        | {aa: _na_to_null(np.ascontiguousarray(fractions[:, i])) for i, aa in enumerate(STANDARD_AAS)}
+    )
+    aa_fraction = wide.unpivot(
+        index="entity_key",
+        on=list(STANDARD_AAS),
+        variable_name="aminoAcid",
+        value_name="value",
     )
     # Quantize at the output boundary like the property columns, through the
     # shared CID-stability helper. aaFraction displays as .3f; round to 5 dp (one
@@ -304,7 +324,21 @@ def _median_cdr3_length_by_chain(reads: pl.DataFrame, chains: list[str]) -> dict
         col = f"{ch}_CDR3"
         if col not in reads.columns:
             continue
-        lengths = [effective_length(s) for s in reads[col].to_list() if s]
+        # Effective length = count of standard-AA characters (case-insensitive),
+        # matching properties.effective_length, computed vectorized in polars over
+        # the whole column instead of one Python call per row (the dominant
+        # antibody-mode stats cost). Only non-empty CDR3 cells count (`if s` in the
+        # scalar form drops null and ""). The median formula below is unchanged, so
+        # the emitted stat is byte-identical.
+        lengths = (
+            reads.select(
+                pl.col(col)
+                .filter(pl.col(col).is_not_null() & (pl.col(col) != ""))
+                .str.count_matches(_STANDARD_AA_CLASS)
+            )
+            .to_series()
+            .to_list()
+        )
         if not lengths:
             continue
         lengths.sort()