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