Install TensorFlow 2.1 on Databricks Runtime 6.5 ML GPU clusters
Databricks Runtime ML includes versions of TensorFlow so you can use it without installing any packages.
Databricks Runtime ML Version | TensorFlow Version |
---|---|
7.0 | 2.2.0 |
6.3 - 6.6 | 1.15.0 |
You can install other versions of TensorFlow by using a cluster-scoped init script.
In this article, you learn how to install TensorFlow 2.1 on Databricks Runtime 6.5 ML GPU clusters.
Important
Removing default libraries and installing new versions may cause instability or completely break your Databricks cluster. You should thoroughly test any new library version in your environment before running production jobs.
Install the init script
Install the following cluster-scoped init script on your Databricks Runtime 6.5 ML GPU cluster.
#!/bin/bash set -e apt-get update apt-get install -y --no-install-recommends --allow-downgrades \ libnccl2=2.4.8-1+cuda10.1 \ libnccl-dev=2.4.8-1+cuda10.1 \ cuda-libraries-10-1 \ libcudnn7=7.6.4.38-1+cuda10.1 \ libcudnn7-dev=7.6.4.38-1+cuda10.1 \ libnvinfer6=6.0.1-1+cuda10.1 \ libnvinfer-dev=6.0.1-1+cuda10.1 \ libnvinfer-plugin6=6.0.1-1+cuda10.1 apt-get clean ln -sfn cuda-10.1 /usr/local/cuda pip install tensorflow==2.1.* setuptools==41.* grpcio==1.24.* # This `conda list` is necessary to recognize the pip-installed packages. conda list conda install cudatoolkit=10.1
Restart the cluster.