MLOps Pipeline for Call Forecasting
Build an MLOps pipeline that trains an LSTM and LightGBM model, tracks them with MLflow, and serves predictions via a Dockerized FastAPI endpoint.
Introduction
30 Second Summary
Every call center needs to predict how many calls are coming in the next hour to staff the right number of agents. The tricky part is not just building a forecast model. It is tracking which model actually performs best and making sure anyone on your team can use it without wrestling with your local Python setup.
In this project, you will build a containerized call volume forecasting service that trains two competing models, tracks them with experiment tracking, selects the best performer, and serves predictions through an API your team can run with a single command.
What You'll Build
A Dockerized MLOps pipeline that generates synthetic 15-minute interval call data, trains a LightGBM model and a Keras LSTM side by side, tracks them with MLflow, registers the winner as "champion," and serves predictions through a FastAPI endpoint.
By the end of this project, you'll have:
- Two trained models (lgb_v1 and lstm_v4) tracked and compared in MLflow with full metrics logging.
- A registered "champion" model in MLflow's model registry, selected programmatically by lowest error.
- A FastAPI prediction endpoint that serves call volume forecasts from the best model, packaged in Docker for team-wide reproducibility.
- Secret Mission: Add a /compare endpoint that returns both models' predictions side-by-side with their logged metrics.
Are there any prerequisites?
This project assumes basic Python knowledge (writing functions, running scripts) and basic Docker knowledge (understanding what containers are). You do not need prior experience with MLflow, FastAPI, LightGBM, or TensorFlow/Keras.
Generate Synthetic Call Volume Data
Your environment is ready with all the dependencies installed. Now it's time to create the data that powers your forecasting models.
Real call center data takes weeks to collect and involves privacy considerations. You'll generate synthetic data that mimics realistic patterns like business-hour peaks, quiet weekends, and random noise.
In this step, get ready to:
- Generate synthetic call volume data with realistic daily patterns.
- Run the script and confirm the output has expected hourly and weekly patterns.
- Build feature engineering functions for both the LightGBM model and the LSTM model.
Create the data generation script
Call centers experience predictable patterns where mornings and early afternoons are busier, and weekends are quieter. Your script will model these patterns so your ML models have realistic training data.
- In VS Code, create a new file called generate_data.py in your project folder.
- Paste the following code into generate_data.py:
import numpy as np
import pandas as pd
from pathlib import Path
def generate_call_data(days=30, seed=42):
np.random.seed(seed)
start = pd.Timestamp("2025-01-01")
periods = days * 24 * 4 # 15-min intervals
timestamps = pd.date_range(start, periods=periods, freq="15min")
calls = []
for ts in timestamps:
hour = ts.hour
dow = ts.dayofweek
is_weekend = int(dow >= 5)
# Base call volume with hourly pattern
if 9 <= hour <= 11:
base = 45
elif 13 <= hour <= 15:
base = 40
elif 7 <= hour <= 18:
base = 25
else:
base = 5
# Weekend reduction
if is_weekend:
base = base * 0.3
# Add noise
noise = np.random.poisson(lam=max(base * 0.2, 1))
call_count = max(0, int(base + noise - base * 0.2))
calls.append(call_count)
df = pd.DataFrame({"timestamp": timestamps, "call_count": calls})
return df
if __name__ == "__main__":
Path("data").mkdir(exist_ok=True)
df = generate_call_data()
df.to_csv("data/calls.csv", index=False)
print(f"Generated {len(df)} rows")
print(df.describe())
What does this code do?
generate_call_data() creates 30 days of timestamps at 15-minute intervals, giving you 2880 data points.
The hourly pattern sets base call volume: 45 calls during 9-11am, 40 during 1-3pm, 25 for other business hours, and 5 for off-hours. Weekend volume drops to 30% of the weekday base.
np.random.poisson() adds realistic noise so the data isn't perfectly predictable. The __main__ block saves the output to data/calls.csv.
- Save generate_data.py.
- You won't see output yet. The next substep runs this script to generate the CSV file.
Code not looking right?
- Make sure you copied the entire file including the import statements at the top.
- Check that the if __name__ == "__main__": block is at the bottom with no extra indentation.
- Help me check my generate_data.py for syntax errors.
✔️ Awesome, I've got everything!
Great. Double check you've saved generate_data.py before moving on.
ⓧ I'd like to double check the full code
import numpy as np
import pandas as pd
from pathlib import Path
def generate_call_data(days=30, seed=42):
np.random.seed(seed)
start = pd.Timestamp("2025-01-01")
periods = days * 24 * 4 # 15-min intervals
timestamps = pd.date_range(start, periods=periods, freq="15min")
calls = []
for ts in timestamps:
hour = ts.hour
dow = ts.dayofweek
is_weekend = int(dow >= 5)
# Base call volume with hourly pattern
if 9 <= hour <= 11:
base = 45
elif 13 <= hour <= 15:
base = 40
elif 7 <= hour <= 18:
base = 25
else:
base = 5
# Weekend reduction
if is_weekend:
base = base * 0.3
# Add noise
noise = np.random.poisson(lam=max(base * 0.2, 1))
call_count = max(0, int(base + noise - base * 0.2))
calls.append(call_count)
df = pd.DataFrame({"timestamp": timestamps, "call_count": calls})
return df
if __name__ == "__main__":
Path("data").mkdir(exist_ok=True)
df = generate_call_data()
df.to_csv("data/calls.csv", index=False)
print(f"Generated {len(df)} rows")
print(df.describe())
Run the script and verify the output
- Generate your synthetic dataset by running this command:
python generate_data.py
- You should see Generated 2880 rows followed by summary statistics showing count, mean, standard deviation, min, max, and percentiles for the call_count column.
- The max value should be in the 40-55 range, reflecting the peak business-hour volumes.
Seeing an error?
- If you see ModuleNotFoundError, make sure your virtual environment is activated.
- If you see FileNotFoundError, make sure you're running the command from your project folder (not a subfolder).
- Help me debug an error when running generate_data.py.
Build the feature engineering module
Raw timestamps aren't useful for ML models. Feature engineering transforms them into numerical signals that models can learn from. Your LightGBM model needs tabular features, while the LSTM needs fixed-length sequences of past observations.
- In VS Code, create a new file called features.py in your project folder.
- Start with the imports and the base feature function by pasting this code:
import pandas as pd
import numpy as np
def create_base_features(df):
"""Extract time features from raw timestamps."""
df = df.copy()
df["timestamp"] = pd.to_datetime(df["timestamp"])
df["hour"] = df["timestamp"].dt.hour
df["day_of_week"] = df["timestamp"].dt.dayofweek
df["is_weekend"] = (df["day_of_week"] >= 5).astype(int)
# Cyclical encoding so the model knows hour 23 is close to hour 0
df["hour_sin"] = np.sin(2 * np.pi * df["hour"] / 24)
df["hour_cos"] = np.cos(2 * np.pi * df["hour"] / 24)
return df
What does this code do?
- df.copy() protects the original DataFrame from being modified.
- dt.hour and dt.dayofweek extract time components that directly correlate with call volume patterns.
- hour_sin and hour_cos encode the hour as a point on a circle so the model knows that 11pm and midnight are close together, not 23 units apart.
- is_weekend converts the day-of-week number into a simple binary flag (0 for weekday, 1 for weekend).
- Save features.py.
- You won't see output yet. The next function builds on this one.
Seeing a red underline in VS Code?
- If VS Code shows an import error for pandas or numpy, your Python interpreter might not be set to the virtual environment.
- Press Cmd+Shift+P (macOS) or Ctrl+Shift+P (Windows) and select Python: Select Interpreter, then choose the .venv option.
- VS Code shows import errors in my features.py file.
- Add the LightGBM feature function below create_base_features in features.py:
def create_all_features(df):
"""Tabular features for LightGBM: base features plus lags and rolling stats."""
df = create_base_features(df)
df["minute_of_day"] = df["timestamp"].dt.hour * 60 + df["timestamp"].dt.minute
df["lag_1"] = df["call_count"].shift(1)
df["lag_4"] = df["call_count"].shift(4) # 1 hour ago
df["lag_96"] = df["call_count"].shift(96) # 1 day ago
df["rolling_mean_4"] = df["call_count"].rolling(window=4).mean()
df = df.dropna().reset_index(drop=True)
return df
What does this code do?
create_all_features calls create_base_features first, then adds advanced signals on top.
- Lag features (lag_1, lag_4, lag_96) give the model access to recent history. lag_4 is the call count from 1 hour ago (four 15-minute intervals). lag_96 is 1 day ago (ninety-six 15-minute intervals).
- rolling_mean_4 smooths out noise by averaging the last hour of data.
- dropna() removes rows where lag values don't exist yet (the first 96 rows have no day-ago data).
- Add the LSTM sequence function below create_all_features in features.py:
def create_sequences(df, lookback=96):
"""Build 3D input arrays for the LSTM (samples, timesteps, features)."""
df = create_base_features(df)
feature_cols = ["call_count", "hour_sin", "hour_cos", "day_of_week"]
data = df[feature_cols].values
target = df["call_count"].values
X, y = [], []
for i in range(lookback, len(data)):
X.append(data[i - lookback:i])
y.append(target[i])
return np.array(X), np.array(y)
What does this code do?
- The LSTM needs a 3D array shaped (samples, timesteps, features). Each sample is a window of 96 consecutive timesteps (24 hours of 15-minute data).
- For each window, the target y is the call count at the very next timestep. The model learns to predict one step ahead from the previous 24 hours.
- The features per timestep include the raw call count plus cyclical hour encoding and day of week. This gives the LSTM both the raw signal and time context.
- Save features.py.
- Confirm the module works by running this command:
python -c "from features import create_all_features, create_sequences; print('Features module ready')"
- You should see Features module ready printed in your terminal.
Seeing an ImportError?
- Make sure features.py is saved in the same folder where you're running the command.
- Check there are no typos in the function name create_all_features.
- Help me fix an ImportError with my features.py module.
✔️ Awesome, I've got everything!
Great. Double check you've saved features.py before moving on.
ⓧ I'd like to double check the full code
import pandas as pd
import numpy as np
def create_base_features(df):
"""Extract time features from raw timestamps."""
df = df.copy()
df["timestamp"] = pd.to_datetime(df["timestamp"])
df["hour"] = df["timestamp"].dt.hour
df["day_of_week"] = df["timestamp"].dt.dayofweek
df["is_weekend"] = (df["day_of_week"] >= 5).astype(int)
# Cyclical encoding so the model knows hour 23 is close to hour 0
df["hour_sin"] = np.sin(2 * np.pi * df["hour"] / 24)
df["hour_cos"] = np.cos(2 * np.pi * df["hour"] / 24)
return df
def create_all_features(df):
"""Tabular features for LightGBM: base features plus lags and rolling stats."""
df = create_base_features(df)
df["minute_of_day"] = df["timestamp"].dt.hour * 60 + df["timestamp"].dt.minute
df["lag_1"] = df["call_count"].shift(1)
df["lag_4"] = df["call_count"].shift(4) # 1 hour ago
df["lag_96"] = df["call_count"].shift(96) # 1 day ago
df["rolling_mean_4"] = df["call_count"].rolling(window=4).mean()
df = df.dropna().reset_index(drop=True)
return df
def create_sequences(df, lookback=96):
"""Build 3D input arrays for the LSTM (samples, timesteps, features)."""
df = create_base_features(df)
feature_cols = ["call_count", "hour_sin", "hour_cos", "day_of_week"]
data = df[feature_cols].values
target = df["call_count"].values
X, y = [], []
for i in range(lookback, len(data)):
X.append(data[i - lookback:i])
y.append(target[i])
return np.array(X), np.array(y)
Your data pipeline is ready with realistic call volume data and feature engineering for both model types. Next up, you'll train lgb_v1 and lstm_v4 side by side and track them with MLflow.
Train Two Models with MLflow Tracking
You have synthetic call volume data sitting in data/calls.csv and feature engineering functions ready for both tabular and sequence-based models. Now it's time to answer the real question: which architecture predicts call volume better?
In this step, you will train two competing models — a LightGBM regressor (lgb_v1) and a Keras LSTM (lstm_v4) — and use MLflow to track every detail of each run. Then you will register the winner as your "champion" model so the rest of your pipeline knows which one to serve.
In this step, get ready to:
- Train a LightGBM model and a Keras LSTM, logging their metrics to MLflow.
- Compare both runs side-by-side in the MLflow UI.
- Register the best model in the MLflow model registry with a "champion" alias.
Create the training script
The training script needs to handle two very different model types: LightGBM trains on flat tabular features, while the LSTM trains on 3D sequence arrays. Both models log their hyperparameters, evaluation metrics, and model artifacts to MLflow.
- Create a new file called train.py in your project folder.
- Add the imports and the reusable training function by pasting this code:
import mlflow
import pandas as pd
import numpy as np
from lightgbm import LGBMRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error, root_mean_squared_error
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense
from features import create_all_features, create_sequences
def train_lgb(df, experiment_name):
"""Train LightGBM on tabular features and log to MLflow."""
mlflow.set_experiment(experiment_name)
df = create_all_features(df)
feature_cols = ["hour", "day_of_week", "is_weekend", "hour_sin", "hour_cos",
"minute_of_day", "lag_1", "lag_4", "lag_96", "rolling_mean_4"]
X = df[feature_cols]
y = df["call_count"]
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, shuffle=False
)
params = {
"n_estimators": 200,
"max_depth": 6,
"learning_rate": 0.1,
"num_leaves": 31,
"random_state": 42,
"verbosity": -1,
}
with mlflow.start_run(run_name="lgb_v1"):
model = LGBMRegressor(**params)
model.fit(X_train, y_train)
predictions = model.predict(X_test)
mae = mean_absolute_error(y_test, predictions)
rmse = root_mean_squared_error(y_test, predictions)
mlflow.log_params(params)
mlflow.log_param("model_tag", "lgb_v1")
mlflow.log_param("features", feature_cols)
mlflow.log_metric("mae", mae)
mlflow.log_metric("rmse", rmse)
mlflow.lightgbm.log_model(model, artifact_path="model")
print(f"lgb_v1 - MAE: {mae:.2f}, RMSE: {rmse:.2f}")
return mae, mlflow.active_run().info.run_id
What does this code do?
The train_lgb function handles the full LightGBM training cycle inside a single MLflow run.
- mlflow.set_experiment() creates or selects an experiment to group related runs together.
- mlflow.start_run() opens a tracked run where everything you log gets recorded.
- mlflow.lightgbm.log_model() saves the trained model as an artifact you can reload later for serving.
You won't see a change yet. The next chunk adds the LSTM training function and the main() that calls both.
- Add the LSTM training function and main() below train_lgb:
def train_lstm(df, experiment_name, lookback=96):
"""Train a 2-layer LSTM on sequences and log to MLflow."""
mlflow.set_experiment(experiment_name)
X, y = create_sequences(df, lookback=lookback)
split = int(len(X) * 0.8)
X_train, X_test = X[:split], X[split:]
y_train, y_test = y[:split], y[split:]
params = {
"lookback": lookback,
"units_1": 64,
"units_2": 32,
"epochs": 20,
"batch_size": 32,
}
with mlflow.start_run(run_name="lstm_v4"):
model = Sequential([
LSTM(params["units_1"], return_sequences=True,
input_shape=(lookback, X.shape[2])),
LSTM(params["units_2"]),
Dense(1),
])
model.compile(optimizer="adam", loss="mse")
model.fit(X_train, y_train, epochs=params["epochs"],
batch_size=params["batch_size"], verbose=0)
predictions = model.predict(X_test, verbose=0).flatten()
mae = mean_absolute_error(y_test, predictions)
rmse = root_mean_squared_error(y_test, predictions)
mlflow.log_params(params)
mlflow.log_param("model_tag", "lstm_v4")
mlflow.log_metric("mae", mae)
mlflow.log_metric("rmse", rmse)
mlflow.keras.log_model(model, artifact_path="model")
print(f"lstm_v4 - MAE: {mae:.2f}, RMSE: {rmse:.2f}")
return mae, mlflow.active_run().info.run_id
def main():
experiment_name = "call-volume-forecasting"
df = pd.read_csv("data/calls.csv")
# Train LightGBM on tabular features
mae_lgb, run_id_lgb = train_lgb(df, experiment_name)
# Train LSTM on sequences
mae_lstm, run_id_lstm = train_lstm(df, experiment_name)
if __name__ == "__main__":
main()
What does this code do?
The train_lstm function builds a 2-layer LSTM using Keras and trains it on the sequence data from create_sequences().
- The first LSTM layer (64 units) with return_sequences=True passes its full output sequence to the second layer.
- The second LSTM layer (32 units) outputs a single vector that feeds into Dense(1) for the final call count prediction.
- mlflow.keras.log_model() saves the Keras model in MLflow's format so it can be loaded later with mlflow.pyfunc.load_model().
- The main() function orchestrates both training runs under the same experiment name so MLflow groups them for comparison.
- Save train.py.
✔️ Awesome, I've got everything!
Great. Double check you've saved your file before moving on.
ⓧ I'd like to double check the full code
import mlflow
import pandas as pd
import numpy as np
from lightgbm import LGBMRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error, root_mean_squared_error
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense
from features import create_all_features, create_sequences
def train_lgb(df, experiment_name):
"""Train LightGBM on tabular features and log to MLflow."""
mlflow.set_experiment(experiment_name)
df = create_all_features(df)
feature_cols = ["hour", "day_of_week", "is_weekend", "hour_sin", "hour_cos",
"minute_of_day", "lag_1", "lag_4", "lag_96", "rolling_mean_4"]
X = df[feature_cols]
y = df["call_count"]
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, shuffle=False
)
params = {
"n_estimators": 200,
"max_depth": 6,
"learning_rate": 0.1,
"num_leaves": 31,
"random_state": 42,
"verbosity": -1,
}
with mlflow.start_run(run_name="lgb_v1"):
model = LGBMRegressor(**params)
model.fit(X_train, y_train)
predictions = model.predict(X_test)
mae = mean_absolute_error(y_test, predictions)
rmse = root_mean_squared_error(y_test, predictions)
mlflow.log_params(params)
mlflow.log_param("model_tag", "lgb_v1")
mlflow.log_param("features", feature_cols)
mlflow.log_metric("mae", mae)
mlflow.log_metric("rmse", rmse)
mlflow.lightgbm.log_model(model, artifact_path="model")
print(f"lgb_v1 - MAE: {mae:.2f}, RMSE: {rmse:.2f}")
return mae, mlflow.active_run().info.run_id
def train_lstm(df, experiment_name, lookback=96):
"""Train a 2-layer LSTM on sequences and log to MLflow."""
mlflow.set_experiment(experiment_name)
X, y = create_sequences(df, lookback=lookback)
split = int(len(X) * 0.8)
X_train, X_test = X[:split], X[split:]
y_train, y_test = y[:split], y[split:]
params = {
"lookback": lookback,
"units_1": 64,
"units_2": 32,
"epochs": 20,
"batch_size": 32,
}
with mlflow.start_run(run_name="lstm_v4"):
model = Sequential([
LSTM(params["units_1"], return_sequences=True,
input_shape=(lookback, X.shape[2])),
LSTM(params["units_2"]),
Dense(1),
])
model.compile(optimizer="adam", loss="mse")
model.fit(X_train, y_train, epochs=params["epochs"],
batch_size=params["batch_size"], verbose=0)
predictions = model.predict(X_test, verbose=0).flatten()
mae = mean_absolute_error(y_test, predictions)
rmse = root_mean_squared_error(y_test, predictions)
mlflow.log_params(params)
mlflow.log_param("model_tag", "lstm_v4")
mlflow.log_metric("mae", mae)
mlflow.log_metric("rmse", rmse)
mlflow.keras.log_model(model, artifact_path="model")
print(f"lstm_v4 - MAE: {mae:.2f}, RMSE: {rmse:.2f}")
return mae, mlflow.active_run().info.run_id
def main():
experiment_name = "call-volume-forecasting"
df = pd.read_csv("data/calls.csv")
# Train LightGBM on tabular features
mae_lgb, run_id_lgb = train_lgb(df, experiment_name)
# Train LSTM on sequences
mae_lstm, run_id_lstm = train_lstm(df, experiment_name)
if __name__ == "__main__":
main()
Run the training script
Time to see which model wins. Running the script trains both models and logs everything to MLflow automatically.
- Run the training script by executing this command in your terminal:
python train.py
- You should see output showing the MAE and RMSE for each model, like this:
lgb_v1 - MAE: X.XX, RMSE: X.XX
lstm_v4 - MAE: X.XX, RMSE: X.XX
What should I see?
The exact numbers will vary. The LSTM takes longer to train (about 30 seconds for 20 epochs) so be patient. The LightGBM model typically trains in under a second.
Either model could win depending on the random seed. The important thing is that both runs complete and their metrics appear.
🙋♀️ Seeing a ModuleNotFoundError?
Make sure your virtual environment is active. You should see (.venv) at the start of your terminal prompt.
Also confirm that data/calls.csv exists. If not, run python generate_data.py first.
help me debug a ModuleNotFoundError when running train.py
Now verify the runs are tracked properly by launching the MLflow UI.
- Start the MLflow UI by running:
mlflow ui
- Open your browser and navigate to http://localhost:5000.
- You should see an experiment called call-volume-forecasting in the left sidebar.
- Click on it to see both runs: lgb_v1 and lstm_v4.
- Select both runs and click Compare to see the MAE and RMSE metrics side-by-side.
What is experiment tracking?
Without experiment tracking, your team ends up with a folder full of scripts and no idea which model version produced which results. MLflow records every run with its parameters, metrics, and artifacts so you can always trace back to exactly how a model was trained.
🙋♀️ MLflow UI not loading?
Make sure no other process is using port 5000. If the port is busy, stop the other process or run mlflow ui --port 5001 instead.
help me troubleshoot the MLflow UI not loading on localhost
Register the best model as champion
Tracking runs is great for comparison, but your API needs to know which model to actually serve. The model registry solves this by letting you assign an alias like "champion" to a specific model version.
Why aliases instead of stages?
Older versions of MLflow used stages (Staging, Production, Archived) to manage model lifecycle. MLflow 3 replaced stages with aliases, which are more flexible. You can name aliases anything: "champion", "challenger", "shadow". Your code loads a model by alias, so promoting a new version is just reassigning the alias.
- In train.py, find the import section at the top of the file.
- Find this line:
from features import create_base_features, create_all_features
- Add this new import directly above it, so you have both lines:
from mlflow import MlflowClient
from features import create_base_features, create_all_features
- Next, find the beginning of main() where you set the experiment name:
def main():
experiment_name = "call-volume-forecasting"
- Add a registry_model_name variable directly below it:
def main():
experiment_name = "call-volume-forecasting"
registry_model_name = "call-volume-model"
- Finally, find the end of main() where the LSTM finishes training. Add the following registration code below the LSTM training call, but above if __name__ == "__main__"::
# Register the better model as champion
client = MlflowClient()
best_run_id = run_id_lstm if mae_lstm < mae_lgb else run_id_lgb
model_uri = f"runs:/{best_run_id}/model"
try:
client.create_registered_model(registry_model_name)
except mlflow.exceptions.MlflowException:
pass # Already exists
mv = client.create_model_version(
name=registry_model_name,
source=model_uri,
run_id=best_run_id,
)
client.set_registered_model_alias(
name=registry_model_name,
alias="champion",
version=mv.version,
)
print(f"\nRegistered version {mv.version} as 'champion' (run: {best_run_id})")
What does this registration code do?
- It compares the MAE from both runs and picks the one with the lower error.
- client.create_registered_model() creates a named entry in the registry (wrapped in try/except so re-running won't crash).
- client.create_model_version() links the trained model artifact to that registry entry.
- client.set_registered_model_alias() assigns the "champion" alias so your API can load it by name instead of by run ID.
- Save train.py.
- Stop the MLflow UI if it's still running (press Ctrl+C in that terminal).
- Re-run the training script to train fresh models and register the champion:
python train.py
- You should see the MAE/RMSE output for both models followed by a registration confirmation like:
model-A-base - MAE: X.XX, RMSE: X.XX
model-B-lags - MAE: X.XX, RMSE: X.XX
Registered version 1 as 'champion' (run: <run_id>)
Seeing a PermissionError or database lock?
This usually means the MLflow UI is still running and holding a lock on the tracking database. Make sure you stopped the MLflow UI process before re-running the script.
help me fix a database lock error when running train.py with MLflow
✔️ Awesome, I've got everything!
Great. Double check you've saved train.py with all the registration code included.
ⓧ I'd like to double check the full code
import mlflow
import pandas as pd
import numpy as np
from lightgbm import LGBMRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error, root_mean_squared_error
from mlflow import MlflowClient
from features import create_base_features, create_all_features
def train_and_log(df, feature_cols, model_name_tag, experiment_name):
mlflow.set_experiment(experiment_name)
X = df[feature_cols]
y = df["call_count"]
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, shuffle=False
)
params = {
"n_estimators": 200,
"max_depth": 6,
"learning_rate": 0.1,
"num_leaves": 31,
"random_state": 42,
"verbosity": -1,
}
with mlflow.start_run(run_name=model_name_tag):
model = LGBMRegressor(**params)
model.fit(X_train, y_train)
predictions = model.predict(X_test)
mae = mean_absolute_error(y_test, predictions)
rmse = root_mean_squared_error(y_test, predictions)
mlflow.log_params(params)
mlflow.log_param("features", feature_cols)
mlflow.log_param("model_tag", model_name_tag)
mlflow.log_metric("mae", mae)
mlflow.log_metric("rmse", rmse)
mlflow.lightgbm.log_model(model, artifact_path="model")
print(f"{model_name_tag} - MAE: {mae:.2f}, RMSE: {rmse:.2f}")
return mae, mlflow.active_run().info.run_id
def main():
experiment_name = "call-volume-forecasting"
registry_model_name = "call-volume-model"
df = pd.read_csv("data/calls.csv")
# Model A: base features only
df_base = create_base_features(df)
base_features = ["hour", "day_of_week", "is_weekend"]
mae_a, run_id_a = train_and_log(df_base, base_features, "model-A-base", experiment_name)
# Model B: all features including lags
df_all = create_all_features(df)
all_features = ["hour", "day_of_week", "is_weekend", "minute_of_day", "lag_1", "lag_4", "lag_96", "rolling_mean_4"]
mae_b, run_id_b = train_and_log(df_all, all_features, "model-B-lags", experiment_name)
# Register the better model as champion
client = MlflowClient()
best_run_id = run_id_b if mae_b < mae_a else run_id_a
model_uri = f"runs:/{best_run_id}/model"
try:
client.create_registered_model(registry_model_name)
except mlflow.exceptions.MlflowException:
pass # Already exists
mv = client.create_model_version(
name=registry_model_name,
source=model_uri,
run_id=best_run_id,
)
client.set_registered_model_alias(
name=registry_model_name,
alias="champion",
version=mv.version,
)
print(f"\nRegistered version {mv.version} as 'champion' (run: {best_run_id})")
if __name__ == "__main__":
main()
Your experiment tracking is locked in and the best model is registered. Next up, you will build a FastAPI endpoint that loads the champion model and serves real-time call volume predictions.
Build the FastAPI Prediction Endpoint
Your champion model is registered in MLflow and ready to serve predictions. Right now, the only way to get a forecast is to load the model manually in a Python script.
In this step, you will build a FastAPI prediction endpoint that loads the champion model at startup and responds to HTTP requests with call volume forecasts. Anyone on your team can query it without touching the training code.
In this step, get ready to:
- Create the FastAPI app with automatic model loading at startup.
- Add a /predict endpoint that returns call volume forecasts for 15-minute intervals.
- Test the API locally and confirm predictions come back.
Create the FastAPI app with model loading
The app needs to load the champion model once when the server starts, then hold it in memory for every incoming request. FastAPI's startup event handler is perfect for this.
- Create a new file called app.py in your project folder.
- Add the following code to set up the app skeleton and model loading:
import mlflow
import pandas as pd
import numpy as np
from fastapi import FastAPI
from pydantic import BaseModel
from features import create_all_features
app = FastAPI(title="Call Volume Forecast API")
model = None
@app.on_event("startup")
def load_model():
global model
model_uri = "models:/call-volume-model@champion"
model = mlflow.pyfunc.load_model(model_uri)
print(f"Loaded champion model from {model_uri}")
What does this code do?
The @app.on_event("startup") decorator tells FastAPI to run load_model() once when the server starts. This loads the champion model into memory so it is ready for every prediction request.
The model URI models:/call-volume-model@champion tells MLflow to fetch whichever model version currently holds the "champion" alias. If you retrain and re-register a better model later, the API automatically picks it up on next restart.
You will not see a change yet. The next chunk adds the endpoints that make this app callable.
Add the prediction and health endpoints
The API needs two things: a data contract defining what requests and responses look like, and the endpoint logic that transforms incoming timestamps into predictions.
- Add the following code below the load_model() function in app.py:
class PredictionRequest(BaseModel):
timestamps: list[str]
class PredictionResponse(BaseModel):
predictions: list[dict]
@app.post("/predict", response_model=PredictionResponse)
def predict(request: PredictionRequest):
df = pd.DataFrame({"timestamp": request.timestamps, "call_count": 0})
df = create_all_features(df)
feature_cols = ["hour", "day_of_week", "is_weekend", "minute_of_day", "lag_1", "lag_4", "lag_96", "rolling_mean_4"]
available_cols = [col for col in feature_cols if col in df.columns]
X = df[available_cols]
preds = model.predict(X)
results = []
for i, ts in enumerate(df["timestamp"]):
results.append({
"timestamp": str(ts),
"predicted_calls": max(0, round(float(preds[i]))),
})
return PredictionResponse(predictions=results)
@app.get("/health")
def health():
return {"status": "healthy", "model_loaded": model is not None}
What does this code do?
Pydantic models PredictionRequest and PredictionResponse define the shape of incoming JSON and outgoing JSON. FastAPI uses these to validate requests and auto-generate documentation.
The predict() function takes a list of timestamp strings, runs them through create_all_features() (the same feature engineering from training), and feeds the resulting DataFrame to the champion model. It returns predicted call counts clamped to zero.
The /health endpoint is a lightweight check that confirms the server is running and the model loaded successfully.
- Save app.py.
You will not see a visible change yet. The next substep starts the server and tests everything end-to-end.
Seeing a red underline on imports?
If your editor shows import errors for mlflow or fastapi, make sure your virtual environment is active and selected as the Python interpreter in VS Code.
help me fix import errors in my FastAPI app
✔️ Awesome, I've got everything!
Great. Make sure you have saved app.py before moving on.
ⓧ I'd like to double check the full code
import mlflow
import pandas as pd
import numpy as np
from fastapi import FastAPI
from pydantic import BaseModel
from features import create_all_features
app = FastAPI(title="Call Volume Forecast API")
model = None
@app.on_event("startup")
def load_model():
global model
model_uri = "models:/call-volume-model@champion"
model = mlflow.pyfunc.load_model(model_uri)
print(f"Loaded champion model from {model_uri}")
class PredictionRequest(BaseModel):
timestamps: list[str]
class PredictionResponse(BaseModel):
predictions: list[dict]
@app.post("/predict", response_model=PredictionResponse)
def predict(request: PredictionRequest):
df = pd.DataFrame({"timestamp": request.timestamps, "call_count": 0})
df = create_all_features(df)
feature_cols = ["hour", "day_of_week", "is_weekend", "minute_of_day", "lag_1", "lag_4", "lag_96", "rolling_mean_4"]
available_cols = [col for col in feature_cols if col in df.columns]
X = df[available_cols]
preds = model.predict(X)
results = []
for i, ts in enumerate(df["timestamp"]):
results.append({
"timestamp": str(ts),
"predicted_calls": max(0, round(float(preds[i]))),
})
return PredictionResponse(predictions=results)
@app.get("/health")
def health():
return {"status": "healthy", "model_loaded": model is not None}
Test the API locally
Time to start the server and confirm predictions flow end-to-end. FastAPI ships with auto-generated interactive docs that let you test endpoints directly in your browser.
- Start the FastAPI development server by running this command:
uvicorn app:app --reload --port 8000
What does this command do?
uvicorn is the ASGI server that runs your FastAPI app. app:app means "in the file app.py, find the object called app." The --reload flag restarts the server whenever you edit code. --port 8000 sets the port to listen on.
You should see output that includes Loaded champion model from models:/call-volume-model@champion followed by a line showing the server is running.
Server fails to start?
If you see an error about the model not being found, make sure you ran train.py successfully in the previous step. The champion alias must exist in the local MLflow tracking store.
If you see a port-in-use error, stop whatever is using port 8000 (possibly an old mlflow ui session if you used port 8000, or stop it with Ctrl+C in that terminal).
help me debug why my FastAPI server won't start
- Open your browser and go to http://localhost:8000/docs.
- Click the POST /predict row to expand it.
- Click Try it out.
- Replace the request body with this JSON:
{
"timestamps": [
"2025-01-15 09:00:00",
"2025-01-15 09:15:00",
"2025-01-15 09:30:00",
"2025-01-15 09:45:00",
"2025-01-15 10:00:00"
]
}
- Click Execute.
You should see a JSON response with a predictions array containing five objects. Each object has a timestamp and a predicted_calls value (an integer representing expected call volume for that 15-minute window).
Why do some predictions show zero?
The create_all_features() function calculates lag features (what happened 15 minutes ago, 1 hour ago, 1 day ago). When you send only 5 timestamps with no historical context, some lag values are missing and get dropped. The model predicts on whatever rows survive.
In production, you would feed real historical data so all lag features resolve properly.
🙋♀️ Getting an empty predictions array?
If the response returns an empty predictions list, all rows were dropped during feature engineering because lag features could not be computed. Try sending more timestamps (at least 97 consecutive 15-minute intervals) to give the lag features enough history.
help me understand why my /predict endpoint returns empty predictions
Your prediction API is live and serving forecasts from the champion model. Next up, you will package this entire service into a Docker container so any team member can run it with a single command.
Dockerize the Service
Your prediction API works on your machine, but what happens when a teammate tries to run it? Different Python versions, missing packages, or conflicting dependencies can break the setup entirely.
Docker solves this by packaging your code, trained models, and dependencies into a single container. In this step, you'll containerize the entire pipeline so any team member can launch the forecasting service with one command.
In this step, get ready to:
- Create a requirements.txt and Dockerfile to package the service.
- Build and run the Docker image.
- Verify the containerized API returns call volume predictions.
Create the requirements file and Dockerfile
Docker needs to know exactly which packages to install inside the container. A requirements.txt file lists every dependency with its version, so the container installs the same environment every time.
- In VS Code, create a new file called requirements.txt in your project root folder.
- Add the following dependencies:
mlflow==3.13.0
fastapi[standard]==0.136.3
lightgbm==4.6.0
pandas
numpy
scikit-learn
What does this file do?
Each line declares a package that pip will install inside the container.
The pinned versions (mlflow==3.13.0, fastapi[standard]==0.136.3, lightgbm==4.6.0) guarantee the container always installs the exact same versions you tested locally. The unpinned packages (pandas, numpy, scikit-learn) will resolve to compatible versions automatically.
- Save requirements.txt.
Now you need the Dockerfile itself. This file tells Docker how to build your container image step by step.
- Create a new file called Dockerfile (no extension) in your project root folder.
- Add the following build instructions:
FROM python:3.12-slim
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY . .
RUN python generate_data.py && python train.py
EXPOSE 8000
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]
What does each instruction do?
- FROM python:3.12-slim starts from a lightweight Python 3.12 base image.
- COPY requirements.txt . copies the requirements file first, so Docker can cache the install layer separately from your source code.
- RUN pip install --no-cache-dir -r requirements.txt installs all dependencies without storing download caches (keeps the image small).
- COPY . . copies all project files into the container.
- RUN python generate_data.py && python train.py generates fresh data and trains both models inside the container during the build. This means the container ships with a ready-to-serve champion model.
- EXPOSE 8000 documents which port the app listens on. It does not publish the port on its own.
- CMD launches the FastAPI server when the container starts, binding to all interfaces so Docker can route traffic in.
💡 Why train inside the container?
Running generate_data.py and train.py at build time means every container starts with a trained model already registered in MLflow. A teammate can pull the image and immediately get predictions without running any setup scripts.
- Save Dockerfile.
- Confirm both files exist by running:
ls
You should see requirements.txt and Dockerfile alongside your existing Python files.
Don't see both files?
- Make sure Dockerfile has no file extension. Some editors add .txt automatically.
- Confirm you saved both files in the same folder as app.py and train.py.
Help me figure out why my Dockerfile or requirements.txt isn't showing up.
✔️ Awesome, I've got everything!
Great. Make sure both files are saved before moving on.
ⓧ I'd like to double check the full code
mlflow==3.13.0
fastapi[standard]==0.136.3
lightgbm==4.6.0
pandas
numpy
scikit-learn
FROM python:3.12-slim
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY . .
RUN python generate_data.py && python train.py
EXPOSE 8000
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]
Build and run the Docker image
With both files in place, you can now build a Docker image that contains your entire forecasting pipeline. The build process installs packages, generates data, and trains models so the resulting image is ready to serve predictions immediately.
- If your local uvicorn server is still running from the previous step, stop it with Ctrl+C so port 8000 is free.
- Build the Docker image by running this command:
docker build -t call-forecast-api .
What does this command do?
- docker build reads the Dockerfile and executes each instruction to create an image.
- -t call-forecast-api tags the image with the name call-forecast-api so you can reference it later.
- . tells Docker to use the current directory as the build context (all your project files).
This build takes 2-3 minutes because it installs packages and trains both models inside the container.
- Wait for the build to finish. You should see output ending with a line like Successfully tagged call-forecast-api:latest.
- Run the container by executing this command:
docker run -p 8000:8000 call-forecast-api
What does this command do?
- docker run starts a new container from the image you just built.
- -p 8000:8000 maps port 8000 on your machine to port 8000 inside the container. Without this flag, you cannot reach the API from your browser or terminal.
You should see uvicorn startup logs in your terminal, ending with a line that says the application is running on 0.0.0.0:8000.
Seeing a port conflict error?
- If you see a message about port 8000 being in use, your local uvicorn server is still running. Stop it in the other terminal with Ctrl+C, then run the docker run command again.
- On macOS, AirPlay Receiver sometimes occupies port 5000. This project uses port 8000, so it should not conflict. If it does, check for other services on 8000.
Help me fix a port conflict when running my Docker container.
Verify the containerized API
The container is running, but you need to confirm the prediction endpoint actually works from inside Docker. This proves the entire pipeline (data generation, model training, and serving) works in the container.
- Open a new terminal window. Your current terminal is attached to the running container.
- Test the /predict endpoint by running this command:
curl -X POST http://localhost:8000/predict -H "Content-Type: application/json" -d '{"timestamps": ["2025-01-15 09:00:00", "2025-01-15 09:15:00", "2025-01-15 09:30:00"]}'
What should I see?
You should get a JSON response with a predictions array. Each item contains a timestamp and a predicted_calls value (a non-negative integer representing predicted call volume for that 15-minute interval).
The exact numbers depend on your trained model, but morning timestamps (9:00-9:30) should show higher predicted volumes than night hours.
💡 Use the interactive docs on Windows
If the curl command gives you quoting issues on Windows, open http://localhost:8000/docs in your browser and use the interactive POST /predict form to send the same test request.
- Confirm the response includes predicted call volumes for each of the three timestamps.
Getting a connection refused error?
- Make sure the Docker container is still running in your other terminal. Look for the uvicorn startup log.
- Confirm you used -p 8000:8000 when starting the container. Without this flag, the port is not published to your host machine.
- If the container exited immediately, run docker logs $(docker ps -lq) to see what went wrong during startup.
Help me debug why my containerized API isn't responding to requests.
Secret mission
Add a Model Comparison Endpoint
Your API serves the champion model's predictions, but what if a stakeholder asks why it was chosen? Add a /compare endpoint that loads both models, runs them on the same sample data, and returns their predictions side-by-side with training metrics. This turns your API into a model evaluation tool, not just a prediction service.
Clean Up Your Resources
Clean Up Your Resources
Decide whether to keep your resources running, pause them to come back later, or delete them entirely. This project runs entirely locally, so there are no ongoing costs.
Resources you used:
- Docker image call-forecast-api (local only, no cost).
- MLflow tracking data in the local mlruns/ directory and mlflow.db SQLite database (no cost).
- Python virtual environment and project files on your Desktop (no cost).
Keep everything running
No action needed. Choose this if you're still actively building or experimenting with different model features and want to keep everything in place.
- Your Docker image stays available for docker run whenever you need it.
- Your MLflow experiment data in mlruns/ stays intact for future training runs and comparisons.
- Your virtual environment keeps all installed packages ready to use.
Pause - I'll come back to this later
Stop running containers to free up system resources. Your files and Docker image stay on disk for when you return.
- If your container is still running, stop it by running:
- If you started it in the foreground (no -d flag), press Ctrl+C in the terminal where it is running.
docker stop $(docker ps -q --filter ancestor=call-forecast-api)
- If the MLflow UI is still running, press Ctrl+C in its terminal to stop it.
- Your project files, Docker image, and MLflow data all remain on disk.
Delete - I don't want to use this again
Remove all project resources and start fresh. This deletes your trained models, experiment history, and the Docker image.
- Stop any running containers from this image:
docker stop $(docker ps -q --filter ancestor=call-forecast-api)
- Remove the Docker image:
docker rmi call-forecast-api
- Delete your project folder from the Desktop. On macOS or Linux:
cd ~/Desktop
rm -rf call-forecast-api
- On Windows:
cd ~\Desktop
rmdir /s /q call-forecast-api
- This removes all Python files, the virtual environment, trained models, and MLflow tracking data.
Nice Work!
Nice Work!
Nice work! You just built a full MLOps pipeline that takes call volume forecasting from a loose script to a versioned, containerized prediction service.
You've learned how to:
- Train and compare two competing LightGBM models using MLflow experiment tracking, logging parameters, metrics, and model artifacts for full reproducibility.
- Register the best-performing model in MLflow's model registry with a champion alias, then serve it through a FastAPI prediction endpoint that returns call volume forecasts on demand.
- Containerize the entire service with Docker so any team member can run a single command and get a working prediction API, regardless of their local environment.
- Secret Mission: Added a /compare endpoint that returns both models' predictions side-by-side with their training metrics for transparent model evaluation.
Ready to quiz yourself?