Azure Databricks AutoML Python API reference

This article describes the Azure Databricks AutoML Python API, which provides methods to start classification, regression, and forecasting AutoML runs. Each method call trains a set of models and generates a trial notebook for each model.

For more information on AutoML, including a low-code UI option, see What is AutoML?.

Classify

The databricks.automl.classify method configures an AutoML run to train a classification model.

databricks.automl.classify(
  dataset: Union[pyspark.sql.DataFrame, pandas.DataFrame, pyspark.pandas.DataFrame, str],
  *,
  target_col: str,
  primary_metric: str = "f1",
  data_dir: Optional[str] = None,
  experiment_dir: Optional[str] = None,                             # <DBR> 10.4 LTS ML and above
  experiment_name: Optional[str] = None,                            # <DBR> 12.1 ML and above
  exclude_cols: Optional[List[str]] = None,                         # <DBR> 10.3 ML and above
  exclude_frameworks: Optional[List[str]] = None,                   # <DBR> 10.3 ML and above
  feature_store_lookups: Optional[List[Dict]] = None,               # <DBR> 11.3 LTS ML and above
  imputers: Optional[Dict[str, Union[str, Dict[str, Any]]]] = None, # <DBR> 10.4 LTS ML and above
  pos_label: Optional[Union[int, bool, str]] = None,                 # <DBR> 11.1 ML and above
  time_col: Optional[str] = None,
  split_col: Optional[str] = None,                                  # <DBR> 15.3 ML and above
  sample_weight_col: Optional[str] = None                           # <DBR> 15.4 ML and above
  max_trials: Optional[int] = None,                                 # <DBR> 10.5 ML and below
  timeout_minutes: Optional[int] = None,
) -> AutoMLSummary

Classify parameters

Parameter name	Type	Description
dataset	str, pandas.DataFrame, pyspark.DataFrame, pyspark.sql.DataFrame	Input table name or DataFrame that contains training features and target. Table name can be in format "<database_name>.<table_name>" or "<schema_name>.<table_name>" for non Unity Catalog tables.
target_col	str	Column name for the target label.
primary_metric	str	Metric used to evaluate and rank model performance. Supported metrics for regression: "r2" (default), "mae", "rmse", "mse" Supported metrics for classification: "f1" (default), "log_loss", "precision", "accuracy", "roc_auc"
data_dir	str of format dbfs:/<folder-name>	Optional. DBFS path used to store the training dataset. This path is visible to both driver and worker nodes. Databricks recommends leaving this field empty, so AutoML can save the training dataset as an MLflow artifact. If a custom path is specified, the dataset does not inherit the AutoML experiment's access permissions.
experiment_dir	str	Optional. Path to the directory in the workspace to save the generated notebooks and experiments. Default: /Users/<username>/databricks_automl/
experiment_name	str	Optional. Name for the MLflow experiment that AutoML creates. Default: Name is automatically generated.
exclude_cols	List[str]	Optional. List of columns to ignore during AutoML calculations. Default: []
exclude_frameworks	List[str]	Optional. List of algorithm frameworks that AutoML should not consider as it develops models. Possible values: empty list, or one or more of "sklearn", "lightgbm", "xgboost". Default: [] (all frameworks are considered
feature_store_lookups	List[Dict]	Optional. List of dictionaries that represent features from Feature Store for data augmentation. Valid keys in each dictionary are: - table_name (str): Required. Name of the feature table. - lookup_key (list or str): Required. Column name(s) to use as key when joining the feature table with the data passed in the dataset param. The order of the column names must match the order of the primary keys of the feature table. - timestamp_lookup_key (str): Required if the specified table is a time series feature table. The column name to use when performing point-in-time lookup on the feature table with the data passed in the dataset param. Default: []
imputers	Dict[str, Union[str, Dict[str, Any]]]	Optional. Dictionary where each key is a column name, and each value is a string or dictionary describing the imputation strategy. If specified as a string, the value must be one of "mean", "median", or "most_frequent". To impute with a known value, specify the value as a dictionary {"strategy": "constant", "fill_value": <desired value>}. You can also specify string options as dictionaries, for example {"strategy": "mean"}. If no imputation strategy is provided for a column, AutoML selects a default strategy based on column type and content. If you specify a non-default imputation method, AutoML does not perform semantic type detection. Default: {}
pos_label	Union[int, bool, str]	(Classification only) The positive class. This is useful for calculating metrics such as precision and recall. Should only be specified for binary classification problems.
time_col	str	Available in Databricks Runtime 10.1 ML and above. Optional. Column name for a time column. If provided, AutoML tries to split the dataset into training, validation, and test sets chronologically, using the earliest points as training data and the latest points as a test set. Accepted column types are timestamp and integer. With Databricks Runtime 10.2 ML and above, string columns are also supported. If column type is string, AutoML tries to convert it to timestamp using semantic detection. If the conversion fails, the AutoML run fails.
split_col	str	Optional. Column name for a split column. Only available in Databricks Runtime 15.3 ML and above for API workflows. If provided, AutoML tries to split train/validate/test sets by user-specified values, and this column is automatically excluded from training features. Accepted column type is string. The value of each entry in this column must be one of the following: "train", "validate", or "test".
sample_weight_col	str	Available in Databricks Runtime 15.4 ML and above for classification API workflows. Optional. Column name in the dataset that contains the sample weights for each row. Classification supports per-class sample weights. These weights adjust the importance of each class during model training. Each sample within a class must have the same sample weight and weights must be non-negative decimal or integer values, ranging from 0 to 10,000. Classes with higher sample weights are considered more important, and have a greater influence on the learning algorithm. If this column is not specified, all classes are assumed to have equal weight.
max_trials	int	Optional. Maximum number of trials to run. This parameter is available in Databricks Runtime 10.5 ML and below, but is deprecated starting in Databricks Runtime 10.3 ML. In Databricks Runtime 11.0 ML and above, this parameter is not supported. Default: 20 If timeout_minutes=None, AutoML runs the maximum number of trials.
timeout_minutes	int	Optional. Maximum time to wait for AutoML trials to complete. Longer timeouts allow AutoML to run more trials and identify a model with better accuracy. Default: 120 minutes Minimum value: 5 minutes An error is reported if the timeout is too short to allow at least one trial to complete.

Regress

The databricks.automl.regress method configures an AutoML run to train a regression model. This method returns an AutoMLSummary.

databricks.automl.regress(
  dataset: Union[pyspark.sql.DataFrame, pandas.DataFrame, pyspark.pandas.DataFrame, str],
  *,
  target_col: str,
  primary_metric: str = "r2",
  data_dir: Optional[str] = None,
  experiment_dir: Optional[str] = None,                             # <DBR> 10.4 LTS ML and above
  experiment_name: Optional[str] = None,                            # <DBR> 12.1 ML and above
  exclude_cols: Optional[List[str]] = None,                         # <DBR> 10.3 ML and above
  exclude_frameworks: Optional[List[str]] = None,                   # <DBR> 10.3 ML and above
  feature_store_lookups: Optional[List[Dict]] = None,               # <DBR> 11.3 LTS ML and above
  imputers: Optional[Dict[str, Union[str, Dict[str, Any]]]] = None, # <DBR> 10.4 LTS ML and above
  time_col: Optional[str] = None,
  split_col: Optional[str] = None,                                  # <DBR> 15.3 ML and above
  sample_weight_col: Optional[str] = None,                          # <DBR> 15.3 ML and above
  max_trials: Optional[int] = None,                                 # <DBR> 10.5 ML and below
  timeout_minutes: Optional[int] = None,
) -> AutoMLSummary

Regress parameters

Parameter name	Type	Description
dataset	str, pandas.DataFrame, pyspark.DataFrame, pyspark.sql.DataFrame	Input table name or DataFrame that contains training features and target. Table name can be in format "<database_name>.<table_name>" or "<schema_name>.<table_name>" for non Unity Catalog tables.
target_col	str	Column name for the target label.
primary_metric	str	Metric used to evaluate and rank model performance. Supported metrics for regression: "r2" (default), "mae", "rmse", "mse" Supported metrics for classification: "f1" (default), "log_loss", "precision", "accuracy", "roc_auc"
data_dir	str of format dbfs:/<folder-name>	Optional. DBFS path used to store the training dataset. This path is visible to both driver and worker nodes. Databricks recommends leaving this field empty, so AutoML can save the training dataset as an MLflow artifact. If a custom path is specified, the dataset does not inherit the AutoML experiment's access permissions.
experiment_dir	str	Optional. Path to the directory in the workspace to save the generated notebooks and experiments. Default: /Users/<username>/databricks_automl/
experiment_name	str	Optional. Name for the MLflow experiment that AutoML creates. Default: Name is automatically generated.
exclude_cols	List[str]	Optional. List of columns to ignore during AutoML calculations. Default: []
exclude_frameworks	List[str]	Optional. List of algorithm frameworks that AutoML should not consider as it develops models. Possible values: empty list, or one or more of "sklearn", "lightgbm", "xgboost". Default: [] (all frameworks are considered
feature_store_lookups	List[Dict]	Optional. List of dictionaries that represent features from Feature Store for data augmentation. Valid keys in each dictionary are: - table_name (str): Required. Name of the feature table. - lookup_key (list or str): Required. Column name(s) to use as key when joining the feature table with the data passed in the dataset param. The order of the column names must match the order of the primary keys of the feature table. - timestamp_lookup_key (str): Required if the specified table is a time series feature table. The column name to use when performing point-in-time lookup on the feature table with the data passed in the dataset param. Default: []
imputers	Dict[str, Union[str, Dict[str, Any]]]	Optional. Dictionary where each key is a column name, and each value is a string or dictionary describing the imputation strategy. If specified as a string, the value must be one of "mean", "median", or "most_frequent". To impute with a known value, specify the value as a dictionary {"strategy": "constant", "fill_value": <desired value>}. You can also specify string options as dictionaries, for example {"strategy": "mean"}. If no imputation strategy is provided for a column, AutoML selects a default strategy based on column type and content. If you specify a non-default imputation method, AutoML does not perform semantic type detection. Default: {}
time_col	str	Available in Databricks Runtime 10.1 ML and above. Optional. Column name for a time column. If provided, AutoML tries to split the dataset into training, validation, and test sets chronologically, using the earliest points as training data and the latest points as a test set. Accepted column types are timestamp and integer. With Databricks Runtime 10.2 ML and above, string columns are also supported. If column type is string, AutoML tries to convert it to timestamp using semantic detection. If the conversion fails, the AutoML run fails.
split_col	str	Optional. Column name for a split column. Only available in Databricks Runtime 15.3 ML and above for API workflows. If provided, AutoML tries to split train/validate/test sets by user-specified values, and this column is automatically excluded from training features. Accepted column type is string. The value of each entry in this column must be one of the following: "train", "validate", or "test".
sample_weight_col	str	Available in Databricks Runtime 15.3 ML and above for regression API workflows. Optional. Column name in the dataset that contains the sample weights for each row. These weights adjust the importance of each row during model training. Weights must be non-negative decimal or integer values, ranging from 0 to 10,000. Rows with higher sample weights are considered more important, and have a greater influence on the learning algorithm. If this column is not specified, all rows are assumed to have equal weight.
max_trials	int	Optional. Maximum number of trials to run. This parameter is available in Databricks Runtime 10.5 ML and below, but is deprecated starting in Databricks Runtime 10.3 ML. In Databricks Runtime 11.0 ML and above, this parameter is not supported. Default: 20 If timeout_minutes=None, AutoML runs the maximum number of trials.
timeout_minutes	int	Optional. Maximum time to wait for AutoML trials to complete. Longer timeouts allow AutoML to run more trials and identify a model with better accuracy. Default: 120 minutes Minimum value: 5 minutes An error is reported if the timeout is too short to allow at least one trial to complete.

Forecast

The databricks.automl.forecast method configures an AutoML run for training a forecasting model. This method returns an AutoMLSummary. To use Auto-ARIMA, the time series must have a regular frequency (that is, the interval between any two points must be the same throughout the time series). The frequency must match the frequency unit specified in the API call. AutoML handles missing time steps by filling in those values with the previous value.

databricks.automl.forecast(
  dataset: Union[pyspark.sql.DataFrame, pandas.DataFrame, pyspark.pandas.DataFrame, str],
  *,
  target_col: str,
  time_col: str,
  primary_metric: str = "smape",
  country_code: str = "US",                                         # <DBR> 12.0 ML and above
  frequency: str = "D",
  horizon: int = 1,
  data_dir: Optional[str] = None,
  experiment_dir: Optional[str] = None,
  experiment_name: Optional[str] = None,                            # <DBR> 12.1 ML and above
  exclude_frameworks: Optional[List[str]] = None,
  feature_store_lookups: Optional[List[Dict]] = None,               # <DBR> 12.2 LTS ML and above
  identity_col: Optional[Union[str, List[str]]] = None,
  sample_weight_col: Optional[str] = None,                          # <DBR> 16.0 ML and above
  output_database: Optional[str] = None,                            # <DBR> 10.5 ML and above
  timeout_minutes: Optional[int] = None,
) -> AutoMLSummary

Forecasting parameters

Parameter name	Type	Description
dataset	str, pandas.DataFrame, pyspark.DataFrame, pyspark.sql.DataFrame	Input table name or DataFrame that contains training features and target. Table name can be in format ".." or "." for non Unity Catalog tables
target_col	str	Column name for the target label.
time_col	str	Name of the time column for forecasting.
primary_metric	str	Metric used to evaluate and rank model performance. Supported metrics: "smape" (default), "mse", "rmse", "mae", or "mdape".
country_code	str	Available in Databricks Runtime 12.0 ML and above. Only supported by the Prophet forecasting model. Optional. Two-letter country code that indicates which country's holidays the forecasting model should use. To ignore holidays, set this parameter to an empty string (""). Supported countries. Default: US (United States holidays).
frequency	str	Frequency of the time series for forecasting. This is the period with which events are expected to occur. The default setting is "D" or daily data. Be sure to change the setting if your data has a different frequency. Possible values: "W" (weeks) "D" / "days" / "day" "hours" / "hour" / "hr" / "h" "m" / "minute" / "min" / "minutes" / "T" "S" / "seconds" / "sec" / "second" The following are only available with Databricks Runtime 12.0 ML and above: "M" / "month" / "months" "Q" / "quarter" / "quarters" "Y" / "year" / "years" Default: "D"
horizon	int	Number of periods into the future for which forecasts should be returned. The units are the time series frequency. Default: 1
data_dir	str of format dbfs:/<folder-name>	Optional. DBFS path used to store the training dataset. This path is visible to both driver and worker nodes. Databricks recommends leaving this field empty, so AutoML can save the training dataset as an MLflow artifact. If a custom path is specified, the dataset does not inherit the AutoML experiment's access permissions.
experiment_dir	str	Optional. Path to the directory in the workspace to save the generated notebooks and experiments. Default: /Users/<username>/databricks_automl/
experiment_name	str	Optional. Name for the MLflow experiment that AutoML creates. Default: Name is automatically generated.
exclude_frameworks	List[str]	Optional. List of algorithm frameworks that AutoML should not consider as it develops models. Possible values: empty list, or one or more of "prophet", "arima". Default: [] (all frameworks are considered
feature_store_lookups	List[Dict]	Optional. List of dictionaries that represent features from Feature Store for covariate data augmentation. Valid keys in each dictionary are: - table_name (str): Required. Name of the feature table. - lookup_key (list or str): Required. Column name(s) to use as key when joining the feature table with the data passed in the dataset param. The order of the column names must match the order of the primary keys of the feature table. - timestamp_lookup_key (str): Required if the specified table is a time series feature table. The column name to use when performing point-in-time lookup on the feature table with the data passed in the dataset param. Default: []
identity_col	Union[str, list]	Optional. Column(s) that identify the time series for multi-series forecasting. AutoML groups by these column(s) and the time column for forecasting.
sample_weight_col	str	Available in Databricks Runtime 16.0 ML and above. Only for multi-time-series workflows. Optional. Specifies the column in the dataset that contains sample weights. These weights indicate the relative importance of each time series during model training and evaluation. Time series with higher weights have a greater influence on the model. If not provided, all time series are treated with equal weight. All rows belonging to the same time series must have the same weight. Weights must be non-negative values, either decimals or integers, and be between 0 and 10,000.
output_database	str	Optional. If provided, AutoML saves predictions of the best model to a new table in the specified database. Default: Predictions are not saved.
timeout_minutes	int	Optional. Maximum time to wait for AutoML trials to complete. Longer timeouts allow AutoML to run more trials and identify a model with better accuracy. Default: 120 minutes Minimum value: 5 minutes An error is reported if the timeout is too short to allow at least one trial to complete.

Import notebook

The databricks.automl.import_notebook method imports a notebook that has been saved as an MLflow artifact. This method returns an ImportNotebookResult.

databricks.automl.import_notebook(
  artifact_uri: str,
  path: str,
  overwrite: bool = False
) -> ImportNotebookResult:

Parameters	Type	Description
artifact_uri	str	The URI of the MLflow artifact that contains the trial notebook.
path	str	The path in the Databricks workspace where the notebook should be imported. This must be an absolute path. The directory will be created if it does not exist.
overwrite	bool	Whether to overwrite the notebook if it already exists. It is False by default.

Import notebook example

summary = databricks.automl.classify(...)
result = databricks.automl.import_notebook(summary.trials[5].artifact_uri, "/Users/you@yourcompany.com/path/to/directory")
print(result.path)
print(result.url)

AutoMLSummary

Summary object for an AutoML run that describes the metrics, parameters, and other details for each of the trials. You also use this object to load the model trained by a specific trial.

Property	Type	Description
experiment	mlflow.entities.Experiment	The MLflow experiment used to log the trials.
trials	List[TrialInfo]	A list of TrialInfo objects containing information about all the trials that were run.
best_trial	TrialInfo	A TrialInfo object containing information about the trial that resulted in the best weighted score for the primary metric.
metric_distribution	str	The distribution of weighted scores for the primary metric across all trials.
output_table_name	str	Used with forecasting only and only if output_database is provided. Name of the table in output_database containing the model's predictions.

TrialInfo

Summary object for each individual trial.

Property	Type	Description
notebook_path	Optional[str]	The path to the generated notebook for this trial in the workspace. For classification and regression, this value is set only for the best trial, while all other trials have the value set to None. For forecasting, this value is present for all trials.
notebook_url	Optional[str]	The URL of the generated notebook for this trial. For classification and regression, this value is set only for the best trial, while all other trials have the value set to None. For forecasting, this value is present for all trials.
artifact_uri	Optional[str]	The MLflow artifact URI for the generated notebook.
mlflow_run_id	str	The MLflow run ID associated with this trial run.
metrics	Dict[str, float]	The metrics logged in MLflow for this trial.
params	Dict[str, str]	The parameters logged in MLflow that were used for this trial.
model_path	str	The MLflow artifact URL of the model trained in this trial.
model_description	str	Short description of the model and the hyperparameters used for training this model.
duration	str	Training duration in minutes.
preprocessors	str	Description of the preprocessors run before training the model.
evaluation_metric_score	float	Score of primary metric, evaluated for the validation dataset.

TrialInfo has a method to load the model generated for the trial.

Method	Description
load_model()	Load the model generated in this trial, logged as an MLflow artifact.

ImportNotebookResult

Property	Type	Description
path	str	The path in the Databricks workspace where the notebook should be imported. This must be an absolute path. The directory will be created if it does not exist.
url	str	The URI of the MLflow artifact that contains the trial notebook.