feat: add mlflow logger in RD loop to log (#815)

* add mlflow logger in DS loop

* fix CI

* fix CI

---------

Co-authored-by: Xu Yang <[email protected]>

Author: peteryang1

rdagent/core/conf.py

@@ -80,5 +80,7 @@ class RDAgentSettings(ExtendedBaseSettings):
     stdout_context_len: int = 400
     stdout_line_len: int = 10000
 
+    enable_mlflow: bool = False
+
 
 RD_AGENT_SETTINGS = RDAgentSettings()

rdagent/utils/workflow.py

@@ -9,18 +9,24 @@
 """
 
 import datetime
+import os
 import pickle
 import time
 from collections import defaultdict
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Callable, Optional, TypeVar, Union, cast
 
+import pytz
 from tqdm.auto import tqdm
 
+from rdagent.core.conf import RD_AGENT_SETTINGS
 from rdagent.log import rdagent_logger as logger
 from rdagent.log.timer import RD_Agent_TIMER_wrapper, RDAgentTimer
 
+if RD_AGENT_SETTINGS.enable_mlflow:
+    import mlflow
+
 
 class LoopMeta(type):
     @staticmethod
@@ -119,7 +125,27 @@ def run(self, step_n: int | None = None, loop_n: int | None = None, all_duration
                     if loop_n <= 0:
                         break
 
+                if RD_AGENT_SETTINGS.enable_mlflow:
+                    mlflow.log_metric("loop_index", self.loop_idx)
+                    mlflow.log_metric("step_index", self.step_idx)
+                    current_local_datetime = datetime.datetime.now(pytz.timezone("Asia/Shanghai"))
+                    float_like_datetime = (
+                        current_local_datetime.second
+                        + current_local_datetime.minute * 1e2
+                        + current_local_datetime.hour * 1e4
+                        + current_local_datetime.day * 1e6
+                        + current_local_datetime.month * 1e8
+                        + current_local_datetime.year * 1e10
+                    )
+                    mlflow.log_metric("current_datetime", float_like_datetime)
+
                 if self.timer.started:
+                    if RD_AGENT_SETTINGS.enable_mlflow:
+                        mlflow.log_metric("remain_time", self.timer.remain_time().seconds)  # type: ignore[union-attr]
+                        mlflow.log_metric(
+                            "remain_percent", self.timer.remain_time() / self.timer.all_duration * 100  # type: ignore[operator]
+                        )
+
                     if self.timer.is_timeout():
                         logger.warning("Timeout, exiting the loop.")
                         break

requirements.txt

@@ -57,3 +57,8 @@ azure.ai.inference
 # data folder desc
 humanize
 genson
+
+# mlflow
+mlflow
+azureml-mlflow
+types-pytz