并行化 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 没有被并行化？