1.11 Manually Logging Metrics, Artifacts, and Models in an MLflow Run

1.11 Manually Logging Metrics, Artifacts, and Models in an MLflow Run

Autologging captures what a framework already knows how to report. Everything else — a custom business metric, a calibration plot, a data-dictionary file, a model wrapped in your own preprocessing — has to be logged deliberately. The exam objective is phrased as manually log metrics, artifacts, and models in an MLflow Run, and it tests whether you reach for the right call.

---

Automated Tracking with mlflow.autolog()

For standard machine learning frameworks, Databricks MLflow provides turnkey automatic logging. Calling mlflow.autolog() before training automatically captures parameters, performance metrics, model signatures, and artifacts without requiring manual logging statements.

import mlflow
import xgboost as xgb
from sklearn.model_selection import train_test_split

# Enable zero-code automatic tracking across supported frameworks
mlflow.autolog(
    log_input_examples=True,
    log_model_signatures=True,
    log_models=True
)

# Standard model training workflow
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2)

with mlflow.start_run(run_name="xgboost_autolog_baseline"):
    model = xgb.XGBClassifier(n_estimators=100, max_depth=5, learning_rate=0.05)
    model.fit(X_train, y_train, eval_set=[(X_val, y_val)], verbose=False)

Framework-Specific Autologging Capabilities

• mlflow.sklearn.autolog(): Logs estimator parameters, train/test scores, confusion matrices, and serialized scikit-learn models.
• mlflow.xgboost.autolog() / mlflow.lightgbm.autolog(): Logs boosting parameters, per-iteration loss curves, feature importance charts, and best iteration index.
• mlflow.pyspark.ml.autolog(): Captures Spark MLlib Pipeline parameters, evaluator metrics, and distributed pipeline models.
• mlflow.pytorch.autolog() / mlflow.tensorflow.autolog(): Logs epoch-by-epoch loss/accuracy metrics, learning rate schedules, and model checkpoints.

---

Granular Manual Logging APIs

When developing custom algorithms, non-standard loss functions, or complex diagnostic visual reports, data scientists leverage MLflow's explicit logging methods.

import mlflow
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

with mlflow.start_run(run_name="custom_rf_tuning", tags={"framework": "scikit-learn", "tier": "gold"}) as run:
    # 1. Log individual and dictionary parameters
    mlflow.log_param("n_estimators", 250)
    mlflow.log_params({"max_depth": 12, "min_samples_split": 4, "criterion": "gini"})
    
    # 2. Log step-wise training and validation metrics
    for epoch in range(1, 11):
        train_loss = compute_train_loss(epoch)
        val_loss = compute_val_loss(epoch)
        mlflow.log_metric("train_loss", train_loss, step=epoch)
        mlflow.log_metric("val_loss", val_loss, step=epoch)
        
    # 3. Log diagnostic plots directly as artifacts
    fig, ax = plt.subplots(figsize=(6, 6))
    cm = confusion_matrix(y_true, y_pred)
    disp = ConfusionMatrixDisplay(confusion_matrix=cm)
    disp.plot(ax=ax)
    mlflow.log_figure(fig, "evaluation/confusion_matrix.png")
    plt.close(fig)
    
    # 4. Log arbitrary files (data dictionary, schema JSON)
    mlflow.log_artifact("config/feature_columns.json", artifact_path="metadata")

Summary of Core MLflow Logging Functions

---

Parameter or Metric? The Distinction That Decides Answers

Two consequences follow directly:

• Logging a loss with log_param produces a static string that cannot be charted or compared numerically. Any answer that tracks a per-epoch value as a parameter is wrong.
• Re-logging the same metric key without a step overwrites the plotted point rather than extending the curve; the step argument is what builds the series.

Logging Models Manually

import mlflow
from mlflow.models import infer_signature

with mlflow.start_run(run_name="rf_manual") as run:
    model.fit(X_train, y_train)
    preds = model.predict(X_train)

    signature = infer_signature(X_train, preds)

    mlflow.sklearn.log_model(
        sk_model=model,
        artifact_path="model",
        signature=signature,           # input/output schema enforced at serving time
        input_example=X_train.iloc[:5],
    )

• The signature records the input and output schema. Serving endpoints validate requests against it, so a logged signature is what turns a schema mistake into a clear 400 response instead of a garbage prediction.
• artifact_path names the folder inside the run; the resulting model URI is runs:/<run_id>/model.
• Adding registered_model_name="catalog.schema.name" logs and registers in one call (Section 1.14).

---

Nested Runs for Hyperparameter Tuning

When performing hyperparameter searches (such as Grid Search, Random Search, or Hyperopt Bayesian optimization), logging all candidate trials as flat top-level runs clutters the experiment UI. Nested Runs organize sweeps hierarchically by establishing a parent container run with child trial runs:

with mlflow.start_run(run_name="parent_hyperopt_sweep") as parent_run:
    mlflow.log_param("search_algorithm", "tpe.suggest")
    mlflow.log_param("total_evals", 20)
    
    for trial_idx, params in enumerate(hyperparameter_candidates):
        # Launch child run nested under parent
        with mlflow.start_run(run_name=f"trial_{trial_idx}", nested=True) as child_run:
            mlflow.log_params(params)
            model = train_model(params)
            val_acc = evaluate(model)
            mlflow.log_metric("val_accuracy", val_acc)