Importing sentence-transformers no longer works on Databricks runtime 17.2 ML

excavator-matt
Contributor III

In Databricks Runtime 16.4 LTS for Machine Learning, I am used to be able to import sentence-transformers without any installation as it is part of the runtime with from sentence_transformers import SentenceTransformer.

In this case I am running on a personal compute cluster and although it generates some advanced warnings that I don't know how to interpret, it runs

2025-09-03 09:39:48.853218: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. 2025-09-03 09:39:48.869748: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered 2025-09-03 09:39:48.890311: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered 2025-09-03 09:39:48.896580: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered 2025-09-03 09:39:48.911486: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags. 2025-09-03 09:39:50.080574: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT

However, if I upgrade my cluster to run the latest 17.2 machine learning (Beta). I instead get this crash.

Failed to import transformers.modeling_utils because of the following error (look up to see its traceback):
/databricks/python/lib/python3.12/site-packages/flash_attn_2_cuda.cpython-312-x86_64-linux-gnu.so: undefined symbol: _ZN3c105ErrorC2ENS_14SourceLocationESs
[Trace ID: 00-91c7cd652996226fe8747ad97efc53e6-f7a274cbb6f58e49-00]

Both runtimes include sentence-transformers although version 3.4.1 versus 4.0.1, so it should work. Is this no longer supported? Is it known to be broken in the beta?