[blog] It’s beaver time! Don’t get logged down with mlflow logging.

logging_many_models.py

# Databricks notebook source
import pandas as pd 
import numpy as np 
import cloudpickle 
from datetime import datetime 
import mlflow
import os 
import json

# COMMAND ----------

# MAGIC %md
# MAGIC ## Creating dummy data for forecasting

# COMMAND ----------

CATALOG = "" 
SCHEMA = ""

# COMMAND ----------

def create_sample_data():
    """Create sample time series data for demonstration"""
    np.random.seed(42)
    
    data = []
    for group_id in ['A', 'B', 'C']:
        dates = pd.date_range('2023-01-01', periods=100, freq='D')
        trend = np.linspace(10, 50, 100)
        noise = np.random.normal(0, 5, 100)
        values = trend + noise + np.random.uniform(0, 20)
        
        for date, value in zip(dates, values):
            data.append({
                'group_id': group_id,
                'date': date,
                'target': max(0, value),
                'feature1': np.random.uniform(-1, 1),
                'feature2': np.random.uniform(0, 10)
            })
    
    return pd.DataFrame(data)
  

main_df = create_sample_data()
main_df.head()

# COMMAND ----------

# MAGIC %md
# MAGIC ## Training multiple models
# MAGIC
# MAGIC Logging many models (in the thousands) using `mlflow.log_model` can slow the code down by a lot. In order to speed this process up, we recommend writing the model to a Delta table. 

# COMMAND ----------

RANDOM_STATE = 42

GROUP_IDS = ['A', 'B', 'C']

# COMMAND ----------

from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error

from pyspark.sql.types import (
    StructType, StructField, StringType, TimestampType, 
    BinaryType, DoubleType, ArrayType
)

from mlflow.models.signature import infer_signature

# COMMAND ----------

###  Create the Delta Table ### 
# this should only be created once
# will hold all of your runs + models

schema = StructType([
    StructField("group_id", StringType(), True),
    StructField("model_type", StringType(), True),
    StructField("model_version", StringType(), True),  
    StructField("model_binary", BinaryType(), True),
    StructField("run_id", StringType(), True),
    StructField("run_date", TimestampType(), True),
    StructField("mse", DoubleType(), True),
    StructField("forecast", ArrayType(DoubleType()), True),
    StructField("actual", ArrayType(DoubleType()), True),
    StructField("is_latest", StringType(), True)  
])

spark.createDataFrame([], schema).write.format("delta").option("overwriteSchema", "true").saveAsTable(f"{CATALOG}.{SCHEMA}.mlflow_runs")

# COMMAND ----------

# creating a dummy wrapper so I can log a Model -> get model artifacts

class DummyWrapper(mlflow.pyfunc.PythonModel): 
  def __init__(self):
    self.dummy = 'dummy'
  
  def predict(self, context, model_input): 
    return model_input

# COMMAND ----------

def train_model(group_df, group_id, latest_model_version, run_id, run_date): 
  features = ['feature1', 'feature2']
  X = group_df[features].values
  y = group_df['target'].values 

  model = RandomForestRegressor(n_estimators=10, random_state=RANDOM_STATE)
  model.fit(X, y)
  predictions = model.predict(X)
  mse = mean_squared_error(y, predictions)

  # get metadata to add to Delta table
  return {
    'group_id': group_id,
    'model_type': 'RandomForestRegressor', 
    'model_version': str(latest_model_version + 1), 
    'model_binary': cloudpickle.dumps(model), 
    'run_id': run_id,
    'run_date': run_date, 
    'mse': mse, 
    'forecast': predictions.tolist(),
    'actual': y.tolist(), 
    'is_latest': "True"
  }

def get_latest_model_versions(table_name):
  spark.sql(f"update {table_name} set is_latest='false'")

  where_clauses = " OR ".join([f"group_id = '{group_id}'" for group_id in GROUP_IDS])
  query = f"SELECT group_id, MAX(CAST(model_version AS int)) AS max_version FROM {table_name} WHERE {where_clauses} GROUP BY group_id"
  version_df = spark.sql(query).toPandas()
  
  # set up the key -- for example, if you are training multiple types of models per group, add that to the key
  version_df['key'] = version_df['group_id'] # + '_' + version_df['model_type']
  version_dict = version_df.set_index('key')['max_version'].fillna(0).to_dict()
  return version_dict

def save_to_delta(model_results, table_name): 
  df = spark.createDataFrame(model_results)
  df.write.format("delta").mode("append").saveAsTable(table_name)

def log_models_to_mlflow(data, table_name=f"{CATALOG}.{SCHEMA}.mlflow_runs"): 
  with mlflow.start_run() as run: 
    run_id = run.info.run_id
    run_date = datetime.now()

    mlflow.log_param("num_groups", len(data['group_id'].unique()))
    mlflow.log_param("delta_table_name", table_name)

    current_model_versions = get_latest_model_versions(table_name) # {'group_id': version #}
    all_model_results = []

    for group_id in GROUP_IDS: 
      group_df = data[data['group_id'] == group_id] 

      latest_version = current_model_versions.get(group_id, 0)
      model_result = train_model(group_df, group_id, latest_version, run_id, run_date)
      all_model_results.append(model_result)

    save_to_delta(all_model_results, table_name)

    mlflow.pyfunc.log_model(
        "dummy_model",
        input_example=main_df,
        python_model=DummyWrapper()
    )



# COMMAND ----------

# Run this a few times-- model_versions get incremented + is_latest is set to the latest models trained
log_models_to_mlflow(main_df)
display(spark.table(f"{CATALOG}.{SCHEMA}.mlflow_runs"))

# COMMAND ----------

# MAGIC %md
# MAGIC ## Load models from Delta tables

# COMMAND ----------

class MultiModelWrapper(): 
  def __init__(self, table_name): 
    self.table = table_name

  def load_model_from_delta(self, group_id, table_name, model_type=None, run_id=None, version=None): 
    query = f"select * from {table_name} where group_id = '{group_id}'"
    if model_type: 
      query += f" and model_type = '{model_type}'"
    if run_id: 
      query += f" and run_id = '{run_id}'"
    if version: 
      query += f" and model_version = '{version}'"
    else: 
      query += f" and is_latest = 'True'" 
    model_df = spark.sql(query).collect()
    if model_df: 
      model = cloudpickle.loads(model_df[0]['model_binary'])
      metadata = model_df[0].asDict(True)
      metadata.pop("model_binary")
      return model, metadata
    else: 
      return None, None
    
  def predict(self, model_input, group_id, model_type=None, run_id=None, version=None): 
    model, _ = self.load_model_from_delta(group_id=group_id, model_type=model_type, run_id=run_id, version=version, table_name=self.table)

    # TODO: make sure the model_input can be ingested by the models!
    return model.predict(model_input.values)

# COMMAND ----------

# instantiate
wrapper_model = MultiModelWrapper(table_name=f"{CATALOG}.{SCHEMA}.mlflow_runs")

test_df = main_df.head(1).drop(columns=["group_id", "target", "date"])
wrapper_model.predict(test_df, 'A', version=2)