diff --git a/sdk/ml/azure-ai-ml/CHANGELOG.md b/sdk/ml/azure-ai-ml/CHANGELOG.md index d3a3c3401922..97c584f2795b 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(name=..., output_name=...)` silently downloading nothing for named data outputs by accepting both RunHistory and ARM asset-type spellings (issue [#48941](https://github.com/Azure/azure-sdk-for-python/issues/48941)). Model output types are also compared case- and separator-insensitively. +- Fixed `MLClient.jobs.stream()` failing for identity-based and SAS-authenticated datastores by using service-provided RunHistory log URLs instead of attempting to sign a new SAS without an account key. +- Fixed datastore-backed log streaming to resolve canonical datastore URI paths, handle trailing slashes, and generate valid Azure Data Lake Storage Gen2 file SAS URLs. ## 1.35.0 (2026-09-08) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_artifacts/_artifact_utilities.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_artifacts/_artifact_utilities.py index 7cd3c243b8e0..74f6ba4524eb 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/_artifacts/_artifact_utilities.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/_artifacts/_artifact_utilities.py @@ -166,13 +166,15 @@ def list_logs_in_datastore( storage_type=ds_info["storage_type"], ) - items = storage_client.list(starts_with=prefix + "/user_logs/") + prefix = prefix.rstrip("/") + log_prefix = prefix + "/" if prefix else "" + items = storage_client.list(starts_with=log_prefix + "user_logs/") # Append legacy log files if present - items.extend(storage_client.list(starts_with=prefix + legacy_log_folder_name)) + items.extend(storage_client.list(starts_with=log_prefix + legacy_log_folder_name.lstrip("/"))) log_dict = {} for item_name in items: - sub_name = item_name.split(prefix + "/")[1] + sub_name = item_name[len(log_prefix) :] if isinstance(storage_client, BlobStorageClient): token = generate_blob_sas( account_name=ds_info["storage_account"], @@ -183,10 +185,12 @@ def list_logs_in_datastore( expiry=datetime.utcnow() + timedelta(minutes=30), ) elif isinstance(storage_client, Gen2StorageClient): - token = generate_file_sas( # pylint: disable=no-value-for-parameter + directory_name, _, file_name = item_name.rpartition("/") + token = generate_file_sas( account_name=ds_info["storage_account"], file_system_name=ds_info["container_name"], - file_name=item_name, + directory_name=directory_name, + file_name=file_name, credential=ds_info["credential"], permission=FileSasPermissions(read=True), expiry=datetime.utcnow() + timedelta(minutes=30), 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..073a7548bb60 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 @@ -12,13 +12,14 @@ import sys import time from typing import Any, Dict, Iterable, List, Optional, TextIO, Union +from urllib.parse import parse_qs from azure.ai.ml._artifacts._artifact_utilities import get_datastore_info, list_logs_in_datastore -from azure.ai.ml._restclient.runhistory.models import Run, RunDetails, TypedAssetReference -from azure.ai.ml._restclient.arm_ml_service.models import DataType +from azure.ai.ml._restclient.arm_ml_service.models import DataType, JobBase from azure.ai.ml._restclient.arm_ml_service.models import JobType as RestJobType -from azure.ai.ml._restclient.arm_ml_service.models import JobBase +from azure.ai.ml._restclient.runhistory.models import Run, RunDetails, TypedAssetReference from azure.ai.ml._utils._http_utils import HttpPipeline +from azure.ai.ml._utils._storage_utils import AzureMLDatastorePathUri from azure.ai.ml._utils.utils import create_requests_pipeline_with_retry, download_text_from_url from azure.ai.ml.constants._common import GitProperties from azure.ai.ml.constants._job.job import JobLogPattern, JobType @@ -34,6 +35,27 @@ 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, @@ -247,11 +269,23 @@ def stream_logs_until_completion( ) is_uri_folder = default_output and default_output.job_output_type == DataType.URI_FOLDER if is_uri_folder: - output_uri = default_output.uri # type: ignore - # Parse the uri format - output_uri = output_uri.split("datastores/")[1] - datastore_name, prefix = output_uri.split("/", 1) + output_uri = AzureMLDatastorePathUri(default_output.uri) # type: ignore + datastore_name = output_uri.datastore + prefix = output_uri.path ds_properties = get_datastore_info(datastore_operations, datastore_name) + credential = ds_properties.get("credential") + # A SAS token authorizes access but cannot sign another SAS. + if ( + not isinstance(credential, str) + or not credential + or "sig" in parse_qs(credential.lstrip("?"), keep_blank_values=True) + ): + module_logger.debug( + "Datastore '%s' has no account key; streaming logs from RunHistory instead.", + datastore_name, + ) + ds_properties = None + prefix = None try: file_handle.write("RunId: {}\n".format(job_name)) @@ -495,14 +529,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..3fada2f868af 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 @@ -1,23 +1,34 @@ import re import time from collections import OrderedDict +from datetime import datetime, timedelta, timezone from io import StringIO from typing import Dict from unittest.mock import Mock +from urllib.parse import parse_qs, urlsplit import pytest +from devtools_testutils.fake_credentials import FakeTokenCredential from mock import mock_open, patch +from azure.ai.ml._artifacts._artifact_utilities import list_logs_in_datastore +from azure.ai.ml._artifacts._blob_storage_helper import BlobStorageClient +from azure.ai.ml._artifacts._gen2_storage_helper import Gen2StorageClient +from azure.ai.ml._restclient.arm_ml_service.models import DatastoreType, JobBase from azure.ai.ml._restclient.runhistory.models import RunDetails, RunDetailsWarning from azure.ai.ml._scope_dependent_operations import OperationScope +from azure.ai.ml.exceptions import ValidationException from azure.ai.ml.operations._job_ops_helper import ( _get_sorted_filtered_logs, - has_pat_token, _incremental_print, + get_job_output_uris_from_dataplane, + has_pat_token, list_logs, stream_logs_until_completion, ) from azure.ai.ml.operations._run_operations import RunOperations +from azure.identity import ChainedTokenCredential +from azure.storage.blob import ContainerSasPermissions, generate_container_sas from .test_vcr_utils import before_record_cb @@ -65,6 +76,37 @@ def mock_run_operations(mock_workspace_scope: OperationScope, mock_aml_services_ yield RunOperations(mock_workspace_scope, mock_aml_services_run_history) +@pytest.fixture +def streaming_job() -> JobBase: + return JobBase( + { + "name": "job-name", + "properties": { + "jobType": "Command", + "properties": {}, + "services": {}, + "outputs": { + "default": { + "jobOutputType": "uri_folder", + "uri": "azureml://datastores/workspaceblobstore/paths/azureml/job-name/", + } + }, + }, + } + ) + + +@pytest.fixture +def datastore_sas_token(fake_datastore_key: str) -> str: + return generate_container_sas( + account_name="teststorage", + container_name="testcontainer", + account_key=fake_datastore_key, + permission=ContainerSasPermissions(read=True, list=True), + expiry=datetime.now(timezone.utc) + timedelta(hours=1), + ) + + @pytest.mark.unittest @pytest.mark.training_experiences_test class TestJobOpsHelper: @@ -74,6 +116,220 @@ 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"), + ("UriFile", "CustomModel"), + ("uri_file", "custom_model"), + ("MLTable", "TritonModel"), + ("mltable", "triton_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( + "credential_kind,expects_datastore_logs", + [ + ("account_key", True), + ("sas_token", False), + ("sas_with_question_mark", False), + ("sas_with_empty_signature", False), + ("identity", False), + ("missing", False), + ("empty", False), + ], + ) + def test_stream_logs_falls_back_to_run_history_for_unsignable_datastore( + self, + credential_kind, + expects_datastore_logs, + streaming_job: JobBase, + fake_datastore_key: str, + datastore_sas_token: str, + capsys, + caplog, + ) -> None: + credentials = { + "account_key": fake_datastore_key, + "sas_token": datastore_sas_token, + "sas_with_question_mark": "?" + datastore_sas_token, + "sas_with_empty_signature": "sp=rl&sig=", + "identity": ChainedTokenCredential(FakeTokenCredential()), + "missing": None, + "empty": "", + } + log_name = "user_logs/std_log.txt" + run_history_url = "https://test.invalid/run-history/log" + datastore_url = "https://test.invalid/datastore/log" + run_operations = Mock() + run_operations.get_run_details.side_effect = [ + RunDetails(status="Running", log_files={}), + RunDetails(status="Running", log_files={log_name: run_history_url}), + RunDetails(status="Completed", log_files={log_name: run_history_url}), + ] + requests_pipeline = Mock() + requests_pipeline.with_policies.return_value = requests_pipeline + requests_pipeline.get.side_effect = [ + Mock(status_code=200, text=Mock(return_value="first line\n")), + Mock(status_code=200, text=Mock(return_value="first line\nsecond line\n")), + ] + + ds_info = {"credential": credentials[credential_kind], "storage_type": DatastoreType.AZURE_BLOB} + caplog.set_level("DEBUG", logger="azure.ai.ml.operations._job_ops_helper") + + 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={log_name: datastore_url}, + ) as mock_list_logs_in_datastore, patch("azure.ai.ml.operations._job_ops_helper.time.sleep"): + stream_logs_until_completion( + run_operations, + streaming_job, + datastore_operations=Mock(), + requests_pipeline=requests_pipeline, + ) + + assert mock_list_logs_in_datastore.called is expects_datastore_logs + expected_url = datastore_url if expects_datastore_logs else run_history_url + assert [call.args[0] for call in requests_pipeline.get.call_args_list] == [expected_url, expected_url] + output = capsys.readouterr().out + assert output.count("\nfirst line\n") == 1 + assert output.count("\nsecond line\n") == 1 + assert "Execution Summary" in output + assert fake_datastore_key not in output + caplog.text + assert datastore_sas_token not in output + caplog.text + + @pytest.mark.parametrize( + "storage_type,storage_class,endpoint", + [ + (DatastoreType.AZURE_BLOB, BlobStorageClient, "blob"), + (DatastoreType.AZURE_DATA_LAKE_GEN2, Gen2StorageClient, "dfs"), + ], + ) + @pytest.mark.parametrize("path", ["azureml/job-name", "azureml/job-name/"]) + @pytest.mark.parametrize("long_uri", [False, True]) + @pytest.mark.parametrize("log_name", ["user_logs/std_log.txt", "azureml-logs/70_driver_log.txt"]) + def test_stream_logs_from_key_datastore_uses_storage_path( + self, storage_type, storage_class, endpoint, path, long_uri, log_name, streaming_job, fake_datastore_key, capsys + ) -> None: + uri_prefix = "azureml://subscriptions/sub/resourcegroups/rg/workspaces/ws/" if long_uri else "azureml://" + streaming_job.properties.outputs["default"].uri = f"{uri_prefix}datastores/workspaceblobstore/paths/{path}" + blob_prefix = "azureml/job-name/" + item_name = blob_prefix + log_name + storage_client = Mock(spec=storage_class) + storage_client.list.side_effect = lambda starts_with: [item_name] if item_name.startswith(starts_with) else [] + ds_info = { + "storage_type": storage_type, + "storage_account": "teststorage", + "account_url": f"https://teststorage.{endpoint}.core.windows.net", + "container_name": "testcontainer", + "credential": fake_datastore_key, + } + run_operations = Mock() + run_operations.get_run_details.side_effect = [ + RunDetails(status="Running", log_files={}), + RunDetails(status="Completed", log_files={}), + ] + requests_pipeline = Mock() + requests_pipeline.with_policies.return_value = requests_pipeline + requests_pipeline.get.return_value = Mock(status_code=200, text=Mock(return_value="datastore log\n")) + + with patch("azure.ai.ml.operations._job_ops_helper.get_datastore_info", return_value=ds_info), patch( + "azure.ai.ml._artifacts._artifact_utilities.get_storage_client", return_value=storage_client + ), patch("azure.ai.ml.operations._job_ops_helper.time.sleep"): + stream_logs_until_completion( + run_operations, streaming_job, datastore_operations=Mock(), requests_pipeline=requests_pipeline + ) + + assert [call.kwargs["starts_with"] for call in storage_client.list.call_args_list] == [ + blob_prefix + "user_logs/", + blob_prefix + "azureml-logs/", + ] + requests_pipeline.get.assert_called_once() + log_url = urlsplit(requests_pipeline.get.call_args.args[0]) + assert log_url.netloc == f"teststorage.{endpoint}.core.windows.net" + assert log_url.path == f"/testcontainer/{item_name}" + assert parse_qs(log_url.query)["sp"] == ["r"] + assert parse_qs(log_url.query)["sig"] + assert "\ndatastore log\n" in capsys.readouterr().out + + @pytest.mark.parametrize( + "uri", + [ + "azureml://datastores/workspaceblobstore/paths/", + "azureml://subscriptions/sub/resourcegroups/rg/workspaces/ws/datastores/workspaceblobstore/paths/", + ], + ) + def test_stream_logs_rejects_empty_datastore_path(self, uri, streaming_job) -> None: + streaming_job.properties.outputs["default"].uri = uri + datastore_operations = Mock() + with pytest.raises(ValidationException, match="Invalid AzureML datastore path URI"): + stream_logs_until_completion( + Mock(), streaming_job, datastore_operations=datastore_operations, requests_pipeline=Mock() + ) + datastore_operations.get.assert_not_called() + + @pytest.mark.parametrize( + "storage_type,storage_class", + [ + (DatastoreType.AZURE_BLOB, BlobStorageClient), + (DatastoreType.AZURE_DATA_LAKE_GEN2, Gen2StorageClient), + ], + ) + def test_list_logs_in_datastore_with_empty_prefix(self, storage_type, storage_class, fake_datastore_key) -> None: + log_name = "user_logs/std_log.txt" + storage_client = Mock(spec=storage_class) + storage_client.list.side_effect = [[log_name], []] + ds_info = { + "storage_type": storage_type, + "storage_account": "teststorage", + "account_url": "https://teststorage.blob.core.windows.net", + "container_name": "testcontainer", + "credential": fake_datastore_key, + } + + with patch("azure.ai.ml._artifacts._artifact_utilities.get_storage_client", return_value=storage_client): + logs = list_logs_in_datastore(ds_info, prefix="", legacy_log_folder_name="/azureml-logs/") + + assert [call.kwargs["starts_with"] for call in storage_client.list.call_args_list] == [ + "user_logs/", + "azureml-logs/", + ] + assert set(logs) == {log_name} + assert urlsplit(logs[log_name]).path == f"/testcontainer/{log_name}" + @pytest.mark.skip("TODO 1907352: Relies on a missing VCR.py recording + test suite needs to be reworked") @pytest.mark.unittest