3.7 Parallelizing Single-Node Models for Hyperparameter Tuning

3.7 Parallelizing Single-Node Models for Hyperparameter Tuning

Hyperparameter optimization and model validation on enterprise datasets require substantial computational resources. On Databricks, machine learning engineers have access to two powerful distributed tuning paradigms: SparkTrials (which distributes independent single-node model evaluations across Spark cluster executors) and pyspark.ml.tuning.CrossValidator (which distributes K-fold data partitions for distributed Spark ML pipelines). Selecting the correct distributed tuning abstraction and validation strategy is critical for avoiding resource contention, preventing data leakage, and optimizing compute expenditure.

---

Scaling Hyperopt with SparkTrials

Standard Hyperopt Trials() executes every trial sequentially on a single machine (the Databricks driver node). SparkTrials overrides this execution engine by broadcasting single-node training code and data to Spark worker nodes, executing multiple hyperparameter trials in parallel as separate Spark tasks.

+-----------------------------------------------------------------------------+
|                   SPARKTRILAS DISTRIBUTED ARCHITECTURE                      |
|                                                                             |
|   DATABRICKS DRIVER NODE                                                    |
|   +---------------------------------------------------------------------+   |
|   | Hyperopt fmin() + TPE Bayesian Engine                               |   |
|   | - Proposes candidate hyperparameter configurations                  |
|   | - SparkContext broadcasts dataset X_train, y_train to cluster       |
|   | - Instantiates SparkTrials(parallelism=4)                           |
|   +---------------------------------------------------------------------+   |
|                                     |
|         +---------------------------+---------------------------+           |
|         |                           |                           |           |
|         v                           v                           v           |
|   SPARK WORKER 1              SPARK WORKER 2              SPARK WORKER 3    |
|   +-------------------+       +-------------------+       +-------------+   |
|   | Task Slot 1       |       | Task Slot 2       |       | Task Slot 3 |   |
|   | Trial 001: Fit    |       | Trial 002: Fit    |       | Trial 003:  |   |
|   | Single-Node Model |       | Single-Node Model |       | Fit Model   |   |
|   | (Scikit / XGBoost)|       | (Scikit / XGBoost)|       |             |   |
|   +-------------------+       +-------------------+       +-------------+   |
|         |                           |                           |           |
|         +---------------------------+---------------------------+           |
|                                     v                                       |
|   All Worker Trial Results (Loss, Metrics) Aggregated back to Driver        |
+-----------------------------------------------------------------------------+

SparkTrials Parameters

from hyperopt import SparkTrials

spark_trials = SparkTrials(
    parallelism=8,           # Number of concurrent trials
    timeout=3600,            # Maximum total runtime in seconds (1 hour)
    max_trials_to_fail=5     # Max consecutive failures before aborting fmin
)

• parallelism: Sets the maximum number of trials evaluated concurrently across the cluster. If set to 8, Spark launches 8 independent Spark tasks across executor core slots.
• timeout: Hard execution time ceiling. If the timeout expires, fmin() terminates gracefully and returns the best trial identified up to that timestamp.
• max_trials_to_fail: Safeguard parameter (default 3, or configurable percentage). If $N$ consecutive trials throw unhandled runtime exceptions (e.g., OOM on specific hyperparameter depths), fmin() halts to prevent burning cluster budget on broken code.

The Parallelism vs. Bayesian Optimization Tradeoff

A critical theoretical concept tested on the Databricks ML Associate exam is the tradeoff between parallelism and TPE optimization quality:

+-----------------------------------------------------------------------------+
|                 PARALLELISM VS. BAYESIAN LEARNING TRADEOFF                  |
|                                                                             |
|   LOW PARALLELISM (e.g., parallelism = 1 to 2)                              |
|   [+] Max Sequential Learning: Every trial learns from all previous results.|
|   [-] High Wall-Clock Duration: Cluster worker cores sit idle.              |
|                                                                             |
|   EXCESSIVE PARALLELISM (e.g., parallelism = 50 with max_evals = 50)        |
|   [-] Zero Sequential Learning: All 50 trials launch simultaneously.        |
|   [-] Degrades to Pure Random Search: TPE has no completed trials to model. |
|   [+] Fast Wall-Clock Duration: Entire cluster saturated immediately.       |
|                                                                             |
|   RECOMMENDED BALANCED HEURISTIC:                                           |
|   Set parallelism = min(Total Worker Core Slots, 4 to 8)                   |
|   Set max_evals >= 4 * parallelism (ensures multiple sequential waves)      |
+-----------------------------------------------------------------------------+

---

Architectural Rule: Single-Node vs. Distributed Models with Hyperopt

A fundamental architectural rule on Databricks dictates which trial runner to use based on the underlying algorithm library:

---

Worked Implementation: SparkTrials with Single-Node XGBoost

Single-Node XGBoost with SparkTrials

import numpy as np
import xgboost as xgb
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score
from hyperopt import fmin, tpe, hp, STATUS_OK, SparkTrials
import mlflow

pdf = spark.table("lakehouse_gold.churn_features").toPandas()
X = pdf.drop(columns=["customer_id", "churned"])
y = pdf["churned"]
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)

search_space = {
    "max_depth": hp.quniform("max_depth", 3, 8, 1),
    "learning_rate": hp.loguniform("learning_rate", np.log(0.01), np.log(0.2)),
    "n_estimators": hp.quniform("n_estimators", 50, 200, 25)
}

def objective(params):
    params["max_depth"] = int(params["max_depth"])
    params["n_estimators"] = int(params["n_estimators"])
    
    clf = xgb.XGBClassifier(**params, random_state=42, eval_metric="logloss")
    clf.fit(X_train, y_train)
    
    preds = clf.predict_proba(X_val)[:, 1]
    auc = roc_auc_score(y_val, preds)
    return {"loss": -auc, "status": STATUS_OK}

# Configure SparkTrials to run 4 concurrent trials across workers
spark_trials = SparkTrials(parallelism=4, timeout=1800, max_trials_to_fail=3)

with mlflow.start_run(run_name="spark_trials_distributed_tuning"):
    best_params = fmin(
        fn=objective,
        space=search_space,
        algo=tpe.suggest,
        max_evals=20,
        trials=spark_trials
    )

print("Best Distributed Tuning Parameters:", best_params)

---

Two Kinds of Parallelism, One Rule

The rule to memorise is that exactly one layer may be distributed.

Running a pyspark.ml model inside a SparkTrials task asks a Spark task to launch Spark jobs, which is not a supported execution pattern and results in contention or outright failure. Conversely, running scikit-learn under plain Trials() leaves every executor idle while the driver trains one model at a time.

Data movement

SparkTrials broadcasts the training data referenced by the objective function to the executors. That works well when the training set fits comfortably in executor memory — which is the same precondition that made a single-node library appropriate in the first place. If the data is too large to broadcast, the answer is not more parallelism; it is a distributed pyspark.ml model.