[azure-ai-ml] Fix job download and streaming regressions

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)
 

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 = {}

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