为什么 Pandas UDF 没有被并行化？

Question

我有来自许多物联网传感器的数据。对于每个特定的传感器，数据框中只有大约 100 行：数据没有倾斜。我正在为每个传感器训练一个单独的机器学习模型。

按照here 的教导，我成功地使用pandas udf 并行训练和记录不同模型的mlflow 指标（假设）。

在 Azure 上使用 Databricks 和单节点集群（Standard_DS3_v2 - 14GB 内存 - 4 核）我能够在大约 23 分钟内完成所有训练。

由于pandas udf 应该为每个组并行计算，我认为我可以通过使用具有更多内核的单节点集群或使用具有更多工作人员的集群来更快地完成训练。所以我尝试运行同一个笔记本：

1. 计算机集群：1 台主机 + 3 台工作机，全部（Standard_DS3_v2 - 14GB 内存 - 4 核）
2. 具有（Standard_DS5_v2 - 56GB 内存 - 16 核）的单节点集群

令我惊讶的是，训练时间并没有减少：选项 1 23min，选项 2 26.5min

我尝试使用newerapplyInPandas，但结果大致相同。

注意：@Chris 回答后，检查 Web UI 上的 Stage Detail 页面（对于具有 1 个 master + 3 个 worker 的集群），我看到我只有一个 stage 负责 udf pandas 培训。花了20分钟。访问这个阶段的细节，我看到它只有一个任务，{'Locality Level': 'PROCESS_LOCAL'}，花了整整 20 分钟。截图如下。

所以@Chris 发现了问题：训练没有被并行化。

为了理解为什么applyInPandas（或udf pandas）没有被并行化，我把我的代码放在下面（applyInPandas 版本）。请注意，我的目标只是使用mlflow 记录训练后的模型指标，因此函数返回的只是它接收到的原始 df。

另请注意，代码按预期工作。 mlflow 正在成功记录培训。我唯一的问题是为什么它没有被并行化。

我感觉问题出在for loop，因为它与tutorial 不同。

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("Princípio_Ativo", PA)
        mlflow.log_param("mae_list", str(mae_list))
        mlflow.log_param("mse_list", str(mse_list))
        mlflow.log_param("status_sucesso", 'sim')
        mlflow.log_artifact('df_predictions.csv')
      except:
        mlflow.log_param("status_falha", 'sim')

  return original_df.drop('run_id', axis=1) 

with mlflow.start_run(run_name="SARIMA", experiment_id=1333367041812290) as run:
  run_id = run.info.run_uuid

  modelDirectoriesDF = (df
    .withColumn("run_id", f.lit(run_id)) # Add run_id
    .groupby("Princípio_Ativo")
    .applyInPandas(train_model, schema=df.schema)
  )
  
combinedDF = (df
  .join(modelDirectoriesDF, on="Princípio_Ativo", how="left")
)

display(combinedDF)

Spark UI 的屏幕截图：

Answer 1

我认为您可以通过在执行最后一个显示命令之前对combinedDF 执行sort() 来改进并行化。比如：

import pyspark.sql.functions as F
display(combinedDF.sort(F.col("Princípio_Ativo").asc()))

上面的代码 sn-p 将阻止 display 函数推测仅评估 combinedDF 中的某些行。 这个想法是，通过对display 之前的行进行排序，spark 必须对它们全部进行评估，以了解display 函数的行顺序。因此，一切都安排为并行执行。 在没有先排序的情况下执行display，spark 采用某种形式的阶段性指数增加计划任务的数量，直到评估足够的行以执行display。

Answer 2

正确。只有一个任务的阶段没有并行化。这可以解释为什么在集群中添加更多核心或节点不会减少运行时间。

您的输入数据集很小 (69KB)，因此除非您明确 repartition，否则如果数据框分区大小保留为默认 128MB（由 spark.sql.files.maxPartitionBytes 参数决定），Spark 会将其写入单个分区。因此，它将被分配给单个任务。

按设备列重新分区输入应该提供您正在寻找的并行训练。