CUBLAS_STATUS_NOT_INITIALIZED when importing Datashader *before* running tsnecuda

[rafaelmessias]

I'm using datashader to visualize tsnecuda's results and came across a weird behavior. If I import datashader at the top of the file, before running tsnecuda, I get the following error when calling fit_transform:

Initializing cuda handles... CUBLAS error in file '/home/conda/feedstock_root/build_artifacts/tsnecuda_1638486315160/work/src/util/cuda_utils.cu', line 101, error: CUBLAS_STATUS_NOT_INITIALIZED
terminating!
python: /home/conda/feedstock_root/build_artifacts/tsnecuda_1638486315160/work/src/util/cuda_utils.cu:103: void __CublasSafeCall(cublasStatus_t, const char*, int): Assertion `0' failed.
Aborted

However, if I run tsnecuda first, and then import datashader afterwards (just before creating the visualization), it works fine!

This is on Windows + WSL2 + Ubuntu 22.04, CUDA 12.1, tsnecuda 3.0.0 (from conda-forge, not the latest version due to libfaiss bug), datashader 0.14.4 (also from conda-forge).

This is not critical since it works, but it's still weird and I guess it would be good to know the cause of the problem so we can import datashader normally at the top. I haven't had the time to look into the datashader code to check if it's doing something strange, so any ideas on what could be the problem here are appreciated. Datashader is a popular way to visualize large scale scatterplots so it fits nicely with tsnecuda's applications.

Simple example for reproduction (also uses Holoviews Panel so it runs outside Jupyter Notebooks):

import pandas as pd
# Does not work
import datashader as ds
from tsnecuda import TSNE
import numpy as np
import panel as pn

# create sample data
X = np.random.rand(100000, 10)
df = pd.DataFrame(TSNE(verbose=9999).fit_transform(X), columns=['x', 'y'])

# Works
#import datashader as ds

fig = ds.tf.shade(ds.Canvas().points(df, 'x', 'y'))

# Point a browser at localhost:12345 to view the figure
bokeh_server = pn.Row(fig).show(port=12345)

bokeh_server.stop()

[DavidMChan]

Hmm - I don't really have the bandwidth to take a look at this, but for anyone out there who is interested, It's likely that we have some code that doesn't use best practices when initializing the CUBLAS handle, meaning that if there's already a CUBLAS context, we aren't fetching it properly. Thus, allowing tsnecuda to create the context, then importing datashader later fixes the problem, but if the context is created with datashader, there are underlying issues.