don't know but changed predict_parquet

Author: EliHei2

src/segger/prediction/predict_parquet.py

@@ -26,7 +26,7 @@
 from segger.prediction.boundary import generate_boundaries
 
 from scipy.sparse.csgraph import connected_components as cc
-from typing import Union, Dict
+from typing import Union, Dict, Tuple
 import dask.dataframe as dd
 from dask import delayed
 from dask.diagnostics import ProgressBar
@@ -287,14 +287,19 @@ def sparse_multiply(embeddings, edge_index, shape) -> coo_matrix:
             # shape = batch[from_type].x.shape[0], batch[to_type].x.shape[0]
             indices = torch.argwhere(edge_index != -1).T
             indices[1] = edge_index[edge_index != -1]
-            rows = cp.fromDlpack(to_dlpack(indices[0, :].to("cuda")))
-            columns = cp.fromDlpack(to_dlpack(indices[1, :].to("cuda")))
+            indices_gpu = indices.to("cuda")  # Keep reference
+            rows = cp.fromDlpack(to_dlpack(indices_gpu[0, :]))
+            columns = cp.fromDlpack(to_dlpack(indices_gpu[1, :]))
+            del indices_gpu  # Delete only after CuPy arrays exist
+            stream = cp.cuda.get_current_stream()
+            stream.synchronize()  # <-- ADD THIS
             # print(rows)
             del indices
             values = similarity[edge_index != -1].flatten()
             sparse_result = coo_matrix(
                 (cp.fromDlpack(to_dlpack(values)), (rows, columns)), shape=shape
             )
+            stream.synchronize()
             return sparse_result
             # Free GPU memory after computation
 
@@ -364,21 +369,23 @@ def _get_id():
             # Convert sparse matrix to dense format (on GPU)
             dense_scores = scores.toarray()  # Convert to dense NumPy array
             del scores  # Remove from memory
-            cp.get_default_memory_pool().free_all_blocks()  # Free CuPy memory
+            cp.cuda.Stream.null.synchronize()
+            # cp.get_default_memory_pool().free_all_blocks()  # Free CuPy memory
 
             # Step 2: Maximize score and assign transcripts based on score threshold
             belongs = cp.max(dense_scores, axis=1)  # Max score per transcript
             assignments["score"] = cp.asnumpy(belongs)  # Move back to CPU
 
             mask = assignments["score"] >= score_cut  # Mask for assigned transcripts
+            cp.cuda.Stream.null.synchronize()
             all_ids = np.concatenate(batch["bd"].id)  # Boundary IDs as NumPy array
             assignments["segger_cell_id"] = np.where(
                 mask, all_ids[cp.argmax(dense_scores, axis=1).get()], None
             )
 
             # Clear memory after score processing
             del dense_scores
-            cp.get_default_memory_pool().free_all_blocks()  # Free CuPy memory
+            # cp.get_default_memory_pool().free_all_blocks()  # Free CuPy memory
             torch.cuda.empty_cache()
 
             assignments["bound"] = np.where(
@@ -470,7 +477,7 @@ def _get_id():
             delayed_write_output_ddf.persist()  # Schedule writing
 
             # Free memory after computation
-            cp.get_default_memory_pool().free_all_blocks()  # Free CuPy memory
+            # cp.get_default_memory_pool().free_all_blocks()  # Free CuPy memory
             torch.cuda.empty_cache()