feat(train): parallelize RF SHAP explainability across CPU cores#259
Conversation
The RF-explainability step ran shap.TreeExplainer single-threaded (shap's C TreeSHAP has no OpenMP / no n_jobs; the RF's own n_jobs doesn't apply — shap re-walks the trees itself), so on a 1000-tree forest the interaction pass alone took ~76h (~183 s/sample x 1500), stalling the end of every training run for days. Add aetherscan.shap_parallel: a TF-free process-pool wrapper that chunks the samples across all cores (manager.n_processes = cpu_count() by default), each worker rebuilding a stock TreeExplainer (rebuild-per-worker avoids the shap #1204 segfault from pickling a prebuilt explainer) with BLAS threads pinned. Byte-identical to the serial result; ~40-45x on a 96-core node. Drives all three passes (summary, interaction, log-loss). select_positive_class_shap moves to shap_parallel (shared with the workers); its tests move to the new TF-free tests/unit/test_shap_parallel.py alongside the MP correctness test. docs/TRAINING_PIPELINE.md documents the CPU-vs-GPU comparison: GPUTreeExplainer is faster on interaction but needs a from-source CUDA build + a depth<32 (distinct-features-per-path) constraint + a broken interventional log-loss path — not worth the maintenance for a few minutes at the end of a multi-day run. max_depth stays unbounded; the constraint is noted only as GPU context, with a switchover runbook. Closes #258
|
Claude finished @zachtheyek's task in 3m 56s —— View job PR review —
|
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| rf_path = os.path.join(self.config.model_path, f"random_forest_{tag}.joblib") | ||
| if not os.path.exists(rf_path): | ||
| joblib.dump(self.rf_model.model, rf_path) |
There was a problem hiding this comment.
Fallback dump could hide a stale-RF divergence. Under the current flow the file should already exist and match self.rf_model.model (either just saved by train_random_forest at L2591, or loaded from disk by try_load_rf_for_resume), so this branch usually doesn't fire. But if it ever does — e.g. the joblib is present from an earlier run under the same tag but was produced by a different model / different config — then:
- workers load the stale on-disk RF for summary / interaction / log-loss
- while
expected_valueat L2789 is computed fromself.rf_model.model(the current in-process RF)
…and every downstream plot silently inherits the mismatch. The mismatch is exactly the case this defensive check was written to catch, but the if not os.path.exists guard turns "the file is stale" into a silent no-op instead of an overwrite.
Two safer options — either is a one-liner:
- Always dump (idempotent, guarantees the file matches
self.rf_model.model):joblib.dump(self.rf_model.model, rf_path)
- Or assert-and-fail-fast:
assert os.path.exists(rf_path), f"expected RF at {rf_path}"— trusting the upstream contract and turning the resume-path bug into a loud failure instead of a silent one.
Given the caller writes the RF unconditionally 190 lines earlier, "always dump" costs one extra disk write per training run (a few hundred MB, dwarfed by the SHAP work) in exchange for eliminating a whole class of mismatch bugs. Recommend that.
Adjacent nit: the comment on L2779-2781 is worth keeping as-is — it's exactly the kind of "why this file has to exist for workers to work" WHY that the CLAUDE.md style guide asks for.
| n = len(x) | ||
| n_workers = max(1, min(n_workers, n)) | ||
|
|
||
| # n_workers == 1: run in-process (no pool) — used on single-core hosts or tiny inputs. | ||
| if n_workers == 1: | ||
| _worker_init(rf_path, background) | ||
| _, out = _worker((0, kind, x, y)) | ||
| return out |
There was a problem hiding this comment.
Empty-input edge case falls into the in-process branch and calls shap on a zero-row array. If x is empty (n == 0), min(n_workers, 0) == 0, then max(1, 0) == 1, so we run the fallback path with x empty. shap's behaviour on a zero-row input isn't well-defined (probably returns an empty array of the right rank, but has been known to raise on some versions).
Not reachable through _compute_or_load_shap_values today (the caller guarantees n_summary/n_interact ≥ 1 via min(..., n_val) with n_val > 0 by the time RF eval runs), so this is a defence-in-depth suggestion rather than a live bug. A cheap guard here would make the module safe to reuse in a context that doesn't have that upstream invariant:
n = len(x)
if n == 0:
# Preserve dtype/rank so callers can concatenate without a special case.
return np.empty_like(x, shape=(0, *x.shape[1:]))
n_workers = max(1, min(n_workers, n))Or, if you'd rather stay strict, raise: if n == 0: raise ValueError("cannot explain zero samples").
| ) as pool: | ||
| results = sorted(pool.map(_worker, tasks), key=lambda r: r[0]) | ||
| return np.concatenate([r[1] for r in results], axis=0) |
There was a problem hiding this comment.
Nit (safe to ignore): sorted(...) is redundant. ProcessPoolExecutor.map is documented to yield results in input order, so results will already be ordered by chunk_index when tasks is emitted in order (which it is — enumerate on L152). The sort is a harmless defensive belt-and-braces, but it also mildly obscures the invariant.
If you want to keep the defensive posture, an explicit assert reads clearer than a re-sort:
results = list(pool.map(_worker, tasks))
assert [r[0] for r in results] == list(range(len(results))), "map lost input order"
return np.concatenate([r[1] for r in results], axis=0)Or just drop the sort and trust the contract. Either way, not blocking.
| 1000-tree forest at production sample counts runs for hours-to-days — dominated by the interaction | ||
| pass. SHAP values are per-sample independent, though, so we chunk the samples across processes and | ||
| call a **stock** ``shap.TreeExplainer`` in each worker (not a fork of the algorithm — just | ||
| parallelism). Measured ~40-45x on a 96-core node, byte-identical to the single-threaded result. |
There was a problem hiding this comment.
Wording nit: "byte-identical" is a stronger claim than the test enforces. The correctness test uses np.testing.assert_allclose(pooled, serial, atol=1e-8) (tests/unit/test_shap_parallel.py:88), which is tolerance-based, not bitwise. In practice TreeSHAP is per-sample independent with no cross-sample reduction, so results should be bitwise identical — but the module docstring and PR body both promise "byte-identical" while the test only enforces allclose.
Two ways to make the wording and the test agree:
- Tighten the test to
np.testing.assert_array_equal(pooled, serial)(cheap, and would fail loudly if a future shap release starts doing something non-deterministic). - Soften the doc: replace "byte-identical" with something like "numerically identical (all-close to machine precision)" here and in the
parallel_shapdocstring on L137. (Also worth updating the PR body if you go this way.)
I'd lean toward the first — if the claim is true today and cheap to enforce, locking it in with assert_array_equal protects against silent drift.
spawn re-imports the parent's __main__ (aetherscan.main -> TensorFlow -> the whole training stack: matplotlib/umap/train.py) into every worker, so the "TF-free workers" goal was not actually met and each worker would pay a heavy import cost. Use a forkserver context with an empty preload: the fork server is spawned once, clean (no __main__ re-import), and workers fork from it importing only shap_parallel + shap. Docstring updated to match.
- train.py: dump the RF unconditionally so the workers load exactly the in-process model — a stale on-disk RF under a reused tag would otherwise silently diverge from the expected_value. - shap_parallel.py: raise on zero-sample input (defence-in-depth); drop the redundant post-map sort (ProcessPoolExecutor.map preserves input order) and the now-dead chunk index. - docstrings/test: TreeSHAP is per-sample deterministic, so the result is bitwise-identical — the correctness test now uses assert_array_equal (was allclose), backing that claim.
|
Thanks — all four addressed in
CI was green on the prior push; this commit is lint-clean locally (the full suite can't run on the |
|
Second-pass check: all four first-pass notes are resolved in |
|
LGTM on Forkserver + empty preload ( Always-dump RF ( Other resolutions confirmed:
CI: green across tests (3.10/3.11/3.12), CodeQL (actions + python), and pre-commit — all successful at No further concerns from my side. Ship it. |
|
Release note draft — Improved Training runs now finish dramatically faster. The final SHAP explainability step — which previously ran single-threaded and could stall a production 1000-tree Random Forest run for roughly 76 hours — is now parallelized across all available CPU cores, delivering an ~40–45× speedup on a 96-core node. Results are byte-identical to the previous serial computation, so existing outputs and downstream analyses are unaffected. Draft for the changelog; edit or drop as needed. |
1 participant
Summary
Closes #258.
train.py::_compute_or_load_shap_valuesranshap.TreeExplainersingle-threaded — shap's C TreeSHAP has no OpenMP /n_jobs, and the RF's ownn_jobsdoesn't apply (shap re-walks the trees itself) — so on a production 1000-tree forest the interaction pass alone took ~76 h (measured ~183 s/sample × 1500), stalling the end of every training run for days. This parallelizes it across all CPU cores.What changed
aetherscan/shap_parallel.py(new, TF-free): a process-pool wrapper. Chunks samples acrossmanager.n_processes(=cpu_count()by default) workers, each rebuilding a stockTreeExplaineron its chunk — rebuild-per-worker avoids the shap #1204 segfault from pickling a pre-built explainer; workers load the RF from its persisted joblib; BLAS threads pinned. Byte-identical to the serial result; ~40–45× on a 96-core node. Runs under aforkservercontext with an empty preload so workers stay lightweight:train.pyimports TF at module level, andspawnwould re-import the parent's__main__(aetherscan.main→ TF → the whole training stack) into every worker — the fork server is spawned once, clean, and workers fork from it importing onlyshap_parallel+shap. (This was corrected from an initialspawnin a follow-up commit during self-review.)train.py:_compute_or_load_shap_valuesdrives summary + interaction + log-loss throughparallel_shap;select_positive_class_shapmoved toshap_parallel(shared with the workers). Expected-value uses one lightweight in-process explainer.tests/unit/test_shap_parallel.py(new, TF-free) — the shape-normalization tests (relocated fromtest_train_utils.py) plus a parametrized MP-vs-single-processallclosecorrectness test for all three passes.docs/TRAINING_PIPELINE.md: a SHAP-performance subsection — why CPU MP, the GPU findings (GPUTreeExplainer is faster on interaction but needs a from-source CUDA build + a< 32distinct-features-per-path warp-lane constraint + a broken interventional log-loss path, shap #4270/#3936/#1726), and a switchover runbook.max_depthstays unbounded; the constraint is documented only as GPU context.Validation
allclose, all three passes) — CI runs it (shap is a runtime dep)..nosyncTF-import deadlock blocks the full pytest suite locally, so the shape-normalization +ValueErrorpaths were checked standalone; the MP path is covered by CI.Notes
# NOTE: come back to this laterquestions (now answered in the docs).max_depthintentionally unchanged (unbounded, sklearn default).