Add mlflow tracker and unit testing for the same.

Also add mlflow docs and add mlflow to docker file and as optional requirement

Signed-off-by: Dushyant Behl <[email protected]>

README.md

@@ -823,12 +823,13 @@ For details about how you can use set a custom stopping criteria and perform cus
 
 ## Experiment Tracking
 
-Experiment tracking in fms-hf-tuning allows users to track their experiments with known trackers like [Aimstack](https://aimstack.io/) or custom trackers built into the code like
+Experiment tracking in fms-hf-tuning allows users to track their experiments with known trackers like [Aimstack](https://aimstack.io/), [MLflow Tracking](https://mlflow.org/docs/latest/tracking.html) or custom trackers built into the code like
 [FileLoggingTracker](./tuning/trackers/filelogging_tracker.py)
 
 The code supports currently two trackers out of the box, 
 * `FileLoggingTracker` : A built in tracker which supports logging training loss to a file.
 * `Aimstack` : A popular opensource tracker which can be used to track any metrics or metadata from the experiments.
+* `MLflow Tracking` : Another popular opensource tracker which stores metrics, metadata or even artifacts from experiments.
 
 Further details on enabling and using the trackers mentioned above can be found [here](docs/experiment-tracking.md).  
 

build/Dockerfile

@@ -19,8 +19,9 @@ ARG USER=tuning
 ARG USER_UID=1000
 ARG PYTHON_VERSION=3.11
 ARG WHEEL_VERSION=""
-## Enable Aimstack if requested via ENABLE_AIM set to "true"
+## Enable Aimstack or MLflow if requested via ENABLE_AIM/MLFLOW set to "true"
 ARG ENABLE_AIM=false
+ARG ENABLE_MLFLOW=false
 ARG ENABLE_FMS_ACCELERATION=true
 
 ## Base Layer ##################################################################
@@ -151,6 +152,10 @@ RUN if [[ "${ENABLE_AIM}" == "true" ]]; then \
         python -m pip install --user "$(head bdist_name)[aim]"; \
     fi
 
+RUN if [[ "${ENABLE_MLFLOW}" == "true" ]]; then \
+    python -m pip install --user "$(head bdist_name)[mlflow]"; \
+fi
+
     # Clean up the wheel module. It's only needed by flash-attn install
 RUN python -m pip uninstall wheel build -y && \
     # Cleanup the bdist whl file

docs/experiment-tracking.md

@@ -115,6 +115,34 @@ sft_trainer.train(train_args=training_args, tracker_configs=tracker_configs,....
 The code expects either the `local` or `remote` repo to be specified and will result in a `ValueError` otherwise.
 See [AimConfig](https://github.com/foundation-model-stack/fms-hf-tuning/blob/a9b8ec8d1d50211873e63fa4641054f704be8712/tuning/config/tracker_configs.py#L25) for more details.
 
+## MLflow Tracker
+
+To enable [MLflow Tracking](https://mlflow.org/docs/latest/tracking.html) users need to pass `"mlflow"` as the requested tracker as part of the [training argument](https://github.com/foundation-model-stack/fms-hf-tuning/blob/a9b8ec8d1d50211873e63fa4641054f704be8712/tuning/config/configs.py#L131).
+
+
+When using MLflow, users need to specify additional arguments which specify [mlflow tracking uri](https://mlflow.org/docs/latest/tracking.html#common-setups) location where either a [mlflow supported database](https://mlflow.org/docs/latest/tracking/backend-stores.html#supported-store-types) or [mlflow remote tracking server](https://mlflow.org/docs/latest/tracking/server.html) is running.
+
+Example
+```
+from tuning import sft_trainer
+from tuning.config.tracker_configs import MLflowConfig, TrackerConfigFactory
+
+training_args = TrainingArguments(
+    ...,
+    trackers = ["mlflow"],
+)
+
+tracker_configs = TrackerConfigFactory(
+    mlflow_config=MLflowConfig(
+        mlflow_experiment="experiment-name",
+        mlflow_tracking_uri=<tracking uri>
+        )
+    )
+
+sft_trainer.train(train_args=training_args, tracker_configs=tracker_configs,....)
+```
+
+The code expects a valid uri to be specified and will result in a `ValueError` otherwise.
 
 ## Running the code via command line `tuning/sft_trainer::main` function
 
@@ -123,10 +151,10 @@ If running the code via main function of [sft_trainer.py](../tuning/sft_trainer.
 To enable tracking please pass
 
 ```
---tracker <aim/file_logger>
+--tracker <aim/file_logger/mlflow>
 ```
 
-To further customise tracking you can specify additional arguments needed by the tracker like
+To further customise tracking you can specify additional arguments needed by the tracker like (example shows aim follow similarly for mlflow)
 
 ```
 --tracker aim --aim_repo <path-to-aimrepo> --experiment <experiment-name>

pyproject.toml

@@ -44,6 +44,7 @@ dependencies = [
 dev = ["wheel>=0.42.0,<1.0", "packaging>=23.2,<25", "ninja>=1.11.1.1,<2.0", "scikit-learn>=1.0, <2.0", "boto3>=1.34, <2.0"]
 flash-attn = ["flash-attn>=2.5.3,<3.0"]
 aim = ["aim>=3.19.0,<4.0"]
+mlflow = ["mlflow"]
 fms-accel = ["fms-acceleration>=0.1"]
 gptq-dev = ["auto_gptq>0.4.2", "optimum>=1.15.0"]
 

tests/test_sft_trainer.py

@@ -359,6 +359,7 @@ def test_parse_arguments(job_config):
         _,
         _,
         _,
+        _,
     ) = sft_trainer.parse_arguments(parser, job_config_copy)
     assert str(model_args.torch_dtype) == "torch.bfloat16"
     assert data_args.dataset_text_field == "output"
@@ -384,6 +385,7 @@ def test_parse_arguments_defaults(job_config):
         _,
         _,
         _,
+        _,
     ) = sft_trainer.parse_arguments(parser, job_config_defaults)
     assert str(model_args.torch_dtype) == "torch.bfloat16"
     assert model_args.use_flash_attn is False
@@ -394,14 +396,14 @@ def test_parse_arguments_peft_method(job_config):
     parser = sft_trainer.get_parser()
     job_config_pt = copy.deepcopy(job_config)
     job_config_pt["peft_method"] = "pt"
-    _, _, _, _, tune_config, _, _, _, _, _, _ = sft_trainer.parse_arguments(
+    _, _, _, _, tune_config, _, _, _, _, _, _, _ = sft_trainer.parse_arguments(
         parser, job_config_pt
     )
     assert isinstance(tune_config, peft_config.PromptTuningConfig)
 
     job_config_lora = copy.deepcopy(job_config)
     job_config_lora["peft_method"] = "lora"
-    _, _, _, _, tune_config, _, _, _, _, _, _ = sft_trainer.parse_arguments(
+    _, _, _, _, tune_config, _, _, _, _, _, _, _ = sft_trainer.parse_arguments(
         parser, job_config_lora
     )
     assert isinstance(tune_config, peft_config.LoraConfig)

tests/trackers/test_aim_tracker.py

@@ -58,7 +58,7 @@ def fixture_aimrepo():
 
 
 @pytest.mark.skipif(aim_not_available, reason="Requires aimstack to be installed")
-def test_run_with_good_tracker_name_but_no_args():
+def test_run_with_aim_tracker_name_but_no_args():
     """Ensure that train() raises error with aim tracker name but no args"""
 
     with tempfile.TemporaryDirectory() as tempdir:

tests/trackers/test_mlflow_tracker.py

@@ -0,0 +1,130 @@
+# Copyright The FMS HF Tuning Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# SPDX-License-Identifier: Apache-2.0
+# https://spdx.dev/learn/handling-license-info/
+
+# Standard
+import copy
+import json
+import os
+import tempfile
+import contextlib
+
+# Third Party
+from transformers.utils.import_utils import _is_package_available
+import pytest
+
+# First Party
+from tests.test_sft_trainer import (
+    DATA_ARGS,
+    MODEL_ARGS,
+    TRAIN_ARGS,
+    _get_checkpoint_path,
+    _test_run_inference,
+    _validate_training,
+)
+
+# Local
+from tuning import sft_trainer
+from tuning.config.tracker_configs import MLflowConfig, TrackerConfigFactory
+
+mlflow_not_available = not _is_package_available("mlflow")
+
+@pytest.mark.skipif(mlflow_not_available, reason="Requires mlflow to be installed")
+def test_run_with_mlflow_tracker_name_but_no_args():
+    """Ensure that train() raises error with mlflow tracker name but no args"""
+
+    with tempfile.TemporaryDirectory() as tempdir:
+        train_args = copy.deepcopy(TRAIN_ARGS)
+        train_args.output_dir = tempdir
+
+        train_args.trackers = ["mlflow"]
+
+        with pytest.raises(
+            ValueError,
+            match="mlflow tracker requested but mlflow_uri is not specified.",
+        ):
+            sft_trainer.train(MODEL_ARGS, DATA_ARGS, train_args)
+
+
+@pytest.mark.skipif(mlflow_not_available, reason="Requires mlflow to be installed")
+def test_e2e_run_with_mlflow_tracker():
+    """Ensure that training succeeds with mlflow tracker"""
+
+    # mlflow performs a cleanup at callback close time which happens post the
+    # delete of this directory so we run into two issues
+    # 1. the temp directory cannot be cleared as it has open pointer by mlflow
+    # 2. mlflow complaints that it cannot find a run which it just created.
+    # this is a race condition which is fixed with mkdtemp() which doesn't delete
+    tempdir = tempfile.mkdtemp()
+
+    train_args = copy.deepcopy(TRAIN_ARGS)
+    train_args.output_dir = tempdir
+
+    # This should not mean file logger is not present.
+    # code will add it by default
+    # The below validate_training check will test for that too.
+    train_args.trackers = ["mlflow"]
+
+    tracker_configs = TrackerConfigFactory(
+        mlflow_config = MLflowConfig(
+            mlflow_experiment="unit_test",
+            mlflow_tracking_uri=os.path.join(tempdir, "mlflowdb.sqlite"),
+        )
+    )
+
+    sft_trainer.train(
+        MODEL_ARGS, DATA_ARGS, train_args, tracker_configs=tracker_configs
+    )
+
+    # validate ft tuning configs
+    _validate_training(tempdir)
+
+    # validate inference
+    _test_run_inference(checkpoint_path=_get_checkpoint_path(tempdir))
+
+
+@pytest.mark.skipif(mlflow_not_available, reason="Requires mlflow to be installed")
+def test_e2e_run_with_mlflow_runuri_export_default_path():
+    """Ensure that mlflow outputs run uri in the output dir by default"""
+
+    tempdir = tempfile.mkdtemp()
+    train_args = copy.deepcopy(TRAIN_ARGS)
+    train_args.output_dir = tempdir
+
+    train_args.trackers = ["mlflow"]
+
+    tracker_configs = TrackerConfigFactory(
+        mlflow_config = MLflowConfig(
+            mlflow_experiment="unit_test",
+            mlflow_tracking_uri=os.path.join(tempdir, "mlflowdb.sqlite"),
+        )
+    )
+
+    sft_trainer.train(
+        MODEL_ARGS, DATA_ARGS, train_args, tracker_configs=tracker_configs
+    )
+
+    # validate ft tuning configs
+    _validate_training(tempdir)
+
+    run_uri_file = os.path.join(tempdir, "mlflow_tracker.json")
+
+    assert os.path.exists(run_uri_file) is True
+    assert os.path.getsize(run_uri_file) > 0
+
+    with open(run_uri_file, "r", encoding="utf-8") as f:
+        content = json.loads(f.read())
+        assert "run_uri" in content

tuning/config/tracker_configs.py

@@ -63,7 +63,27 @@ def __post_init__(self):
             )
 
 
+@dataclass
+class MLflowConfig:
+    # Name of the experiment
+    mlflow_experiment: str = None
+    mlflow_tracking_uri: str = None
+    # Location of where mlflow's run uri is to be exported.
+    # If mlflow_run_uri_export_path is set the run uri will be output in a json format
+    # to the location pointed to by `mlflow_run_uri_export_path/mlflow_tracker.json`
+    # If this is not set then the default location where run uri will be exported
+    # is training_args.output_dir/mlflow_tracker.json
+    # Run uri is not exported if mlflow_run_uri_export_path variable is not set
+    # and output_dir is not specified.
+    mlflow_run_uri_export_path: str = None
+
+    def __post_init__(self):
+        if self.mlflow_experiment is None:
+            self.mlflow_experiment = "fms-hf-tuning"
+
+
 @dataclass
 class TrackerConfigFactory:
     file_logger_config: FileLoggingTrackerConfig = None
     aim_config: AimConfig = None
+    mlflow_config: MLflowConfig = None

tuning/sft_trainer.py

@@ -51,6 +51,7 @@
 from tuning.config.tracker_configs import (
     AimConfig,
     FileLoggingTrackerConfig,
+    MLflowConfig,
     TrackerConfigFactory,
 )
 from tuning.data.setup_dataprocessor import process_dataargs
@@ -444,6 +445,7 @@ def get_parser():
             QuantizedLoraConfig,
             FusedOpsAndKernelsConfig,
             AttentionAndDistributedPackingConfig,
+            MLflowConfig,
         )
     )
     parser.add_argument(
@@ -493,8 +495,10 @@ def parse_arguments(parser, json_config=None):
             Configuration for fused operations and kernels.
         AttentionAndDistributedPackingConfig
             Configuration for padding free and packing.
+        MLflowConfig
+            Configuration for mlflow tracker.
         dict[str, str]
-            Extra AIM metadata.
+            Extra tracker metadata.
     """
     if json_config:
         (
@@ -509,6 +513,7 @@ def parse_arguments(parser, json_config=None):
             quantized_lora_config,
             fusedops_kernels_config,
             attention_and_distributed_packing_config,
+            mlflow_config,
         ) = parser.parse_dict(json_config, allow_extra_keys=True)
         peft_method = json_config.get("peft_method")
         exp_metadata = json_config.get("exp_metadata")
@@ -525,6 +530,7 @@ def parse_arguments(parser, json_config=None):
             quantized_lora_config,
             fusedops_kernels_config,
             attention_and_distributed_packing_config,
+            mlflow_config,
             additional,
             _,
         ) = parser.parse_args_into_dataclasses(return_remaining_strings=True)
@@ -550,6 +556,7 @@ def parse_arguments(parser, json_config=None):
         quantized_lora_config,
         fusedops_kernels_config,
         attention_and_distributed_packing_config,
+        mlflow_config,
         exp_metadata,
     )
 
@@ -571,6 +578,7 @@ def main():
             quantized_lora_config,
             fusedops_kernels_config,
             attention_and_distributed_packing_config,
+            mlflow_config,
             exp_metadata,
         ) = parse_arguments(parser, job_config)
 
@@ -582,7 +590,8 @@ def main():
             model_args %s, data_args %s, training_args %s, trainer_controller_args %s, \
             tune_config %s, file_logger_config, %s aim_config %s, \
             quantized_lora_config %s, fusedops_kernels_config %s, \
-            attention_and_distributed_packing_config %s exp_metadata %s",
+            attention_and_distributed_packing_config %s,\
+            mlflow_config %s, exp_metadata %s",
             model_args,
             data_args,
             training_args,
@@ -593,6 +602,7 @@ def main():
             quantized_lora_config,
             fusedops_kernels_config,
             attention_and_distributed_packing_config,
+            mlflow_config,
             exp_metadata,
         )
     except Exception as e:  # pylint: disable=broad-except
@@ -617,10 +627,11 @@ def main():
                 "failed while parsing extra metadata. pass a valid json %s", repr(e)
             )
 
-    combined_tracker_configs = TrackerConfigFactory()
-
-    combined_tracker_configs.file_logger_config = file_logger_config
-    combined_tracker_configs.aim_config = aim_config
+    tracker_configs = TrackerConfigFactory(
+        file_logger_config=file_logger_config,
+        aim_config=aim_config,
+        mlflow_config=mlflow_config,
+    )
 
     if training_args.output_dir:
         os.makedirs(training_args.output_dir, exist_ok=True)
@@ -632,7 +643,7 @@ def main():
             train_args=training_args,
             peft_config=tune_config,
             trainer_controller_args=trainer_controller_args,
-            tracker_configs=combined_tracker_configs,
+            tracker_configs=tracker_configs,
             additional_callbacks=None,
             exp_metadata=metadata,
             quantized_lora_config=quantized_lora_config,