Add MLflow support and expose logging configuration in TrainingArgs

README.md

@@ -22,6 +22,7 @@ The library now supports reasoning traces through the `reasoning_content` field
 
 - [Installing](#installing-the-library)
   - [Additional Nvidia packages](#additional-nvidia-packages)
+  - [Optional logging dependencies](#optional-logging-dependencies)
 - [Using the library](#using-the-library)
 - [Data format](#data-format)
   - [Reasoning content support](#reasoning-content-support-1)
@@ -68,6 +69,23 @@ Editable install (development)
 pip install -e .[cuda]
 ```
 
+### Optional Logging Dependencies
+
+The library supports optional logging backends for experiment tracking. Install the ones you need:
+
+```bash
+# MLflow logging
+pip install mlflow
+
+# Weights & Biases logging
+pip install wandb
+
+# TensorBoard logging
+pip install tensorboard
+```
+
+For more details on configuring logging, see the [Logging Documentation](docs/logging.md).
+
 ## Using the library
 
 See the `examples` dir for guided sample notebooks on library usage. Below provides some added details on library options:

pyproject.toml

@@ -48,6 +48,9 @@ optional-dependencies.cuda = { file = ["requirements-cuda.txt"] }
 optional-dependencies.rocm = { file = ["requirements-rocm.txt"] }
 optional-dependencies.hpu = { file = ["requirements-hpu.txt"] }
 optional-dependencies.deepspeed = { file = ["requirements-deepspeed.txt"] }
+optional-dependencies.mlflow = { file = ["requirements-mlflow.txt"] }
+optional-dependencies.wandb = { file = ["requirements-wandb.txt"] }
+optional-dependencies.tensorboard = { file = ["requirements-tensorboard.txt"] }
 
 [tool.setuptools.packages.find]
 where = ["src"]

requirements-mlflow.txt

@@ -0,0 +1 @@
+mlflow

requirements-tensorboard.txt

@@ -0,0 +1 @@
+tensorboard

requirements-wandb.txt

@@ -0,0 +1 @@
+wandb

src/instructlab/training/config.py

@@ -305,3 +305,38 @@ class TrainingArgs(BaseModel):
     log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = Field(
         default="INFO"
     )
+
+    mlflow_tracking_uri: str | None = Field(
+        default=None,
+        description="MLflow tracking server URI (e.g., 'http://localhost:5000'). Falls back to MLFLOW_TRACKING_URI env var.",
+    )
+
+    mlflow_experiment_name: str | None = Field(
+        default=None,
+        description="MLflow experiment name. Falls back to MLFLOW_EXPERIMENT_NAME env var.",
+    )
+
+    mlflow_run_name: str | None = Field(
+        default=None,
+        description="MLflow run name. Supports placeholders: {time}, {rank}, {utc_time}, {local_rank}",
+    )
+
+    wandb_project: str | None = Field(
+        default=None,
+        description="Weights & Biases project name.",
+    )
+
+    wandb_entity: str | None = Field(
+        default=None,
+        description="Weights & Biases team/entity name.",
+    )
+
+    wandb_run_name: str | None = Field(
+        default=None,
+        description="Weights & Biases run name. Supports placeholders: {time}, {rank}, {utc_time}, {local_rank}",
+    )
+
+    tensorboard_log_dir: str | None = Field(
+        default=None,
+        description="Directory for TensorBoard logs. Defaults to ckpt_output_dir if not specified.",
+    )