在 databricks 上使用 MLFlow 和 Spark 时出现运行时错误

Question

这是我创建的一些模型：

class SomeModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, input):
        # do fancy ML stuff
        # log results
        pandas_df = pd.DataFrame(...insert predictions here...)
        spark_df = spark.createDataFrame(pandas_df)
        spark_df.write.saveAsTable('tablename', mode='append')

我正在尝试通过稍后在代码中调用它来以这种方式记录我的模型：

with mlflow.start_run(run_name="SomeModel_run"):
    model = SomeModel()
    mlflow.pyfunc.log_model("somemodel", python_model=model)

不幸的是，它给了我这个错误消息：

RuntimeError: It appears that you are attempting to reference SparkContext from a broadcast variable, action, or transformation. SparkContext can only be used on the driver, not in code that it run on workers. For more information, see SPARK-5063.

该错误是由

mlflow.pyfunc.log_model("somemodel", python_model=model)

行引起的，如果我将其注释掉，我的模型将做出预测并将结果记录在我的表中。

或者，删除预测函数中调用 Spark 创建数据帧并保存表的行，我就可以记录我的模型。

我该如何解决这个问题？我需要我的模型不仅要写入表还要被记录

Answer 1

我遇到了类似的错误，我可以通过在类之外使用“spark.createDataFrame(pandas_df)”等 Spark 函数来解决此问题。如果您想使用 Spark 读写数据，请在 main 函数中进行。

    class SomeModel(mlflow.pyfunc.PythonModel):
       def predict(self, context, input):
       # do fancy ML stuff
       # log results
       return predictions

    with mlflow.start_run(run_name="SomeModel_run"):
      model = SomeModel()
      pandas_df = pd.DataFrame(...insert predictions here...)
      mlflow.pyfunc.log_model("somemodel", python_model=model)

Answer 2

如果您的目标是在端点中记录模型的输出，Databricks 发布了一项名为“推理表”的监控功能。激活它们时，它将记录端点的请求和响应。然后，您可以对这些表执行转换以监督模型性能、生成统计数据、检查模型/数据偏差，甚至创建仪表板，Databricks 文档中有一些示例。