Databricks AI environment

The Databricks AI environment is a curated GPU-enabled runtime tailored for AI development. It is supported in serverless GPU environment 4 and above.

This new environment streamlines development by delivering a fully pre-configured library stack for machine learning (including frameworks like PyTorch, HuggingFace Transformers, etc.) and native support for GPUs. It integrates with Databricks notebooks, Unity Catalog, and MLflow, providing an integrated experience. With the AI environment, ML teams can simply select a ready-to-run GPU cluster and begin training models immediately, rather than spending days on setup and troubleshooting.

Connect to the AI environment

To use the Databricks AI environment from a Databricks notebook connected to serverless GPU compute:

1. From a notebook, click the Connect drop-down menu at the top and select Serverless GPU.
2. Click the to open the Environment side panel.
3. Select A10 from the Accelerator field.
4. Select AI v4 for the AI environment from the Base environment field.
5. If you chose None from the Base environment field, select the Environment version.
6. Click Apply and then Confirm that you want to apply the serverless GPU compute to your notebook environment.

To setup Databricks AI environment for a notebook job on serverless GPU:

1. From the notebook job setting, click Edit the notebook’s environment from the Environment and libraries section.
2. Click the to open the Environment side panel.
3. Select AI v4 for the AI environment from the Base environment field.

The new job runs will be able to pick up the Databricks AI environment.

FAQ

What is the difference between Databricks AI environment and the environment in Databricks Runtime for Machine Learning?

Both Databricks Runtime for Machine Learning and the Databricks AI environment provide a pre-configured compute environment tailored for AI/ML use cases. While Databricks Runtime for Machine Learning is used with classic compute resources, the Databricks AI environment is for serverless GPU compute.

Both the environment in Databricks Runtime for Machine Learning and the Databricks AI environment include common machine learning packages, with some differences. Most notably, the Databricks AI environment includes more updated packages, but does not include Tensorflow and GraphFrames. For more information on what’s included in the Databricks AI environment and Databricks Runtime for Machine Learning, see the Serverless environment release notes and Databricks runtime release notes.

Known issues

• The Databricks AI environment does not work with Environment and Libraries field in the task configuration for notebook jobs. If you create a new jobs environment from that field, you may not be able to select Databricks AI environment.
• The Databricks AI environment does not support environment export.