feat: mlflow logging integration (#1773)

docs/mlflow_integration.md

@@ -0,0 +1,65 @@
+## MLFlow Integration
+YOLOX now supports MLFlow integration. MLFlow is an open-source platform for managing the end-to-end machine learning lifecycle. It is designed to work with any ML library, algorithm, deployment tool, or language. MLFlow can be used to track experiments, metrics, and parameters, and to log and visualize model artifacts. \
+For more information, please refer to: [MLFlow Documentation](https://www.mlflow.org/docs/latest/index.html)
+
+## Follow these steps to start logging your experiments to MLFlow:
+### Step-1: Install MLFlow via pip 
+```bash
+pip install mlflow python-dotenv
+```
+
+### Step-2: Set up MLFlow Tracking Server
+Start or connect to a MLFlow tracking server like databricks. You can start a local tracking server by running the following command:
+```bash
+mlflow server --host 127.0.0.1 --port 8080
+```
+Read more about setting up MLFlow tracking server [here](https://mlflow.org/docs/latest/tracking/server.html#mlflow-tracking-server)
+
+### Step-3: Set up MLFlow Environment Variables
+Set the following environment variables in your `.env` file:
+```bash
+MLFLOW_TRACKING_URI="127.0.0.1:5000"  # set to your mlflow server URI
+MLFLOW_EXPERIMENT_NAME="/path/to/experiment"  # set to your experiment name
+MLFLOW_TAGS={"release.candidate": "DEV1", "release.version": "0.0.0"}
+# config related to logging model to mlflow as pyfunc
+YOLOX_MLFLOW_LOG_MODEL_ARTIFACTS="True" # whether to log model (best or historical) or not 
+YOLOX_MLFLOW_LOG_MODEL_PER_n_EPOCHS=30 # try logging model only after every n epochs
+YOLOX_MLFLOW_LOG_Nth_EPOCH_MODELS="False" # whether to log step model along with best_model or not
+YOLOX_MLFLOW_RUN_NAME="" # give a custom name to your run, otherwise a random name is assign by mlflow
+YOLOX_MLFLOW_FLATTEN_PARAMS="True" # flatten any sub sub params of dict to be logged as simple key value pair
+
+
+MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING=True # log system gpu usage and other metrices
+MLFLOW_NESTED_RUN="False" #whether to run as a nested run of given run_id
+MLFLOW_RUN_ID="" # continue training from a given run_id
+```
+### Step-5: Provide --logger "mlflow" to the training script
+```bash
+python tools/train.py -l mlflow -f exps/path/to/exp.py -d 1 -b 8 --fp16 -o -c 
+pre_trained_model/<model>.pth
+# note the -l mlflow flag
+# one working example is this
+python tools/train.py -l mlflow -f exps/example/custom/yolox_s.py -d 1 -b 8 --fp16 -o -c pre_trained_model/yolox_s.pth
+```
+### Step-4: optional; start the mlflow ui and track your experiments
+If you log runs to a local mlruns directory, run the following command in the directory above it, then access http://127.0.0.1:5000 in your browser.
+
+```bash
+mlflow ui --port 5000
+```
+
+## Optional Databricks Integration
+
+### Step-1: Install Databricks sdk
+```bash
+pip install databricks-sdk
+```
+
+### Step-2: Set up Databricks Environment Variables
+Set the following environment variables in your `.env` file:
+```bash
+MLFLOW_TRACKING_URI="databricks"  # set to databricks
+MLFLOW_EXPERIMENT_NAME="/Users/<user>/<experiment_name>/"
+DATABRICKS_HOST = "https://dbc-1234567890123456.cloud.databricks.com" # set to your server URI
+DATABRICKS_TOKEN = "dapixxxxxxxxxxxxx"
+```

tools/train.py

@@ -85,7 +85,7 @@ def make_parser():
         "--logger",
         type=str,
         help="Logger to be used for metrics. \
-        Implemented loggers include `tensorboard` and `wandb`.",
+                Implemented loggers include `tensorboard`, `mlflow` and `wandb`.",
         default="tensorboard"
     )
     parser.add_argument(

yolox/core/trainer.py

@@ -14,6 +14,7 @@ from yolox.data import DataPrefetcher
 from yolox.exp import Exp
 from yolox.utils import (
     MeterBuffer,
+    MlflowLogger,
     ModelEMA,
     WandbLogger,
     adjust_status,
@@ -74,7 +75,8 @@ class Trainer:
         self.before_train()
         try:
             self.train_in_epoch()
-        except Exception:
+        except Exception as e:
+            logger.error("Exception in training: ", e)
             raise
         finally:
             self.after_train()
@@ -185,8 +187,11 @@ class Trainer:
                     self.exp,
                     self.evaluator.dataloader.dataset
                 )
+            elif self.args.logger == "mlflow":
+                self.mlflow_logger = MlflowLogger()
+                self.mlflow_logger.setup(args=self.args, exp=self.exp)
             else:
-                raise ValueError("logger must be either 'tensorboard' or 'wandb'")
+                raise ValueError("logger must be either 'tensorboard', 'mlflow' or 'wandb'")
 
         logger.info("Training start...")
         logger.info("\n{}".format(model))
@@ -198,6 +203,16 @@ class Trainer:
         if self.rank == 0:
             if self.args.logger == "wandb":
                 self.wandb_logger.finish()
+            elif self.args.logger == "mlflow":
+                metadata = {
+                    "epoch": self.epoch + 1,
+                    "input_size": self.input_size,
+                    'start_ckpt': self.args.ckpt,
+                    'exp_file': self.args.exp_file,
+                    "best_ap": float(self.best_ap)
+                }
+                self.mlflow_logger.on_train_end(self.args, file_name=self.file_name,
+                                                metadata=metadata)
 
     def before_epoch(self):
         logger.info("---> start train epoch{}".format(self.epoch + 1))
@@ -276,6 +291,10 @@ class Trainer:
                         "train/lr": self.meter["lr"].latest
                     })
                     self.wandb_logger.log_metrics(metrics, step=self.progress_in_iter)
+                if self.args.logger == 'mlflow':
+                    logs = {"train/" + k: v.latest for k, v in loss_meter.items()}
+                    logs.update({"train/lr": self.meter["lr"].latest})
+                    self.mlflow_logger.on_log(self.args, self.exp, self.epoch+1, logs)
 
             self.meter.clear_meters()
 
@@ -351,6 +370,14 @@ class Trainer:
                     "train/epoch": self.epoch + 1,
                 })
                 self.wandb_logger.log_images(predictions)
+            if self.args.logger == "mlflow":
+                logs = {
+                    "val/COCOAP50": ap50,
+                    "val/COCOAP50_95": ap50_95,
+                    "val/best_ap": round(self.best_ap, 3),
+                    "train/epoch": self.epoch + 1,
+                }
+                self.mlflow_logger.on_log(self.args, self.exp, self.epoch+1, logs)
             logger.info("\n" + summary)
         synchronize()
 
@@ -358,6 +385,17 @@ class Trainer:
         if self.save_history_ckpt:
             self.save_ckpt(f"epoch_{self.epoch + 1}", ap=ap50_95)
 
+        if self.args.logger == "mlflow":
+            metadata = {
+                    "epoch": self.epoch + 1,
+                    "input_size": self.input_size,
+                    'start_ckpt': self.args.ckpt,
+                    'exp_file': self.args.exp_file,
+                    "best_ap": float(self.best_ap)
+                }
+            self.mlflow_logger.save_checkpoints(self.args, self.exp, self.file_name, self.epoch,
+                                                metadata, update_best_ckpt)
+
     def save_ckpt(self, ckpt_name, update_best_ckpt=False, ap=None):
         if self.rank == 0:
             save_model = self.ema_model.ema if self.use_model_ema else self.model

yolox/utils/__init__.py

@@ -11,6 +11,7 @@ from .ema import *
 from .logger import WandbLogger, setup_logger
 from .lr_scheduler import LRScheduler
 from .metric import *
+from .mlflow_logger import MlflowLogger
 from .model_utils import *
 from .setup_env import *
 from .visualize import *

yolox/utils/mlflow_logger.py

@@ -0,0 +1,420 @@
+#!/usr/bin/env python3
+# Copyright (c) Megvii Inc. All rights reserved.
+# Please read docs/mlflow_integration.md for more details.
+"""
+Logging training runs with hyperparameter, datasets and trained models to MlFlow.
+Mlflow support Model Tracking, Experiment Tracking, and Model Registry.
+It can be hosted on-premises or in all the major cloud provider or with databricks also.
+Please read docs/mlflow_integration.md for more details.
+
+For changing default logging Behaviour you can change mlflow environment variables:
+    https://mlflow.org/docs/latest/python_api/mlflow.environment_variables.html
+
+For more information, please refer to:
+https://mlflow.org/docs/latest/introduction/index.html
+"""
+import importlib.metadata
+import importlib.util
+import json
+import os
+from collections.abc import MutableMapping
+import packaging.version
+from loguru import logger
+
+import torch
+
+from yolox.utils import is_main_process
+
+
+class MlflowLogger:
+    """
+    Main Mlflow logging class to log hyperparameters, metrics, and models to Mlflow.
+    """
+    def __init__(self):
+        if not self.is_required_library_available():
+            raise RuntimeError(
+                "MLflow Logging requires mlflow and python-dotenv to be installed. "
+                "Run `pip install mlflow python-dotenv`.")
+
+        import mlflow
+        from dotenv import find_dotenv, load_dotenv
+        load_dotenv(find_dotenv())
+        self.ENV_VARS_TRUE_VALUES = {"1", "ON", "YES", "TRUE"}
+        self._MAX_PARAM_VAL_LENGTH = mlflow.utils.validation.MAX_PARAM_VAL_LENGTH
+        self._MAX_PARAMS_TAGS_PER_BATCH = mlflow.utils.validation.MAX_PARAMS_TAGS_PER_BATCH
+        self._initialized = False
+        self._auto_end_run = False
+        self.best_ckpt_upload_pending = False
+        self._tracking_uri = None
+        self._experiment_name = None
+        self._mlflow_log_artifacts = None
+        self._mlflow_log_model_per_n_epochs = None
+        self._mlflow_log_nth_epoch_models = None
+        self.run_name = None
+        self._flatten_params = None
+        self._nested_run = None
+        self._run_id = None
+        self._async_log = None
+        self._ml_flow = mlflow
+
+    def is_required_library_available(self):
+        """
+        check if required libraries are available.
+
+        Args: None
+
+        Returns:
+            bool: True if required libraries are available, False otherwise.
+        """
+        dotenv_availaible = importlib.util.find_spec("dotenv") is not None
+        mlflow_available = importlib.util.find_spec("mlflow") is not None
+        return dotenv_availaible and mlflow_available
+
+    def flatten_dict(self, d: MutableMapping, parent_key: str = "", delimiter: str = "."):
+        """
+        Flatten a nested dict into a single level dict.
+
+        Args:
+            d(MutableMapping): nested dictionary
+            parent_key(str): parent key
+            delimiter(str): delimiter to use
+
+        Returns:
+            flattened_dict(dict): flattened dictionary
+
+        """
+
+        def _flatten_dict(d, parent_key="", delimiter="."):
+            for k, v in d.items():
+                key = str(parent_key) + delimiter + str(k) if parent_key else k
+                if v and isinstance(v, MutableMapping):
+                    yield from self.flatten_dict(v, key, delimiter=delimiter).items()
+                else:
+                    yield key, v
+
+        return dict(_flatten_dict(d, parent_key, delimiter))
+
+    def setup(self, args, exp):
+        """
+        Set up the optional MLflow integration.
+
+        Args:
+            args(dict): training args dictionary
+            exp(dict): Experiment related hyperparameters
+
+        Returns:
+            None
+
+        Environment:
+        - **YOLOX_MLFLOW_LOG_MODEL_ARTIFACTS** (`str`, *optional*, defaults to `False`):
+            Whether to use MLflow `.log_artifact()` facility to log artifacts. This only makes
+            sense if logging to a remote server, e.g. s3 or GCS. If set to `True` or *1*,
+            will copy each check-points on each save in [`TrainingArguments`]'s `output_dir` to the
+            local or remote artifact storage. Using it without a remote storage will just copy the
+            files to your artifact location.
+        - **YOLOX_MLFLOW_LOG_MODEL_PER_n_EPOCHS** (`int`, *optional*, defaults to 30):
+            If ``YOLOX_MLFLOW_LOG_MODEL_ARTIFACTS`` is enabled then Log model checkpoints after
+            every n epochs. Default is 30. ``best_ckpt.pth`` will be updated after `n` epochs if
+            it has been updated during last `n`  epochs.
+        - **YOLOX_MLFLOW_LOG_Nth_EPOCH_MODELS** (`str`, *optional*, defaults to `False`):
+            Whether to log the ``epoch_n_ckpt.pth`` models along with best_ckpt.pth model after
+             every `n` epoch as per YOLOX_MLFLOW_LOG_MODEL_PER_n_EPOCHS.
+             If set to `True` or *1*, will log ``epoch_n_ckpt.pth`` along with
+             ``best_ckpt.pth`` and as mlflow artifacts in different folders.
+        - **YOLOX_MLFLOW_RUN_NAME** (`str`, *optional*, defaults to random name):
+            Name of new run. Used only when ``run_id`` is unspecified. If a new run is
+            created and ``run_name`` is not specified, a random name will be generated for the run.
+        - **YOLOX_MLFLOW_FLATTEN_PARAMS** (`str`, *optional*, defaults to `False`):
+            Whether to flatten the parameters dictionary before logging.
+        - **MLFLOW_TRACKING_URI** (`str`, *optional*):
+            Whether to store runs at a specific path or remote server. Unset by default, which
+            skips setting the tracking URI entirely.
+        - **MLFLOW_EXPERIMENT_NAME** (`str`, *optional*, defaults to `None`):
+            Whether to use an MLflow experiment_name under which to launch the run. Default to
+            `None` which will point to the `Default` experiment in MLflow. Otherwise, it is a
+            case-sensitive name of the experiment to be activated. If an experiment with this
+            name does not exist, a new experiment with this name is created.
+        - **MLFLOW_TAGS** (`str`, *optional*):
+            A string dump of a dictionary of key/value pair to be added to the MLflow run as tags.
+             Example: `os.environ['MLFLOW_TAGS']=
+             '{"release.candidate": "RC1", "release.version": "2.2.0"}'`.
+        - **MLFLOW_NESTED_RUN** (`str`, *optional*):
+            Whether to use MLflow nested runs. If set to `True` or *1*, will create a nested run
+            inside the current run.
+        - **MLFLOW_RUN_ID** (`str`, *optional*):
+            Allow to reattach to an existing run which can be useful when resuming training from a
+             checkpoint. When `MLFLOW_RUN_ID` environment variable is set, `start_run` attempts
+             to resume a run with the specified run ID and other parameters are ignored.
+        - Other MLflow environment variables: For changing default logging Behaviour refer mlflow
+            environment variables:
+        https://mlflow.org/docs/latest/python_api/mlflow.environment_variables.html
+        - Setup ``Databricks`` integration with MLflow: Provide these two environment variables:
+            DATABRICKS_HOST="https://adb-4273978218682429.9.azuredatabricks.net"
+            DATABRICKS_TOKEN="dapixxxxxxxxxxxxx"
+        """
+        self._tracking_uri = os.getenv("MLFLOW_TRACKING_URI", None)
+        self._experiment_name = os.getenv("MLFLOW_EXPERIMENT_NAME", None)
+        self._mlflow_log_artifacts = os.getenv("YOLOX_MLFLOW_LOG_MODEL_ARTIFACTS",
+                                               "False").upper() in self.ENV_VARS_TRUE_VALUES
+        self._mlflow_log_model_per_n_epochs = int(os.getenv(
+            "YOLOX_MLFLOW_LOG_MODEL_PER_n_EPOCHS", 30))
+
+        self._mlflow_log_nth_epoch_models = os.getenv("YOLOX_MLFLOW_LOG_Nth_EPOCH_MODELS",
+                                                      "False").upper() in self.ENV_VARS_TRUE_VALUES
+        self.run_name = os.getenv("YOLOX_MLFLOW_RUN_NAME", None)
+        self.run_name = None if len(self.run_name.strip()) == 0 else self.run_name
+        self._flatten_params = os.getenv("YOLOX_MLFLOW_FLATTEN_PARAMS",
+                                         "FALSE").upper() in self.ENV_VARS_TRUE_VALUES
+        self._nested_run = os.getenv("MLFLOW_NESTED_RUN",
+                                     "FALSE").upper() in self.ENV_VARS_TRUE_VALUES
+        self._run_id = os.getenv("MLFLOW_RUN_ID", None)
+
+        # "synchronous" flag is only available with mlflow version >= 2.8.0
+        # https://github.com/mlflow/mlflow/pull/9705
+        # https://github.com/mlflow/mlflow/releases/tag/v2.8.0
+        self._async_log = packaging.version.parse(
+            self._ml_flow.__version__) >= packaging.version.parse("2.8.0")
+
+        logger.debug(
+            f"MLflow experiment_name={self._experiment_name}, run_name={self.run_name}, "
+            f"nested={self._nested_run}, tags={self._nested_run}, tracking_uri={self._tracking_uri}"
+        )
+        if is_main_process():
+            if not self._ml_flow.is_tracking_uri_set():
+                if self._tracking_uri:
+                    self._ml_flow.set_tracking_uri(self._tracking_uri)
+                    logger.debug(f"MLflow tracking URI is set to {self._tracking_uri}")
+                else:
+                    logger.debug(
+                        "Environment variable `MLFLOW_TRACKING_URI` is not provided and therefore"
+                        " will not be explicitly set."
+                    )
+            else:
+                logger.debug(f"MLflow tracking URI is set to {self._ml_flow.get_tracking_uri()}")
+
+            if self._ml_flow.active_run() is None or self._nested_run or self._run_id:
+                if self._experiment_name:
+                    # Use of set_experiment() ensure that Experiment is created if not exists
+                    self._ml_flow.set_experiment(self._experiment_name)
+                self._ml_flow.start_run(run_name=self.run_name, nested=self._nested_run)
+                logger.debug(
+                    f"MLflow run started with run_id={self._ml_flow.active_run().info.run_id}")
+                self._auto_end_run = True
+                self._initialized = True
+            # filters these params from args
+            keys = ['experiment_name', 'batch_size', 'exp_file', 'resume', 'ckpt', 'start_epoch',
+                    'num_machines', 'fp16', 'logger']
+            combined_dict = {k: v for k, v in vars(args).items() if k in keys}
+            if exp is not None:
+                exp_dict = self.convert_exp_todict(exp)
+                combined_dict = {**exp_dict, **combined_dict}
+            self.log_params_mlflow(combined_dict)
+            mlflow_tags = os.getenv("MLFLOW_TAGS", None)
+            if mlflow_tags:
+                mlflow_tags = json.loads(mlflow_tags)
+                self._ml_flow.set_tags(mlflow_tags)
+
+    def log_params_mlflow(self, params_dict):
+        """
+        Log hyperparameters to MLflow.
+        MLflow's log_param() only accepts values no longer than 250 characters.
+        No overwriting of existing parameters is allowed by default from mlflow.
+
+        Args:
+            params_dict(dict): dict of hyperparameters
+
+        Returns:
+            None
+        """
+        if is_main_process():
+            params_dict = self.flatten_dict(params_dict) if self._flatten_params else params_dict
+            # remove params that are too long for MLflow
+            for name, value in list(params_dict.items()):
+                # internally, all values are converted to str in MLflow
+                if len(str(value)) > self._MAX_PARAM_VAL_LENGTH:
+                    logger.warning(
+                        f'Trainer is attempting to log a value of "{value}" for key "{name}" as a '
+                        f'parameter. MLflow\'s log_param() only accepts values no longer than 250 '
+                        f'characters so we dropped this attribute. You can use '
+                        f'`MLFLOW_FLATTEN_PARAMS` environment variable to flatten the parameters '
+                        f'and avoid this message.'
+                    )
+                    del params_dict[name]
+            # MLflow cannot log more than 100 values in one go, so we have to split it
+            combined_dict_items = list(params_dict.items())
+            for i in range(0, len(combined_dict_items), self._MAX_PARAMS_TAGS_PER_BATCH):
+                if self._async_log:
+                    self._ml_flow.log_params(
+                        dict(combined_dict_items[i: i + self._MAX_PARAMS_TAGS_PER_BATCH]),
+                        synchronous=False
+                    )
+                else:
+                    self._ml_flow.log_params(
+                        dict(combined_dict_items[i: i + self._MAX_PARAMS_TAGS_PER_BATCH])
+                    )
+
+    def convert_exp_todict(self, exp):
+        """
+        Convert the experiment object to dictionary for required parameter only
+
+        Args:
+            exp(dict): Experiment object
+
+        Returns:
+            exp_dict(dict): dict of experiment parameters
+
+        """
+        filter_keys = ['max_epoch', 'num_classes', 'input_size', 'output_dir',
+                       'data_dir', 'train_ann', 'val_ann', 'test_ann',
+                       'test_conf', 'nmsthre']
+        exp_dict = {k: v for k, v in exp.__dict__.items()
+                    if not k.startswith("__") and k in filter_keys}
+        return exp_dict
+
+    def on_log(self, args, exp, step, logs):
+        """
+        Log metrics to MLflow.
+
+        Args:
+            args(dict): training args dictionary
+            exp(dict): Experiment related hyperparameters
+            step(int): current training step
+            logs(dict): dictionary of logs to be logged
+
+        Returns:
+            None
+        """
+        # step = trainer.progress_in_iter
+        if not self._initialized:
+            self.setup(args, exp)
+        if is_main_process():  # master thread only
+            metrics = {}
+            for k, v in logs.items():
+                if isinstance(v, (int, float)):
+                    metrics[k] = v
+                elif isinstance(v, torch.Tensor) and v.numel() == 1:
+                    metrics[k] = v.item()
+                else:
+                    logger.warning(
+                        f'Trainer is attempting to log a value of "{v}" of type {type(v)} for key '
+                        f'"{k}" as a metric. MLflow log_metric() only accepts float and int types '
+                        f'so we dropped this attribute.'
+                    )
+
+            if self._async_log:
+                self._ml_flow.log_metrics(metrics=metrics, step=step, synchronous=False)
+            else:
+                self._ml_flow.log_metrics(metrics=metrics, step=step)
+
+    def on_train_end(self, args, file_name, metadata):
+        """
+        Mlflow logging action to take when training ends:
+            1. log the training log file
+            2. publish the latest best model to model_registry if it is allowed in config file
+            3. close the mlfow run
+
+        Args:
+            args(dict): training args dictionary
+            file_name(str): output directory
+            metadata(dict): model related metadata
+
+        Returns:
+            None
+        """
+        if is_main_process() and self._initialized:
+            self.save_log_file(args, file_name)
+            if self.best_ckpt_upload_pending:
+                model_file_name = "best_ckpt"
+                mlflow_out_dir = f"{args.experiment_name}/{model_file_name}"
+                artifact_path = os.path.join(file_name, f"{model_file_name}.pth")
+                self.mlflow_save_pyfunc_model(metadata, artifact_path, mlflow_out_dir)
+            if self._auto_end_run and self._ml_flow.active_run():
+                self._ml_flow.end_run()
+
+    def save_log_file(self, args, file_name):
+        """
+        Save the training log file to mlflow artifact path
+        Args:
+            args(dict): training args dictionary
+            file_name(str): output directory
+
+        Returns:
+            None
+        """
+        log_file_path = os.path.join(file_name, "train_log.txt")
+        mlflow_out_dir = f"{args.experiment_name}"
+        logger.info(f"Logging logfile: {log_file_path} in mlflow artifact path: {mlflow_out_dir}.")
+        self._ml_flow.log_artifact(log_file_path, mlflow_out_dir)
+
+    def save_checkpoints(self, args, exp, file_name, epoch, metadata, update_best_ckpt):
+        """
+        Save the model checkpoints to mlflow artifact path
+        if save_history_ckpt is enabled then
+
+        Args:
+            args(dict): training args dictionary
+            exp(dict): Experiment related hyperparameters
+            file_name(str): output directory
+            epoch(int): current epoch
+            metadata(dict): model related metadata
+            update_best_ckpt(bool): bool to show if best_ckpt was updated
+
+        Returns:
+            None
+        """
+        if is_main_process() and self._mlflow_log_artifacts:
+            if update_best_ckpt:
+                self.best_ckpt_upload_pending = True
+            if ((epoch + 1) % self._mlflow_log_model_per_n_epochs) == 0:
+                self.save_log_file(args, file_name)
+                if self.best_ckpt_upload_pending:
+                    model_file_name = "best_ckpt"
+                    mlflow_out_dir = f"{args.experiment_name}/{model_file_name}"
+                    artifact_path = os.path.join(file_name, f"{model_file_name}.pth")
+                    self.mlflow_save_pyfunc_model(metadata, artifact_path, mlflow_out_dir)
+                    self.best_ckpt_upload_pending = False
+                if self._mlflow_log_nth_epoch_models and exp.save_history_ckpt:
+                    model_file_name = f"epoch_{epoch + 1}_ckpt"
+                    mlflow_out_dir = f"{args.experiment_name}/hist_epochs/{model_file_name}"
+                    artifact_path = os.path.join(file_name, f"{model_file_name}.pth")
+                    self.mlflow_save_pyfunc_model(metadata, artifact_path, mlflow_out_dir)
+
+    def mlflow_save_pyfunc_model(self, metadata, artifact_path, mlflow_out_dir):
+        """
+        This will send the given model to mlflow server if HF_MLFLOW_LOG_ARTIFACTS is true
+            - optionally publish to model registry if allowed in config file
+
+        Args:
+            metadata(dict): model related metadata
+            artifact_path(str): model checkpoint path
+            mlflow_out_dir(str): mlflow artifact path
+
+        Returns:
+            None
+        """
+        if is_main_process() and self._initialized and self._mlflow_log_artifacts:
+            logger.info(
+                f"Logging checkpoint {artifact_path} artifacts in mlflow artifact path: "
+                f"{mlflow_out_dir}. This may take time.")
+            if os.path.exists(artifact_path):
+                self._ml_flow.pyfunc.log_model(
+                    mlflow_out_dir,
+                    artifacts={"model_path": artifact_path},
+                    python_model=self._ml_flow.pyfunc.PythonModel(),
+                    metadata=metadata
+                )
+
+    def __del__(self):
+        """
+        if the previous run is not terminated correctly, the fluent API will
+        not let you start a new run before the previous one is killed
+
+        Args: None
+        Return: None
+        """
+        if (
+                self._auto_end_run
+                and callable(getattr(self._ml_flow, "active_run", None))
+                and self._ml_flow.active_run() is not None
+        ):
+            self._ml_flow.end_run()