Modelo KNN mediante pyfunc devuelve ModuleNotFoundError o FileNotFoundError

Problema

Ha creado un modelo de Sklearn mediante KNeighborsClassifier y está usando pyfunc para ejecutar una predicción.

Por ejemplo:

import mlflow.pyfunc
pyfunc_udf = mlflow.pyfunc.spark_udf(spark, model_uri=model_uri, result_type='string')
predicted_df = merge.withColumn("prediction", pyfunc_udf(*merge.columns[1:]))
predicted_df.collect()

La predicción devuelve un ModuleNotFoundError: No module named 'sklearn.neighbors._classification' mensaje de error.

La predicción también puede devolver un FileNotFoundError: [Errno 2] No usable temporary directory found mensaje de error.

Causa

Cuando se registra un modelo de KNN, todos los puntos de datos usados para el entrenamiento se guardan como parte del archivo Pickle.

Si el modelo se entrena con millones de registros, todos los datos se agregan al modelo, lo que puede aumentar drásticamente su tamaño. Un modelo entrenado en millones de registros puede totalizar fácilmente varios GB.

pyfunc intenta cargar el modelo completo en la memoria caché del Ejecutor cuando se ejecuta una predicción.

Si el modelo es demasiado grande para caber en la memoria, se produce uno de los mensajes de error anteriores.

Solución

Debe usar un algoritmo basado en árbol, como Random Forest o XGBoost para reducir la resolución de los datos en un modelo KNN.

Si tiene datos desequilibrados, intente un método de muestreo como SMOTE al entrenar un algoritmo basado en árbol.