Hi 

I have created a model and pipeline using xgboost.spark's sparkXGBregressor and pyspark.ml's Pipeline instance. However, i run into a "RuntimeError: _get_spark_session should not be invoked from executor side." when i try to save the predictions i generate with the model pipeline to a featurestore table. Code:

import mlflow
from mlflow.tracking import MlflowClient
from datetime import datetime
from pyspark.sql.functions import col, struct, lit
from databricks.feature_engineering.entities.feature_lookup import FeatureLookup

from databricks.feature_engineering import FeatureEngineeringClient
fs = FeatureEngineeringClient()

def create_training_data(training_days):
    return (
    fs.read_table(name='dev_fs.xxx_featurestore_experiment.label')
    .where(col("calendardate").isin(training_days))
    )

training_data = create_training_data(training_days)

# Create or get existing feature table
aacc_features_lookup = [
        FeatureLookup(
            table_name='dev_fs.xxx_featurestore_experiment.xxx',
            lookup_key=['xxxid'],
            timestamp_lookup_key=['calendardate'],
            feature_names=['xxx4', 'xxx5', 'xxx6']
        )
    ]

# Create or get existing feature table
apcp_features_lookup = [
        FeatureLookup(
            table_name='dev_fs.xxx_featurestore_experiment.xxx',
            lookup_key=['xxxid'],
            timestamp_lookup_key=['calendardate'],
            feature_names=['xxx1', 'xxx2', 'xxx3']
        )
    ]

feature_lookups = aacc_features_lookup + apcp_features_lookup

# ---------------------
# the model pipeline

from pyspark.ml.feature import VectorAssembler, Imputer
from pyspark.ml import Pipeline
from xgboost.spark import SparkXGBRegressor
from typing import List

def create_new_xgb_model(features: List[str], **kwargs) -> Pipeline:
    """Create a new untrained XGBoost model using PySpark's built-in Imputer"""

    # Use PySpark's built-in Imputer
    imputer = Imputer(
        inputCols=features,
        outputCols=[f"{col}_imputed" for col in features]
    )

    # Assemble the features into a single vector column
    assembler = VectorAssembler(
        inputCols=[f"{col}_imputed" for col in features],
        outputCol="features"
    )

    # Define the Spark XGBoost Regressor
    xgb_classifier = SparkXGBRegressor(
        enable_sparse_data_optim=True,
        features_col="features",
        label_col="age_truncated",
        prediction_col="prediction",
        num_workers=4,
        missing=0.0,
        **kwargs
    )

    # Create and return the pipeline
    return Pipeline(stages=[imputer, assembler, xgb_classifier])

# ---------------------
# mlflow experiment

# Start MLflow run with feature store
with mlflow.start_run(run_name='feature_store_model') as run:
        # Create training set with features from Feature Store
    training_set = fs.create_training_set(
            df=training_data,
            feature_lookups=feature_lookups,
            label='age_truncated'
        )
        
    # Get training data as pandas DataFrame
    training_data = training_set.load_df()
        
    # Train model
    model = create_new_xgb_model(features=list_features, **optim_params)
    pipeline = model.fit(training_data)

    # Log model with feature store
    fs.log_model(
            model=pipeline,
            artifact_path="feature_store_experiment_model",
            flavor=mlflow.spark,
            training_set=training_set,
            registered_model_name="feature_store_experiment-model"
        )

# ----------------------
# batch scoring

batch_df = spark.table("dev_fs.xxx_featurestore_experiment.label")
predictions = fs.score_batch(model_uri=model_uri, df=batch_df.sample(fraction=0.001))
fs.create_table(
    name="dev_fs.xxx_featurestore_experiment.predictions", 
    primary_keys=["xxxid","calendardate"], df=predictions
    )

the full error:

PythonException: 
  An exception was thrown from the Python worker. Please see the stack trace below.
Traceback (most recent call last):
  File "/databricks/python/lib/python3.11/site-packages/mlflow/pyfunc/__init__.py", line 2109, in udf
    loaded_model = mlflow.pyfunc.load_model(local_model_path)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/databricks/python/lib/python3.11/site-packages/mlflow/tracing/provider.py", line 244, in wrapper
    except MlflowTracingException as e:
  File "/databricks/python/lib/python3.11/site-packages/mlflow/pyfunc/__init__.py", line 1046, in load_model
    _clear_dependencies_schemas()
  File "/databricks/python/lib/python3.11/site-packages/mlflow/spark/__init__.py", line 957, in _load_pyfunc
    spark_model = _load_model(model_uri=path)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/databricks/python/lib/python3.11/site-packages/mlflow/spark/__init__.py", line 838, in _load_model
    return PipelineModel.load(model_uri)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/databricks/spark/python/pyspark/ml/util.py", line 465, in load
    return cls.read().load(path)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/databricks/spark/python/pyspark/ml/pipeline.py", line 288, in load
    uid, stages = PipelineSharedReadWrite.load(metadata, self.sc, path)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/databricks/spark/python/pyspark/ml/pipeline.py", line 442, in load
    stage: "PipelineStage" = DefaultParamsReader.loadParamsInstance(stagePath, sc)
                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/databricks/spark/python/pyspark/ml/util.py", line 750, in loadParamsInstance
    instance = py_type.load(path)
               ^^^^^^^^^^^^^^^^^^
  File "/databricks/spark/python/pyspark/ml/util.py", line 465, in load
    return cls.read().load(path)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/databricks/python/lib/python3.11/site-packages/xgboost/spark/core.py", line 1733, in load
    _get_spark_session().sparkContext.textFile(model_load_path).collect()[0]
    ^^^^^^^^^^^^^^^^^^^^
  File "/databricks/python/lib/python3.11/site-packages/xgboost/spark/utils.py", line 94, in _get_spark_session
    raise RuntimeError(
RuntimeError: _get_spark_session should not be invoked from executor side.

any help would be much appreciated. I am trying to test out databricks featurestore clients functions (like fs.log_model, fs.batch_score, etc.), so i would like suggestions where these functions are used.