title: Why is Pandas UDF not being parallelized?
body: <p>I have data from many IoT sensors. For each particular sensor, there's only  about 100 rows in the dataframe: the data is not skewed. I'm training an individual machine learning model for each sensor.</p>
  <p>I'm using <code>pandas udf</code> successfully to train and log mlflow metrics of different models in parallel (supposedly), as taught <a href="https://databricks.com/blog/2020/05/19/manage-and-scale-machine-learning-models-for-iot-devices.html" rel="nofollow noreferrer">here</a>.</p>
  <p>Using Databricks on Azure with a single node cluster (Standard_DS3_v2 - 14GB of memory - 4 cores) I was able to finish all the training in about <strong>23min</strong>.</p>
  <p>Since <code>pandas udf</code>, supposedly, does the computing in parallel for each group, I thought that I could finish the training faster by using a single node cluster with more cores, or by using a cluster with more workers. So I tried to run the same notebook with:</p>
  <ol>
  <li>A cluster of computers: 1 master + 3 workers, all (Standard_DS3_v2 - 14GB of memory - 4 cores)</li>
  <li>A single node cluster with (Standard_DS5_v2 - 56GB of memory - 16 cores)</li>
  </ol>
  <p>For my surprise, the training time did not decreased: <strong>23min</strong> for option 1, and <strong>26.5min</strong> for option 2</p>
  <p>I tried to use the <a href="https://spark.apache.org/docs/3.0.1/sql-pyspark-pandas-with-arrow.html#pandas-function-apis" rel="nofollow noreferrer">newer</a> <code>applyInPandas</code>, but the result was roughly the same.</p>
  <p>NB: After @Chris answer, checking the Stage Detail page on the Web UI (for the cluster with 1 master + 3 workers), I see that I have only one stage responsible for the udf pandas training. It took 20 min. Acessing the details of this stage, I see it had only one task, with <code>{'Locality Level': 'PROCESS_LOCAL'}</code>, that took the whole 20min. Screenshots below.</p>
  <p>So @Chris has identified the problem: the training is not being parallelized.</p>
  <p>In order to understand why the <code>applyInPandas</code> (or <code>udf pandas</code>)  is not being parallelized, I put my code below (with the <code>applyInPandas</code> version). Note that my goal is only to log the trained models metrics with <code>mlflow</code>, so what the function returns is just the original df it receives.</p>
  <p>Note also that the code is working as expected. mlflow is logging the training successfully. My only issue is why it's not being parallelized.</p>
  <p>I have a feeling that the issue is with the <code>for loop</code>, since it differs from the <a href="https://databricks.com/blog/2020/05/19/manage-and-scale-machine-learning-models-for-iot-devices.html" rel="nofollow noreferrer">tutorial</a>.</p>
  <pre><code>import pyspark.sql.functions as f
  import mlflow
  import pandas as pd
  import numpy as np
  from sklearn.metrics import mean_squared_error, mean_absolute_error
  import pmdarima as pm
  from statsmodels.tsa.statespace.sarimax import SARIMAX
  
  def train_model(df_pandas):
    '''
    Trains a model on grouped instances
    '''
    original_df = df_pandas.copy() #the original df will be returned in the end
    PA = df_pandas['Princípio_Ativo'].iloc[0]
    run_id = df_pandas['run_id'].iloc[0] # Pulls run ID to do a nested run
    
    
    observacoes_no_teste = 12
    horizonte = 1
    observacoes_total = len(df_pandas.index)
    observacoes_no_train = len(df_pandas.index) - observacoes_no_teste
    
    try:
      #train test split
      X = df_pandas[:observacoes_no_train]['Demanda']
      y = df_pandas[observacoes_no_train:]['Demanda']
  
      # Train the model
      model = pm.auto_arima(X, seasonal=True, m=12)
  
      order = model.get_params()['order']
      seasonal_order = model.get_params()['seasonal_order']
  
  
    except:
      pass
   
    # Resume the top-level training
    with mlflow.start_run(run_id=run_id, experiment_id=1333367041812290):
      # Create a nested run for the specific device
      with mlflow.start_run(run_name=str(PA), nested=True, experiment_id=1333367041812290) as run:
        
        mae_list = []
        mse_list = []
        previsoes_list = []
        teste_list = []
        predictions_list = []
  
      
        try:
          #the purpose of the following loop is to do backtesting: the model is trained with n observations, and the (n+1)th is predicted. n is increased by 1 on each iteration.
          for i in range(observacoes_total-observacoes_no_train-horizonte+1):
            #train test split
            X = df_pandas[:observacoes_no_train+i]['Demanda']
            y = df_pandas[observacoes_no_train+i:observacoes_no_train+i+horizonte]['Demanda']
            #train model
            model = SARIMAX(X, order=order, seasonal_order=seasonal_order)
            model = model.fit()
            #make predictions
            predictions = model.predict(start=observacoes_no_train + i, end=(observacoes_no_train + i + horizonte-1))
  
            predictions_list.append(predictions)
  
            mse = round(mean_squared_error(y, predictions),2)
            mae = round(mean_absolute_error(y, predictions),2)
  
            mse_list.append(mse)
            mae_list.append(mae)
  
          #series with predictions
          in_sample_predictions = pd.concat(predictions_list)
          in_sample_predictions.name = 'in_sample'
          #out of sample predictions
          hp = 3
          out_of_sample_predictions = model.predict(start=observacoes_total, end=(observacoes_total + hp - 1))
          out_of_sample_predictions.name = 'out_sample'
          #in sample + out of sample predictions
          df_predictions = pd.concat([df_pandas.drop('run_id',axis=1), in_sample_predictions,out_of_sample_predictions], axis=1)
          #save df with predictions to be logged as an artifact my mlflow.
          df_predictions.to_csv('df_predictions.csv')
  
          #mlflow logging
          mlflow.log_param(&quot;Princípio_Ativo&quot;, PA)
          mlflow.log_param(&quot;mae_list&quot;, str(mae_list))
          mlflow.log_param(&quot;mse_list&quot;, str(mse_list))
          mlflow.log_param(&quot;status_sucesso&quot;, 'sim')
          mlflow.log_artifact('df_predictions.csv')
        except:
          mlflow.log_param(&quot;status_falha&quot;, 'sim')
  
    return original_df.drop('run_id', axis=1) 
  
  </code></pre>
  <pre><code>with mlflow.start_run(run_name=&quot;SARIMA&quot;, experiment_id=1333367041812290) as run:
    run_id = run.info.run_uuid
  
    modelDirectoriesDF = (df
      .withColumn(&quot;run_id&quot;, f.lit(run_id)) # Add run_id
      .groupby(&quot;Princípio_Ativo&quot;)
      .applyInPandas(train_model, schema=df.schema)
    )
    
  combinedDF = (df
    .join(modelDirectoriesDF, on=&quot;Princípio_Ativo&quot;, how=&quot;left&quot;)
  )
  
  display(combinedDF)
  </code></pre>
  <p>Screenshots from the Spark UI:
  <a href="https://i.stack.imgur.com/9bMNu.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/9bMNu.png" alt="Picture1" /></a>
  <a href="https://i.stack.imgur.com/qOOza.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/qOOza.png" alt="picture2" /></a>
  <a href="https://i.stack.imgur.com/Yc0V5.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/Yc0V5.png" alt="picture3" /></a></p>