Spaces:
Sleeping
Sleeping
File size: 2,008 Bytes
c1b1f3a | 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 57 58 59 60 | 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"
|