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.