databricks-sdk-py icon indicating copy to clipboard operation
databricks-sdk-py copied to clipboard

[ISSUE] RegisteredModelInfo.aliases is empty despite registered aliases existing

Open j-lootens opened this issue 1 year ago • 0 comments

Description RegisteredModelInfo does not return model aliases as expected. It always contains an empty list instead of a list of all the aliases of all model versions under this registered model.

Reproduction

import mlflow

mlflow.set_registry_uri('databricks-uc')

from databricks.sdk import WorkspaceClient
from mlflow.models.signature import infer_signature
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

# Load dataset
data = load_iris()
X_train, X_test, y_train, y_test = train_test_split(data.data, data.target, test_size=0.2, random_state=42)

# Train model
model = LogisticRegression(max_iter=200)
model.fit(X_train, y_train)

# Evaluate model
predictions = model.predict(X_test)
accuracy = accuracy_score(y_test, predictions)
signature = infer_signature(X_train, model.predict(X_train))

catalog = "INSERT_CATALOG"
schema = "INSERT_SCHEMA"
# Log model and metrics to MLflow
with mlflow.start_run() as run:
    mlflow.sklearn.log_model(model, "logistic_regression_model", registered_model_name=f"{catalog}.{schema}.logistic_regression_model_uc", signature=signature)
    mlflow.log_metric("accuracy", accuracy)


w = WorkspaceClient()
w.registered_models.set_alias(f"{catalog}.{schema}.logistic_regression_model_uc", "test",1)

model = MlflowClient().get_model_version_by_alias(f"{catalog}.{schema}.logistic_regression_model_uc", "test")
assert len(model.aliases) == 1 # True

registered_models = list(w.registered_models.list(catalog_name=catalog, schema_name=schema))
# assumes only one model has been registered
assert len(registered_models[0].aliases) > 0 # False

Expected behavior RegisteredModelInfo.aliases contains all aliases of all model versions under this registered model instead of always being an empty list

Is it a regression? /

Debug Logs /

Other Information /

Additional context

  • DBR 16.1 ML (includes Apache Spark 3.5.0, Scala 2.12)
  • databricks-sdk version: 0.30.0

j-lootens avatar Jan 21 '25 16:01 j-lootens