Let make_serving_endpoint reference a valid model version (#106)

## Changes
The `make_serving_endpoint` fixture [stopped working as of
yesterday](https://github.com/databrickslabs/pytester/actions/runs/13405748534).
Apparently, the created models from the `make_model` fixture do not come
with a model version anymore (which used to be `'1'`).

The `make_serving_endpoint` is updated to fallback on a UC model
version. Also, it allows users to provide input parameters in case they
want to use another model.

### Linked issues

Resolves #databrickslabs/ucx#3714
Resolves #databrickslabs/ucx#3715

### Tests

- [x] added unit tests
- [ ] fixed integration test

Author: JCZuurmond

README.md

@@ -1065,7 +1065,7 @@ def test_models(make_group, make_model, make_registered_model_permissions):
     )
 ```
 
-See also [`make_serving_endpoint`](#make_serving_endpoint-fixture), [`ws`](#ws-fixture), [`make_random`](#make_random-fixture), [`watchdog_remove_after`](#watchdog_remove_after-fixture).
+See also [`ws`](#ws-fixture), [`make_random`](#make_random-fixture), [`watchdog_remove_after`](#watchdog_remove_after-fixture).
 
 
 [[back to top](#python-testing-for-databricks)]
@@ -1204,6 +1204,13 @@ The function returns a [`ServingEndpointDetailed`](https://databricks-sdk-py.rea
 
 Under the covers, this fixture also creates a model to serve on a small workload size.
 
+Keyword arguments:
+* `endpoint_name` (str, optional): The name of the endpoint. Defaults to `dummy-*`.
+* `model_name` (str, optional): The name of the model to serve on the endpoint.
+    Defaults to system model `system.ai.llama_v3_2_1b_instruct`.
+* `model_version` (str, optional): The model version to serve. If None, tries to get the latest version for
+    workspace local models. Otherwise, defaults to version `1`.
+
 Usage:
 ```python
 def test_endpoints(make_group, make_serving_endpoint, make_serving_endpoint_permissions):
@@ -1216,7 +1223,7 @@ def test_endpoints(make_group, make_serving_endpoint, make_serving_endpoint_perm
     )
 ```
 
-See also [`ws`](#ws-fixture), [`make_random`](#make_random-fixture), [`make_model`](#make_model-fixture), [`watchdog_remove_after`](#watchdog_remove_after-fixture).
+See also [`ws`](#ws-fixture), [`make_random`](#make_random-fixture), [`watchdog_remove_after`](#watchdog_remove_after-fixture).
 
 
 [[back to top](#python-testing-for-databricks)]

src/databricks/labs/pytester/fixtures/ml.py

@@ -1,19 +1,27 @@
+import logging
 from collections.abc import Callable, Generator
+from unittest.mock import Mock
 
 from pytest import fixture
+from databricks.sdk.errors import BadRequest
 from databricks.sdk.service._internal import Wait
 from databricks.sdk.service.serving import (
-    ServingEndpointDetailed,
     EndpointCoreConfigInput,
+    EndpointPendingConfig,
+    EndpointTag,
     ServedModelInput,
     ServedModelInputWorkloadSize,
-    EndpointTag,
+    ServedModelOutput,
+    ServingEndpointDetailed,
 )
 from databricks.sdk.service.ml import CreateExperimentResponse, ModelDatabricks, ModelTag
 
 from databricks.labs.pytester.fixtures.baseline import factory
 
 
+logger = logging.getLogger(__name__)
+
+
 @fixture
 def make_experiment(
     ws,
@@ -103,13 +111,20 @@ def create(*, model_name: str | None = None, **kwargs) -> ModelDatabricks:
 
 
 @fixture
-def make_serving_endpoint(ws, make_random, make_model, watchdog_remove_after):
+def make_serving_endpoint(ws, make_random, watchdog_remove_after):
     """
     Returns a function to create Databricks Serving Endpoints and clean them up after the test.
     The function returns a `databricks.sdk.service.serving.ServingEndpointDetailed` object.
 
     Under the covers, this fixture also creates a model to serve on a small workload size.
 
+    Keyword arguments:
+    * `endpoint_name` (str, optional): The name of the endpoint. Defaults to `dummy-*`.
+    * `model_name` (str, optional): The name of the model to serve on the endpoint.
+        Defaults to system model `system.ai.llama_v3_2_1b_instruct`.
+    * `model_version` (str, optional): The model version to serve. If None, tries to get the latest version for
+        workspace local models. Otherwise, defaults to version `1`.
+
     Usage:
     ```python
     def test_endpoints(make_group, make_serving_endpoint, make_serving_endpoint_permissions):
@@ -123,27 +138,51 @@ def test_endpoints(make_group, make_serving_endpoint, make_serving_endpoint_perm
     ```
     """
 
-    def create() -> Wait[ServingEndpointDetailed]:
-        endpoint_name = make_random(4)
-        model = make_model()
+    def create(
+        *,
+        endpoint_name: str | None = None,
+        model_name: str | None = None,
+        model_version: str | None = None,
+    ) -> Wait[ServingEndpointDetailed]:
+        endpoint_name = endpoint_name or make_random(4)
+        model_name = model_name or "system.ai.llama_v3_2_1b_instruct"
+        if not model_version and "." not in model_name:  # The period in the name signals it is NOT workspace local
+            try:
+                model_version = ws.model_registry.get_latest_versions(model_name).version
+            except BadRequest as e:
+                logger.warning(
+                    f"Cannot get latest version for model: {model_name}. Fallback to version '1'.", exc_info=e
+                )
+        model_version = model_version or "1"
+        tags = [EndpointTag(key="RemoveAfter", value=watchdog_remove_after)]
+        served_model_input = ServedModelInput(
+            model_name=model_name,
+            model_version=model_version,
+            scale_to_zero_enabled=True,
+            workload_size=ServedModelInputWorkloadSize.SMALL,
+        )
         endpoint = ws.serving_endpoints.create(
             endpoint_name,
-            config=EndpointCoreConfigInput(
-                served_models=[
-                    ServedModelInput(
-                        model_name=model.name,
-                        model_version="1",
-                        scale_to_zero_enabled=True,
-                        workload_size=ServedModelInputWorkloadSize.SMALL,
-                    )
-                ]
-            ),
-            tags=[EndpointTag(key="RemoveAfter", value=watchdog_remove_after)],
+            EndpointCoreConfigInput(served_models=[served_model_input]),
+            tags=tags,
         )
+        if isinstance(endpoint, Mock):  # For testing
+            served_model_output = ServedModelOutput(
+                model_name=model_name,
+                model_version=model_version,
+                scale_to_zero_enabled=True,
+                workload_size=ServedModelInputWorkloadSize.SMALL.value,
+            )
+            endpoint = ServingEndpointDetailed(
+                name=endpoint_name,
+                pending_config=EndpointPendingConfig(served_models=[served_model_output]),
+                tags=tags,
+            )
         return endpoint
 
-    def remove(endpoint_name: str):
-        ws.serving_endpoints.delete(endpoint_name)
+    def remove(endpoint: ServingEndpointDetailed) -> None:
+        if endpoint.name:
+            ws.serving_endpoints.delete(endpoint.name)
 
     yield from factory("Serving endpoint", create, remove)
 

tests/unit/fixtures/test_ml.py

@@ -1,5 +1,10 @@
+import pytest
+
+from databricks.sdk.errors import InvalidParameterValue
+from databricks.sdk.service.ml import ModelVersion
+
 from databricks.labs.pytester.fixtures.ml import make_experiment, make_model, make_serving_endpoint
-from databricks.labs.pytester.fixtures.unwrap import call_stateful
+from databricks.labs.pytester.fixtures.unwrap import CallContext, call_stateful
 
 
 def test_make_experiment_no_args():
@@ -14,7 +19,66 @@ def test_make_model_no_args():
     assert model is not None
 
 
-def test_make_serving_endpoint_no_args():
+def test_make_serving_endpoint_no_args() -> None:
     ctx, serving_endpoint = call_stateful(make_serving_endpoint)
     assert ctx is not None
     assert serving_endpoint is not None
+
+
+def test_make_serving_endpoint_sets_default_endpoint_name() -> None:
+    """Default endpoint name should be random."""
+    _, serving_endpoint = call_stateful(make_serving_endpoint)
+    assert serving_endpoint.name == "RANDOM"  # Mocked value in unit tests
+
+
+def test_make_serving_endpoint_sets_endpoint_name() -> None:
+    _, serving_endpoint = call_stateful(make_serving_endpoint, endpoint_name="test")
+    assert serving_endpoint.name == "test"
+
+
+def test_make_serving_endpoint_sets_default_model_name() -> None:
+    """The default model name should be 'system.ai.llama_v3_2_1b_instruct'."""
+    _, serving_endpoint = call_stateful(make_serving_endpoint)
+    assert serving_endpoint.pending_config.served_models[0].model_name == "system.ai.llama_v3_2_1b_instruct"
+
+
+def test_make_serving_endpoint_sets_model_name() -> None:
+    _, serving_endpoint = call_stateful(make_serving_endpoint, model_name="test")
+    assert serving_endpoint.pending_config.served_models[0].model_name == "test"
+
+
+@pytest.mark.parametrize("model_name", [None, "test"])
+def test_make_serving_endpoint_sets_default_model_version_to_one(model_name: str | None) -> None:
+    """The default model version should be '1' independent
+
+    Independent of the model name, if the latest version cannot be retrieved.
+    """
+
+    def _setup_model_registry_api(call_context: CallContext) -> CallContext:
+        """Set up the model registry api for unit testing."""
+        call_context["ws"].model_registry.get_latest_versions.side_effect = InvalidParameterValue("test")
+        return call_context
+
+    _, serving_endpoint = call_stateful(
+        make_serving_endpoint, model_name=model_name, call_context_setup=_setup_model_registry_api
+    )
+    assert serving_endpoint.pending_config.served_models[0].model_version == "1"
+
+
+@pytest.mark.parametrize("model_name", [None, "test"])
+def test_make_serving_endpoint_sets_model_version(model_name: str | None) -> None:
+    """The model version should be the value passed into the fixture.
+
+    Independent of the model name, also if the latest version can be retrieved.
+    """
+
+    def _setup_model_registry_api(call_context: CallContext) -> CallContext:
+        """Set up the model registry api for unit testing."""
+        model_version = ModelVersion(version="3")  # Latest version is higher than the expected version
+        call_context["ws"].model_registry.get_latest_versions.return_value = model_version
+        return call_context
+
+    _, serving_endpoint = call_stateful(
+        make_serving_endpoint, model_name=model_name, model_version="2", call_context_setup=_setup_model_registry_api
+    )
+    assert serving_endpoint.pending_config.served_models[0].model_version == "2"