Added mlflow system metrics.

Author: fmind

poetry.lock

@@ -23,17 +23,19 @@ bikes = 'bikes.scripts:main'
 [tool.poetry.dependencies]
 python = "^3.12"
 cloudpathlib = "^0.17.0"
+codecarbon = "^2.3.4"
 loguru = "^0.7.2"
+mlflow-skinny = "^2.10.2"
 omegaconf = "^2.3.0"
 pandas = "^2.2.0"
 pandera = "^0.18.0"
 plotly = "^5.18.0"
+psutil = "^5.9.8"
 pyarrow = "^15.0.0"
 pydantic = "^2.5.3"
 pydantic-settings = "^2.1.0"
+pynvml = "^11.5.0"
 scikit-learn = "^1.4.0"
-mlflow-skinny = "^2.10.2"
-codecarbon = "^2.3.4"
 
 [tool.poetry.group.dev.dependencies]
 invoke = "^2.2.0"

pyproject.toml

@@ -132,6 +132,7 @@ class MLflowService(Service):
         autolog_log_models (bool): If True, enables logging of models during autologging.
         autolog_log_datasets (bool): If True, logs datasets used during autologging.
         autolog_silent (bool): If True, suppresses all MLflow warnings during autologging.
+        enable_system_metrics (bool): enable system metrics logging.
         tracking_uri (str): The URI for the MLflow tracking server.
         experiment_name (str): The name of the experiment to log runs under.
         registry_uri (str): The URI for the MLflow model registry.
@@ -147,6 +148,8 @@ class MLflowService(Service):
     autolog_log_models: bool = False
     autolog_log_datasets: bool = True
     autolog_silent: bool = False
+    # system
+    enable_system_metrics: bool = True
     # tracking
     tracking_uri: str = "http://localhost:5000"
     experiment_name: str = "bikes"
@@ -171,6 +174,9 @@ def start(self):
             log_models=self.autolog_log_models,
             silent=self.autolog_silent,
         )
+        # system metrics
+        if self.enable_system_metrics:
+            mlflow.enable_system_metrics_logging()
 
     def client(self) -> MlflowClient:
         """Get an instance of MLflow client."""

src/bikes/services.py

@@ -3,6 +3,7 @@
 import os
 
 import mlflow
+import pytest
 from bikes import services
 from loguru import logger
 
@@ -37,9 +38,11 @@ def test_carbon_service(tmp_carbon_path: str):
     assert os.stat(output_path).st_size > 0, "Output path should not be empty!"
 
 
-def test_mlflow_service(mlflow_service: services.MLflowService):
+@pytest.mark.parametrize("enable_system_metrics", [True, False])
+def test_mlflow_service(enable_system_metrics: bool, mlflow_service: services.MLflowService):
     # given
     service = mlflow_service
+    service.enable_system_metrics = enable_system_metrics
     # when
     service.start()
     client = service.client()