Load features from Unity Catalog for model training

from databricks.feature_engineering import FeatureLookup
import mlflow

pickup_features_table = "ml.taxi_example.trip_pickup_time_series_features"

pickup_features_lookups = [
	FeatureLookup(
		table_name=pickup_features_table
		feature_names=[
			"mean_fare_window_1h_pickup_zip",
      "count_trips_window_1h_pickup_zip",
		],
		lookup_key=["pickup_zip"],
    timestamp_lookup_key="tpep_pickup_datetime",
	)
]

Configure MLflow client to access models in Unity Catalog

import mlflow
mlflow.set_registry_uri("databricks-uc")

# End any existing runs (in the case this notebook is being run for a second time)
mlflow.end_run()

# Start an mlflow run, which is needed to log the model
mlflow.start_run()

# Since the timestamp columns would likely cause the model to overfit the data
# unless additional feature engineering was performed, exclude them to avoid training on them.
exclude_columns = ["tpep_pickup_datetime", "tpep_dropoff_datetime"]

# Create the training set that includes the raw input data merged with corresponding features from both feature tables
training_set = fe.create_training_set(
    df=raw_data,
    feature_lookups=pickup_feature_lookups + dropoff_feature_lookups,
    label="fare_amount",
    exclude_columns=exclude_columns,
)

# Load the TrainingSet into a dataframe which can be passed into sklearn for training a model
training_df = training_set.load_df()

Train Model

from sklearn.model_selection import train_test_split
from mlflow.tracking import MlflowClient
import lightgbm as lgb
import mlflow.lightgbm
from mlflow.models.signature import infer_signature

features_and_label = training_df.columns

# Collect data into a Pandas array for training
data = training_df.toPandas()[features_and_label]

train, test = train_test_split(data, random_state=123)
X_train = train.drop(["fare_amount"], axis=1)
X_test = test.drop(["fare_amount"], axis=1)
y_train = train.fare_amount
y_test = test.fare_amount

mlflow.lightgbm.autolog()
train_lgb_dataset = lgb.Dataset(X_train, label=y_train.values)
test_lgb_dataset = lgb.Dataset(X_test, label=y_test.values)

param = {"num_leaves": 32, "objective": "regression", "metric": "rmse"}
num_rounds = 100

# Train a lightGBM model
model = lgb.train(param, train_lgb_dataset, num_rounds)

# Log the trained model with MLflow and package it with feature lookup information.
fe.log_model(
    model=model,
    artifact_path="model_packaged",
    flavor=mlflow.lightgbm,
    training_set=training_set,
    registered_model_name="ml.taxi_example.taxi_example_fare_time_series_packaged",
)