From a1df5a7d2bc852722327c42a8c320566a5daf861 Mon Sep 17 00:00:00 2001 From: Net Zhang Date: Tue, 6 Oct 2026 10:41:42 -0400 Subject: [PATCH 1/4] Pick sklearn or cuML per method & dataset size under auto mode resolve_backend("auto") resolved to cuML whenever a GPU was present, which made the size heuristics inside reduce_dim and run_kmeans unreachable from the apps and sent small datasets through cuML start-up for no gain. Replace both with one shared rule, CUML_AUTO_RANGE: cuML for PCA from 500 samples and KMeans from 1000, UMAP only from 10000 (subprocess start-up has to be amortized), and t-SNE only below 3000 (the collapse-safe exact solver is O(N^2)). The sidebars now pass the dataset size and method; explicit backend choices are honored unchanged. Co-Authored-By: Claude Fable 5.1 --- README.md | 2 +- apps/embed_explore/components/sidebar.py | 4 +- apps/precalculated/components/sidebar.py | 4 +- shared/components/clustering_controls.py | 12 ++-- shared/utils/backend.py | 73 ++++++++++++++++++++---- shared/utils/clustering.py | 8 +-- 6 files changed, 79 insertions(+), 24 deletions(-) diff --git a/README.md b/README.md index 2cff7ad..4ee17e9 100644 --- a/README.md +++ b/README.md @@ -55,7 +55,7 @@ uv pip install -e ".[gpu-cu12]" uv pip install -e ".[gpu-cu13]" ``` -The app auto-detects GPU availability at runtime and falls back to CPU if anything goes wrong — no configuration needed. The CPU sklearn path is auto-accelerated by [scikit-learn-intelex](https://github.com/uxlfoundation/scikit-learn-intelex)[^1]. You can also manually select backends (`cuML`, `sklearn`) in the sidebar. +The app auto-detects GPU availability at runtime and falls back to CPU if anything goes wrong; no configuration needed. With the default `auto` backend, the faster library is chosen per method and dataset size from measured crossovers, even when a GPU is present: cuML for PCA from 500 samples and KMeans from 1000; UMAP only from 10000 samples; t-SNE only below 3000 samples (the collapse-safe exact solver is O(N^2)). **Select `cuml` explicitly to force the GPU**. The CPU sklearn path is auto-accelerated by [scikit-learn-intelex](https://github.com/uxlfoundation/scikit-learn-intelex)[^1]. You can also manually select backends (`cuML`, `sklearn`) in the sidebar. ## Usage diff --git a/apps/embed_explore/components/sidebar.py b/apps/embed_explore/components/sidebar.py index 511d88b..5735205 100644 --- a/apps/embed_explore/components/sidebar.py +++ b/apps/embed_explore/components/sidebar.py @@ -147,7 +147,7 @@ def _run_projection(embeddings, valid_paths, reduction_method, dim_reduction_bac """Run dim reduction and create the 2D scatter plot dataframe.""" try: cuda_available, device_info = check_cuda_available() - actual_backend = resolve_backend(dim_reduction_backend, "reduction") + actual_backend = resolve_backend(dim_reduction_backend, "reduction", n_samples=embeddings.shape[0], method=reduction_method) logger.info("=" * 60) logger.info("PROJECTION START") @@ -210,7 +210,7 @@ def _run_projection(embeddings, valid_paths, reduction_method, dim_reduction_bac def _run_kmeans(embeddings, n_clusters, clustering_backend, n_workers, seed): """Run KMeans on already-extracted embeddings and add labels to df_plot.""" try: - actual_backend = resolve_backend(clustering_backend, "clustering") + actual_backend = resolve_backend(clustering_backend, "clustering", n_samples=embeddings.shape[0]) logger.info(f"KMeans: k={n_clusters}, backend={actual_backend}") with st.spinner(f"Running KMeans (k={n_clusters})..."): diff --git a/apps/precalculated/components/sidebar.py b/apps/precalculated/components/sidebar.py index 162f6cf..a4a2a44 100644 --- a/apps/precalculated/components/sidebar.py +++ b/apps/precalculated/components/sidebar.py @@ -615,7 +615,7 @@ def _run_projection(filtered_df, reduction_method, dim_reduction_backend, seed): """Run dim reduction and create the 2D scatter plot dataframe.""" try: cuda_available, device_info = check_cuda_available() - actual_backend = resolve_backend(dim_reduction_backend, "reduction") + actual_backend = resolve_backend(dim_reduction_backend, "reduction", n_samples=len(filtered_df), method=reduction_method) logger.info("=" * 60) logger.info("PROJECTION START") @@ -678,7 +678,7 @@ def _run_projection(filtered_df, reduction_method, dim_reduction_backend, seed): def _run_kmeans(embeddings, n_clusters, clustering_backend, n_workers, seed): """Run KMeans on already-extracted embeddings and add labels to df_plot.""" try: - actual_backend = resolve_backend(clustering_backend, "clustering") + actual_backend = resolve_backend(clustering_backend, "clustering", n_samples=embeddings.shape[0]) logger.info(f"KMeans: k={n_clusters}, backend={actual_backend}") with st.spinner(f"Running KMeans (k={n_clusters})..."): diff --git a/shared/components/clustering_controls.py b/shared/components/clustering_controls.py index 34d3e1c..99243d9 100644 --- a/shared/components/clustering_controls.py +++ b/shared/components/clustering_controls.py @@ -52,7 +52,8 @@ def render_clustering_backend_controls(): "Dimensionality Reduction Backend", options=dim_reduction_options, index=0, - help="Backend for PCA/t-SNE/UMAP computation" + help="Backend for PCA/t-SNE/UMAP. auto picks the faster library per method and " + "dataset size from measured crossovers; choose cuml to force the GPU." ) with col2: @@ -60,7 +61,8 @@ def render_clustering_backend_controls(): "Clustering Backend", options=clustering_options, index=0, - help="Backend for K-means clustering computation" + help="Backend for K-means. auto: sklearn below 1000 samples, cuML on GPU from " + "there (measured crossover); choose cuml to force the GPU." ) # Performance and reproducibility settings @@ -100,7 +102,8 @@ def render_projection_controls(): dim_reduction_backend = st.selectbox( "Dim Reduction Backend", options=dim_reduction_options, index=0, - help="Backend for PCA/t-SNE/UMAP computation", + help="Backend for PCA/t-SNE/UMAP. auto picks the faster library per method and " + "dataset size from measured crossovers; choose cuml to force the GPU.", key="proj_backend" ) @@ -130,7 +133,8 @@ def render_kmeans_controls(): clustering_backend = st.selectbox( "Clustering Backend", options=clustering_options, index=0, - help="Backend for K-means computation", + help="Backend for K-means. auto: sklearn below 1000 samples, cuML on GPU from " + "there (measured crossover); choose cuml to force the GPU.", key="km_backend" ) n_workers = st.number_input( diff --git a/shared/utils/backend.py b/shared/utils/backend.py index 7af37a3..8d5c65f 100644 --- a/shared/utils/backend.py +++ b/shared/utils/backend.py @@ -29,6 +29,39 @@ # Cache CUDA availability to avoid repeated checks _cuda_check_cache: Optional[Tuple[bool, str]] = None +# Dataset-size range (samples) in which "auto" picks cuML over sklearn, per +# method. Chosen from measured crossovers of the app's own reduce_dim / +# run_kmeans paths (scripts/bench_backend_threshold.py); the exact numbers +# depend on the GPU, CPU count and library versions, so treat them as +# defaults to re-measure, not constants. The shape of the rule is stable: +# - PCA, KMEANS: cheap either way; cuML wins once transfer cost is amortized. +# - UMAP: cuML runs in a subprocess, so a fixed start-up cost must be +# amortized over a large enough dataset. +# - TSNE: the app forces cuML's exact O(N^2) solver (#40), so cuML only +# wins on small data and loses increasingly above the bound. +# Entries are (min_samples, max_samples); None is unbounded. An explicit +# "cuml" request is always honored. reduce_dim / run_kmeans use this same +# rule so the layers agree (#50). +CUML_AUTO_RANGE = { + "PCA": (500, None), + "TSNE": (None, 3000), + "UMAP": (10000, None), + "KMEANS": (1000, None), +} + + +def auto_prefers_cuml(method: str, n_samples: Optional[int]) -> bool: + """Whether "auto" should pick cuML for ``method`` at this dataset size, + hardware aside. Unknown size or method keeps the hardware-only rule.""" + if n_samples is None: + return True + lo, hi = CUML_AUTO_RANGE.get(method.upper(), (None, None)) + if lo is not None and n_samples < lo: + return False + if hi is not None and n_samples >= hi: + return False + return True + def check_cuda_available() -> Tuple[bool, str]: """ @@ -83,16 +116,27 @@ def check_cuml_available() -> bool: return False -def resolve_backend(backend: str, operation: str = "general") -> str: +def resolve_backend( + backend: str, + operation: str = "general", + n_samples: Optional[int] = None, + method: Optional[str] = None, +) -> str: """ - Resolve 'auto' backend to actual backend based on available hardware. + Resolve 'auto' backend to actual backend based on hardware, method and data size. Args: backend: Requested backend ("auto", "sklearn", "cuml") - operation: Operation type for logging ("clustering", "reduction", "general") + operation: Operation type ("clustering", "reduction", "general") + n_samples: Dataset size. With "auto", sizes outside the method's + CUML_AUTO_RANGE resolve to sklearn even when a GPU is available. + None keeps the hardware-only rule. + method: "PCA" / "TSNE" / "UMAP" for reduction; defaults to "KMEANS" + for clustering. Needed for the size rule. Returns: - Resolved backend name. CPU paths always go through sklearn. + Resolved backend name. CPU paths always go through sklearn; an + explicit "cuml" or "sklearn" request is returned unchanged. """ if backend != "auto": logger.debug(f"Using explicitly requested backend: {backend}") @@ -100,14 +144,21 @@ def resolve_backend(backend: str, operation: str = "general") -> str: cuda_available, device_info = check_cuda_available() # Only probe for cuML when CUDA is actually available. - if cuda_available and check_cuml_available(): - resolved = "cuml" - logger.info(f"Auto-resolved {operation} backend to cuML (GPU: {device_info})") - else: - resolved = "sklearn" + if not (cuda_available and check_cuml_available()): logger.info(f"Auto-resolved {operation} backend to sklearn (CPU)") - - return resolved + return "sklearn" + + method_key = (method or ("KMEANS" if operation == "clustering" else "")).upper() + if n_samples is not None and method_key and not auto_prefers_cuml(method_key, n_samples): + lo, hi = CUML_AUTO_RANGE[method_key] + where = f"below {lo}" if lo is not None and n_samples < lo else f"at or above {hi}" + logger.info(f"Auto-resolved {operation} backend to sklearn for {method_key}: {n_samples} " + f"samples is {where}, where sklearn is faster (GPU {device_info} available; " + f"choose 'cuml' explicitly to force it)") + return "sklearn" + + logger.info(f"Auto-resolved {operation} backend to cuML (GPU: {device_info})") + return "cuml" def get_backend_info() -> dict: diff --git a/shared/utils/clustering.py b/shared/utils/clustering.py index 4115ef5..69af649 100644 --- a/shared/utils/clustering.py +++ b/shared/utils/clustering.py @@ -8,7 +8,7 @@ from shared.utils.logging_config import get_logger from shared.utils.backend import ( - HAS_CUML_PACKAGE, HAS_CUPY_PACKAGE, + HAS_CUML_PACKAGE, HAS_CUPY_PACKAGE, auto_prefers_cuml, check_cuda_available, check_cuml_available, ) @@ -172,8 +172,8 @@ def reduce_dim(embeddings: np.ndarray, method: str = "PCA", seed: Optional[int] use_cuml = False if backend == "cuml" and HAS_CUML and cuda_available: use_cuml = True - elif backend == "auto" and HAS_CUML and cuda_available and n_samples > 5000: - # Use cuML automatically for large datasets on GPU + elif backend == "auto" and HAS_CUML and cuda_available and auto_prefers_cuml(method, n_samples): + # Same measured per-method size rule as resolve_backend use_cuml = True start_time = time.time() @@ -376,7 +376,7 @@ def run_kmeans(embeddings: np.ndarray, n_clusters: int, seed: Optional[int] = No if backend == "cuml" and HAS_CUML and cuda_available: logger.info("Using cuML backend for KMeans") result = _run_kmeans_cuml(embeddings, n_clusters, seed, n_workers) - elif backend == "auto" and HAS_CUML and cuda_available and n_samples > 500: + elif backend == "auto" and HAS_CUML and cuda_available and auto_prefers_cuml("KMEANS", n_samples): logger.info("Auto-selected cuML backend for KMeans (GPU available, large dataset)") result = _run_kmeans_cuml(embeddings, n_clusters, seed, n_workers) else: From cf71f614c67dc984ab13bd63f61e52eb85e6708a Mon Sep 17 00:00:00 2001 From: Net Zhang Date: Tue, 6 Oct 2026 11:01:21 -0400 Subject: [PATCH 2/4] Report failure instead of success in the progress context Co-Authored-By: Claude Fable 5.1 --- apps/embed_explore/components/sidebar.py | 8 +-- shared/lib/progress.py | 65 ++++++++++++++++++------ 2 files changed, 53 insertions(+), 20 deletions(-) diff --git a/apps/embed_explore/components/sidebar.py b/apps/embed_explore/components/sidebar.py index 5735205..499ceec 100644 --- a/apps/embed_explore/components/sidebar.py +++ b/apps/embed_explore/components/sidebar.py @@ -67,7 +67,7 @@ def render_embedding_section() -> Tuple[bool, Optional[str], Optional[str], int, ) if embeddings.shape[0] == 0: - st.error("No valid image embeddings found.") + progress.fail("No images found in this folder. Check the path.") logger.warning("Embedding generation returned 0 embeddings") st.session_state.embeddings = None st.session_state.valid_paths = None @@ -88,7 +88,7 @@ def render_embedding_section() -> Tuple[bool, Optional[str], Optional[str], int, st.session_state.selected_image_idx = None except Exception as e: - st.error(f"Error during embedding: {e}") + progress.fail(f"Error during embedding: {e}") logger.exception("Embedding generation failed") elif embed_button: @@ -327,7 +327,7 @@ def render_save_section(): ) st.info(f"Summary CSV saved at {csv_path}") except Exception as e: - save_status_placeholder.error(f"Error saving images: {e}") + progress.fail(f"Error saving images: {e}") elif save_cluster_button: save_status_placeholder.warning("Please select at least one cluster.") @@ -363,7 +363,7 @@ def render_save_section(): ) st.info(f"Summary CSV saved at {csv_path}") except Exception as e: - repartition_status_placeholder.error(f"Error repartitioning images: {e}") + progress.fail(f"Error repartitioning images: {e}") def render_clustering_sidebar(): diff --git a/shared/lib/progress.py b/shared/lib/progress.py index 3c49bd9..8f3757d 100644 --- a/shared/lib/progress.py +++ b/shared/lib/progress.py @@ -3,15 +3,30 @@ """ from abc import ABC, abstractmethod -from typing import Optional, Callable +from typing import Optional import streamlit as st class ProgressContext(ABC): - """Base class for different progress UI patterns""" - + """Base class for different progress UI patterns. + + Entering the context returns an object that is callable as + ``progress(fraction, text)`` and also offers ``progress.fail(message)`` + for a step that finishes without raising but did not succeed (for + example "no images found"), so the context reports an error instead of + its success message. + """ + + @abstractmethod + def __enter__(self) -> "ProgressContext": + pass + + @abstractmethod + def __call__(self, progress: float, text: str) -> None: + pass + @abstractmethod - def __enter__(self) -> Callable[[float, str], None]: + def fail(self, message: str) -> None: pass @abstractmethod @@ -26,23 +41,34 @@ def __init__(self, placeholder, success_message: Optional[str] = None): self.placeholder = placeholder self.success_message = success_message self.progress_bar = None - + self.failure_message: Optional[str] = None + def __enter__(self): self.progress_bar = self.placeholder.progress(0, text="Starting...") - return self.update_progress - + return self + + def __call__(self, progress: float, text: str) -> None: + self.update_progress(progress, text) + def update_progress(self, progress: float, text: str): if self.progress_bar: self.progress_bar.progress(progress, text=text) - + + def fail(self, message: str) -> None: + """Mark the step as failed: on exit the placeholder shows this error + instead of the success message.""" + self.failure_message = message + def __exit__(self, exc_type, exc_val, exc_tb): if self.progress_bar: self.progress_bar.empty() - - if exc_type is None and self.success_message: - self.placeholder.success(self.success_message) - elif exc_type is not None: + + if exc_type is not None: self.placeholder.error(f"Error: {exc_val}") + elif self.failure_message: + self.placeholder.error(self.failure_message) + elif self.success_message: + self.placeholder.success(self.success_message) class MockProgressContext(ProgressContext): @@ -50,12 +76,19 @@ class MockProgressContext(ProgressContext): def __init__(self): self.updates = [] - + self.failure_message: Optional[str] = None + def __enter__(self): - return self.capture_progress - + return self + + def __call__(self, progress: float, text: str) -> None: + self.capture_progress(progress, text) + def capture_progress(self, progress: float, text: str): self.updates.append((progress, text)) - + + def fail(self, message: str) -> None: + self.failure_message = message + def __exit__(self, *args): pass From 87db0009a358015f7bc9e31e32539db33d3f98ef Mon Sep 17 00:00:00 2001 From: Net Zhang Date: Tue, 6 Oct 2026 13:59:04 -0400 Subject: [PATCH 3/4] Let the Embed section include subfolders EmbeddingService.generate_embeddings already supported recursive listing, but the sidebar never passed it, so a folder organized b, for example, species subfolders reported 0 images. Add an Include subfolders checkbox (default on) and pass it through. --- apps/embed_explore/components/sidebar.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/apps/embed_explore/components/sidebar.py b/apps/embed_explore/components/sidebar.py index 499ceec..ca318be 100644 --- a/apps/embed_explore/components/sidebar.py +++ b/apps/embed_explore/components/sidebar.py @@ -55,6 +55,12 @@ def render_embedding_section() -> Tuple[bool, Optional[str], Optional[str], int, value=32, step=1 ) + recursive = st.checkbox( + "Include subfolders", + value=True, + key="embed_recursive", + help="Also embed images in nested folders (e.g. one folder per species).", + ) embed_button = st.button("Run Embedding") # Handle embedding execution @@ -63,7 +69,7 @@ def render_embedding_section() -> Tuple[bool, Optional[str], Optional[str], int, try: embeddings, valid_paths = EmbeddingService.generate_embeddings( image_dir, model_name, batch_size, n_workers, - progress_callback=progress + progress_callback=progress, recursive=recursive, ) if embeddings.shape[0] == 0: From bdb71e1f03aa5baf82d20c800a147a20b0f2ac17 Mon Sep 17 00:00:00 2001 From: Net Zhang Date: Tue, 6 Oct 2026 14:15:02 -0400 Subject: [PATCH 4/4] Return before loading the model when no images are found --- shared/services/embedding_service.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/shared/services/embedding_service.py b/shared/services/embedding_service.py index 0eb2399..2782dd5 100644 --- a/shared/services/embedding_service.py +++ b/shared/services/embedding_service.py @@ -173,6 +173,13 @@ def generate_embeddings( total = len(image_paths) logger.info(f"Found {total} images in {image_dir}") + if total == 0: + # Nothing to embed: return before paying for the model load (#50). + logger.warning(f"No images found in {image_dir} (recursive={recursive})") + if progress_callback: + progress_callback(1.0, "No images found") + return np.empty((0, 0), dtype=np.float32), [] + if progress_callback: progress_callback(0.05, f"Found {total} images. Loading model...")