Switch to using temporary data reference SAS uri for non-registry snapshot uploads (Azure#28957)

* Refactor code create_or_update to use temporary data reference for uploads

Author: diondrapeck

sdk/ml/azure-ai-ml/azure/ai/ml/_artifacts/_artifact_utilities.py

@@ -4,6 +4,7 @@
 
 # pylint: disable=protected-access
 
+import json
 import logging
 import os
 import uuid
@@ -13,7 +14,8 @@
 
 from azure.ai.ml._artifacts._blob_storage_helper import BlobStorageClient
 from azure.ai.ml._artifacts._gen2_storage_helper import Gen2StorageClient
-from azure.ai.ml._azure_environments import _get_storage_endpoint_from_metadata
+from azure.ai.ml._azure_environments import _get_storage_endpoint_from_metadata, _get_cloud_details
+from azure.ai.ml._restclient.v2022_05_01.models import Workspace
 from azure.ai.ml._restclient.v2022_10_01.models import DatastoreType
 from azure.ai.ml._scope_dependent_operations import OperationScope
 from azure.ai.ml._utils._arm_id_utils import (
@@ -29,20 +31,29 @@
     _validate_path,
     get_ignore_file,
     get_object_hash,
+    get_content_hash,
+    get_content_hash_version,
 )
+from azure.ai.ml._utils._http_utils import HttpPipeline
 from azure.ai.ml._utils._storage_utils import (
     AzureMLDatastorePathUri,
     get_artifact_path_from_storage_url,
     get_storage_client,
 )
-from azure.ai.ml._utils.utils import is_mlflow_uri, is_url
-from azure.ai.ml.constants._common import SHORT_URI_FORMAT, STORAGE_ACCOUNT_URLS
+from azure.ai.ml._utils.utils import is_mlflow_uri, is_url, retry, replace_between
+from azure.ai.ml.constants._common import (
+    SHORT_URI_FORMAT,
+    STORAGE_ACCOUNT_URLS,
+    MAX_ASSET_STORE_API_CALL_RETRIES,
+    HTTPS_PREFIX,
+)
 from azure.ai.ml.entities import Environment
 from azure.ai.ml.entities._assets._artifacts.artifact import Artifact, ArtifactStorageInfo
 from azure.ai.ml.entities._credentials import AccountKeyConfiguration
 from azure.ai.ml.entities._datastore._constants import WORKSPACE_BLOB_STORE
 from azure.ai.ml.exceptions import ErrorTarget, ValidationException
 from azure.ai.ml.operations._datastore_operations import DatastoreOperations
+from azure.core.exceptions import HttpResponseError
 from azure.storage.blob import BlobSasPermissions, generate_blob_sas
 from azure.storage.filedatalake import FileSasPermissions, generate_file_sas
 
@@ -359,6 +370,126 @@ def _update_gen2_metadata(name, version, indicator_file, storage_client) -> None
     artifact_directory_client.set_metadata(_build_metadata_dict(name=name, version=version))
 
 
+def _generate_temporary_data_reference_id() -> str:
+    """Generate a temporary data reference id."""
+    return str(uuid.uuid4())
+
+
+@retry(
+    exceptions=HttpResponseError,
+    failure_msg="Artifact upload exceeded maximum retries. Try again.",
+    logger=module_logger,
+    max_attempts=MAX_ASSET_STORE_API_CALL_RETRIES,
+)
+def _get_snapshot_temporary_data_reference(
+    asset_name: str,
+    asset_version: str,
+    request_headers: Dict[str, str],
+    workspace: Workspace,
+    requests_pipeline: HttpPipeline,
+) -> Tuple[str, str]:
+    """
+    Make a temporary data reference for an asset and return SAS uri and blob storage uri.
+    :param asset_name: Name of the asset to be created
+    :type asset_name: str
+    :param asset_version: Version of the asset to be created
+    :type asset_version: str
+    :param request_headers: Request headers for API call
+    :type request_headers: Dict[str, str]
+    :param workspace: Workspace object
+    :type workspace: azure.ai.ml._restclient.v2022_05_01.models.Workspace
+    :param requests_pipeline: Proxy for sending HTTP requests
+    :type requests_pipeline: azure.ai.ml._utils._http_utils.HttpPipeline
+    :return: Existing asset's name and version, if found
+    :rtype: Tuple[str, str]
+    """
+
+    # create temporary data reference
+    temporary_data_reference_id = _generate_temporary_data_reference_id()
+
+    # build and send request
+    asset_id = (
+        f"azureml://locations/{workspace.location}/workspaces/{workspace.workspace_id}/"
+        f"codes/{asset_name}/versions/{asset_version}"
+    )
+    data = {
+        "assetId": asset_id,
+        "temporaryDataReferenceId": temporary_data_reference_id,
+        "temporaryDataReferenceType": "TemporaryBlobReference",
+    }
+    data_encoded = json.dumps(data).encode("utf-8")
+    serialized_data = json.loads(data_encoded)
+
+    # make sure correct cloud endpoint is used
+    cloud_endpoint = _get_cloud_details()["registry_discovery_endpoint"]
+    service_url = replace_between(cloud_endpoint, HTTPS_PREFIX, ".", workspace.location)
+
+    # send request
+    request_url = f"{service_url}assetstore/v1.0/temporaryDataReference/createOrGet"
+    response = requests_pipeline.post(request_url, json=serialized_data, headers=request_headers)
+
+    if response.status_code != 200:
+        raise HttpResponseError(response=response)
+
+    response_json = json.loads(response.text())
+
+    # get SAS uri for upload and blob uri for asset creation
+    blob_uri = response_json["blobReferenceForConsumption"]["blobUri"]
+    sas_uri = response_json["blobReferenceForConsumption"]["credential"]["sasUri"]
+
+    return sas_uri, blob_uri
+
+
+def _get_asset_by_hash(
+    operations: "DatastoreOperations",
+    hash_str: str,
+    request_headers: Dict[str, str],
+    workspace: Workspace,
+    requests_pipeline: HttpPipeline,
+) -> Dict[str, str]:
+    """
+    Check if an asset with the same hash already exists in the workspace. If so, return the asset name and version.
+    :param operations: Datastore Operations object from MLClient
+    :type operations: azure.ai.ml.operations._datastore_operations.DatastoreOperations
+    :param hash_str: The hash of the specified local upload
+    :type hash_str: str
+    :param request_headers: Request headers for API call
+    :type request_headers: Dict[str, str]
+    :param workspace: Workspace object
+    :type workspace: azure.ai.ml._restclient.v2022_05_01.models.Workspace
+    :param requests_pipeline: Proxy for sending HTTP requests
+    :type requests_pipeline: azure.ai.ml._utils._http_utils.HttpPipeline
+    :return: Existing asset's name and version, if found
+    :rtype: Optional[Dict[str, str]]
+    """
+    existing_asset = {}
+    hash_version = get_content_hash_version()
+
+    # get workspace credentials
+    subscription_id = operations._subscription_id
+    resource_group_name = operations._resource_group_name
+
+    # make sure correct cloud endpoint is used
+    cloud_endpoint = _get_cloud_details()["registry_discovery_endpoint"]
+    service_url = replace_between(cloud_endpoint, HTTPS_PREFIX, ".", workspace.location)
+    request_url = (
+        f"{service_url}content/v2.0/subscriptions/{subscription_id}/"
+        f"resourceGroups/{resource_group_name}/providers/Microsoft.MachineLearningServices/workspaces/"
+        f"{workspace.name}/snapshots/getByHash?hash={hash_str}&hashVersion={hash_version}"
+    )
+
+    response = requests_pipeline.get(request_url, headers=request_headers)
+    if response.status_code != 200:
+        # If API is unresponsive, create new asset
+        return None
+
+    response_json = json.loads(response.text())
+    existing_asset["name"] = response_json["name"]
+    existing_asset["version"] = response_json["version"]
+
+    return existing_asset
+
+
 T = TypeVar("T", bound=Artifact)
 
 
@@ -422,6 +553,193 @@ def _check_and_upload_path(
     return artifact, indicator_file
 
 
+def _get_snapshot_path_info(artifact) -> Tuple[str, str, str]:
+    """
+    Validate an Artifact's local path and get its resolved path, ignore file, and hash
+    :param artifact: Artifact object
+    :type artifact: azure.ai.ml.entities._assets._artifacts.artifact.Artifact
+    :return: Artifact's path, ignorefile, and hash
+    :rtype: Tuple[str, str, str]
+    """
+    if (
+        hasattr(artifact, "local_path")
+        and artifact.local_path is not None
+        or (
+            hasattr(artifact, "path")
+            and artifact.path is not None
+            and not (is_url(artifact.path) or is_mlflow_uri(artifact.path))
+        )
+    ):
+        path = (
+            Path(artifact.path)
+            if hasattr(artifact, "path") and artifact.path is not None
+            else Path(artifact.local_path)
+        )
+        if not path.is_absolute():
+            path = Path(artifact.base_path, path).resolve()
+
+    _validate_path(path, _type=ErrorTarget.CODE)
+
+    ignore_file = get_ignore_file(path)
+    asset_hash = get_content_hash(path, ignore_file)
+
+    return path, ignore_file, asset_hash
+
+
+def _get_existing_snapshot_by_hash(
+    datastore_operation,
+    asset_hash,
+    workspace: Workspace,
+    requests_pipeline: HttpPipeline,
+) -> Dict[str, str]:
+    """
+    Check if an asset with the same hash already exists in the workspace. If so, return the asset name and version.
+    :param datastore_operation: Datastore Operations object from MLClient
+    :type operations: azure.ai.ml.operations._datastore_operations.DatastoreOperations
+    :param asset_hash: The hash of the specified local upload
+    :type asset_hash: str
+    :param workspace: Workspace object
+    :type workspace: azure.ai.ml._restclient.v2022_05_01.models.Workspace
+    :param requests_pipeline: Proxy for sending HTTP requests
+    :type requests_pipeline: azure.ai.ml._utils._http_utils.HttpPipeline
+    :return: Existing asset's name and version, if found
+    :rtype: Optional[Dict[str, str]]
+    """
+    ws_base_url = datastore_operation._operation._client._base_url
+    token = datastore_operation._credential.get_token(ws_base_url + "/.default").token
+    request_headers = {"Authorization": "Bearer " + token}
+    request_headers["Content-Type"] = "application/json; charset=UTF-8"
+
+    existing_asset = _get_asset_by_hash(
+        operations=datastore_operation,
+        hash_str=asset_hash,
+        request_headers=request_headers,
+        workspace=workspace,
+        requests_pipeline=requests_pipeline,
+    )
+
+    return existing_asset
+
+
+def _upload_snapshot_to_datastore(
+    operation_scope: OperationScope,
+    datastore_operation: DatastoreOperations,
+    path: Union[str, Path, os.PathLike],
+    workspace: Workspace,
+    requests_pipeline: HttpPipeline,
+    datastore_name: str = None,
+    show_progress: bool = True,
+    asset_name: str = None,
+    asset_version: str = None,
+    asset_hash: str = None,
+    ignore_file: IgnoreFile = IgnoreFile(),
+    sas_uri: str = None,  # contains registry sas url
+) -> ArtifactStorageInfo:
+    """
+    Upload a code snapshot to workspace datastore.
+    :param operation_scope: Workspace scope information
+    :type operation_scope: azure.ai.ml._scope_dependent_operations.OperationScope
+    :param datastore_operation: Datastore Operations object from MLClient
+    :type datastore_operation: azure.ai.ml.operations._datastore_operations.DatastoreOperations
+    :param path: The local path of the artifact
+    :type path: Union[str, Path, os.PathLike]
+    :param workspace: Workspace object
+    :type workspace: azure.ai.ml._restclient.v2022_05_01.models.Workspace
+    :param requests_pipeline: Proxy for sending HTTP requests
+    :type requests_pipeline: azure.ai.ml._utils._http_utils.HttpPipeline
+    :param datastore_name: Name of the datastore to upload to
+    :type datastore_name: str
+    :param show_progress: Whether or not to show progress bar during upload, defaults to True
+    :type show_progress: bool
+    :param asset_name: Name of the asset to be created
+    :type asset_name: str
+    :param asset_version: Version of the asset to be created
+    :type asset_version: str
+    :param asset_hash: The hash of the specified local upload
+    :type asset_hash: str
+    :param ignore_file: Information about the path's .gitignore or .amlignore file, if exists
+    :type ignore_file: azure.ai.ml._utils._asset_utils.IgnoreFile
+    :param sas_uri: SAS uri for uploading to datastore
+    :type sas_uri: str
+    :return: Uploaded artifact's storage information
+    :rtype: azure.ai.ml.entities._assets._artifacts.artifact.ArtifactStorageInfo
+    """
+    ws_base_url = datastore_operation._operation._client._base_url
+    token = datastore_operation._credential.get_token(ws_base_url + "/.default").token
+    request_headers = {"Authorization": "Bearer " + token}
+    request_headers["Content-Type"] = "application/json; charset=UTF-8"
+
+    if not sas_uri:
+        sas_uri, blob_uri = _get_snapshot_temporary_data_reference(
+            requests_pipeline=requests_pipeline,
+            asset_name=asset_name,
+            asset_version=asset_version,
+            request_headers=request_headers,
+            workspace=workspace,
+        )
+
+    artifact = upload_artifact(
+        str(path),
+        datastore_operation,
+        operation_scope,
+        datastore_name,
+        show_progress=show_progress,
+        asset_hash=asset_hash,
+        asset_name=asset_name,
+        asset_version=asset_version,
+        ignore_file=ignore_file,
+        sas_uri=sas_uri,
+    )
+    artifact.storage_account_url = blob_uri
+
+    return artifact
+
+
+def _check_and_upload_snapshot(
+    artifact: T,
+    asset_operations: Union["DataOperations", "ModelOperations", "CodeOperations"],
+    path: Union[str, Path, os.PathLike],
+    workspace: Workspace,
+    requests_pipeline: HttpPipeline,
+    ignore_file: IgnoreFile = None,
+    datastore_name: str = WORKSPACE_BLOB_STORE,
+    sas_uri: str = None,
+    show_progress: bool = True,
+) -> Tuple[T, str]:
+    """Checks whether `artifact` is a path or a uri and uploads it to the
+    datastore if necessary.
+    param T artifact: artifact to check and upload param
+    Union["DataOperations", "ModelOperations", "CodeOperations"]
+    asset_operations:     the asset operations to use for uploading
+    param str datastore_name: the name of the datastore to upload to
+    param str sas_uri: the sas uri to use for uploading
+    """
+    uploaded_artifact = _upload_snapshot_to_datastore(
+        operation_scope=asset_operations._operation_scope,
+        datastore_operation=asset_operations._datastore_operation,
+        path=path,
+        workspace=workspace,
+        requests_pipeline=requests_pipeline,
+        datastore_name=datastore_name,
+        asset_name=artifact.name,
+        asset_version=str(artifact.version),
+        asset_hash=artifact._upload_hash if hasattr(artifact, "_upload_hash") else None,
+        show_progress=show_progress,
+        sas_uri=sas_uri,
+        ignore_file=ignore_file,
+    )
+
+    if artifact._is_anonymous:
+        artifact.name, artifact.version = (
+            uploaded_artifact.name,
+            uploaded_artifact.version,
+        )
+    # Pass all of the upload information to the assets, and they will each construct the URLs that they support
+    artifact._update_path(uploaded_artifact)
+
+    return artifact
+
+
 def _check_and_upload_env_build_context(
     environment: Environment,
     operations: "EnvironmentOperations",

sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py

@@ -350,6 +350,7 @@ def __init__(
             self._operation_config,
             self._service_client_10_2021_dataplanepreview if registry_name else self._service_client_05_2022,
             self._datastores,
+            requests_pipeline=self._requests_pipeline,
             **ops_kwargs,
         )
         self._operation_container.add(AzureMLResourceType.CODE, self._code)

sdk/ml/azure-ai-ml/azure/ai/ml/_utils/utils.py

@@ -2,7 +2,7 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
 
-# pylint: disable=protected-access
+# pylint: disable=protected-access,too-many-lines
 import copy
 import decimal
 import hashlib
@@ -996,3 +996,8 @@ def get_valid_dot_keys_with_wildcard(
     """
     left_reversed_parts = dot_key_wildcard.split(".")[::-1]
     return _get_valid_dot_keys_with_wildcard_impl(left_reversed_parts, root, validate_func=validate_func)
+
+
+def replace_between(s: str, start: str, end: str, replace: str) -> str:
+    """Replace string between two substrings."""
+    return start + replace + s[s.find(end) - 1 + len(end) :]