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

1. 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
   

2. Restart the cluster.