Traceability lets a production ML team follow a result back through its hyperparameters, metrics, model artifacts, inputs, outputs, code, and processing run. It replaces manual experiment memory with centrally recorded evidence that supports repeatable training, evaluation, and deployment. section 9.3.4 section 9.0
A traceable ML workflow connects:
- Training inputs and configuration to logged losses, optimizer settings, model summaries, and serialized models. section 9.3.4
- A validation result to the exact model run, allowing evaluation to load the model through its run ID rather than reconstructing weights manually. section 9.3.4
- A deployed model to registered versions, aliases, quality gates, promotion decisions, and rollback targets. section 9.3.5
Listing 9.13 Model training code with MLflow tracking
def train_model(
mlflow_experiment_name='recommender',
mlflow_run_id=None,
mlflow_tags={},
hot_reload_model_run_id=None,
model_embedding_factors=20,
model_learning_rate=1e-3,
model_hidden_dims=256,
model_dropout_rate=0.2,
optimizer_step_size=10,
optimizer_gamma=0.1,
training_epochs=30,
train_batch_size=64,
test_batch_size=64,
shuffle_training_data=True,
shuffle_testing_data=True):
import torch
from torch.autograd import Variable
from torch.utils.data import DataLoader
import mlflow #1
from torchinfo import summary #1
from mlflow.models import infer_signature #1
input_params = {}
for k, v in locals().items(): #1
if k == 'input_params':
continue
input_params[k] = v
class MatrixFactorization(torch.nn.Module):
. . .
if hot_reload_model_run_id is not None: #2
model_uri = f"runs:/{hot_reload_model_run_id}/model" #2
model = mlflow.pytorch.load_model(model_uri) #2
else:
model = MatrixFactorization( . . .
. . .
mlflow.set_experiment(mlflow_experiment_name) #1
with mlflow.start_run(run_id=mlflow_run_id): #1
for k,v in input_params.items():
if 'mlflow_' not in k:
mlflow.log_param(k, v) #1
mlflow.log_param("loss_function",
loss_func.__class__.__name__) #1
mlflow.log_param("optimizer", "SGD") #1
mlflow.log_params({'n_user':
dataset_map['n_users'], 'n_items':
dataset_map['n_items']}) #1
for k,v in mlflow_tags.items():
mlflow.set_tag(k, v) #1
with open("model_summary.txt", "w") as f:
f.write(str(summary(model)))
mlflow.log_artifact("model_summary.txt") #3
model_signature = None
for train_iter in range(training_epochs):
print(train_iter)
model.train()
t_loss = 0
t_count = 0
for row, col, rating in train_dataloader:
prediction = model(row, col)
if model_signature is None:
model_signature = infer_signature(
{'user': row.cpu().detach().numpy(),
'movie': col.cpu().detach().numpy()},
prediction.cpu().detach().numpy()) #3
loss = loss_func(prediction, rating.unsqueeze(1))
. . .
mlflow.log_metric(
"avg_training_loss",
f"{(t_loss/t_count):3f}",
step=train_iter) #1
scheduler.step()
model.eval()
. . .
with torch.no_grad():
for row, col,rating in test_dataloader:
prediction = model(row, col)
loss = loss_func(prediction, rating.unsqueeze(1))
te_loss += loss
te_count += 1
mlflow.log_metric(
"avg_testing_loss",
f"{(te_loss/te_count):3f}",
step=train_iter) #1
print(f"Test loss: {te_loss/te_count}")
print(f"Train loss: {t_loss/t_count}")
mlflow.pytorch.log_model(model,
"model", signature=model_signature) #3
Inspect the training implementation that collects parameters, logs metrics and artifacts, infers the model signature, and stores the model in MLflow. section 9.3.4
Listing 9.15 Promoting model to staging
def promote_model_to_staging(
new_model_run_id,
registered_model_name="recommender_production",
rms_threshold=0.0,
precision_threshold=-0.3,
recall_threshold=-0.2,
):
import mlflow.pytorc
import mlflow
from mlflow import MlflowClient
from mlflow.exceptions import RestException
client = MlflowClient()
current_staging = None
try:
current_staging = client.get_model_version_by_alias(
registered_model_name, "staging"
)
except RestException:
print("No staging model found. Upgrade current run to staging.")
if current_staging.run_id == new_model_run_id:
print("Input run is already the current staging.")
return
if current_staging is not None:
current_staging_model_data = client.get_run(
current_staging.run_id
).data.to_dictionary()
staging_metrics = current_staging_model_data["metrics"]
new_model_data = client.get_run(new_model_run_id)
new_metrics = new_model_data.data.to_dictionary()["metrics"]
rms_diff = new_metrics["rms"] - staging_metrics["rms"]
prec_key = "precision_50"
rec_key = "recall_50"
prec_diff = new_metrics[prec_key] - staging_metrics[prec_key]
rec_diff = new_metrics[rec_key] - staging_metrics[rec_key]
if rms_diff > rms_threshold:
return
if prec_diff < precision_threshold:
return
if rec_diff < recall_threshold:
return
result = mlflow.register_model(
f"runs:/{new_model_run_id}/model", "recommender_production"
) #1
client.set_registered_model_alias(
"recommender_production", "staging", result.version
) #1
Inspect the promotion workflow that compares candidate runs with the current staging model and applies operational quality thresholds. section 9.3.5
