Spaces:
Sleeping
Sleeping
| import mlflow | |
| from mlflow.tracking import MlflowClient | |
| def register_and_promote_model( | |
| model_uri: str, | |
| model_name: str, | |
| val_f1: float, | |
| robustness_drop: float = 0.0, | |
| description: str = "", | |
| tags: dict = None | |
| ) -> str: | |
| """ | |
| Registers model in the MLflow Model Registry and promotes alias to 'champion' | |
| if it surpasses previous performance or if no champion exists. | |
| """ | |
| client = MlflowClient() | |
| # Ensure registered model exists | |
| try: | |
| client.create_registered_model(model_name) | |
| except Exception: | |
| pass | |
| # Create model version | |
| mv = client.create_model_version( | |
| name=model_name, | |
| source=model_uri, | |
| run_id=model_uri.split("/")[1] if "runs:/" in model_uri else "", | |
| description=description, | |
| tags=tags or {} | |
| ) | |
| version = mv.version | |
| # Check existing champion | |
| try: | |
| champ_version = client.get_model_version_by_alias(model_name, "champion") | |
| prev_f1 = float(champ_version.tags.get("val_f1", 0.0)) | |
| if val_f1 > prev_f1 and robustness_drop < 0.25: | |
| client.set_registered_model_alias(model_name, "champion", version) | |
| client.set_registered_model_alias(model_name, "challenger", champ_version.version) | |
| action = "PROMOTED_CHAMPION" | |
| else: | |
| client.set_registered_model_alias(model_name, "challenger", version) | |
| action = "SET_CHALLENGER" | |
| except Exception: | |
| # No champion yet, promote current | |
| client.set_registered_model_alias(model_name, "champion", version) | |
| action = "FIRST_CHAMPION" | |
| # Save validation metrics as tags | |
| client.set_model_version_tag(model_name, version, "val_f1", str(val_f1)) | |
| client.set_model_version_tag(model_name, version, "robustness_drop", str(robustness_drop)) | |
| return action | |
| def get_champion_model_uri(model_name: str) -> str: | |
| """Returns the MLflow models URI for the current champion alias.""" | |
| return f"models:/{model_name}@champion" | |