Re: spark_session invocation from executor side er...

Walter_C · ‎11-14-2024

The error you're encountering is due to attempting to access the Spark session on the executor side, which is not allowed in Spark's distributed computing model. This typically happens when trying to use Spark-specific functionality within a UDF or during model inference on executors.To resolve this issue and use Databricks Feature Store functions with your XGBoost model, you need to make some adjustments to your approach. Here are some suggestions:

Use MLflow's pyfunc flavor instead of Spark flavor:

When logging the model, use mlflow.pyfunc instead of mlflow.spark. This will allow you to use the model for batch scoring without encountering Spark session issues on executors.

python

fs.log_model(
    model=pipeline,
    artifact_path="feature_store_experiment_model",
    flavor=mlflow.pyfunc,
    training_set=training_set,
    registered_model_name="feature_store_experiment-model"
)

Use fs.score_batch() for batch scoring:

Instead of loading the model and applying it directly, use the Feature Store's score_batch() function, which is designed to handle this scenario:

python

batch_df = spark.table("dev_fs.xxx_featurestore_experiment.label")
predictions = fs.score_batch(
    model_uri="models:/feature_store_experiment-model/latest",
    df=batch_df.sample(fraction=0.001)
)

Create the feature store table with predictions:

After getting the predictions, you can create or update the feature store table:

python

fs.create_table(
    name="dev_fs.xxx_featurestore_experiment.predictions", 
    primary_keys=["xxxid", "calendardate"],
    df=predictions
)

Ensure consistent feature names:

Make sure that the feature names used during training match those in your feature lookups and batch scoring dataframe.

Consider using Koalas or Pandas UDFs:

If you need to perform operations that require access to the Spark session on executors, consider using Koalas or Pandas UDFs, which are designed to work in a distributed environment.

Optimize for large datasets:

If you're working with large datasets, consider using Spark's built-in ML library (MLlib) instead of XGBoost, as it's designed to work natively with Spark's distributed computing model.