Spaces:
Running
Running
File size: 2,071 Bytes
ba003d8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | 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") |