並行化 MLflow 項目在 Azure Databricks Spark 上使用 Pandas UDF 運行

Question

我正在嘗試在 Azure Databricks 上使用 Spark 並行化多個時間序列的訓練。
除了培訓，我還想使用 MLflow 記錄指標和模型。

代碼結構很簡單（基本上改編了這個例子）。

1. Databricks 筆記本觸發 MLflow 項目

mlflow.run( uri="/dbfs/mlflow-project", parameters={"data_path": "dbfs:/data/", "experiment_name": "test"}, experiment_id=575501044793272, use_conda=False, backend="databricks", backend_config={ "new_cluster": { "spark_version": "9.1.x-cpu-ml-scala2.12", "num_workers": 8, "node_type_id": "Standard_DS4_v2", }, "libraries": [{"pypi": {"package": "pyarrow"}}] }, synchronous=False )

2. 主叫function 。 它基本上執行三個步驟：

   1. 讀取提供的data_path指示的增量表
   2. 定義一個 function 觸發 MLflow 項目的“火車入口”
   3. 在 Spark DataFrame 上將此 function 作為 Pandas UDF 應用

這里的代碼：

sc = sparkContext('local')
spark = SparkSession(sc)

@click.argument("data_path")
@click.argument("experiment_name")
def run(data_path: str, experiment_name: str):
            
    df = spark.read.format("delta").load(f"{data_path}")
    result_schema = StructType([StructField("key", StringType())])

    def forecast(data: pd.DataFrame) -> pd.DataFrame:
        child_run = client.create_run(
            experiment_id=experiment,
            tags={MLFLOW_PARENT_RUN_ID: parent_run_id},
        )
        p = mlflow.projects.run(
            run_id=child_run.info.run_id, 
            uri=".",
            entry_points="train",
            parameters={"data": data.to_json(), "run_id": child_run.info.run_id}, 
            experiment_id=experiment,
            backend="local",
            usa_conda=False,
            synchronous=False,
        )

        # Just a placeholder to use pandas UDF
        out = pd.DataFrame(data={"key": ["1"]})
        return out

    client = MLflowClient()
    experiment_path = f"/mlflow/experiments/{experiment_name}"
    experiment = client.create_experiment(experiment_path)

    parent_run = client.create_run(experiment_id=experiment)
    parent_run_id = parent_run.run_id

    # Apply pandas UDF (count() used just to avoid lazy evaluation)
    df.groupBy("key").applyInPandas(forecast, result_schema).count()

3. 火車 function 在每個 key 上被調用。
   這基本上為每個時間序列（即每個鍵）訓練了一個 Prophet model，它記錄了參數和 model。

從集群 stderr 和 stdout 我可以看到 pandas UDF 被正確應用，因為它根據“鍵”列正確划分了整個數據，即一次處理一個時間序列。

問題是監控集群使用只使用一個節點，驅動程序節點：工作沒有分布在可用的工作人員上，盡管 pandas UDF 似乎已正確應用。

這里可能是什么問題？ 我可以提供更多細節嗎？

非常感謝你，馬特奧

Answer 1

看來您需要重新分區輸入 dataframe。否則 spark 將看到單個分區 dataframe 並進行相應處理。

Answer 2

代碼中，spark-config定義'local'為master節點，所以在partitions上有限制，最多是CPU上的core數。 這降低了使用並行化的好處，從而降低了速度。 為了充分利用並行化提高速度，應調用“yarn”或類似的多節點配置，並相應地對輸入 dataframe 進行重新分區。 在代碼中，每個“點擊”裝飾器都將由驅動程序作為單一作業圖進行處理。

一個類似的問題在

為什么 Pandas UDF 沒有被並行化？