Tuning with Optuna and MlflowSparkStudy

[AdamIH123]

I am following the guide for tuning a model with Optuna and MlflowSparkStudy. My compute is configured with autoscaling enabled, with 1–2 Spark workers, each with 8 cores and 32 GB of memory. I set n_jobs=2 and trials=100, in mlflow_study.optimize().

I have a few questions about how MlflowSparkStudy distributes the workload:

1. How should I think about n_jobs in mlflow_study.optimize() vs. num_threads/n_jobs in LightGBM?

   For example, with 1–2 Spark workers and 8 cores per worker, would it make sense to set mlflow_study.optimize(n_jobs=2) and LightGBM n_jobs=7? My dataset is large, so ideally I would like to run only one model-training job per Spark worker and use the remaining cores on that worker for the LightGBM training.

2. Does MlflowSparkStudy automatically trigger or make use of Spark autoscaling?

   If I configure my cluster with a minimum of 1 worker and a maximum of 2 workers, will MlflowSparkStudy cause Spark to scale up to 2 workers as needed when running multiple trials in parallel?

3. What is the recommended way to make a large DataFrame available to each Spark worker?

   Currently, I have a Pandas DataFrame that I pass to the objective function used by mlflow_study.optimize(). Should I broadcast the DataFrame, cache it in Spark, or use another approach to avoid repeatedly transferring the data to each worker for every trial?

Any guidance on the recommended configuration or best practices would be greatly appreciated.

#mlflow #optuna #MlflowSparkStudy

https://learn.microsoft.com/en-us/azure/databricks/machine-learning/automl-hyperparam-tuning/optuna 

[kunduruanil]

@AdamIH123 Did you try this?

from Synapse. ml.lightgbm import LightGBMClassifier

Convert a Spark Pandas DataFrame to a standard PySpark DataFrame.

spark_df = spark_pandas_df.to_spark()

Train distributed LightGBM

lgbm = LightGBMClassifier(learningRate=0.1, numIterations=100)

model = lgbm.fit(spark_df)

[ThiamLee]

Great questions—especially the distinction between Optuna’s trial-level parallelism and LightGBM’s intra-trial threading. The interaction with Spark autoscaling and data locality is also something I’d love to see documented with a concrete example. Curious what configuration others have found most efficient for large datasets.