This message was deleted.

Hey. Thanks for using statsforecast. • I believe the df is repartitioned again after you pass it to the forecast method, so I think the partitioning wouldn't make a difference, this is something I'm going to work on this week. • The

n_jobs

controlls the parallelism within each partition, but since I think each partition holds a single serie it wouldn't make a difference (I'll also verify this). • There's an

nmodels

argument that controls how many combinations are tried, so that's an easy one to limit the search space.

Hi @José Morales, are you talking about the default partition config in spark, or the DF partition after reading? I don't know if the data is distributed differently later in the code, but at the time of reading it has only one partition (because it's read from a single file, written with pandas)

Hey. I meant the DF partitioning, so exactly what you did is what I expect would help reducing the time. How many CPUs does each executor have? I'm not sure if two tasks can run concurrently on the same executor

Hello José, I don't find the precise answer, but during the ARIMA fit, I can see that the maximum number of task is 200 (20/32 per executor = compute workers on databricks)

I think it's doing a shuffle first and then the training, so that's probably the shuffle partitions. Can you increase them to a multiple of your cores? e.g.

spark.conf.set("spark.sql.shuffle.partitions", 640)

It seems to work, up to 98% cpu usage with this settings:

num_worker_cores = 320
spark.conf.set("spark.sql.shuffle.partitions", num_worker_cores*4)
spark.conf.set("spark.sql.adaptive.enabled", "false")

Nice!

I retested,

spark.conf.set("spark.sql.shuffle.partitions", num_worker_cores)

is enough, the problem was simply to overcome the default limit of 200. Thanks again for your help @José Morales!