|
12 | 12 | ModelVersionProperties as ModelVersionDetails, |
13 | 13 | ) |
14 | 14 | from azure.ai.ml._scope_dependent_operations import OperationConfig, OperationScope |
| 15 | +from azure.ai.ml._restclient.v2021_10_01_dataplanepreview.models import ( |
| 16 | + ModelVersionData as RegistryModelVersionData, |
| 17 | + ModelVersionDetails as RegistryModelVersionDetails, |
| 18 | +) |
15 | 19 | from azure.ai.ml.entities._assets import Model |
16 | 20 | from azure.ai.ml.entities._assets._artifacts.artifact import ArtifactStorageInfo |
17 | 21 | from azure.ai.ml.exceptions import ErrorTarget, ValidationException |
@@ -64,6 +68,49 @@ def mock_model_operation_reg( |
64 | 68 | ) |
65 | 69 |
|
66 | 70 |
|
| 71 | +def _make_registry_model_rest_object( |
| 72 | + version: str, |
| 73 | + default_dt_asset_id: Optional[str] = None, |
| 74 | + allowed_dt_asset_ids: Optional[list] = None, |
| 75 | +) -> Mock: |
| 76 | + """Build a registry ModelVersionData mock mirroring the real GET/LIST response shapes. |
| 77 | +
|
| 78 | + The GET response carries default_deployment_template / allowed_deployment_templates; the |
| 79 | + LIST (top=1) response used for label resolution omits them, which is what drops the |
| 80 | + references on the label path (bug 5423568). |
| 81 | + """ |
| 82 | + rest_properties = Mock(spec=RegistryModelVersionDetails) |
| 83 | + rest_properties.description = "Test model" |
| 84 | + rest_properties.tags = {} |
| 85 | + rest_properties.properties = {} |
| 86 | + rest_properties.flavors = {} |
| 87 | + rest_properties.model_uri = "azureml://locations/test/artifacts/model" |
| 88 | + rest_properties.model_type = "custom_model" |
| 89 | + rest_properties.stage = None |
| 90 | + rest_properties.job_name = None |
| 91 | + rest_properties.intellectual_property = None |
| 92 | + rest_properties.system_metadata = None |
| 93 | + rest_properties.default_deployment_template = {"asset_id": default_dt_asset_id} if default_dt_asset_id else None |
| 94 | + rest_properties.allowed_deployment_templates = ( |
| 95 | + [{"asset_id": asset_id} for asset_id in allowed_dt_asset_ids] if allowed_dt_asset_ids else None |
| 96 | + ) |
| 97 | + |
| 98 | + rest_object = Mock(spec=RegistryModelVersionData) |
| 99 | + rest_object.id = ( |
| 100 | + "/subscriptions/sub/resourceGroups/rg/providers/Microsoft.MachineLearningServices" |
| 101 | + f"/workspaces/ws/models/test-model/versions/{version}" |
| 102 | + ) |
| 103 | + rest_object.properties = rest_properties |
| 104 | + |
| 105 | + system_data = Mock() |
| 106 | + system_data.created_by = "test_user" |
| 107 | + system_data.created_at = None |
| 108 | + system_data.last_modified_by = None |
| 109 | + system_data.last_modified_at = None |
| 110 | + rest_object.system_data = system_data |
| 111 | + return rest_object |
| 112 | + |
| 113 | + |
67 | 114 | @pytest.mark.unittest |
68 | 115 | @pytest.mark.production_experiences_test |
69 | 116 | class TestModelOperations: |
@@ -172,6 +219,62 @@ def test_get_no_version(self, mock_model_operation: ModelOperations) -> None: |
172 | 219 | with pytest.raises(Exception): |
173 | 220 | mock_model_operation.get(name=name) |
174 | 221 |
|
| 222 | + def test_get_with_label_rehydrates_deployment_template_references( |
| 223 | + self, mock_model_operation_reg: ModelOperations |
| 224 | + ) -> None: |
| 225 | + """Bug 5423568: models.get(name, label='latest') must hydrate deployment-template |
| 226 | + references identically to the explicit version= path. |
| 227 | +
|
| 228 | + Label resolution goes through the version LIST endpoint (top=1), whose items omit |
| 229 | + default_deployment_template / allowed_deployment_templates. The fix re-fetches the |
| 230 | + resolved version through the GET endpoint, which carries those references. |
| 231 | + """ |
| 232 | + default_dt = "azureml://registries/azure-huggingface/deploymenttemplates/dt/versions/5" |
| 233 | + allowed_dt = "azureml://registries/azure-huggingface/deploymenttemplates/dt/labels/latest" |
| 234 | + |
| 235 | + # LIST shape (used by label resolution) -> deployment-template references dropped. |
| 236 | + list_pager = Mock() |
| 237 | + list_pager.next.return_value = _make_registry_model_rest_object(version="5") |
| 238 | + mock_model_operation_reg._model_versions_operation.list.return_value = list_pager |
| 239 | + |
| 240 | + # GET shape (used by the re-fetch) -> deployment-template references populated. |
| 241 | + mock_model_operation_reg._model_versions_operation.get.return_value = _make_registry_model_rest_object( |
| 242 | + version="5", default_dt_asset_id=default_dt, allowed_dt_asset_ids=[allowed_dt] |
| 243 | + ) |
| 244 | + |
| 245 | + result = mock_model_operation_reg.get(name="test-model", label="latest") |
| 246 | + |
| 247 | + # The resolved version was re-fetched through GET (not returned straight from the LIST shape). |
| 248 | + mock_model_operation_reg._model_versions_operation.get.assert_called_once_with( |
| 249 | + name="test-model", |
| 250 | + version="5", |
| 251 | + registry_name=mock_model_operation_reg._registry_name, |
| 252 | + **mock_model_operation_reg._scope_kwargs, |
| 253 | + ) |
| 254 | + # Deployment-template references are hydrated, matching the explicit version= path. |
| 255 | + assert result.version == "5" |
| 256 | + assert result.default_deployment_template is not None |
| 257 | + assert result.default_deployment_template.asset_id == default_dt |
| 258 | + assert result.allowed_deployment_templates is not None |
| 259 | + assert len(result.allowed_deployment_templates) == 1 |
| 260 | + assert result.allowed_deployment_templates[0].asset_id == allowed_dt |
| 261 | + |
| 262 | + def test_get_with_label_workspace_does_not_refetch(self, mock_model_operation: ModelOperations) -> None: |
| 263 | + """Bug 5423568: the deployment-template re-fetch is scoped to the registry. |
| 264 | +
|
| 265 | + Workspace models carry no deployment-template references, so the label path must not issue |
| 266 | + the extra GET; this keeps workspace behaviour (and its recorded e2e sessions) unchanged. |
| 267 | + """ |
| 268 | + resolved = Model(name="test-model", version="4") |
| 269 | + with patch( |
| 270 | + "azure.ai.ml.operations._model_operations._resolve_label_to_asset", |
| 271 | + return_value=resolved, |
| 272 | + ): |
| 273 | + result = mock_model_operation.get(name="test-model", label="latest") |
| 274 | + |
| 275 | + assert result is resolved |
| 276 | + assert mock_model_operation._model_versions_operation.get.call_count == 0 |
| 277 | + |
175 | 278 | @patch.object(Model, "_from_rest_object", new=Mock()) |
176 | 279 | @patch.object(Model, "_from_container_rest_object", new=Mock()) |
177 | 280 | def test_list(self, mock_model_operation: ModelOperations) -> None: |
|
0 commit comments