---
title: Train forecasting models with AutoML Python API | Databricks on AWS
source: https://docs.databricks.com/aws/en/machine-learning/automl/forecasting-train-api
ingestedAt: "2026-06-18T08:09:35.361Z"
---

This example notebook shows how to train a time-series forecasting model on Databricks using the AutoML Python API. Using a COVID-19 case-count dataset, you call `automl.forecast()` with a 30-day daily horizon to project future case counts, then load the best model with MLflow to generate and plot forecasts.

## Requirements[​](#requirements "Direct link to Requirements")

Databricks Runtime for Machine Learning 10.0 or above.  
To save model predictions, Databricks Runtime for Machine Learning 10.5 or above.

## COVID-19 dataset[​](#covid-19-dataset "Direct link to COVID-19 dataset")

The dataset contains records for the number of cases of the COVID-19 virus by date in the US, with additional geographical information. The goal is to forecast how many cases of the virus will occur over the next 30 days in the US.

Python

    import pyspark.pandas as psdf = ps.read_csv("/databricks-datasets/COVID/covid-19-data")df["date"] = ps.to_datetime(df['date'], errors='coerce')df["cases"] = df["cases"].astype(int)display(df)

## AutoML training[​](#automl-training "Direct link to AutoML training")

The following command starts an AutoML run. You must provide the column that the model should predict in the `target_col` argument and the time column. When the run completes, you can follow the link to the best trial notebook to examine the training code.

This example also specifies:

*   `horizon=30` to specify that AutoML should forecast 30 days into the future.
*   `frequency="d"` to specify that a forecast should be provided for each day.
*   `primary_metric="mdape"` to specify the metric to optimize for during training.

Python

    import databricks.automlimport logging# Disable informational messages from fbprophetlogging.getLogger("py4j").setLevel(logging.WARNING)# Note: If you are running Databricks Runtime for Machine Learning 10.4 or below, use this line instead:# summary = databricks.automl.forecast(df, target_col="cases", time_col="date", horizon=30, frequency="d",  primary_metric="mdape")summary = databricks.automl.forecast(df, target_col="cases", time_col="date", horizon=30, frequency="d",  primary_metric="mdape", output_database="default")

## Iterate on the model[​](#iterate-on-the-model "Direct link to Iterate on the model")

*   Explore the notebooks and experiments linked above.
*   If the metrics for the best trial notebook look good, you can continue with the next cell.
*   If you want to improve on the model generated by the best trial:
    *   Go to the notebook with the best trial and clone it.
    *   Edit the notebook as necessary to improve the model.
    *   When you are satisfied with the model, note the URI where the artifact for the trained model is logged. Assign this URI to the `model_uri` variable in the next cell.

## Show the predicted results from the best model[​](#show-the-predicted-results-from-the-best-model "Direct link to Show the predicted results from the best model")

**Note:** This section requires Databricks Runtime for Machine Learning 10.5 or above.

### Load predictions from the best model[​](#load-predictions-from-the-best-model "Direct link to Load predictions from the best model")

In Databricks Runtime for Machine Learning 10.5 or above, if `output_database` is provided, AutoML saves the predictions from the best model.

Python

    # Load the saved predictions.forecast_pd = spark.table(summary.output_table_name)display(forecast_pd)

## Use the model for forecasting[​](#use-the-model-for-forecasting "Direct link to Use the model for forecasting")

You can use the commands in this section with Databricks Runtime for Machine Learning 10.0 or above.

### Load the model with MLflow[​](#load-the-model-with-mlflow "Direct link to Load the model with MLflow")

MLflow allows you to easily import models back into Python by using the AutoML `trial_id` .

Python

    import mlflow.pyfuncfrom mlflow.tracking import MlflowClientrun_id = MlflowClient()trial_id = summary.best_trial.mlflow_run_idmodel_uri = "runs:/{run_id}/model".format(run_id=trial_id)pyfunc_model = mlflow.pyfunc.load_model(model_uri)

### Use the model to make forecasts[​](#use-the-model-to-make-forecasts "Direct link to Use the model to make forecasts")

Call the `predict_timeseries` model method to generate forecasts.  
In Databricks Runtime for Machine Learning 10.5 or above, you can set `include_history=False` to get the predicted data only.

Python

    forecasts = pyfunc_model._model_impl.python_model.predict_timeseries()display(forecasts)# Option for Databricks Runtime for Machine Learning 10.5 or above# forecasts = pyfunc_model._model_impl.python_model.predict_timeseries(include_history=False)

### Plot the forecasted points[​](#plot-the-forecasted-points "Direct link to Plot the forecasted points")

In the plot below, the thick black line shows the time series dataset, and the blue line is the forecast created by the model.

Python

    df_true = df.groupby("date").agg(y=("cases", "avg")).reset_index().to_pandas()

Python

    import matplotlib.pyplot as pltfig = plt.figure(facecolor='w', figsize=(10, 6))ax = fig.add_subplot(111)forecasts = pyfunc_model._model_impl.python_model.predict_timeseries(include_history=True)fcst_t = forecasts['ds'].dt.to_pydatetime()ax.plot(df_true['date'].dt.to_pydatetime(), df_true['y'], 'k.', label='Observed data points')ax.plot(fcst_t, forecasts['yhat'], ls='-', c='#0072B2', label='Forecasts')ax.fill_between(fcst_t, forecasts['yhat_lower'], forecasts['yhat_upper'],                color='#0072B2', alpha=0.2, label='Uncertainty interval')ax.legend()plt.show()

## Register and deploy the model[​](#register-and-deploy-the-model "Direct link to Register and deploy the model")

You can register and deploy a model trained by AutoML like any other model in the MLflow Model Registry. See [Log, load, and register MLflow models](https://docs.databricks.com/aws/en/mlflow/models).

### Troubleshooting: `No module named pandas.core.indexes.numeric`[​](#troubleshooting-no-module-named-pandascoreindexesnumeric "Direct link to troubleshooting-no-module-named-pandascoreindexesnumeric")

When serving an AutoML\-trained model with Model Serving, you may see the error `No module named pandas.core.indexes.numeric`. This happens when the `pandas` version used by AutoML differs from the one in the model serving endpoint environment. To resolve:

1.  Download the [add-pandas-dependency.py script](https://docs.databricks.com/aws/en/assets/files/add-pandas-dependency-4808dc9dcdfb035bdca6ebce6b86d719.py). The script edits `requirements.txt` and `conda.yaml` for the logged model to pin `pandas==1.5.3`.
2.  Edit the script to include the `run_id` of the MLflow run where the model was logged.
3.  Re-register the model.
4.  Serve the new model version.

## Example notebook[​](#example-notebook "Direct link to Example notebook")

#### Train forecasting models with AutoML Python API

## Next steps[​](#next-steps "Direct link to Next steps")

[AutoML Python API reference](https://docs.databricks.com/aws/en/machine-learning/automl/automl-api-reference).
