aetherscan 1.0.0__py3-none-any.whl

aetherscan/run_state.py

@@ -0,0 +1,277 @@
+"""
+Persisted training-run state for fault-tolerant resume.
+
+A TrainingRunState manifest lives at {output_path}/run_state_{save_tag}.json and carries
+everything a retry (in-process or a full relaunch of the identical command) needs to resume
+where the previous attempt died:
+
+- run_start_time: wall clock of attempt 1. TrainingPipeline.__init__ seeds self.start_time
+  from it, so every DB query/plot spans the whole run rather than just the current attempt.
+- completed_rounds: beta-VAE rounds whose checkpoint was saved; resume starts at max + 1.
+- stages_done / stages_failed: drive run_training_pipeline's stage machine — done stages are
+  skipped, failed non-critical stages (plots, hf_upload) are retried on the next run and force a nonzero
+  exit if they never succeed.
+
+Writes are atomic (tmp -> os.replace, mirroring round_data's .done manifest protocol) so a
+crash mid-write can never leave a truncated manifest — the file is either the previous state
+or the new one.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import hashlib
+import json
+import logging
+import os
+from dataclasses import asdict, dataclass, field
+
+logger = logging.getLogger(__name__)
+
+# Ordered stage names for run_training_pipeline's stage machine.
+STAGE_VAE_ROUNDS = "vae_rounds"
+STAGE_VAE_PLOTS = "vae_plots"
+STAGE_RF_TRAIN = "rf_train"
+STAGE_RF_PLOTS = "rf_plots"
+STAGE_FINAL_SAVE = "final_save"
+STAGE_HF_UPLOAD = "hf_upload"  # Opt-in (config.hf.upload_after_training); non-critical
+TRAINING_STAGES = (
+    STAGE_VAE_ROUNDS,
+    STAGE_VAE_PLOTS,
+    STAGE_RF_TRAIN,
+    STAGE_RF_PLOTS,
+    STAGE_FINAL_SAVE,
+    STAGE_HF_UPLOAD,
+)
+
+
+def run_state_path(output_path: str, tag: str) -> str:
+    """Manifest location for one training run: {output_path}/run_state_{save_tag}.json."""
+    return os.path.join(output_path, f"run_state_{tag}.json")
+
+
+# Config sections excluded from a run's fingerprint: pure infra/runtime (db, manager, monitor,
+# logger), resume-control (checkpoint), environment paths, the inference config, and the HF
+# upload config — none of which change the training result, so a change to any of them must NOT
+# force a fresh run. In particular, toggling --hf-upload or changing its repo/revision must never
+# be read as training drift (that would discard the manifest and overwrite a completed model).
+_FINGERPRINT_EXCLUDE_SECTIONS = frozenset(
+    {"db", "manager", "monitor", "logger", "inference", "paths", "checkpoint", "hf"}
+)
+# Retry knobs live in the training section but only control the retry loop, not the result.
+_FINGERPRINT_EXCLUDE_TRAINING_KEYS = frozenset({"max_retries", "retry_delay"})
+
+
+def config_fingerprint(config_dict: dict) -> str:
+    """
+    Stable hash of the training-result-affecting config (from Config.to_dict()), used to detect
+    that a run's config changed under a reused --save-tag. Excludes infra/runtime, resume-control,
+    environment paths, the inference config, and the retry knobs (see the constants above) so a
+    change to any of those does not spuriously force a fresh run.
+    """
+    relevant = {
+        section: values
+        for section, values in config_dict.items()
+        if section not in _FINGERPRINT_EXCLUDE_SECTIONS
+    }
+    training = relevant.get("training")
+    if isinstance(training, dict):
+        relevant["training"] = {
+            k: v for k, v in training.items() if k not in _FINGERPRINT_EXCLUDE_TRAINING_KEYS
+        }
+    canonical = json.dumps(relevant, sort_keys=True, default=str)
+    return hashlib.sha256(canonical.encode()).hexdigest()
+
+
+def config_changed(state: TrainingRunState | None, current_fingerprint: str) -> bool:
+    """
+    True when a persisted manifest exists but was written under a different config fingerprint —
+    i.e. the resolved training config changed under a reused save tag. The caller downgrades to a
+    fresh run so a stale model is never silently reused/reported as success (a manifest predating
+    fingerprinting carries an empty string, which mismatches any real fingerprint -> fresh run).
+    """
+    return state is not None and state.config_fingerprint != current_fingerprint
+
+
+# Inference-result-affecting config for inference_config_fingerprint(). We DENYLIST the inference
+# keys that don't change per-cadence results (I/O paths, batching, retry knobs, and viz/debug
+# toggles) rather than allowlisting the ones that do — a denylist can only over-invalidate (force a
+# harmless re-inference), never under-invalidate (the actual stale-reuse footgun). The data section
+# is the reverse: it is mostly training params, so we ALLOWLIST only the geometry fields that drive
+# stamp extraction and the encoder input.
+_INFERENCE_FINGERPRINT_EXCLUDE_INFERENCE_KEYS = frozenset(
+    {
+        "config_path",  # source-config file path, not a result input
+        "per_replica_batch_size",  # batching only; #120's pad+truncate makes results batch-invariant
+        "coarse_channel_log_interval",  # progress-logging chunk size only (inert)
+        "bandpass_debug_plot",  # opt-in debug artifact
+        "preprocess_output_dir",  # where stamps are written (folded into npy_path, not values)
+        "inference_viz_enabled",  # viz toggle
+        "inference_viz_scope",  # viz only (#301)
+        "stamp_gallery_top_k",  # viz only
+        "max_candidate_plots",  # viz only
+        "max_retries",  # retry loop only
+        "retry_delay",  # retry loop only
+        "prefetch_depth",  # scheduling only; per-cadence results are depth-invariant (#298 N2)
+        # Stamp-cache retention only (#302): deletes already-scored artifacts, never
+        # changes what a cadence scores. MUST stay excluded — as a new inference key it
+        # would otherwise enter BOTH fingerprints, staling every 'inferred' resume row
+        # and renaming every ED cache directory on upgrade.
+        "prune_stamps",
+    }
+)
+_INFERENCE_FINGERPRINT_DATA_KEYS = frozenset(
+    {"downsample_factor", "width_bin", "num_observations", "time_bins"}
+)
+# Scoring/model keys excluded ON TOP of the inference denylist for the PREPROCESSING
+# fingerprint (#298 I3): energy detection is deterministic given (csv files, h5 files, ED
+# config), so a changed encoder/RF/threshold must REUSE stamps — that is the whole point of
+# the fingerprint-scoped stamp cache. Everything else in the inference section (cadence
+# grouping columns, channelization, bandpass method/taps, detection windows/threshold, stamp
+# geometry, overlap, the downsample toggle) stays IN the hash. This too is a DENYLIST, never
+# an ED-key allowlist: an allowlist that missed a key (say coarse_channel_width) would
+# silently REUSE WRONG STAMPS when that key changes; a denylist can only over-invalidate.
+_PREPROCESSING_FINGERPRINT_EXTRA_EXCLUDE_KEYS = frozenset(
+    {
+        "encoder_path",
+        "rf_path",
+        "classification_threshold",
+        "screening_threshold",
+        "mc_draws",
+        "reference_cloud_size",
+    }
+)
+
+
+def preprocessing_config_fingerprint(config_dict: dict) -> str:
+    """
+    Stable hash of the PREPROCESSING-result-affecting config — everything that changes what
+    energy detection writes into a stamp .npy (#298 I3). Keys the default stamp cache
+    directory ({data_path}/inference/preprocessed/<csv_stem>_ed<hash12>/), is persisted into
+    each cadence's metadata JSON as ed_config_fingerprint, and is verified by the resume
+    guard — so runs sharing an ED config share stamps, and any ED-config change lands in a
+    different directory by construction.
+    """
+    inference = config_dict.get("inference") or {}
+    data = config_dict.get("data") or {}
+    excluded = (
+        _INFERENCE_FINGERPRINT_EXCLUDE_INFERENCE_KEYS
+        | _PREPROCESSING_FINGERPRINT_EXTRA_EXCLUDE_KEYS
+    )
+    relevant = {
+        "inference": {k: v for k, v in inference.items() if k not in excluded},
+        "data": {k: data[k] for k in _INFERENCE_FINGERPRINT_DATA_KEYS if k in data},
+    }
+    canonical = json.dumps(relevant, sort_keys=True, default=str)
+    return hashlib.sha256(canonical.encode()).hexdigest()
+
+
+def inference_config_fingerprint(config_dict: dict) -> str:
+    """
+    Stable hash of the inference-result-affecting config (from Config.to_dict()) — the inference
+    counterpart of config_fingerprint. Used to detect that inference was re-run under a reused
+    --save-tag with a changed config so a stale 'inferred' manifest row is not silently reused as
+    a skip. Covers the inference section minus non-result knobs (see the denylist above) plus the
+    data-section geometry that drives stamp extraction and the encoder input.
+    """
+    inference = config_dict.get("inference") or {}
+    data = config_dict.get("data") or {}
+    relevant = {
+        "inference": {
+            k: v
+            for k, v in inference.items()
+            if k not in _INFERENCE_FINGERPRINT_EXCLUDE_INFERENCE_KEYS
+        },
+        "data": {k: data[k] for k in _INFERENCE_FINGERPRINT_DATA_KEYS if k in data},
+    }
+    canonical = json.dumps(relevant, sort_keys=True, default=str)
+    return hashlib.sha256(canonical.encode()).hexdigest()
+
+
+@dataclass
+class TrainingRunState:
+    """Persisted state of one training run (identified by its save tag) across attempts."""
+
+    tag: str
+    run_start_time: float  # Wall clock of attempt 1 — used by ALL DB queries/plots
+    attempt: int = 1  # Incremented per retry (each TrainingPipeline rebuild)
+    config_fingerprint: str = ""  # Hash of the result-affecting config (config_fingerprint())
+    completed_rounds: list[int] = field(default_factory=list)
+    stages_done: list[str] = field(default_factory=list)
+    stages_failed: list[str] = field(default_factory=list)
+
+    def is_stage_done(self, stage: str) -> bool:
+        return stage in self.stages_done
+
+    def mark_stage_done(self, stage: str) -> None:
+        """Record a stage success; a previously recorded failure of the same stage is cleared."""
+        if stage not in self.stages_done:
+            self.stages_done.append(stage)
+        if stage in self.stages_failed:
+            self.stages_failed.remove(stage)
+
+    def record_stage_failure(self, stage: str) -> None:
+        """Record a non-critical stage failure; not in stages_done, so relaunches retry it."""
+        if stage not in self.stages_failed:
+            self.stages_failed.append(stage)
+
+    def clear_stage_failure(self, stage: str) -> None:
+        """Drop a recorded failure without marking the stage done — used when the user opts
+        out of an optional stage (e.g. re-running without --hf-upload) so a stale failure
+        can't force a nonzero exit forever."""
+        if stage in self.stages_failed:
+            self.stages_failed.remove(stage)
+
+    def mark_round_completed(self, round_number: int) -> None:
+        if round_number not in self.completed_rounds:
+            self.completed_rounds.append(round_number)
+            self.completed_rounds.sort()
+
+    @property
+    def resume_round(self) -> int:
+        """1-based round the beta-VAE loop should (re)start from."""
+        return max(self.completed_rounds) + 1 if self.completed_rounds else 1
+
+
+def save_run_state(state: TrainingRunState, path: str) -> None:
+    """Persist the manifest atomically (tmp -> os.replace): a crash mid-write leaves either
+    the previous manifest or the new one, never a truncated file."""
+    parent = os.path.dirname(path)
+    if parent:
+        os.makedirs(parent, exist_ok=True)
+    tmp_path = path + ".tmp"
+    try:
+        with open(tmp_path, "w") as f:
+            json.dump(asdict(state), f, indent=2)
+        os.replace(tmp_path, path)
+    except Exception:
+        with contextlib.suppress(OSError):
+            os.remove(tmp_path)
+        raise
+
+
+def load_run_state(path: str) -> TrainingRunState | None:
+    """Load a persisted manifest; returns None when the file is missing or unreadable
+    (a corrupt manifest downgrades to a fresh run rather than blocking training)."""
+    if not os.path.isfile(path):
+        return None
+    try:
+        with open(path) as f:
+            payload = json.load(f)
+        return TrainingRunState(
+            tag=str(payload["tag"]),
+            run_start_time=float(payload["run_start_time"]),
+            attempt=int(payload.get("attempt", 1)),
+            config_fingerprint=str(payload.get("config_fingerprint", "")),
+            completed_rounds=sorted(int(r) for r in payload.get("completed_rounds", [])),
+            stages_done=[str(s) for s in payload.get("stages_done", [])],
+            stages_failed=[str(s) for s in payload.get("stages_failed", [])],
+        )
+    except KeyError as e:
+        logger.warning(
+            f"Run state at {path} is missing required field {e} — treating as a fresh run"
+        )
+        return None
+    except Exception as e:
+        logger.warning(f"Failed to load run state from {path}: {e} — treating as a fresh run")
+        return None

aetherscan/seeding.py

@@ -0,0 +1,160 @@
+"""
+Root-seed stream derivation for reproducible pipeline runs (issues #49, #279).
+
+A single root seed (config.reproducibility.seed, --seed on both subcommands) makes every
+random stream in the pipeline deterministic. Rather than seeding numpy's global RNG once
+(streams would collide and any consumer could perturb every other), each consumer derives its
+own independent stream here via np.random.SeedSequence keyed on (root_seed, stream id,
+...extras) — the same key always reproduces the same stream, and distinct keys are
+statistically independent. Consumers that need a numpy Generator use derive_rng; APIs that
+take an int random_state (sklearn, umap) use derive_seed; TensorFlow's global RNG (weight
+init + the VAE Sampling layer) is seeded via seed_tensorflow, called by BOTH pipeline
+constructors so training and inference cannot drift (#279 — inference used to be entirely
+unseeded, making candidate sets unreproducible).
+
+The Random Forest seed is likewise derived from the root (STREAM_RF); config.rf.seed remains
+only as an explicit override for the deprecated --rf-seed flag.
+
+Deliberately NOT seeded (do not "fix" these): uuid4 temp-file names (must stay real entropy),
+round_data._array_checksum's content-keyed default_rng (keyed on the array, never the root,
+or checksums stop being comparable), and data_generation's per-PID worker init seeding (always
+overwritten by the per-task root-derived reseed before any draw).
+"""
+
+from __future__ import annotations
+
+import logging
+import random
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+# Stream ids: the leading component of every derived stream key. Each consumer of the root seed
+# owns one id (extra key components — e.g. the round index — subdivide it further), so no two
+# consumers can ever end up sharing a stream. Add a new id here for any new consumer.
+STREAM_DATA_GEN = 0  # per-round worker-task seeds (data_generation.generate_round_to_memmap)
+STREAM_DATASET = 1  # stratified split / trim / epoch shuffles (prepare_distributed_train_dataset)
+STREAM_VIZ = 2  # latent-viz batch selection + padding (train.TrainingPipeline)
+STREAM_PLOT = 3  # plot subsampling (train.TrainingPipeline plot helpers)
+STREAM_RF = 4  # RandomForest estimator random_state + pre-fit shuffle (models/random_forest)
+STREAM_UMAP = 5  # every umap.UMAP random_state + fit-pool subsampling (sub-keyed per site)
+STREAM_KMEANS = 6  # SHAP-space KMeans random_state (train.plot_rf_shap_explanation_clustering)
+STREAM_TF = 7  # tf.random.set_seed roots — sub-keyed 0=training (then per round), 1=inference
+STREAM_INFERENCE_VIZ = 8  # inference_viz latent budget-fill subsample
+STREAM_SHAP_SAMPLES = 9  # SHAP summary/interaction row subsampling (feeds cached artifacts)
+STREAM_RF_PLOTS = 10  # RF plot subsamples (sub-keyed: 0=learning curve, 1=decision boundary)
+STREAM_INFERENCE_MC = 11  # pass-2 MC latent draws (sub-keyed by the cadence's catalog index)
+STREAM_REFERENCE_CLOUD = 12  # reservoir subsample of pass-1 rejects for the reference cloud
+STREAM_KERAS_INIT = 13  # Python's global random, which tf_keras initializers seed from (see below)
+
+# One-shot flag so an unseeded run warns exactly once instead of per derived stream.
+# Module-level, so the "once" is PER PROCESS: forkserver/spawn workers each carry their own
+# copy and may repeat the warning in their own logs — intentional (each worker's log is
+# self-contained), not a bug.
+_UNSEEDED_WARNED = False
+
+
+def _warn_if_unseeded(root_seed: int | None) -> None:
+    global _UNSEEDED_WARNED
+    if root_seed is None and not _UNSEEDED_WARNED:
+        logger.warning(
+            "reproducibility.seed is None: every derived random stream falls back to OS "
+            "entropy, so this run is NOT reproducible. Set --seed (or leave the default) "
+            "for reproducible artifacts."
+        )
+        _UNSEEDED_WARNED = True
+
+
+def derive_rng(root_seed: int | None, *stream_key: int) -> np.random.Generator:
+    """
+    Build the numpy Generator for one consumer stream of the pipeline root seed.
+
+    With root_seed None this returns a fresh OS-entropy Generator (the historical
+    non-reproducible behavior), warning once per process. Otherwise the Generator is seeded
+    from SeedSequence([root_seed, *stream_key]), so the same (root_seed, stream_key) always
+    yields an identical stream and different keys yield independent ones. `stream_key` must
+    start with one of the STREAM_* ids above.
+    """
+    _warn_if_unseeded(root_seed)
+    if root_seed is None:
+        return np.random.default_rng()
+    return np.random.default_rng(np.random.SeedSequence([root_seed, *stream_key]))
+
+
+def derive_seed(root_seed: int | None, *stream_key: int) -> int:
+    """
+    The int-valued sibling of derive_rng, for APIs that take an integer random_state
+    (sklearn RandomForestClassifier/KMeans/shuffle, umap.UMAP, tf.random.set_seed).
+
+    Deterministic given (root_seed, stream_key); with root_seed None it returns a fresh
+    OS-entropy int (warning once) so call sites can keep passing a CONCRETE random_state —
+    per #279's constraint, determinism knobs stay set at every site that sets one today,
+    only the value's provenance changes. Returns a value in [0, 2**32).
+    """
+    _warn_if_unseeded(root_seed)
+    if root_seed is None:
+        return int(np.random.SeedSequence().generate_state(1)[0])
+    return int(np.random.SeedSequence([root_seed, *stream_key]).generate_state(1)[0])
+
+
+def seed_tensorflow(root_seed: int | None, deterministic_ops: bool, *stream_key: int) -> int | None:
+    """
+    Seed TensorFlow's global RNG from the root and (optionally) force deterministic op
+    implementations. Called by BOTH TrainingPipeline and InferencePipeline constructors —
+    and again at round/cadence boundaries with extended stream keys, so a resumed or
+    partially-rerun pipeline reproduces an uninterrupted one (#279: a single __init__-time
+    set_seed is not resume-safe, because skipping rounds shifts the stream position).
+
+    ALSO seeds Python's global `random` module, which is NOT incidental: tf_keras's weight
+    initializers (HeNormal/GlorotNormal in models/vae.py) build a
+    `backend.RandomGenerator(seed=None, rng_type="stateless")`, whose `_create_seed()` falls
+    back to `random.randint(1, int(1e9))` when no Keras seed generator is set. So
+    `tf.random.set_seed()` alone does NOT pin weight initialization — measured on the NGC
+    2.17 image: two same-`tf.random.set_seed` builds gave 100% different weights. Seeding
+    Python's `random` closes that hole (verified: identical weights same seed, different
+    weights different seed).
+
+    Do NOT "simplify" this to `tf_keras.utils.set_random_seed()`, the canonical Keras API:
+    it populates the thread-local `_SEED_GENERATOR`, after which `_create_seed()` calls
+    `generator.randint(1, 1e9)` with a FLOAT bound, which Python 3.12's `random.randint`
+    rejects — every subsequent initializer raises `TypeError: 'float' object cannot be
+    interpreted as an integer`. That path is broken on this stack; this one is not.
+
+    The VAE `Sampling` layer calls `tf.random.normal` directly and IS covered by
+    `tf.random.set_seed` alone — only the initializers needed the extra seeding.
+
+    Returns the seed applied, or None when root_seed is None (TF keeps its own entropy;
+    warned once). TF is imported lazily so this module stays importable without the
+    scientific stack.
+    """
+    import tensorflow as tf  # noqa: PLC0415
+
+    _warn_if_unseeded(root_seed)
+    applied: int | None = None
+    if root_seed is not None:
+        applied = derive_seed(root_seed, STREAM_TF, *stream_key)
+        tf.random.set_seed(applied)
+        # NOTE: random.seed() is PROCESS-WIDE, with two consequences worth knowing:
+        # (1) it pins any transitive consumer of Python's `random` in the imported dependency
+        #     tree, not just tf_keras — a feature for a reproducibility fix, but it means a
+        #     library whose behavior varies with RNG state now behaves differently than it did
+        #     before this call existed. Nothing in the current dep set is known to care.
+        # (2) it is not thread-local. data_generation reseeds per task before every draw
+        #     (data_generation.py:795), so today's single-threaded-per-task usage is safe; but
+        #     if a CONCURRENT main-process consumer of `random` is ever added (a monitor
+        #     thread, a plot renderer, a callback), the seed pins the aggregate stream while
+        #     the per-thread interleaving stays nondeterministic.
+        random.seed(derive_seed(root_seed, STREAM_KERAS_INIT, *stream_key))
+
+    if deterministic_ops:
+        # Deterministic kernels only pin op results; without a seed the draws still differ
+        # run to run, so flag the (legal but pointless) combination loudly
+        tf.config.experimental.enable_op_determinism()
+        logger.info("Enabled deterministic TF op implementations (tf_deterministic_ops)")
+        if root_seed is None:
+            logger.warning(
+                "tf_deterministic_ops is enabled but reproducibility.seed is None: "
+                "deterministic kernels alone do not make runs reproducible"
+            )
+    return applied

aetherscan/shap_parallel.py

@@ -0,0 +1,238 @@
+"""
+Process-pool wrapper for the RF SHAP passes.
+
+shap's TreeSHAP C extension is single-threaded (no OpenMP, no ``n_jobs``), so explaining a
+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.
+
+This module is deliberately kept small and off the TF/training import graph. ``train.py`` imports
+TensorFlow at module level (and pulls in matplotlib/umap/…), and multiprocessing's ``spawn`` re-imports
+the parent's ``__main__`` (``aetherscan.main`` → TF → the whole training stack) into every worker. So
+the pool runs under the **forkserver** start method with an **empty preload**: the fork server is
+spawned once, clean (no ``__main__`` re-import), and workers fork from it — importing only this module
++ shap, never TF. See ``_compute_or_load_shap_values`` in ``train.py`` for the caller and
+``docs/TRAINING_PIPELINE.md`` for the CPU-vs-GPU comparison behind this choice.
+
+Design notes:
+  * **Rebuild the explainer inside each worker** from the picklable sklearn RF — never pickle a
+    pre-built ``TreeExplainer`` into workers (segfaults large-model workers; shap #1204). Workers
+    load the RF from its persisted joblib path so the model isn't re-serialised through the pool.
+  * **One pool per session** (``shap_pool``): all three passes map onto the same worker set, so
+    workers start, load the RF, and build their explainers once per session rather than once per
+    pass. Each worker caches one explainer per pass family (plain vs log-loss) — see
+    ``_get_explainer``.
+  * **forkserver + empty preload** keeps workers off the TF/training stack (see above); the fork
+    server is spawned once and forks lightweight workers.
+  * **Pin BLAS/OpenMP threads to 1 per worker** so N workers don't oversubscribe the CPU (best-effort:
+    shap's TreeSHAP is single-threaded, so this only tames incidental numpy/BLAS threads).
+"""
+
+from __future__ import annotations
+
+import contextlib
+import multiprocessing as mp
+import os
+from concurrent.futures import ProcessPoolExecutor
+
+import numpy as np
+
+# Kinds of SHAP pass a worker can run.
+_RAW_KINDS = ("summary", "interaction")
+_LOGLOSS_KIND = "logloss"
+
+# Per-worker globals, populated once by the pool initializer (or in-process for the n_workers==1
+# fallback). Module-level so the forkserver workers can resolve them. _EXPLAINERS caches one built
+# explainer per pass family ("raw" vs "logloss") so a worker serving several chunks — and several
+# passes, now that one pool spans all of them — parses the forest once per family, not per chunk.
+_RF = None
+_BACKGROUND = None
+_EXPLAINERS: dict[str, object] = {}
+
+# Chunks submitted per worker per pass: >1 smooths stragglers (TreeSHAP cost varies per sample, so
+# a worker stuck on a slow chunk no longer idles the rest for the whole tail of the pass).
+_CHUNKS_PER_WORKER = 4
+
+
+def select_positive_class_shap(values, log_loss: bool = False) -> np.ndarray:
+    """
+    Normalize SHAP output across shap versions into a single positive-class ndarray.
+
+    TreeExplainer returns results in several shapes depending on shap version & task:
+    - classic binary classification: list ``[neg, pos]`` of ``(n, F)`` arrays
+    - newer shap, last axis is class: ``(n, F, 2)`` for values, ``(n, F, F, 2)`` for interactions
+    - log-loss (single scalar output): ``(n, F)`` or ``(n, F, F)``
+
+    Picks the positive-class slice in all cases. For log-loss there is only one output, so it is
+    returned as-is (modulo a preserved class axis). Kept here (not in train.py) so both the pipeline
+    and the TF-free workers share one definition.
+    """
+    if log_loss:
+        if isinstance(values, list):
+            return np.asarray(values[0])
+        values = np.asarray(values)
+        # newer shap preserves the class axis even for model_output="log_loss":
+        # (n, F, 2) -> (n, F); (n, F, F, 2) -> (n, F, F)
+        if values.ndim >= 3 and values.shape[-1] == 2:
+            return values[..., 1]
+        return values
+
+    if isinstance(values, list):
+        return np.asarray(values[1])
+
+    values = np.asarray(values)
+    # (n, F, 2) -> (n, F); (n, F, F, 2) -> (n, F, F)
+    if values.ndim >= 3 and values.shape[-1] == 2:
+        return values[..., 1]
+    return values
+
+
+def _pin_worker_threads() -> None:
+    # Must run before numpy/shap import their native thread pools. shap's TreeSHAP is single-threaded
+    # anyway, so this only tames incidental numpy/BLAS threads — but with all cores as workers that
+    # still matters. setdefault so an explicitly-set env is respected.
+    for var in (
+        "OMP_NUM_THREADS",
+        "OPENBLAS_NUM_THREADS",
+        "MKL_NUM_THREADS",
+        "NUMEXPR_NUM_THREADS",
+        "VECLIB_MAXIMUM_THREADS",
+        "BLIS_NUM_THREADS",
+    ):
+        os.environ.setdefault(var, "1")
+
+
+def _worker_init(rf_path: str, background) -> None:
+    """Pool initializer: pin threads, then load the RF once per worker (rebuild-per-worker source)."""
+    _pin_worker_threads()
+    import joblib  # noqa: PLC0415  (deferred so _pin_worker_threads() sets env before numpy/BLAS init)
+
+    global _RF, _BACKGROUND
+    _RF = joblib.load(rf_path)
+    _BACKGROUND = background
+    # In-process (n_workers==1) sessions reuse this module's globals across calls — drop any
+    # explainers built for a previous session's RF.
+    _EXPLAINERS.clear()
+
+
+def _get_explainer(kind: str):
+    """
+    Return the cached per-worker explainer for ``kind``, building it on first use.
+
+    The log-loss pass needs a structurally different explainer (interventional:
+    ``model_output="log_loss"`` + background data), so it gets its own cache slot; the summary
+    and interaction passes share the one plain ``TreeExplainer``.
+    """
+    import shap  # noqa: PLC0415  (deferred so the worker pins threads before shap/BLAS init)
+
+    # Validate here, not just at the top-level API: the in-process (n_workers <= 1) path
+    # reaches _worker without _validate_pass, and silently collapsing an unknown kind into
+    # the "raw" slot would hand a future pass the wrong explainer.
+    if kind != _LOGLOSS_KIND and kind not in _RAW_KINDS:
+        raise ValueError(f"unknown SHAP kind {kind!r}")
+    key = _LOGLOSS_KIND if kind == _LOGLOSS_KIND else "raw"
+    explainer = _EXPLAINERS.get(key)
+    if explainer is None:
+        if key == _LOGLOSS_KIND:
+            explainer = shap.TreeExplainer(_RF, data=_BACKGROUND, model_output="log_loss")
+        else:
+            explainer = shap.TreeExplainer(_RF)
+        _EXPLAINERS[key] = explainer
+    return explainer
+
+
+def _worker(task):
+    """Compute one SHAP pass on one chunk of samples using the module-global RF. Returns the
+    positive-class array for the chunk; ProcessPoolExecutor.map preserves input order."""
+    kind, x_chunk, y_chunk = task
+    explainer = _get_explainer(kind)
+
+    # Suppress shap's tqdm (it writes to stderr, which the pipeline logger forwards to Slack).
+    with open(os.devnull, "w") as devnull, contextlib.redirect_stderr(devnull):
+        if kind == _LOGLOSS_KIND:
+            raw = explainer.shap_values(x_chunk, y=y_chunk)
+        elif kind == "interaction":
+            raw = explainer.shap_interaction_values(x_chunk)
+        else:
+            raw = explainer.shap_values(x_chunk)
+    return select_positive_class_shap(raw, log_loss=(kind == _LOGLOSS_KIND))
+
+
+def _validate_pass(kind: str, n: int) -> None:
+    if kind != _LOGLOSS_KIND and kind not in _RAW_KINDS:
+        raise ValueError(f"unknown SHAP kind {kind!r}")
+    if n == 0:
+        raise ValueError("cannot explain zero samples")
+
+
+@contextlib.contextmanager
+def shap_pool(rf_path: str, n_workers: int, *, background: np.ndarray | None = None):
+    """
+    One worker pool shared by every SHAP pass computed inside the block.
+
+    Yields ``run(kind, x, y=None) -> np.ndarray`` computing the positive-class SHAP values for
+    one pass ``kind`` in ``{"summary", "interaction", "logloss"}``. Each pass historically spun
+    up (and tore down) its own pool; sharing one pool starts the workers — and parses the RF
+    into per-worker explainers (see ``_get_explainer``) — once per session instead of once per
+    pass. ``background`` must be supplied at pool creation when a ``"logloss"`` pass will run
+    (workers receive it through the initializer, not per task). ``n_workers <= 1`` runs every
+    pass in-process (no pool) — single-core hosts or tiny inputs.
+    """
+    if n_workers <= 1:
+        _worker_init(rf_path, background)
+
+        def run(kind: str, x: np.ndarray, y: np.ndarray | None = None) -> np.ndarray:
+            _validate_pass(kind, len(x))
+            return _worker((kind, x, y))
+
+        yield run
+        return
+
+    # forkserver with an EMPTY preload: the fork server is spawned clean and does NOT re-import the
+    # parent's __main__ (aetherscan.main -> TF -> the whole training stack), so workers stay light.
+    ctx = mp.get_context("forkserver")
+    ctx.set_forkserver_preload([])
+    with ProcessPoolExecutor(
+        max_workers=n_workers,
+        mp_context=ctx,
+        initializer=_worker_init,
+        initargs=(rf_path, background),
+    ) as pool:
+
+        def run(kind: str, x: np.ndarray, y: np.ndarray | None = None) -> np.ndarray:
+            _validate_pass(kind, len(x))
+            # ~_CHUNKS_PER_WORKER chunks per worker smooth stragglers. The chunk count cannot
+            # change the numbers: TreeSHAP is per-sample exact and per-sample independent, and
+            # pool.map yields chunk results in submission order, so concatenating them is
+            # bitwise-identical to the serial result for ANY chunking.
+            n_chunks = min(len(x), n_workers * _CHUNKS_PER_WORKER)
+            chunks = np.array_split(np.arange(len(x)), n_chunks)
+            tasks = [(kind, x[c], (None if y is None else y[c])) for c in chunks if len(c)]
+            return np.concatenate(list(pool.map(_worker, tasks)), axis=0)
+
+        yield run
+
+
+def parallel_shap(
+    rf_path: str,
+    x: np.ndarray,
+    kind: str,
+    n_workers: int,
+    *,
+    background: np.ndarray | None = None,
+    y: np.ndarray | None = None,
+) -> np.ndarray:
+    """
+    Compute the positive-class SHAP values for ``x`` under one pass ``kind`` in
+    ``{"summary", "interaction", "logloss"}``, chunked across up to ``n_workers`` worker processes
+    (a dedicated single-pass pool; callers running several passes should hold one ``shap_pool``
+    open instead — ``_compute_or_load_shap_values`` in ``train.py`` does).
+
+    ``rf_path`` is the persisted sklearn RF joblib (each worker loads it). ``background`` and ``y`` are
+    required only for the ``"logloss"`` (interventional) pass. TreeSHAP is per-sample deterministic, so
+    the concatenated result is bitwise-identical to the single-process computation.
+    """
+    _validate_pass(kind, len(x))
+    with shap_pool(rf_path, max(1, min(n_workers, len(x))), background=background) as run:
+        return run(kind, x, y)