mirror of
https://github.com/NicolasBohn/NexQuant.git
synced 2026-07-28 16:07:46 +00:00
TPTBusiness
20428f7d91
- quant.py: guard against empty orch_factors, move strategy_name before try block - quant_proposal.py: fix __init__ return type Tuple[dict,bool] -> None - strategy_orchestrator.py: remove dead rdagent_logger import shadowed by getLogger - factor.py: replace unusual 'not x is None' with idiomatic 'x is not None' - workflow/loop.py: withdraw_loop(0) raises RuntimeError instead of looking for folder -1 - workflow/tracking.py: replace crash-prone AssertionError with logger.warning + skip - factor_from_report.py: fix misleading comment about loop_n/step_n dual use
98 lines
3.6 KiB
Python
98 lines
3.6 KiB
Python
"""
|
|
Tracking module for experiment tracking using MLflow.
|
|
|
|
This module provides a clean interface for tracking metrics and parameters
|
|
while keeping the MLflow dependency optional based on configuration.
|
|
"""
|
|
|
|
import datetime
|
|
from typing import TYPE_CHECKING
|
|
|
|
import pytz
|
|
from rdagent.core.conf import RD_AGENT_SETTINGS
|
|
from rdagent.log.timer import RD_Agent_TIMER_wrapper
|
|
|
|
if TYPE_CHECKING:
|
|
# Import here to avoid circular dependency
|
|
from rdagent.utils.workflow.loop import LoopBase
|
|
|
|
from rdagent.log import rdagent_logger as logger
|
|
|
|
# Define a placeholder for mlflow if it's not available
|
|
mlflow = None
|
|
|
|
# Conditional import to make MLflow optional
|
|
if RD_AGENT_SETTINGS.enable_mlflow:
|
|
try:
|
|
import mlflow # type: ignore[assignment]
|
|
except ImportError:
|
|
logger.warning("MLflow is enabled in settings but could not be imported.")
|
|
RD_AGENT_SETTINGS.enable_mlflow = False
|
|
|
|
|
|
class WorkflowTracker:
|
|
"""
|
|
A workflow-specific tracking system that logs metrics related to workflow execution.
|
|
|
|
This class handles metric logging while keeping the MLflow dependency optional.
|
|
If MLflow is not enabled in settings, tracking calls become no-ops.
|
|
"""
|
|
|
|
def __init__(self, loop_base: "LoopBase"):
|
|
"""
|
|
Initialize a WorkflowTracker with a LoopBase instance.
|
|
|
|
Args:
|
|
loop_base: The LoopBase instance to track metrics for
|
|
"""
|
|
self.loop_base = loop_base
|
|
|
|
@staticmethod
|
|
def is_enabled() -> bool:
|
|
"""Check if tracking is enabled."""
|
|
return RD_AGENT_SETTINGS.enable_mlflow
|
|
|
|
@staticmethod
|
|
def _datetime_to_float(dt: datetime.datetime) -> float:
|
|
"""Convert datetime to a structured float representation."""
|
|
return dt.second + dt.minute * 1e2 + dt.hour * 1e4 + dt.day * 1e6 + dt.month * 1e8 + dt.year * 1e10
|
|
|
|
def log_workflow_state(self) -> None:
|
|
"""
|
|
Log all workflow state metrics from the associated LoopBase instance.
|
|
"""
|
|
if not RD_AGENT_SETTINGS.enable_mlflow or mlflow is None:
|
|
return
|
|
|
|
try:
|
|
# Log workflow progress
|
|
mlflow.log_metric("loop_index", self.loop_base.loop_idx)
|
|
mlflow.log_metric("step_index", self.loop_base.step_idx[self.loop_base.loop_idx])
|
|
|
|
current_local_datetime = datetime.datetime.now(pytz.timezone("Asia/Shanghai"))
|
|
float_like_datetime = self._datetime_to_float(current_local_datetime)
|
|
mlflow.log_metric("current_datetime", float_like_datetime)
|
|
|
|
# Log API status
|
|
mlflow.log_metric("api_fail_count", RD_Agent_TIMER_wrapper.api_fail_count)
|
|
latest_api_fail_time = RD_Agent_TIMER_wrapper.latest_api_fail_time
|
|
if latest_api_fail_time is not None:
|
|
float_like_datetime = self._datetime_to_float(latest_api_fail_time)
|
|
mlflow.log_metric("lastest_api_fail_time", float_like_datetime)
|
|
|
|
# Log timer status if timer is started
|
|
if self.loop_base.timer.started:
|
|
remain_time = self.loop_base.timer.remain_time()
|
|
if remain_time is None:
|
|
logger.warning("remain_time is None despite timer.started, skipping timer metrics")
|
|
else:
|
|
mlflow.log_metric("remain_time", remain_time.total_seconds())
|
|
mlflow.log_metric(
|
|
"remain_percent",
|
|
remain_time / self.loop_base.timer.all_duration * 100,
|
|
)
|
|
|
|
# Keep only the log_workflow_state method as it's the primary entry point now
|
|
except Exception as e:
|
|
logger.warning(f"Error in log_workflow_state: {e}")
|