diff --git a/sdk/ml/azure-ai-ml/CHANGELOG.md b/sdk/ml/azure-ai-ml/CHANGELOG.md index d3a3c3401922..20da452d53a1 100644 --- a/sdk/ml/azure-ai-ml/CHANGELOG.md +++ b/sdk/ml/azure-ai-ml/CHANGELOG.md @@ -5,6 +5,8 @@ ### Features Added ### Bugs Fixed +- Fixed `MLClient.jobs.download(name=..., output_name=...)` silently downloading nothing for named `uri_file` / `uri_folder` / `mltable` outputs (issue [#48941](https://github.com/Azure/azure-sdk-for-python/issues/48941)). `get_job_output_uris_from_dataplane` filtered the RunHistory `run_metadata.outputs` types against the `DataType` enum, whose wire values changed from PascalCase (`"UriFolder"`) to snake_case (`"uri_folder"`) when the REST clients were migrated to the shared `arm_ml_service` client in 1.35.0, while RunHistory still reports PascalCase. Asset types are now compared case- and separator-insensitively, so both spellings resolve. +- Fixed `MLClient.jobs.stream()` raising `TypeError: argument should be a bytes-like object or ASCII string, not 'ChainedTokenCredential'` for workspaces whose default datastore uses identity-based access. The same `DataType` wire-value change in 1.35.0 made the `job_output_type == DataType.URI_FOLDER` check start matching (it had silently evaluated to `False` since the job clients moved to snake_case), which activated the read-logs-directly-from-the-datastore path. That path signs a short-lived SAS and only works with an account key or SAS token; for identity-based datastores the resolved credential is a `TokenCredential`, which cannot sign. Log streaming now falls back to the RunHistory log files when the datastore credential is not signable. ## 1.35.0 (2026-09-08) 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..91e9742b068b 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 @@ -34,6 +34,28 @@ module_logger = logging.getLogger(__name__) +def _normalize_asset_type(asset_type: Optional[str]) -> str: + """Normalizes an asset type so it can be compared across REST contracts. + + The RunHistory dataplane reports asset types in PascalCase (e.g. ``"UriFolder"``, ``"MLFlowModel"``) while the + ARM/Machine Learning Services contract uses snake_case (e.g. ``"uri_folder"``, ``"mlflow_model"``). Stripping + underscores and lower-casing makes both spellings comparable. + + :param asset_type: The asset type reported by a service. + :type asset_type: Optional[str] + :return: The normalized asset type. + :rtype: str + """ + return (asset_type or "").replace("_", "").lower() + + +_DATA_ASSET_TYPES = { + _normalize_asset_type(data_type) + for data_type in (DataType.URI_FILE, DataType.URI_FOLDER, DataType.MLTABLE) +} +_MODEL_ASSET_TYPES = {_normalize_asset_type(t) for t in ("CustomModel", "MLFlowModel", "TritonModel")} + + def _get_sorted_filtered_logs( logs_iterable: Iterable[str], job_type: str, @@ -252,6 +274,17 @@ def stream_logs_until_completion( output_uri = output_uri.split("datastores/")[1] datastore_name, prefix = output_uri.split("/", 1) ds_properties = get_datastore_info(datastore_operations, datastore_name) + # Reading logs straight from the datastore requires signing a short-lived SAS, which is only + # possible with an account key or an existing SAS token (both plain strings). Identity-based + # datastores resolve to a TokenCredential instead, which cannot sign a SAS. In that case fall + # back to the RunHistory log files, which are already SAS-scoped by the service. + if not isinstance(ds_properties.get("credential"), str): + module_logger.debug( + "Datastore '%s' has no signable key or SAS token; streaming logs from RunHistory instead.", + datastore_name, + ) + ds_properties = None + prefix = None try: file_handle.write("RunId: {}\n".format(job_name)) @@ -495,14 +528,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 = {} 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..e6e19e0785e6 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 @@ -12,6 +12,7 @@ from azure.ai.ml._scope_dependent_operations import OperationScope 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, @@ -74,6 +75,97 @@ 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 PascalCase, the ARM contract reports snake_case. Both must resolve. + ("UriFolder", "MLFlowModel"), + ("uri_folder", "mlflow_model"), + ], + ) + def test_get_job_output_uris_from_dataplane_matches_both_type_spellings(self, data_type, model_type) -> None: + run_outputs = { + "forecast_data": Mock(asset_id="data-asset-id", type=data_type), + "trained_model": Mock(asset_id="model-asset-id", type=model_type), + } + run_operations = Mock() + run_operations.get_run_data.return_value.run_metadata.outputs = run_outputs + + dataset_dataplane_operations = Mock() + dataset_dataplane_operations.get_batch_dataset_uris.return_value.values_property = { + "data-asset-id": Mock(uri="azureml://datastores/ds/paths/forecast_data") + } + + model_dataplane_operations = Mock() + model_dataplane_operations.get_batch_model_uris.return_value.values = { + "model-asset-id": Mock(path="azureml://datastores/ds/paths/trained_model") + } + + uris = get_job_output_uris_from_dataplane( + "job-name", + run_operations, + dataset_dataplane_operations, + model_dataplane_operations, + ) + + dataset_dataplane_operations.get_batch_dataset_uris.assert_called_once_with(["data-asset-id"]) + model_dataplane_operations.get_batch_model_uris.assert_called_once_with(["model-asset-id"]) + assert uris == { + "forecast_data": "azureml://datastores/ds/paths/forecast_data", + "trained_model": "azureml://datastores/ds/paths/trained_model", + } + + @pytest.mark.parametrize( + "datastore_credential,expects_datastore_logs", + [ + # Account key / SAS token: a signable string, so the datastore fast path is used. + ("fake-account-key", True), + # Identity-based datastore: a TokenCredential cannot sign a SAS, so fall back to RunHistory. + (Mock(name="ChainedTokenCredential"), False), + (None, False), + ], + ) + def test_stream_logs_falls_back_to_run_history_for_unsignable_datastore( + self, datastore_credential, expects_datastore_logs + ) -> None: + job_resource = Mock() + job_resource.name = "job-name" + job_resource.properties.job_type = "Command" + job_resource.properties.properties = {} + job_resource.properties.services = {} + job_resource.properties.outputs = { + "default": Mock( + job_output_type="uri_folder", + uri="azureml://.../datastores/workspaceblobstore/paths/azureml/job-name/", + ) + } + + run_operations = Mock() + run_operations.get_run_details.side_effect = [ + RunDetails(status="Running", log_files={}), + RunDetails(status="Completed", log_files={}), + ] + + ds_info = {"credential": datastore_credential, "storage_type": "AzureBlob"} + + with patch("azure.ai.ml.operations._job_ops_helper.get_datastore_info", return_value=ds_info), patch( + "azure.ai.ml.operations._job_ops_helper.list_logs_in_datastore", return_value={} + ) as mock_list_logs_in_datastore, patch( + "azure.ai.ml.operations._job_ops_helper.create_requests_pipeline_with_retry" + ), patch( + "azure.ai.ml.operations._job_ops_helper.time.sleep" + ): + stream_logs_until_completion( + run_operations, + job_resource, + datastore_operations=Mock(), + requests_pipeline=Mock(), + ) + + # Regression guard for the 1.35.0 `TypeError: ... not 'ChainedTokenCredential'`: a non-signable + # credential must never reach the SAS-generating datastore log reader. + assert mock_list_logs_in_datastore.called is expects_datastore_logs + @pytest.mark.skip("TODO 1907352: Relies on a missing VCR.py recording + test suite needs to be reworked") @pytest.mark.unittest