How to update libraries in runtime on Databricks?

Viewed 157

I have certain code snippets that need to be run using library versions that are higher than the versions currently shipped within the Databricks Runtime version.

I have an external file logged by mlflow which contains a collection of Python libraries & versions, which are required for running certain code, and I need to ensure that those are properly installed before running the cells in the Databricks notebook.

I'm aware that it's possible to just do %pip install in a cell to upgrade to the required versions, but I want to avoid hardcoding packages and package versions, that's why I am attempting to rely on the external file logged by mlflow.

I am currently using Databricks Runtime 9.1 LTS for Machine Learning which comes with pandas==1.2.4 pre-installed. However, for my workflow I would need a more up-to-date package, pandas==1.3.5.

My approach for ensuring that the version seen by the notebook is indeed the newer 1.3.5 one, and not the default 1.2.4 that comes pre-installed, was to use the Libraries API of Databricks:

import requests

import pandas as pd 
print(f"pandas version BEFORE installation with Libraries API:", pd.__version__)

ctx = dbutils.notebook.entry_point.getDbutils().notebook().getContext()
host_name  = ctx.tags().get("browserHostName").get()
host_token = ctx.apiToken().get()
cluster_id = ctx.tags().get("clusterId").get()

response = requests.post(
    f"https://{host_name}/api/2.0/libraries/install",
    headers = {'Authorization': f'Bearer {host_token}'},
    json = {
        "cluster_id": cluster_id,
        "libraries": [
            {
                "pypi": {
                    "package": "pandas==1.3.5",
                }
            }
        ]
    }
)
assert response.status_code == 200, f"Package installation might have failed, status code: {response.status_code}!"

import pandas as pd 
print(f"pandas version AFTER installation with Libraries API:", pd.__version__)

The output of the above cell is:

pandas version BEFORE installation with Libraries API: 1.2.4
pandas version AFTER installation with Libraries API: 1.2.4

While the intended output which I want to achieve is:

pandas version BEFORE installation with Libraries API: 1.2.4
pandas version AFTER installation with Libraries API: 1.3.5

How could such library upgrade/overwrite over the Databricks Runtime be achieved?

1 Answers

The problem with your approach is that you don't wait until installation finishes. From documentation:

The installation is asynchronous - it completes in the background after the request.

So you need to check that library installation finishes. You can do it using a cluster-status API.

Although I'm not 100% sure that you will achieve your task for the current notebook - Python interpreter already loaded the Pandas, so it should be restarted to get changes into effect.

Related