src/mlpipeline/utils/mlflow_utils.py

import mlflow
from typing import Dict, Any, Optional
from pathlib import Path

from mlpipeline.logging.logger import get_logger

logger = get_logger(__name__)


class MLflowManager:
    def __init__(self, tracking_uri: str, experiment_name: str):
        self.tracking_uri = tracking_uri
        self.experiment_name = experiment_name
        mlflow.set_tracking_uri(tracking_uri)
        mlflow.set_experiment(experiment_name)
        logger.info(f"MLflow tracking URI: {tracking_uri}")
        logger.info(f"MLflow experiment: {experiment_name}")
    
    def start_run(self, run_name: Optional[str] = None):
        mlflow.start_run(run_name=run_name)
        logger.info(f"Started MLflow run: {run_name or 'auto'}")
    
    def end_run(self):
        mlflow.end_run()
        logger.info("Ended MLflow run")
    
    def log_params(self, params: Dict[str, Any]):
        mlflow.log_params(params)
        logger.info(f"Logged {len(params)} parameters")
    
    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None):
        mlflow.log_metrics(metrics, step=step)
        logger.info(f"Logged {len(metrics)} metrics")
    
    def log_metric(self, key: str, value: float, step: Optional[int] = None):
        mlflow.log_metric(key, value, step=step)
    
    def log_artifact(self, local_path: str, artifact_path: Optional[str] = None):
        mlflow.log_artifact(local_path, artifact_path)
        logger.info(f"Logged artifact: {local_path}")
    
    def log_model(self, model: Any, artifact_path: str, **kwargs):
        mlflow.sklearn.log_model(model, artifact_path, **kwargs)
        logger.info(f"Logged model: {artifact_path}")
    
    def register_model(self, model_uri: str, name: str) -> Any:
        result = mlflow.register_model(model_uri, name)
        logger.info(f"Registered model: {name}")
        return result
    
    def set_tag(self, key: str, value: str):
        mlflow.set_tag(key, value)
    
    def set_tags(self, tags: Dict[str, str]):
        mlflow.set_tags(tags)
        logger.info(f"Set {len(tags)} tags")