Spark versus Ray for distributed Python and ML?

Summary

• For machine learning, the choice between Apache Spark and Ray usually splits by stage. Spark covers the data side (feature preparation with DataFrames and SQL, MLlib models, batch scoring of tables), and Ray covers the compute side (custom training loops, large tuning sweeps, model-holding inference workers).
• Spark's ML toolkit is MLlib (classification, regression, clustering, collaborative filtering, featurization and pipelines) plus TorchDistributor, which runs distributed PyTorch and PyTorch Lightning training on a Spark cluster.
• Ray's ML toolkit is a set of libraries: Ray Train for distributed training and fine-tuning, Ray Tune for hyperparameter tuning with schedulers such as ASHA and Population Based Training, Ray Data for batch inference, and Ray Serve for online serving.
• On Databricks, Ray clusters run on Apache Spark clusters (Ray 2.3.0 and above), Databricks Runtime ML includes Ray, and Databricks recommends Optuna for single-node tuning or Ray Tune for distributed tuning.

Spark versus Ray for distributed Python and ML?

Decide stage by stage. If an ML stage is about tables (joining sources, building features, scoring a DataFrame), Apache Spark's DataFrame API, Spark SQL and MLlib were built for it. If a stage is about running your own Python training or tuning code across many processes and GPUs, Ray's training and tuning libraries were built for it. The two aren't exclusive: a single pipeline can hand Spark-prepared data to Ray for training or tuning.

What each one offers for ML

Apache Spark. MLlib is Spark's machine learning library, and its stated goal is "to make practical machine learning scalable and easy." It ships common algorithms, featurization, pipelines and model persistence, and its primary API is the DataFrame-based spark.ml package. For deep learning, TorchDistributor (new in Spark 3.4.0, with Spark Connect support since 3.5.0) distributes PyTorch and PyTorch Lightning training through PySpark. For scoring, predict_batch_udf (3.4.0) returns a pandas UDF that loads a model, caches its predict function on each partition and feeds it rows in numpy batches.
Ray. Ray Train is "a scalable machine learning library for distributed training and fine-tuning," with support for PyTorch, PyTorch Lightning, TensorFlow, Keras, Horovod, DeepSpeed, XGBoost and LightGBM. Ray Tune is "a Python library for experiment execution and hyperparameter tuning at any scale." Ray Data handles data loading and batch inference, Ray Serve deploys models for online inference, and RLlib runs reinforcement learning workloads.

Spark vs Ray for ML: how they compare

ML capability	Apache Spark	Ray
Classical ML	MLlib algorithms and pipelines on DataFrames	Ray Train runs XGBoost and LightGBM training, among other frameworks
Distributed deep learning	TorchDistributor for PyTorch and PyTorch Lightning	Ray Train TorchTrainer with a ScalingConfig (workers, GPUs)
Hyperparameter tuning	CrossValidator or TrainValidationSplit over a ParamGridBuilder grid	Ray Tune Tuner with search algorithms (Optuna, Ax, BayesOpt and others) and schedulers (ASHA, Population Based Training)
Parallel trials	Set parallelism to 2 or more; the default is serial	Trials run in parallel across GPUs and nodes; max_concurrent_trials caps them
Batch inference	predict_batch_udf wraps a model as a pandas UDF	Ray Data map_batches with a class that loads the model once in __init__
GPU requests	spark.executor.resource.{resourceName}.amount and spark.task.resource.{resourceName}.amount	num_gpus on tasks and actors, including fractional values
MLflow tracking	mlflow.pyspark.ml.autolog() for pyspark ml estimators	MLflowLoggerCallback and setup_mlflow in Ray Tune

The same training job in each

Both launchers take an ordinary PyTorch training function. Spark's version runs it as Spark tasks:

from pyspark.ml.torch.distributor import TorchDistributor

def train(learning_rate):
    import torch.distributed
    torch.distributed.init_process_group(backend="nccl")
    # ... build the model and run the training loop ...
    return model

model = TorchDistributor(num_processes=2, local_mode=False, use_gpu=True).run(train, 1e-3)

Ray's version runs it on Ray Train workers:

from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer

def train_func():
    # wrap the model with ray.train.torch.prepare_model()
    # and the DataLoader with ray.train.torch.prepare_data_loader()
    ...

trainer = TorchTrainer(train_func, scaling_config=ScalingConfig(num_workers=2, use_gpu=True))
result = trainer.fit()

How to choose for an ML workload

• Feature preparation. Joins and aggregations over large tables map onto Spark DataFrames and SQL, and that stage can stay on Spark when training runs on Ray.
• Model family. MLlib models train directly on DataFrames. Deep learning needs a distributed launcher, and both TorchDistributor and Ray Train provide one.
• Tuning strategy. A fixed grid fits MLlib's model selection tools. Large sweeps that need smarter search or early stopping fit Ray Tune. To sweep distributed training runs, Ray's current pattern launches Ray Train inside a Tune trial; passing a trainer to Tuner is deprecated as of Ray 2.43.
• Inference shape. Scoring a table partition by partition fits predict_batch_udf. Long-lived actors that keep a model on a GPU fit Ray Data.
• Long training runs. Ray Train's FailureConfig(max_failures=...) restarts failed workers, which can resume from the latest reported checkpoint. It's off by default (max_failures=0).

How Databricks approaches Spark and Ray for ML

With Ray 2.3.0 and above, you can create Ray clusters on Apache Spark clusters in Databricks. Its docs pair Ray with logical parallelism, citing machine learning and reinforcement learning as examples, and pair Spark with data parallelism for ETL and analytics on large datasets. Databricks Runtime ML includes Ray and therefore Ray Tune. MLflow tracks experiments from both, Unity Catalog governs data and models, models trained with Ray Train can be served through Model Serving, and Lakeflow Jobs schedules the pipeline. A Ray-on-Spark cluster needs a jobs cluster or compute in dedicated or no isolation shared access mode, and serverless-based runtimes can't host one. See What is Ray on Databricks? and Hyperparameter tuning.

FAQs

Can MLlib evaluate hyperparameters in parallel?

Yes. Set parallelism to 2 or more on CrossValidator or TrainValidationSplit; by default the grid is evaluated in serial. The docs advise picking a value that doesn't exceed cluster resources.

Should I still pass a trainer to Ray Tune's Tuner?

No. Tuner(trainer) is deprecated as of Ray 2.43. The documented pattern launches Ray Train inside a Tune trial function and sweeps its settings through param_space.

Which tuning library does Databricks recommend?

Optuna for single-node optimization or Ray Tune for distributed tuning. Hyperopt is no longer included in Databricks Runtime ML after 16.4 LTS ML.