assume cupy for sam when returning embeddings in umap

[eordentlich]

No description provided.

[greptile-apps]

Greptile Summary

This PR fixes a bug where UMAP's embedding-return logic didn't account for CUDA System Allocator Mode (SAM, cuda_system_mem_enabled). Without the fix, SAM mode would fall into the code path that calls list(embedding[...]) directly on cupy arrays instead of first calling .get() to move data to host, causing incorrect behaviour or runtime errors.

• Adds or cuda_system_mem_enabled to both the sparse-output and dense-output branches so the .get() conversion path is taken whenever either GPU memory mode is active.
• The two changed conditions now mirror the guard already used for the dense-input preparation at line 1203, making all three sites consistent.

Confidence Score: 5/5

Safe to merge — the two-line change closes a gap where SAM mode would attempt to iterate cupy arrays without calling .get(), and the fix is consistent with the existing guard used for input preparation.

The change is minimal and targeted: both output branches (sparse and dense) now match the guard already applied to the dense-input preparation path at line 1203. There are no new dependencies, no schema changes, and the fix directly resolves the mismatch between SAM array types and the serialization code.

No files require special attention; the only changed file is python/src/spark_rapids_ml/umap.py and the modifications are straightforward.

Important Files Changed

Filename	Overview
python/src/spark_rapids_ml/umap.py	Two conditions updated to include cuda_system_mem_enabled, ensuring cupy .get() is called when SAM is active; fix is symmetric and consistent with the existing input-preparation guard at line 1203.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[cuml_fit_func returns embedding & raw_data as cupy arrays] --> B{use_sparse_array?}
    B -- Yes --> C[slice csr_chunk, extract indices/indptr/data]
    B -- No --> D{cuda_managed_mem_enabled\nor cuda_system_mem_enabled?}
    C --> E{cuda_managed_mem_enabled\nor cuda_system_mem_enabled?\n★ fixed}
    E -- Yes --> F["call .get() on embedding, indices, indptr, data\nyield pd.DataFrame"]
    E -- No --> G["iterate cupy arrays directly\nyield pd.DataFrame"]
    D -- Yes --> H["call .get() on embedding & raw_data\nyield pd.DataFrame"]
    D -- No --> I["iterate cupy arrays directly\nyield pd.DataFrame"]

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
    A[cuml_fit_func returns embedding & raw_data as cupy arrays] --> B{use_sparse_array?}
    B -- Yes --> C[slice csr_chunk, extract indices/indptr/data]
    B -- No --> D{cuda_managed_mem_enabled\nor cuda_system_mem_enabled?}
    C --> E{cuda_managed_mem_enabled\nor cuda_system_mem_enabled?\n★ fixed}
    E -- Yes --> F["call .get() on embedding, indices, indptr, data\nyield pd.DataFrame"]
    E -- No --> G["iterate cupy arrays directly\nyield pd.DataFrame"]
    D -- Yes --> H["call .get() on embedding & raw_data\nyield pd.DataFrame"]
    D -- No --> I["iterate cupy arrays directly\nyield pd.DataFrame"]

Loading

_{Reviews (1): Last reviewed commit: "assume cupy for sam when returning embed..." | Re-trigger Greptile}

[eordentlich]

build