diff --git a/rdagent/core/conf.py b/rdagent/core/conf.py index c6b7bdf8..0824c2fe 100644 --- a/rdagent/core/conf.py +++ b/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() diff --git a/rdagent/utils/workflow.py b/rdagent/utils/workflow.py index 4bcedaf5..2abeacd8 100644 --- a/rdagent/utils/workflow.py +++ b/rdagent/utils/workflow.py @@ -9,6 +9,7 @@ Postscripts: """ import datetime +import os import pickle import time from collections import defaultdict @@ -16,11 +17,16 @@ 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 @@ class LoopBase: 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 diff --git a/requirements.txt b/requirements.txt index 6bb09aba..ebac1ec2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -57,3 +57,8 @@ azure.ai.inference # data folder desc humanize genson + +# mlflow +mlflow +azureml-mlflow +types-pytz \ No newline at end of file