Caching and Memory Optimization

This document describes the runtime caching subsystem introduced to reduce memory usage and wall-clock time during pipeline execution. It covers block-based feature storage, copy-on-write branch snapshots, step-level caching, and observability tooling.

Overview

The caching subsystem addresses three performance problems in pipeline execution:

Generator variant redundancy: When _or_ / _range_ generators expand a pipeline into many variants, shared preprocessing prefixes are re-executed from scratch for every variant.
Branch deep-copy explosion: Each branch in a {"branch": [...]} step deep-copies the entire feature set, even when most branches only read features without modifying them.
Array concatenation overhead: Adding a new processing via add_processing() previously reallocated and copied the entire 3D array, causing transient 2x memory spikes.

The subsystem has four layers, each independently toggleable:

Layer	What it does	Default	Config flag
Block-based storage	Per-processing 2D blocks instead of monolithic 3D array	Always on	N/A (structural)
Copy-on-Write snapshots	Branch snapshots share arrays via refcounting	ON	`use_cow_snapshots`
Step-level caching	Reuse preprocessing results across generator variants	OFF	`step_cache_enabled`
Observability	Per-step RSS logging, cache stats, memory warnings	ON	`log_step_memory`, `log_cache_stats`

Configuration

All cache settings live in a single dataclass:

from nirs4all.config.cache_config import CacheConfig

cache = CacheConfig(
    # Step cache (off by default)
    step_cache_enabled=True,
    step_cache_max_mb=2048,       # 2 GB budget
    step_cache_max_entries=200,

    # Branch snapshots (on by default)
    use_cow_snapshots=True,

    # Observability
    log_cache_stats=True,
    log_step_memory=True,
    memory_warning_threshold_mb=3072,  # 3 GB
)

result = nirs4all.run(
    pipeline=pipeline,
    dataset=dataset,
    cache=cache,
)

When no cache= argument is provided, CacheConfig() defaults are used: CoW snapshots are on, step caching is off, and observability logging is on.

Configuration fields

Field	Type	Default	Description
`step_cache_enabled`	`bool`	`False`	Enable step-level caching for cross-variant reuse
`step_cache_max_mb`	`int`	`2048`	Maximum memory budget for the step cache (MB)
`step_cache_max_entries`	`int`	`200`	Maximum number of cached step entries
`use_cow_snapshots`	`bool`	`True`	Use copy-on-write snapshots for branch features
`log_cache_stats`	`bool`	`True`	Log end-of-run cache statistics at `verbose >= 1`
`log_step_memory`	`bool`	`True`	Log per-step memory stats at `verbose >= 2`
`memory_warning_threshold_mb`	`int`	`3072`	RSS threshold (MB) for memory warnings

Config flow

nirs4all.run(cache=CacheConfig(...))
  -> PipelineRunner.cache_config
    -> PipelineOrchestrator.cache_config
      -> RuntimeContext(cache_config=, step_cache=)
        -> PipelineExecutor (reads from runtime_context)
        -> BranchController (reads use_cow_snapshots)

Block-Based Feature Storage

Location: nirs4all/data/_features/array_storage.py

Problem

The previous implementation stored all processings in a single 3D numpy array of shape (samples, processings, features). Every call to add_processing() allocated a new 3D array and copied all existing data via np.concatenate. Peak memory during the concatenation was ~2x steady-state.

Solution

ArrayStorage now maintains a list of 2D blocks, one per processing:

Before (monolithic):
  _array: np.ndarray  # shape (N, P, F) — one 3D array

After (block-based):
  _blocks: List[np.ndarray]          # P blocks, each (N, F)
  _cached_3d: Optional[np.ndarray]   # lazily computed
  _shared: Optional[SharedBlocks]    # CoW state (Phase 3)

Key operations

Operation	Complexity	Notes
`add_processing(data)`	O(1)	Appends a 2D block, no copy of existing data
`add_processings_batch(blocks)`	O(k)	Appends k blocks, invalidates cache once
`remove_processing(idx)`	O(1)	Drops a block from the list
`.array` (property)	O(NPF) lazy	Builds 3D via `np.stack`, cached until next mutation
`.nbytes` (property)	O(P)	Sum of block `.nbytes` values
`get_data(indices)`	O(NPF)	Indexes into the 3D view

Lazy 3D cache

The .array property builds the 3D array only when accessed and caches it in _cached_3d. Any mutation (add_processing, remove_processing, update_processing, resize_features) invalidates the cache by calling _prepare_for_mutation().

Code paths that only need 2D access (e.g., per-processing iteration) can avoid the 3D build entirely by iterating _blocks directly.

Processing eviction

FeatureSource.evict_oldest_processings(max_processings) removes the oldest non-raw processings (by insertion order, preserving the raw processing at index 0). This prevents unbounded growth during feature augmentation.

# In FeatureSource:
evicted = source.evict_oldest_processings(max_processings=50)

The eviction works with ProcessingManager.remove_processing() to drop the processing ID and reindex, and ArrayStorage.remove_processing() to drop the corresponding 2D block.

Copy-on-Write Branch Snapshots

Location: nirs4all/data/_features/array_storage.py (SharedBlocks), nirs4all/controllers/data/branch.py

Problem

Branch execution deep-copies the entire feature set for each branch. For N duplication branches with a 64 MB feature set, total copies are ~N*128 MB (snapshot + restore copies per branch). Most post-branch steps (model training) only read features without modifying them — the deep-copy is wasted.

Solution: SharedBlocks

SharedBlocks is a reference-counted wrapper around a numpy array:

class SharedBlocks:
    __slots__ = ("_array", "_refcount")

    def acquire(self) -> "SharedBlocks"   # refcount += 1, return self
    def release(self) -> None             # refcount -= 1
    def is_shared(self) -> bool           # refcount > 1
    def detach(self) -> "SharedBlocks"    # copy only if shared

The key insight is detach(): if the array is shared by multiple branches (refcount > 1), it creates an independent copy and decrements the original’s refcount. If only one holder exists, it returns self — no copy.

ArrayStorage CoW integration

ArrayStorage provides three methods for CoW lifecycle:

# 1. Before branching: wrap current data in SharedBlocks
shared = storage.ensure_shared()

# 2. Branch restore: enter shared mode (zero-copy)
storage.restore_from_shared(shared.acquire())

# 3. On mutation: CoW trigger (copy only if shared)
storage._prepare_for_mutation()  # calls _materialize_blocks() if needed

Branch lifecycle

BranchController.execute()
  |
  | 1. _snapshot_features(dataset, use_cow=True)
  |    -> For each source: source._storage.ensure_shared().acquire()
  |    -> Returns lightweight snapshot (SharedBlocks refs + metadata)
  |
  | 2. Per branch:
  |    _restore_features(dataset, snapshot, use_cow=True)
  |      -> source._storage.restore_from_shared(shared.acquire())
  |    Execute branch steps...
  |      -> If step mutates features (transformer):
  |         _prepare_for_mutation() -> _materialize_blocks()
  |           -> detach() copies array only if refcount > 1
  |
  | 3. _release_snapshot(snapshot, use_cow=True)
  |    -> Releases shared references for prompt GC

Memory improvement

For 6 branches with a 64 MB feature set:

Approach	Additional memory
Deep-copy (old)	~768 MB (12 full copies)
CoW (new)	~0 MB if read-only, ~64 MB per branch that mutates

Rollback

Set use_cow_snapshots=False in CacheConfig to revert to deep-copy behavior.

Step-Level Caching

Location: nirs4all/pipeline/execution/step_cache.py, nirs4all/pipeline/execution/executor.py

Problem

Generator sweeps like [{_or_: [SNV, MSC]}, PLS({_range_: [1, 20]})] expand into many variants. Each variant re-executes all preprocessing from scratch, even when variants share a common prefix. For 2 scalers x 20 PLS components = 40 variants, the SNV transform runs 20 times with identical input.

Solution

StepCache stores lightweight feature snapshots (CachedStepState) after each preprocessing step, keyed by (chain_path_hash, data_hash, selector_fingerprint).

When a variant encounters a step that has already been computed with the same input:

Cache hit: Restore the feature state from the snapshot, skip execution.
Cache miss: Execute normally, store the result.

Cache key composition

key = "step:{chain_path_hash}:{data_hash}:{selector_fingerprint}"

Component	What it captures	Why
`chain_path_hash`	MD5 of the step’s serialized config	Identifies the operator and its parameters
`data_hash`	Content hash of the dataset before the step	Ensures input data matches
`selector_fingerprint`	SHA-256 of partition, fold_id, processing, include_augmented, tag_filters, branch_path	Prevents cross-contamination between folds/partitions

CachedStepState

@dataclass
class CachedStepState:
    features_sources: List[Any]       # Deep-copied FeatureSources
    processing_names: List[List[str]] # Processing chain per source
    content_hash: str                 # Post-step hash
    bytes_estimate: int               # numpy nbytes (no pickle)

Only feature sources are deep-copied — targets, metadata, fold assignments, and tags are immutable during preprocessing and are not included. This makes snapshots significantly smaller than full dataset copies.

Which steps are cached

Only controllers that return True from supports_step_cache() are cached:

Controller	Cached	Reason
`TransformerMixinController`	Yes	Preprocessing transforms benefit from cross-variant reuse
`ModelController`	No	Variant-specific — caching defeats the purpose
`SplitterController`	No	Cheap to execute, stateful with fold indices
`BranchController`	No	Complex state, handled by CoW snapshots
`MergeController`	No	Depends on branch results
`FeatureAugmentationController`	No	Produces many processings, full result is large

Eviction

The step cache uses LRU (Least Recently Used) eviction, which is appropriate because early pipeline steps are accessed often during variant expansion, then become cold. The underlying DataCache backend updates timestamps on every access and evicts by oldest timestamp when the budget is exceeded.

Budget is enforced by both entry count (step_cache_max_entries) and total bytes (step_cache_max_mb).

Execution flow

PipelineExecutor._execute_single_step():
  |
  | step_cache = runtime_context.step_cache
  | if step_cache and mode == "train" and _is_step_cacheable(step):
  |   |
  |   | 1. Compute cache key
  |   | 2. step_cache.get(key) -> CachedStepState or None
  |   |
  |   | HIT:
  |   |   step_cache.restore(state, dataset)
  |   |   return  # Skip execution
  |   |
  |   | MISS:
  |   |   Execute step normally
  |   |   step_cache.put(key, dataset)

Observability

Per-step memory logging

At verbose >= 2, each step logs its steady-state memory footprint and process RSS:

[Step 3/8] SNV | shape: 5000x4x2000 | steady: 152.6 MB | RSS: 1.2 GB

At any verbosity level, a warning fires when RSS exceeds memory_warning_threshold_mb:

Warning: Process RSS at 3.8 GB (threshold: 3.0 GB) after step SNV

Cache statistics

At verbose >= 1, after each dataset execution completes (when log_cache_stats=True):

Step cache: 45 hits / 12 misses (78.9% hit rate) | 312.4 MB peak | 8 evictions

Available via step_cache.stats():

{
    "hits": 45,
    "misses": 12,
    "hit_rate": 0.789,
    "evictions": 8,
    "peak_mb": 312.4,
    "entries": 10,
    "max_entries": 200,
    "size_mb": 280.5,
    "max_size_mb": 2048.0,
}

Memory estimation utilities

Location: nirs4all/utils/memory.py

Function	Description
`estimate_dataset_bytes(dataset)`	Sum of numpy `nbytes` across all feature sources and targets
`estimate_cache_entry_bytes(entry)`	numpy `nbytes` if available, `sys.getsizeof` fallback (no pickle)
`get_process_rss_mb()`	Process RSS via `/proc/self/statm` (Linux) or `psutil` (macOS)
`format_bytes(n)`	Human-readable: `"152.6 MB"`, `"1.2 GB"`

All estimation functions are O(1) per array (read shape and dtype only). No serialization or pickle calls.

Architecture Diagram

nirs4all.run(cache=CacheConfig(...))
  |
  v
PipelineRunner -------- cache_config ---------> PipelineOrchestrator
                                                  |
                                                  | Creates StepCache (if enabled)
                                                  | Creates RuntimeContext
                                                  v
                                              PipelineExecutor
                                                  |
                               +------------------+------------------+
                               |                                     |
                         Step cacheable?                        Not cacheable
                               |                                     |
                        cache.get(key)                         Execute normally
                          /        \
                       HIT          MISS
                        |             |
                   Restore from    Execute step
                   CachedStepState    |
                        |          cache.put(key)
                        v             v
                     Next step     Next step


Branch execution with CoW:

  BranchController.execute()
    |
    |  ensure_shared() + acquire()     # O(1) snapshot
    |
    +-- Branch 1: restore_from_shared()
    |     |-> Step modifies features? -> detach() -> copy
    |     |-> Step reads only?        -> no copy
    |
    +-- Branch 2: restore_from_shared()
    |     |-> ...
    |
    |  release_snapshot()              # Free shared refs

Performance Tuning

Step cache budget

The default 2 GB / 200 entries is suitable for most workloads. For large datasets with many generator variants, increase the budget:

cache = CacheConfig(
    step_cache_enabled=True,
    step_cache_max_mb=4096,       # 4 GB
    step_cache_max_entries=500,
)

Monitor eviction counts via the end-of-run stats. If evictions are high, the budget is too small and cache effectiveness is reduced.

Memory warning threshold

The default 3 GB threshold may need adjustment for machines with more or less RAM:

cache = CacheConfig(
    memory_warning_threshold_mb=8192,  # 8 GB for large-memory machines
)

Disabling features for debugging

Each layer can be independently disabled:

# Disable step caching (default)
CacheConfig(step_cache_enabled=False)

# Disable CoW snapshots (fall back to deep-copy)
CacheConfig(use_cow_snapshots=False)

# Disable observability logging
CacheConfig(log_cache_stats=False, log_step_memory=False)

Example

See examples/developer/06_internals/D03_cache_performance.py for a runnable side-by-side comparison of cache configurations measuring wall-clock time, RSS delta, and cache hit statistics.

cd examples
./run.sh -n "D03*"

Source Files

File	Contents
`nirs4all/config/cache_config.py`	`CacheConfig` dataclass
`nirs4all/data/_features/array_storage.py`	`ArrayStorage` (block-based), `SharedBlocks` (CoW)
`nirs4all/data/_features/feature_source.py`	`FeatureSource.evict_oldest_processings()`
`nirs4all/data/_features/processing_manager.py`	`ProcessingManager.remove_processing()`
`nirs4all/data/performance/cache.py`	`DataCache` (LRU/LFU backend), `CacheEntry`
`nirs4all/pipeline/execution/step_cache.py`	`StepCache`, `CachedStepState`
`nirs4all/pipeline/execution/executor.py`	Cache wiring in `_execute_single_step()`
`nirs4all/pipeline/execution/orchestrator.py`	`StepCache` creation and stats logging
`nirs4all/pipeline/config/context.py`	`RuntimeContext` carries cache state
`nirs4all/controllers/controller.py`	`supports_step_cache()` base method
`nirs4all/controllers/transforms/transformer.py`	`supports_step_cache() -> True`
`nirs4all/controllers/data/branch.py`	CoW snapshot/restore methods
`nirs4all/utils/memory.py`	Memory estimation utilities