Log and load models
model_name = "sk-learn-random-forest-reg-model"
model_version = "latest"
# Load the model from the Model Registry
model_uri = f"models:/{model_name}/{model_version}"
model = mlflow.sklearn.load_model(model_uri)
# Alternatively
# load input data table as a Spark DataFrame
input_data = spark.table(input_table_name)
model_udf = mlflow.pyfunc.spark_udf(spark, model_path)
df = input_data.withColumn("prediction", model_udf())
Deploy single-node model in parallel using spark_udf
predict = mlflow.pyfunc.spark_udf(spark, f"runs:/{run.info.run_id}/random_forest_model")
# Apply the model as a standard UDF using the column names as the input to the function.
prediction_df = spark_df.withColumn("prediction", predict(*spark_df.columns))
Write Optimizations
Partition
dbutils.fs.rm(f"{DA.paths.working_dir}/batch-predictions-partitioned.delta", recurse=True)
delta_partitioned_path = f"{DA.paths.working_dir}/batch-predictions-partitioned.delta"
prediction_df.write.partitionBy("neighbourhood_cleansed").mode("OVERWRITE").format("delta").save(delta_partitioned_path)
Z-Ordering
spark.sql(f"OPTIMIZE delta.`{delta_partitioned_path}` ZORDER BY (zipcode)")
Feature Store Batch Score
batch_input_df = inference_df.drop("price") #exclude true label
with_predictions = fs.score_batch(f"models:/{model_name}/1",
batch_input_df, result_type='double')
display(with_predictions)
- Workflow:
- Load dataset: The dataset used for the example is
diamonds dataset from Seaborn library
- Create Feature table in Feature Store
- Push preprocessed features to Feature Store
- Create train and test sets
- Prepare and train models
- Log models to associate them to features in Features Store
- Score models on test set using
score_batch
- Case of new data
# 1. Load dataset
pd_diamonds = sns.load_dataset('diamonds').reset_index()
diamonds_full = spark.createDataFrame(pd_diamonds).withColumnRenamed('x', 'x_r')
display(diamonds_full.limit(5))
# 2. Create Feature table in Feature Store
fs = feature_store.FeatureStoreClient()
# fs.drop_table("default.diamonds_fs")
# create feature table - as only the scema is provided in the command below, it will only create the table structure without populating it with data
result = fs.create_table(name="diamonds_fs", # required
primary_keys=["index"], # required
schema=diamonds_full.drop("price").schema, # need either dataframe schema
#df=diamonds_full, # or dataframe itself
description="seaborn diamonds dataset");
# 3. Push preprocessed features to Feature table
fs.write_table(name="diamonds_fs",
df=diamonds_full.drop('price'),
mode='merge')
# 4. Create train and tst sets
y_test = diamonds_full.select("price", "index").orderBy(rand()).limit(int(33*diamonds_full.count()/100))
y_train = diamonds_full.select("price", "index").subtract(y_test)
display(y_train.limit(5))
# With 4 features: x, y, z, carat
feature_lookups_4_features = [feature_store.FeatureLookup(table_name="diamonds_fs",
feature_names=['x_r', 'y', 'z', 'carat'],
lookup_key="index")]
# create associated training set
train_set_4_features = fs.create_training_set(y_train,
feature_lookups_4_features,
label="price",
exclude_columns="index")
# load training set
train_set_4 = train_set_4_features.load_df()
# display to check
display(train_set_4.limit(5))
# With all numerical features
feature_lookups_all_features = [feature_store.FeatureLookup(table_name="diamonds_fs",
feature_names=[c for c in diamonds_full.columns if c not in ['index', 'cut', 'clarity', 'price', 'color']],
lookup_key="index")]
# create associated training set
train_set_all_features = fs.create_training_set(y_train,
feature_lookups_all_features,
label="price",
exclude_columns="index")
# load training set
train_set_all = train_set_all_features.load_df()
# display to check
display(train_set_all.limit(5))
# 5. Prepare and train models
X_train_4 = train_set_4.drop("price").toPandas()
y_train_4 = train_set_4.toPandas()["price"]
rf_4_model = RandomForestRegressor()
rf_4_model.fit(X_train_4, y_train_4)
X_train_all = train_set_all.drop("price").toPandas()
y_train_all = train_set_all.toPandas()["price"]
rf_all_model = RandomForestRegressor()
rf_all_model.fit(X_train_all, y_train_all)
# 6. Log models to associate them to featurs in Feature Store
model_name_4_features = "trained_with_4_features"
fs.log_model(rf_4_model,
artifact_path=model_name_4_features, # parameter required
flavor=mlflow.sklearn, # parameter required
training_set=train_set_4_features, # either training_set or feature_spec_path parameters required
registered_model_name=model_name_4_features); # not required. However model will not be linked to features in features store
model_name_all_features = "trained_with_all_features"
fs.log_model(rf_all_model,
artifact_path=model_name_all_features, # parameter required
flavor=mlflow.sklearn, # parameter required
training_set=train_set_all_features, # either training_set or feature_spec_path parameters required
registered_model_name=model_name_all_features); # not required. However model will not be linked to features in features store until model is registered
# 7. Score models on test set using score_batch
# latest run id for model named "trained_with_4_features"
for val in mlflow.MlflowClient().get_registered_model(model_name_4_features):
if val[0]=='latest_versions':
run_id_4 = val[1][0].run_id
# uri to latest run
uri_4_features = f"runs:/{run_id_4}/{model_name_4_features}"
print(uri_4_features)
# predict on test set
predictions_df_4_features = fs.score_batch(uri_4_features, y_test).select("price", "prediction");
display(predictions_df_4_features.orderBy(rand()).limit(5))
# Score RMSE on test set for model trained with 4 features
print("RMSE for model trained on 4 features:",
mean_squared_error(predictions_df_4_features.toPandas()['price'], predictions_df_4_features.toPandas()['prediction'], squared=False))
# latest run id for model named "trained_with_all_features"
for val in mlflow.MlflowClient().get_registered_model(model_name_all_features):
if val[0]=='latest_versions':
run_id_all = val[1][0].run_id
# uri to latest run
uri_all_features = f"runs:/{run_id_all}/{model_name_all_features}"
print(uri_all_features)
# predict on test set
predictions_df_all_features = fs.score_batch(uri_all_features, y_test).select("price", "prediction", "carat");
display(predictions_df_all_features.select("price", "prediction").orderBy(rand()).limit(5));
print(mean_squared_error(predictions_df_all_features.toPandas()['price'], predictions_df_all_features.toPandas()['prediction'], squared=False))
# Comparison of actual price and predicted price according to `carat` for model trained using all numerical features, on a sample of 1000 random entries:
plt.figure(figsize=(17, 6))
sample_predictions = predictions_df_all_features.orderBy(rand()).limit(1000)
plt.scatter(list(sample_predictions.toPandas()['carat']), list(sample_predictions.toPandas()['price']), label='Actual Price', color='blue', marker='o')
plt.scatter(list(sample_predictions.toPandas()['carat']), list(sample_predictions.toPandas()['prediction']), label='Prediction', color='orange', marker='s')
# Adding labels and title
plt.xlabel('Carat')
plt.ylabel('Prices')
plt.title('Actual Price vs Prediction')
# Adding grid for better readability
plt.grid(True, linestyle='--', alpha=0.7)
# Adding legend
plt.legend()
# Show plot
plt.show();
# 8.Case of new data
# How to predict newly arriving data if there is no information on it in Feature Store? Need to update first the Feature Store with new data:
new_diamond = (diamonds_full.limit(1).withColumn('index', lit(88887777).cast('long'))
.withColumn('carat', lit(2).cast('double'))
.withColumn('cut', lit('Good').cast('string'))
.withColumn('color', lit('E').cast('string'))
.withColumn('clarity', lit('VS1').cast('string'))
.withColumn('depth', lit(40).cast('double'))
.withColumn('table', lit(64).cast('double'))
.withColumn('x_r', lit(4.14).cast('double'))
.withColumn('y', lit(3.5).cast('double'))
.withColumn('z', lit(2.1).cast('double')))
new_diamond_with_price = spark.createDataFrame(pd.DataFrame({'index': [88887777], 'price': [4500]}))
new_diamond_without_price = spark.createDataFrame(pd.DataFrame({'index': [88887777]}))
diamond_unknown = spark.createDataFrame(pd.DataFrame({'index': [98989898]}))
display(new_diamond)
display(new_diamond_with_price)
display(new_diamond_without_price)
# Update the Feature Store with the new diamond data
fs.write_table(name="diamonds_fs",
df=new_diamond,
mode='merge')
# predict with price
predictions_new_diamond_with_price = fs.score_batch(uri_all_features, new_diamond_with_price)
display(predictions_new_diamond_with_price)
# predict without price
predictions_new_diamond_without_price = fs.score_batch(uri_all_features, new_diamond_without_price)
display(predictions_new_diamond_without_price)