Distributed Hyperparameter Tuning on Databricks: Optuna with MLflow and Ray Tune

1.Why Hyperopt is no longer the answer

For years, distributed tuning on Databricks meant Hyperopt, which worked with both distributed algorithms such as Spark MLlib and single-machine libraries such as scikit-learn. That has changed for two reasons. The open-source Hyperopt project is no longer maintained, and Hyperopt is not included in Databricks Runtime for Machine Learning after 16.4 LTS ML. Databricks now recommends two replacements, and the split between them matters on the exam: Optuna for single-node optimization, and Ray Tune for an experience similar to Hyperopt's distributed tuning.

In practice, a migration starts by looking at the shape of the workload. If one machine is enough, Optuna is the documented choice. If trials must fan out across a cluster in the way Hyperopt did, Ray Tune is the closest match. Older Databricks examples, such as the notebook that compares three model types using Hyperopt with SparkTrials, remain useful for understanding existing code, but they are patterns to replace rather than patterns to build new work on. Keep this in mind when a question describes a legacy job and asks what to do next.

Sources1

2.Optuna on Spark executors with MLflow 3

Optuna is a lightweight framework with dynamic search spaces. In plain Optuna you parallelize with Joblib and track results in MLflow. MLflow 3.0, which comes pre-installed in Databricks Runtime 17.0 ML and above, adds two classes that let Optuna scale horizontally. MlflowStorage makes the MLflow Tracking Server Optuna's storage backend. It batches writes so the trials don't hit REST API throttling. MlflowSparkStudy wraps an Optuna Study and runs its trials on PySpark executors, with all of them reading and writing through that shared storage.

The workflow runs in a fixed order. First, write an objective function that calls a trial's suggest methods to define the search space and returns the metric to minimize. Second, create the shared storage. Third, create the study and call its optimize function.

import mlflow
from mlflow.optuna.storage import MlflowStorage

experiment_id = mlflow.get_experiment_by_name(dbutils.notebook.entry_point.getDbutils().notebook().getContext().notebookPath().get()).experiment_id

mlflow_storage = MlflowStorage(experiment_id=experiment_id)

from mlflow.pyspark.optuna.study import MlflowSparkStudy

mlflow_study = MlflowSparkStudy(
    study_name="spark-mlflow-tuning",
    storage=mlflow_storage,
)

mlflow_study. ? (objective, n_trials=8, n_jobs=4)

Because the sampler and pruner have defaults, a minimal study needs only a name and a storage object. In the documented example the study runs 8 trials with n_jobs=4, so several trials are in flight at once rather than one after another. The storage class has its own tuning knobs for the batching mentioned above: batch_flush_interval (default 1.0 seconds between automatic flushes) and batch_size_threshold (default 100 items before a flush is triggered). You rarely need to change them, but they explain how MLflow avoids throttling when many executors report results at once.

Sources21

3.Ray Tune: distributed search on a Ray-on-Spark cluster

Ray Tune comes with Ray, which is included in Databricks Runtime ML, and uses Ray as its distributed backend. It is the recommended replacement for Hyperopt-style distributed tuning, and it fits trials that are compute-heavy and independent, such as deep learning runs. Ray Tune runs on a Ray cluster that you start inside your Spark cluster with setup_ray_cluster. That call starts a background Spark job. Each Spark task in the job becomes a Ray worker node, and the Ray head node runs on the driver.

from ray.util.spark import setup_ray_cluster, shutdown_ray_cluster

setup_ray_cluster(
  min_worker_nodes=2,
  max_worker_nodes=4,
  num_cpus_per_node=4,
  collect_log_to_path="/Volumes/<catalog>/<schema>/<volume>/path/to/ray_collected_logs"
)

# Pass any custom Ray configuration with ray.init
ray.init(ignore_reinit_error=True)

How you size the cluster affects tuning throughput. Databricks recommends at least 4 CPU cores and 10GB of heap memory per Ray worker node. Setting the per-worker CPU (or GPU) count to match a Spark worker node gives one Ray worker per node that uses all of that node's resources. By default the head node gets 0 CPUs and 0 GPUs, because it only coordinates. In hybrid Spark-and-Ray clusters, reduce spark.executor.memory, because the Spark executor's lazy garbage collection and dataset cache can take memory Ray needs.

For tracking, Ray Tune uses a child-run approach with MLflow. You create a parent run in the driver process, and each Ray Tune task logs a child run under it, so all trials are grouped in one hierarchy.

Autoscaling is controlled by the range between min_worker_nodes and max_worker_nodes. If min_worker_nodes is left undefined, a fixed-size cluster starts with max_worker_nodes workers. A bounded range lets the Ray cluster grow for a search and shrink afterward, which frees resources for Spark tasks. Databricks also documents a Ray Tune example that uses the ASHA scheduler, which stops trials that fall behind instead of training every candidate to completion. Giving each trial a whole GPU through tune.with_resources keeps trials from sharing a device.

Sources134

Exam traps

Each one states something that sounds right. Open it to see what is actually true.

1. 1.Hyperopt is still the recommended tool for distributed tuning on current Databricks Runtime ML.Why is that wrong?

   Hyperopt is not included after 16.4 LTS ML. Databricks recommends Optuna for single-node optimization and Ray Tune for Hyperopt-like distributed tuning.

   Covered in Why Hyperopt is no longer the answer

2. 2.MlflowSparkStudy defaults to random search with no pruning.Why is that wrong?

   The default sampler is TPESampler and the default pruner is MedianPruner, which stops unpromising trials early.

   Covered in Optuna on Spark executors with MLflow 3