Repository: mlflow/mlflow

Issue:
[BUG] unity catalog 
### Issues Policy acknowledgement

- [X] I have read and agree to submit bug reports in accordance with the [issues policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md)

### Where did you encounter this bug?

Local machine

### Willingness to contribute

No. I cannot contribute a bug fix at this time.

### MLflow version

- Client: 2.13.0


### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Ubuntu 22.04
- **Python version**: 3.10


### Describe the problem

Databricks uc backend search_model_versions by run_id does not work.

### Tracking information

DATABRICKS UC

### Code to reproduce issue

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
import mlflow
from mlflow import MlflowClient
from mlflow.legacy_databricks_cli.configure.provider import set_config_provider
from sklearn.datasets import load_iris
from sklearn.linear_model import LogisticRegression

from rmg_lakehouse.session import MlFlowConfigProvider

mlflow.set_tracking_uri("databricks")
mlflow.set_registry_uri("databricks-uc")
mlflow.set_experiment("/experiments/test-experiment")
X, y = load_iris(return_X_y=True)
model = LogisticRegression().fit(X, y)
with mlflow.start_run():
    model_info = mlflow.sklearn.log_model(model, "model")

client = MlflowClient()
model_versions = client.search_model_versions(f"run_id='{model_info.run_id}'")
```


### Stack trace

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/home/noortw01/.conda/envs/rmg-lakehouse/lib/python3.10/site-packages/mlflow/tracking/client.py", line 4190, in search_model_versions
    return self._get_registry_client().search_model_versions(
  File "/home/noortw01/.conda/envs/rmg-lakehouse/lib/python3.10/site-packages/mlflow/tracking/_model_registry/client.py", line 347, in search_model_versions
    return self.store.search_model_versions(filter_string, max_results, order_by, page_token)
  File "/home/noortw01/.conda/envs/rmg-lakehouse/lib/python3.10/site-packages/mlflow/store/_unity_catalog/registry/rest_store.py", line 904, in search_model_versions
    response_proto = self._call_endpoint(SearchModelVersionsRequest, req_body)
  File "/home/noortw01/.conda/envs/rmg-lakehouse/lib/python3.10/site-packages/mlflow/store/model_registry/base_rest_store.py", line 44, in _call_endpoint
    return call_endpoint(
  File "/home/noortw01/.conda/envs/rmg-lakehouse/lib/python3.10/site-packages/mlflow/utils/rest_utils.py", line 304, in call_endpoint
    response = verify_rest_response(response, endpoint)
  File "/home/noortw01/.conda/envs/rmg-lakehouse/lib/python3.10/site-packages/mlflow/utils/rest_utils.py", line 174, in verify_rest_response
    raise RestException(json.loads(response.text))
mlflow.exceptions.RestException: INVALID_PARAMETER_VALUE: Unsupported filter query : `run_id='23bfb54b487445d19ebc5a07b717bf2b'`. Please specify your filter parameter in the format `name = 'model_name'`.
```


### Other info / logs

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
REPLACE_ME
```


### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/deployments`: MLflow Deployments client APIs, server, and third-party Deployments integrations
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [X] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [X] `integrations/databricks`: Databricks integrations


Discussion:
I think databricks-uc backend hasn't supported search_model_versions by run_id,
I pinged our UC team about this issue.
Thought so. I need the version so I can put an alias on it. I can get versions by name and take the last one but seems more deterministic to actually get the model produced by this run_id. Not sure if it's feasible to add the version to the ModelInfo returned from log_model function?
@Willem-J-an thanks for flagging this - support for complex searches in UC is a bit further out, but it should definitely be easier to get the version number of a newly-registered model version when calling `log_model`, @WeichenXu123 I wonder if we could return it as part of the [ModelInfo](https://mlflow.org/docs/latest/python_api/mlflow.sklearn.html?highlight=log_model#mlflow.sklearn.log_model) returned by `log_model`? If the interface change seems sensible and won't break things/abstractions elsewhere I'm happy to review a PR to add this


> @Willem-J-an thanks for flagging this - support for complex searches in UC is a bit further out, but it should definitely be easier to get the version number of a newly-registered model version when calling `log_model`, @WeichenXu123 I wonder if we could return it as part of the [ModelInfo](https://mlflow.org/docs/latest/python_api/mlflow.sklearn.html?highlight=log_model#mlflow.sklearn.log_model) returned by `log_model`? If the interface change seems sensible and won't break things/abstractions elsewhere I'm happy to review a PR to add this

Sounds good to me!

Resolving patch:
diff --git a/mlflow/models/model.py b/mlflow/models/model.py
--- a/mlflow/models/model.py
+++ b/mlflow/models/model.py
@@ -87,6 +87,7 @@ def __init__(
         mlflow_version: str,
         signature_dict: Optional[Dict[str, Any]] = None,
         metadata: Optional[Dict[str, Any]] = None,
+        registered_model_version: Optional[int] = None,
     ):
         self._artifact_path = artifact_path
         self._flavors = flavors
@@ -99,6 +100,7 @@ def __init__(
         self._utc_time_created = utc_time_created
         self._mlflow_version = mlflow_version
         self._metadata = metadata
+        self._registered_model_version = registered_model_version
 
     @property
     def artifact_path(self):
@@ -276,6 +278,21 @@ def metadata(self) -> Optional[Dict[str, Any]]:
         """
         return self._metadata
 
+    @property
+    def registered_model_version(self) -> Optional[int]:
+        """
+        The registered model version, if the model is registered.
+
+        :getter: Gets the registered model version, if the model is registered in Model Registry.
+        :setter: Sets the registered model version.
+        :type: Optional[int]
+        """
+        return self._registered_model_version
+
+    @registered_model_version.setter
+    def registered_model_version(self, value):
+        self._registered_model_version = value
+
 
 class Model:
     """
@@ -656,6 +673,7 @@ def log(
         from mlflow.models.wheeled_model import _ORIGINAL_REQ_FILE_NAME, WheeledModel
         from mlflow.utils.model_utils import _validate_and_get_model_config_from_file
 
+        registered_model = None
         with TempDir() as tmp:
             local_path = tmp.path("model")
             if run_id is None:
@@ -737,14 +755,18 @@ def log(
                 # older tracking servers. Only print out a warning for now.
                 _logger.warning(_LOG_MODEL_METADATA_WARNING_TEMPLATE, mlflow.get_artifact_uri())
                 _logger.debug("", exc_info=True)
+
             if registered_model_name is not None:
-                mlflow.tracking._model_registry.fluent._register_model(
+                registered_model = mlflow.tracking._model_registry.fluent._register_model(
                     f"runs:/{run_id}/{mlflow_model.artifact_path}",
                     registered_model_name,
                     await_registration_for=await_registration_for,
                     local_model_path=local_path,
                 )
-        return mlflow_model.get_model_info()
+        model_info = mlflow_model.get_model_info()
+        if registered_model is not None:
+            model_info.registered_model_version = registered_model.version
+        return model_info
 
 
 def get_model_info(model_uri: str) -> ModelInfo:


Test patch:
diff --git a/tests/models/test_model.py b/tests/models/test_model.py
--- a/tests/models/test_model.py
+++ b/tests/models/test_model.py
@@ -227,6 +227,17 @@ def test_model_info():
         assert model_info_fetched.mlflow_version == loaded_model.mlflow_version
 
 
+def test_model_info_with_model_version(tmp_path):
+    experiment_id = mlflow.create_experiment("test", artifact_location=str(tmp_path))
+    with mlflow.start_run(experiment_id=experiment_id):
+        model_info = Model.log("some/path", TestFlavor, registered_model_name="model_abc")
+        assert model_info.registered_model_version == 1
+        model_info = Model.log("some/path", TestFlavor, registered_model_name="model_abc")
+        assert model_info.registered_model_version == 2
+        model_info = Model.log("some/path", TestFlavor)
+        assert model_info.registered_model_version is None
+
+
 def test_model_metadata():
     with TempDir(chdr=True) as tmp:
         metadata = {"metadata_key": "metadata_value"}