mmm_mlflow_logging.py

# Initiate the MLflow run context
with mlflow.start_run(run_name="your_run_name") as run:

    # Log git hash
    git_commit_hash = get_git_revision_hash()
    if git_commit_hash:
        # Log the git commit hash as a tag
        mlflow.set_tag("git_commit_id", git_commit_hash)


    # Log the pre-processing steps taken (Dictionaries of constants)
    mlflow.log_params(FEATURES)
    mlflow.log_params(DIM_REDUCTION_CONFIG)
    mlflow.log_params(SEASONALITY_CONFIG)
    mlflow.log_params(TRAIN_TEST_CONFIG)
    mlflow.log_params(MODEL_CONFIG)
    

    # Log artifacts
    mlflow.log_figure(prior_plot, "graphs/prior_plot.png")
    mlflow.log_figure(train_plot, "graphs/train_plot.png")
    mlflow.log_figure(test_plot, "graphs/test_plot.png")
    mlflow.log_figure(cpa_plot, "graphs/cpa_plot.png")
    mlflow.log_figure(waterfall_plot, "graphs/waterfall_plot.png")
    # Also log the data output
    mlflow.log_artifact("model_priors.csv", "output_data")
    mlflow.log_artifact("model_coeffs.csv", "output_data")
    mlflow.log_artifact("channel_cpa.csv", "output_data")
    mlflow.log_artifact("cpa_comparison_to_expectations.csv", "output_data")
    

    # Log model metrics
    mlflow.log_metrics(model_metrics)
    mlflow.log_metrics(model_diagnostics)
    mlflow.log_metrics(loocv_metrics)