diff --git a/sdk/ml/azure-ai-ml/CHANGELOG.md b/sdk/ml/azure-ai-ml/CHANGELOG.md index d3a3c3401922..ee44884c29f9 100644 --- a/sdk/ml/azure-ai-ml/CHANGELOG.md +++ b/sdk/ml/azure-ai-ml/CHANGELOG.md @@ -5,6 +5,9 @@ ### Features Added ### Bugs Fixed +- Fixed `MLClient.jobs.download(output_name=...)` silently downloading nothing for a job's `uri_file` / `uri_folder` / `mltable` outputs (issue [#48941](https://github.com/Azure/azure-sdk-for-python/issues/48941)). RunHistory reports a job output's type in PascalCase (for example `UriFolder`), but the filter in `get_job_output_uris_from_dataplane` compared it against the snake_case values of the `arm_ml_service` `DataType` enum (`uri_folder`), so every data output was dropped from the resolution and the download completed without an error. Output types are now compared regardless of casing and separators, so both spellings resolve. +- Fixed downloading a job's model output (`CustomModel` / `MLFlowModel` / `TritonModel`) failing with `TypeError: BatchGetResolvedUrisDto.__init__() got an unexpected keyword argument 'values'`, and then with `AttributeError: 'function' object has no attribute 'items'`. The migrated model dataplane request and response models name that field `values_property` (the `values` attribute is the mapping method inherited by the model), so the request is now built and the resolved paths are now read through it. The on-the-wire request body is unchanged (`{"values": [...]}`). +- `MLClient.jobs.download` now logs a warning, rather than a debug message, when an explicitly requested `output_name` could not be resolved, so a download that produces nothing is no longer silent. ## 1.35.0 (2026-09-08) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_job_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_job_operations.py index 7782d3d1c486..3646674d9e95 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_job_operations.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_job_operations.py @@ -1159,8 +1159,9 @@ def download( artifact_directory_name = "artifacts" output_directory_name = "named-outputs" - def log_missing_uri(what: str) -> None: - module_logger.debug( + def log_missing_uri(what: str, *, warn: bool = False) -> None: + log = module_logger.warning if warn else module_logger.debug + log( 'Could not download %s for job "%s" (job status: %s)', what, job_details.name, @@ -1197,7 +1198,8 @@ def log_missing_uri(what: str) -> None: outputs = self._get_named_output_uri(name, output_name) if output_name not in outputs: - log_missing_uri(what=f'output "{output_name}"') + # The output was explicitly requested, so make the fact that nothing was downloaded visible. + log_missing_uri(what=f'output "{output_name}"', warn=True) elif all: outputs = self._get_named_output_uri(name) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_job_ops_helper.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_job_ops_helper.py index ec31ecf6248e..5e46b5228e3f 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_job_ops_helper.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_job_ops_helper.py @@ -450,6 +450,30 @@ def has_pat_token(url: Optional[str]) -> bool: return re.search(pat_regex, url) is not None +def _normalize_asset_type(asset_type: Optional[str]) -> str: + """Normalize an asset type so that it can be compared regardless of casing and separators. + + A job output's asset type is reported by the RunHistory dataplane in PascalCase (for example + "UriFolder" or "MLFlowModel"), while the control plane enums spell the same types in snake_case + (for example "uri_folder" or "mlflow_model"). Normalizing both sides lets the two spellings of the + same type compare equal. + + :param asset_type: The asset type to normalize. + :type asset_type: Optional[str] + :return: The normalized asset type, or an empty string if no asset type was provided. + :rtype: str + """ + return asset_type.lower().replace("_", "") if asset_type else "" + + +# Asset types of job outputs that are resolved through the dataset dataplane, and the ones resolved +# through the model dataplane. Both are normalized so that either wire spelling is recognized. +_DATA_ASSET_TYPES = frozenset(_normalize_asset_type(data_type.value) for data_type in DataType) +_MODEL_ASSET_TYPES = frozenset( + _normalize_asset_type(model_type) for model_type in ("CustomModel", "MLFlowModel", "TritonModel") +) + + def get_job_output_uris_from_dataplane( job_name: Optional[str], run_operations: RunOperations, @@ -495,14 +519,14 @@ def get_job_output_uris_from_dataplane( dataset_ids = [ run_outputs[output_name].asset_id for output_name in output_names - if run_outputs[output_name].type in [o.value for o in DataType] + if _normalize_asset_type(run_outputs[output_name].type) in _DATA_ASSET_TYPES ] # Collect all output ids that correspond to models model_ids = [ run_outputs[output_name].asset_id for output_name in output_names - if run_outputs[output_name].type in ["CustomModel", "MLFlowModel", "TritonModel"] + if _normalize_asset_type(run_outputs[output_name].type) in _MODEL_ASSET_TYPES ] output_name_to_dataset_uri = {} @@ -525,6 +549,6 @@ def get_job_output_uris_from_dataplane( else None ) output_name_to_model_uri = { - asset_id_to_output_name[k]: v.path for k, v in model_uris.values.items() # type: ignore + asset_id_to_output_name[k]: v.path for k, v in model_uris.values_property.items() # type: ignore } return {**output_name_to_dataset_uri, **output_name_to_model_uri} diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_model_dataplane_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_model_dataplane_operations.py index ad5c4100acbb..6efca5cadd27 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_model_dataplane_operations.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_model_dataplane_operations.py @@ -23,7 +23,7 @@ def __init__( self._operation = service_client.models def get_batch_model_uris(self, model_ids: List[str]) -> BatchModelPathResponseDto: - batch_uri_request = BatchGetResolvedUrisDto(values=model_ids) + batch_uri_request = BatchGetResolvedUrisDto(values_property=model_ids) return self._operation.batch_get_resolved_uris( self._operation_scope.subscription_id, self._operation_scope.resource_group_name, diff --git a/sdk/ml/azure-ai-ml/tests/job_common/unittests/test_job_operations.py b/sdk/ml/azure-ai-ml/tests/job_common/unittests/test_job_operations.py index 6b62706f3c13..0691c293fb34 100644 --- a/sdk/ml/azure-ai-ml/tests/job_common/unittests/test_job_operations.py +++ b/sdk/ml/azure-ai-ml/tests/job_common/unittests/test_job_operations.py @@ -1,6 +1,8 @@ import json +import logging import os import platform +from types import SimpleNamespace from unittest.mock import Mock, patch import jwt @@ -17,6 +19,7 @@ from azure.ai.ml.operations import DatastoreOperations, EnvironmentOperations, JobOperations, WorkspaceOperations from azure.ai.ml.operations._code_operations import CodeOperations from azure.ai.ml.operations._job_ops_helper import get_git_properties +from azure.ai.ml.operations._run_history_constants import JobStatus from azure.ai.ml.operations._run_operations import RunOperations from azure.core.credentials import AccessToken from azure.identity import DefaultAzureCredential @@ -546,3 +549,14 @@ def test_download_with_none(self, mock_job_operation: JobOperations) -> None: with pytest.raises(Exception) as ex: mock_job_operation.download(None) assert "None is a invalid input for client.jobs.get()." in ex.value.message + + def test_download_warns_when_named_output_could_not_be_resolved(self, mock_job_operation: JobOperations, caplog): + job_details = SimpleNamespace(name="mock-job", status=JobStatus.COMPLETED, properties={}, tags={}) + + with patch.object(JobOperations, "get", return_value=job_details), patch.object( + JobOperations, "_get_named_output_uri", return_value={} + ): + with caplog.at_level(logging.WARNING): + mock_job_operation.download("mock-job", output_name="my_output") + + assert 'Could not download output "my_output"' in caplog.text diff --git a/sdk/ml/azure-ai-ml/tests/job_common/unittests/test_job_ops_helper.py b/sdk/ml/azure-ai-ml/tests/job_common/unittests/test_job_ops_helper.py index 66181c161b48..1b05af49cbd8 100644 --- a/sdk/ml/azure-ai-ml/tests/job_common/unittests/test_job_ops_helper.py +++ b/sdk/ml/azure-ai-ml/tests/job_common/unittests/test_job_ops_helper.py @@ -8,20 +8,59 @@ import pytest from mock import mock_open, patch -from azure.ai.ml._restclient.runhistory.models import RunDetails, RunDetailsWarning -from azure.ai.ml._scope_dependent_operations import OperationScope +from azure.ai.ml._restclient.dataset_dataplane.models import BatchDataUriResponse, DataUriV2Response +from azure.ai.ml._restclient.model_dataplane.models import BatchModelPathResponseDto, ModelPathResponseDto +from azure.ai.ml._restclient.runhistory.models import ( + GetRunDataResult, + Run, + RunDetails, + RunDetailsWarning, + TypedAssetReference, +) +from azure.ai.ml._scope_dependent_operations import OperationConfig, OperationScope +from azure.ai.ml.operations._dataset_dataplane_operations import DatasetDataplaneOperations from azure.ai.ml.operations._job_ops_helper import ( _get_sorted_filtered_logs, + get_job_output_uris_from_dataplane, has_pat_token, _incremental_print, list_logs, stream_logs_until_completion, ) +from azure.ai.ml.operations._model_dataplane_operations import ModelDataplaneOperations from azure.ai.ml.operations._run_operations import RunOperations from .test_vcr_utils import before_record_cb +DATA_OUTPUT_ASSET_ID = "azureml://locations/eastus/workspaces/00000/data/azureml_dummy_output_data/versions/1" +MODEL_OUTPUT_ASSET_ID = "azureml://locations/eastus/workspaces/00000/models/azureml_dummy_output_model/versions/1" +DATA_OUTPUT_URI = "azureml://datastores/workspaceblobstore/paths/azureml/dummy/forecast_data/" +MODEL_OUTPUT_URI = "azureml://datastores/workspaceblobstore/paths/azureml/dummy/trained_model/" + + +def build_run_operations(outputs: Dict[str, TypedAssetReference]) -> Mock: + run_operations = Mock() + run_operations.get_run_data.return_value = GetRunDataResult(run_metadata=Run(outputs=outputs)) + return run_operations + + +def build_dataset_dataplane_operations(uris: Dict[str, str]) -> Mock: + dataset_dataplane_operations = Mock() + dataset_dataplane_operations.get_batch_dataset_uris.return_value = BatchDataUriResponse( + values_property={asset_id: DataUriV2Response(uri=uri) for asset_id, uri in uris.items()} + ) + return dataset_dataplane_operations + + +def build_model_dataplane_operations(paths: Dict[str, str]) -> Mock: + model_dataplane_operations = Mock() + model_dataplane_operations.get_batch_model_uris.return_value = BatchModelPathResponseDto( + values_property={asset_id: ModelPathResponseDto(path=path) for asset_id, path in paths.items()} + ) + return model_dataplane_operations + + class DummyJob: class InteractionEndpoint: def __init__(self, **kwargs): @@ -65,6 +104,20 @@ def mock_run_operations(mock_workspace_scope: OperationScope, mock_aml_services_ yield RunOperations(mock_workspace_scope, mock_aml_services_run_history) +@pytest.fixture +def mock_dataset_dataplane_operations( + mock_workspace_scope: OperationScope, mock_operation_config: OperationConfig +) -> DatasetDataplaneOperations: + yield DatasetDataplaneOperations(mock_workspace_scope, mock_operation_config, Mock()) + + +@pytest.fixture +def mock_model_dataplane_operations( + mock_workspace_scope: OperationScope, mock_operation_config: OperationConfig +) -> ModelDataplaneOperations: + yield ModelDataplaneOperations(mock_workspace_scope, mock_operation_config, Mock()) + + @pytest.mark.unittest @pytest.mark.training_experiences_test class TestJobOpsHelper: @@ -74,6 +127,105 @@ def test_has_pat_token(self) -> None: assert not has_pat_token("https://dev.azure.com/organization/project/_apis/pipelines/1/runs") assert not has_pat_token("https://learn.microsoft.com/en-us/ai/?tabs=developer") + @pytest.mark.parametrize( + "data_type,model_type", + [ + # RunHistory reports the type of a job's outputs in PascalCase + ("UriFolder", "CustomModel"), + ("UriFile", "MLFlowModel"), + ("MLTable", "TritonModel"), + # while the control plane spells the same types in snake_case + ("uri_folder", "custom_model"), + ("uri_file", "mlflow_model"), + ("mltable", "triton_model"), + ], + ) + def test_get_job_output_uris_from_dataplane(self, data_type: str, model_type: str) -> None: + run_operations = build_run_operations( + { + "forecast_data": TypedAssetReference(asset_id=DATA_OUTPUT_ASSET_ID, type=data_type), + "trained_model": TypedAssetReference(asset_id=MODEL_OUTPUT_ASSET_ID, type=model_type), + } + ) + + outputs = get_job_output_uris_from_dataplane( + "dummy", + run_operations, + build_dataset_dataplane_operations({DATA_OUTPUT_ASSET_ID: DATA_OUTPUT_URI}), + build_model_dataplane_operations({MODEL_OUTPUT_ASSET_ID: MODEL_OUTPUT_URI}), + ) + + assert outputs == {"forecast_data": DATA_OUTPUT_URI, "trained_model": MODEL_OUTPUT_URI} + + def test_get_job_output_uris_from_dataplane_with_output_name(self) -> None: + run_operations = build_run_operations( + {"forecast_data": TypedAssetReference(asset_id=DATA_OUTPUT_ASSET_ID, type="UriFolder")} + ) + dataset_dataplane_operations = build_dataset_dataplane_operations({DATA_OUTPUT_ASSET_ID: DATA_OUTPUT_URI}) + + outputs = get_job_output_uris_from_dataplane( + "dummy", + run_operations, + dataset_dataplane_operations, + None, + output_names="forecast_data", + ) + + assert outputs == {"forecast_data": DATA_OUTPUT_URI} + dataset_dataplane_operations.get_batch_dataset_uris.assert_called_once_with([DATA_OUTPUT_ASSET_ID]) + + def test_get_job_output_uris_from_dataplane_skips_unknown_types(self) -> None: + run_operations = build_run_operations( + { + "unknown_type": TypedAssetReference(asset_id=DATA_OUTPUT_ASSET_ID, type="SomethingElse"), + "no_type": TypedAssetReference(asset_id=MODEL_OUTPUT_ASSET_ID), + } + ) + dataset_dataplane_operations = build_dataset_dataplane_operations({}) + model_dataplane_operations = build_model_dataplane_operations({}) + + outputs = get_job_output_uris_from_dataplane( + "dummy", + run_operations, + dataset_dataplane_operations, + model_dataplane_operations, + ) + + assert outputs == {} + dataset_dataplane_operations.get_batch_dataset_uris.assert_not_called() + model_dataplane_operations.get_batch_model_uris.assert_not_called() + + def test_get_job_output_uris_from_dataplane_requests_asset_ids( + self, + mock_dataset_dataplane_operations: DatasetDataplaneOperations, + mock_model_dataplane_operations: ModelDataplaneOperations, + ) -> None: + mock_dataset_dataplane_operations._operation.batch_get_resolved_uris.return_value = BatchDataUriResponse( + values_property={DATA_OUTPUT_ASSET_ID: DataUriV2Response(uri=DATA_OUTPUT_URI)} + ) + mock_model_dataplane_operations._operation.batch_get_resolved_uris.return_value = BatchModelPathResponseDto( + values_property={MODEL_OUTPUT_ASSET_ID: ModelPathResponseDto(path=MODEL_OUTPUT_URI)} + ) + run_operations = build_run_operations( + { + "forecast_data": TypedAssetReference(asset_id=DATA_OUTPUT_ASSET_ID, type="UriFolder"), + "trained_model": TypedAssetReference(asset_id=MODEL_OUTPUT_ASSET_ID, type="CustomModel"), + } + ) + + outputs = get_job_output_uris_from_dataplane( + "dummy", + run_operations, + mock_dataset_dataplane_operations, + mock_model_dataplane_operations, + ) + + assert outputs == {"forecast_data": DATA_OUTPUT_URI, "trained_model": MODEL_OUTPUT_URI} + dataset_request = mock_dataset_dataplane_operations._operation.batch_get_resolved_uris.call_args[1]["body"] + assert dataset_request.values_property == [DATA_OUTPUT_ASSET_ID] + model_request = mock_model_dataplane_operations._operation.batch_get_resolved_uris.call_args[1]["body"] + assert model_request.values_property == [MODEL_OUTPUT_ASSET_ID] + @pytest.mark.skip("TODO 1907352: Relies on a missing VCR.py recording + test suite needs to be reworked") @pytest.mark.unittest