Skip to content

BACKEND_PIPELINE.md

Latest commit

 

History

History
132 lines (96 loc) · 5.69 KB

BACKEND_PIPELINE.md

File metadata and controls

132 lines (96 loc) · 5.69 KB
Raw

Backend Pipeline

A quick walkthrough of what happens to your embeddings from the moment you click "Run Clustering" to the scatter plot on screen.

The Pipeline at a Glance

Raw Embeddings (from parquet or model)
  │
  ├─ Validate: check for NaN/Inf, cast to float32
  ├─ L2 Normalize: project onto unit hypersphere
  │
  ├─► Step 1: KMeans Clustering (high-dimensional)
  │     Backend: cuML (GPU) → sklearn (CPU, auto-accelerated by `sklearn-intelex`)
  │
  ├─► Step 2: Dimensionality Reduction to 2D
  │     Method:  PCA / t-SNE / UMAP
  │     Backend: cuML (GPU) → sklearn (CPU, auto-accelerated by `sklearn-intelex` for PCA/TSNE)
  │
  └─► Scatter Plot (Altair)
        Color = cluster, position = 2D projection

Note that sklearn-intelex acceleration is used for CPU operations where available1.

Step 0: Embedding Preparation

Before any computation, every embedding goes through _prepare_embeddings():

  1. Cast to float32 — GPU backends require it; keeps memory predictable.
  2. NaN/Inf check — replaces bad values with 0 and logs a warning.
  3. L2 normalization — divides each vector by its magnitude so every point sits on the unit hypersphere. This is critical for two reasons:
    • Prevents cuML UMAP's NN-descent from crashing with SIGFPE on large-magnitude vectors (see investigation/cuml_umap_sigfpe/).
    • Appropriate for contrastive embeddings (CLIP, BioCLIP) whose training objective is cosine-similarity based — magnitude isn't a learned signal.

Input norms are logged so you can always verify what came in.

Step 1: KMeans Clustering

Clusters the full high-dimensional embeddings (e.g., 768-d for BioCLIP 2). Runs before dimensionality reduction so clusters are based on the full feature space, not a lossy 2D projection.

Backend When It's Used How It Works
cuML GPU available + >500 samples GPU-accelerated KMeans via RAPIDS. Runs on CuPy arrays. Falls back to sklearn on any error.
sklearn CPU path (default on machines without a GPU) Standard scikit-learn KMeans, auto-accelerated by scikit-learn-intelex (Intel oneDAL) when installed1 — typically 10–17× faster than vanilla sklearn on CPU. Disable with EMB_EXPLORER_DISABLE_SKLEARNEX=1.

Auto-selection priority: cuML > sklearn. You can override in the sidebar.

Step 2: Dimensionality Reduction

Projects embeddings from high-dimensional space down to 2D for visualization. This is purely for the scatter plot — clustering uses the full-dimensional data.

PCA (Principal Component Analysis)

The fastest option. Linear projection onto the two directions of maximum variance. Good for getting a quick overview; doesn't capture nonlinear structure.

Backend Notes
cuML GPU-accelerated, near-instant even on large datasets
sklearn CPU-based, still fast since PCA is O(n)

t-SNE

Nonlinear method that preserves local neighborhoods. Good at revealing clusters but slow on large datasets. Perplexity is auto-adjusted based on sample size.

Backend Notes
cuML GPU-accelerated, handles thousands of samples well
sklearn CPU-based, can be slow above ~5k samples

UMAP

The recommended default. Nonlinear like t-SNE but faster and better at preserving global structure. Neighbor count is auto-adjusted.

Backend Notes
cuML Runs in an isolated subprocess so a crash doesn't kill the app. The subprocess verifies L2 normalization as a safety net. Falls back to sklearn on failure.
sklearn CPU-based umap-learn. Slower but numerically stable.

Why the subprocess? cuML UMAP's NN-descent algorithm can occasionally trigger a SIGFPE (floating-point exception) that kills the process instantly — no Python try/except can catch it. The subprocess isolates this risk.

Backend Selection

When you select "auto" (the default), the app picks the fastest available backend:

Operation Auto Logic
KMeans cuML if GPU + >500 samples, else sklearn (auto-accelerated by sklearn-intelex when installed1)
Dim. Reduction cuML if GPU + >5000 samples, else sklearn (auto-accelerated by sklearn-intelex for PCA / t-SNE)

Any GPU error (architecture mismatch, missing libraries, out of memory (OOM)) triggers an automatic retry with sklearn. OOM errors are surfaced to the user with guidance.

Logging

Every step is logged to logs/emb_explorer.log (DEBUG level) and console (INFO):

  • Embedding extraction: shape, dtype
  • Preparation: input norms (min/max/mean), non-finite count, L2 normalization
  • Backend selection: which backend was chosen and why
  • KMeans: cluster count, sample count, elapsed time
  • Reduction: method, sample count, elapsed time
  • Fallbacks: what failed and what we fell back to
  • Visualization: point selection events, density mode changes

Check the log file for the full picture when debugging.

GPU Fallback Chain

cuML (GPU)
  │ error?
  ▼
sklearn (CPU, always works)

The app is designed to always produce a result. GPU acceleration is a nice-to-have, never a hard requirement.

Footnotes

  1. sklearn-intelex is powered by the oneDAL library that provides accelerations on x86_64 Linux and Windows machines, and silently fall back to vanilla sklearn on unsupported architectures like Apple Silicon and ARM Linux. The package is under the UXL Foundation (a Linux Foundation project) so cross-vendor support is a stated goal. ↩ ↩2 ↩3