two version of R2R are here HEAD master

1 files changed, 598 insertions, 0 deletions

diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/_artifacts/_artifact_utilities.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/_artifacts/_artifact_utilities.py
new file mode 100644
index 00000000..c225a73e
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/_artifacts/_artifact_utilities.py
@@ -0,0 +1,598 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+# pylint: disable=protected-access
+
+import logging
+import os
+import uuid
+from datetime import datetime, timedelta
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, Optional, Tuple, TypeVar, Union
+
+from typing_extensions import Literal
+
+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._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 (
+    AMLNamedArmId,
+    get_datastore_arm_id,
+    get_resource_name_from_arm_id,
+    is_ARM_id_for_resource,
+    remove_aml_prefix,
+)
+from azure.ai.ml._utils._asset_utils import (
+    IgnoreFile,
+    _build_metadata_dict,
+    _validate_path,
+    get_content_hash,
+    get_ignore_file,
+    get_object_hash,
+)
+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.entities import Environment
+from azure.ai.ml.entities._assets._artifacts.artifact import Artifact, ArtifactStorageInfo
+from azure.ai.ml.entities._datastore._constants import WORKSPACE_BLOB_STORE
+from azure.ai.ml.exceptions import ErrorTarget, MlException, 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
+
+if TYPE_CHECKING:
+    from azure.ai.ml.operations import (
+        DataOperations,
+        EnvironmentOperations,
+        EvaluatorOperations,
+        FeatureSetOperations,
+        IndexOperations,
+        ModelOperations,
+    )
+    from azure.ai.ml.operations._code_operations import CodeOperations
+
+module_logger = logging.getLogger(__name__)
+
+
+def _get_datastore_name(*, datastore_name: Optional[str] = WORKSPACE_BLOB_STORE) -> str:
+    datastore_name = WORKSPACE_BLOB_STORE if not datastore_name else datastore_name
+    try:
+        datastore_name = get_resource_name_from_arm_id(datastore_name)
+    except (ValueError, AttributeError, ValidationException):
+        module_logger.debug("datastore_name %s is not a full arm id. Proceed with a shortened name.\n", datastore_name)
+    datastore_name = remove_aml_prefix(datastore_name)
+    if is_ARM_id_for_resource(datastore_name):
+        datastore_name = get_resource_name_from_arm_id(datastore_name)
+    return str(datastore_name)
+
+
+def get_datastore_info(
+    operations: DatastoreOperations,
+    name: str,
+    *,
+    credential=None,
+    **kwargs,
+) -> Dict[Literal["storage_type", "storage_account", "account_url", "container_name", "credential"], str]:
+    """Get datastore account, type, and auth information.
+
+    :param operations: DatastoreOperations object
+    :type operations: DatastoreOperations
+    :param name: Name of the datastore. If not provided, the default datastore will be used.
+    :type name: str
+    :keyword credential: Local credential to use for authentication. This argument is no longer used as of 1.18.0.
+        Instead, a SAS token will be requested from the datastore, and the MLClient credential will be used as backup,
+        if necessary.
+    :paramtype credential: str
+    :return: The dictionary with datastore info
+    :rtype: Dict[Literal["storage_type", "storage_account", "account_url", "container_name", "credential"], str]
+    """
+    datastore_info: Dict = {}
+    datastore = operations.get(name) if name else operations.get_default()
+
+    storage_endpoint = _get_storage_endpoint_from_metadata()
+    datastore_info["storage_type"] = datastore.type
+    datastore_info["storage_account"] = datastore.account_name
+    datastore_info["account_url"] = STORAGE_ACCOUNT_URLS[datastore.type].format(
+        datastore.account_name, storage_endpoint
+    )
+
+    try:
+        credential = operations._list_secrets(name=name, expirable_secret=True)
+        datastore_info["credential"] = credential.sas_token
+    except HttpResponseError:
+        datastore_info["credential"] = operations._credential
+
+    if datastore.type == DatastoreType.AZURE_BLOB:
+        datastore_info["container_name"] = str(datastore.container_name)
+    elif datastore.type == DatastoreType.AZURE_DATA_LAKE_GEN2:
+        datastore_info["container_name"] = str(datastore.filesystem)
+    else:
+        msg = (
+            f"Datastore type {datastore.type} is not supported for uploads. "
+            f"Supported types are {DatastoreType.AZURE_BLOB} and {DatastoreType.AZURE_DATA_LAKE_GEN2}."
+        )
+        raise MlException(message=msg, no_personal_data_message=msg)
+
+    for override_param_name, value in kwargs.items():
+        if override_param_name in datastore_info:
+            datastore_info[override_param_name] = value
+
+    return datastore_info
+
+
+def list_logs_in_datastore(
+    ds_info: Dict[Literal["storage_type", "storage_account", "account_url", "container_name", "credential"], str],
+    prefix: str,
+    legacy_log_folder_name: str,
+) -> Dict[str, str]:
+    """Returns a dictionary of file name to blob or data lake uri with SAS token, matching the structure of
+    RunDetails.logFiles.
+
+    :param ds_info: The datastore info
+    :type ds_info: Dict[Literal["storage_type", "storage_account", "account_url", "container_name", "credential"], str]
+    :param prefix: A prefix used to filter logs by path
+    :type prefix: str
+    :param legacy_log_folder_name: the name of the folder in the datastore that contains the logs
+        * /azureml-logs/*.txt is the legacy log structure for commandJob and sweepJob
+        * /logs/azureml/*.txt is the legacy log structure for pipeline parent Job
+    :type legacy_log_folder_name: str
+    :return: A mapping of log file name to the remote URI
+    :rtype: Dict[str, str]
+    """
+    if ds_info["storage_type"] not in [
+        DatastoreType.AZURE_BLOB,
+        DatastoreType.AZURE_DATA_LAKE_GEN2,
+    ]:
+        msg = "Only Blob and Azure DataLake Storage Gen2 datastores are supported."
+        raise MlException(message=msg, no_personal_data_message=msg)
+
+    storage_client = get_storage_client(
+        credential=ds_info["credential"],
+        container_name=ds_info["container_name"],
+        storage_account=ds_info["storage_account"],
+        storage_type=ds_info["storage_type"],
+    )
+
+    items = storage_client.list(starts_with=prefix + "/user_logs/")
+    # Append legacy log files if present
+    items.extend(storage_client.list(starts_with=prefix + legacy_log_folder_name))
+
+    log_dict = {}
+    for item_name in items:
+        sub_name = item_name.split(prefix + "/")[1]
+        if isinstance(storage_client, BlobStorageClient):
+            token = generate_blob_sas(
+                account_name=ds_info["storage_account"],
+                container_name=ds_info["container_name"],
+                blob_name=item_name,
+                account_key=ds_info["credential"],
+                permission=BlobSasPermissions(read=True),
+                expiry=datetime.utcnow() + timedelta(minutes=30),
+            )
+        elif isinstance(storage_client, Gen2StorageClient):
+            token = generate_file_sas(  # pylint: disable=no-value-for-parameter
+                account_name=ds_info["storage_account"],
+                file_system_name=ds_info["container_name"],
+                file_name=item_name,
+                credential=ds_info["credential"],
+                permission=FileSasPermissions(read=True),
+                expiry=datetime.utcnow() + timedelta(minutes=30),
+            )
+
+        log_dict[sub_name] = "{}/{}/{}?{}".format(
+            ds_info["account_url"],
+            ds_info["container_name"],
+            item_name,
+            token,  # pylint: disable=possibly-used-before-assignment
+        )
+    return log_dict
+
+
+def _get_default_datastore_info(datastore_operation):
+    return get_datastore_info(datastore_operation, None)
+
+
+def upload_artifact(
+    local_path: str,
+    datastore_operation: DatastoreOperations,
+    operation_scope: OperationScope,
+    datastore_name: Optional[str],
+    asset_hash: Optional[str] = None,
+    show_progress: bool = True,
+    asset_name: Optional[str] = None,
+    asset_version: Optional[str] = None,
+    ignore_file: IgnoreFile = IgnoreFile(None),
+    sas_uri: Optional[str] = None,
+) -> ArtifactStorageInfo:
+    """Upload local file or directory to datastore.
+
+    :param local_path: The local file or directory to upload
+    :type local_path: str
+    :param datastore_operation: The datastore operation
+    :type datastore_operation: DatastoreOperations
+    :param operation_scope: The operation scope
+    :type operation_scope: OperationScope
+    :param datastore_name: The datastore name
+    :type datastore_name: Optional[str]
+    :param asset_hash: The asset hash
+    :type asset_hash: Optional[str]
+    :param show_progress: Whether to show progress on the console. Defaults to True.
+    :type show_progress: bool
+    :param asset_name: The asset name
+    :type asset_name: Optional[str]
+    :param asset_version: The asset version
+    :type asset_version: Optional[str]
+    :param ignore_file: The IgnoreFile determining which, if any, files to ignore when uploading
+    :type ignore_file: IgnoreFile
+    :param sas_uri: The sas uri to use for uploading
+    :type sas_uri: Optional[str]
+    :return: The artifact storage info
+    :rtype: ArtifactStorageInfo
+    """
+    if sas_uri:
+        storage_client = get_storage_client(credential=None, storage_account=None, account_url=sas_uri)
+    else:
+        datastore_name = _get_datastore_name(datastore_name=datastore_name)
+        datastore_info = get_datastore_info(datastore_operation, datastore_name)
+        storage_client = get_storage_client(**datastore_info)
+
+    artifact_info = storage_client.upload(
+        local_path,
+        asset_hash=asset_hash,
+        show_progress=show_progress,
+        name=asset_name,
+        version=asset_version,
+        ignore_file=ignore_file,
+    )
+
+    artifact = ArtifactStorageInfo(
+        name=artifact_info["name"],
+        version=artifact_info["version"],
+        relative_path=artifact_info["remote path"],
+        datastore_arm_id=get_datastore_arm_id(datastore_name, operation_scope) if not sas_uri else None,
+        container_name=(
+            storage_client.container if isinstance(storage_client, BlobStorageClient) else storage_client.file_system
+        ),
+        storage_account_url=datastore_info.get("account_url") if not sas_uri else sas_uri,
+        indicator_file=artifact_info["indicator file"],
+        is_file=Path(local_path).is_file(),
+    )
+    return artifact
+
+
+def download_artifact(
+    starts_with: Union[str, os.PathLike],
+    destination: str,
+    datastore_operation: DatastoreOperations,
+    datastore_name: Optional[str],
+    datastore_info: Optional[Dict] = None,
+) -> str:
+    """Download datastore path to local file or directory.
+
+    :param starts_with: Prefix of blobs to download
+    :type starts_with: Union[str, os.PathLike]
+    :param destination: Path that files will be written to
+    :type destination: str
+    :param datastore_operation: Datastore operations
+    :type datastore_operation: DatastoreOperations
+    :param datastore_name: name of datastore
+    :type datastore_name: Optional[str]
+    :param datastore_info: the return value of invoking get_datastore_info
+    :type datastore_info: Optional[Dict]
+    :return: Path that files were written to
+    :rtype: str
+    """
+    starts_with = starts_with.as_posix() if isinstance(starts_with, Path) else starts_with
+    datastore_name = _get_datastore_name(datastore_name=datastore_name)
+    if datastore_info is None:
+        datastore_info = get_datastore_info(datastore_operation, datastore_name)
+    storage_client = get_storage_client(**datastore_info)
+    storage_client.download(starts_with=starts_with, destination=destination)
+    return destination
+
+
+def download_artifact_from_storage_url(
+    blob_url: str,
+    destination: str,
+    datastore_operation: DatastoreOperations,
+    datastore_name: Optional[str],
+) -> str:
+    """Download datastore blob URL to local file or directory.
+
+    :param blob_url: The blob url to download
+    :type blob_url: str
+    :param destination: Path that the artifact will be written to
+    :type destination: str
+    :param datastore_operation: The datastore operations
+    :type datastore_operation: DatastoreOperations
+    :param datastore_name: The datastore name
+    :type datastore_name: Optional[str]
+    :return: Path that files were written to
+    :rtype: str
+    """
+    datastore_name = _get_datastore_name(datastore_name=datastore_name)
+    datastore_info = get_datastore_info(datastore_operation, datastore_name)
+    starts_with = get_artifact_path_from_storage_url(
+        blob_url=str(blob_url), container_name=datastore_info.get("container_name")
+    )
+    return download_artifact(
+        starts_with=starts_with,
+        destination=destination,
+        datastore_operation=datastore_operation,
+        datastore_name=datastore_name,
+        datastore_info=datastore_info,
+    )
+
+
+def download_artifact_from_aml_uri(uri: str, destination: str, datastore_operation: DatastoreOperations) -> str:
+    """Downloads artifact pointed to by URI of the form `azureml://...` to destination.
+
+    :param str uri: AzureML uri of artifact to download
+    :param str destination: Path to download artifact to
+    :param DatastoreOperations datastore_operation: datastore operations
+    :return: Path that files were downloaded to
+    :rtype: str
+    """
+    parsed_uri = AzureMLDatastorePathUri(uri)
+    return download_artifact(
+        starts_with=parsed_uri.path,
+        destination=destination,
+        datastore_operation=datastore_operation,
+        datastore_name=parsed_uri.datastore,
+    )
+
+
+def aml_datastore_path_exists(
+    uri: str, datastore_operation: DatastoreOperations, datastore_info: Optional[dict] = None
+) -> bool:
+    """Checks whether `uri` of the form "azureml://" points to either a directory or a file.
+
+    :param str uri: azure ml datastore uri
+    :param DatastoreOperations datastore_operation: Datastore operation
+    :param dict datastore_info: return value of get_datastore_info
+    :return: True if uri exists False otherwise
+    :rtype: bool
+    """
+    parsed_uri = AzureMLDatastorePathUri(uri)
+    datastore_info = datastore_info or get_datastore_info(datastore_operation, parsed_uri.datastore)
+    return get_storage_client(**datastore_info).exists(parsed_uri.path)
+
+
+def _upload_to_datastore(
+    operation_scope: OperationScope,
+    datastore_operation: DatastoreOperations,
+    path: Union[str, Path, os.PathLike],
+    artifact_type: str,
+    datastore_name: Optional[str] = None,
+    show_progress: bool = True,
+    asset_name: Optional[str] = None,
+    asset_version: Optional[str] = None,
+    asset_hash: Optional[str] = None,
+    ignore_file: Optional[IgnoreFile] = None,
+    sas_uri: Optional[str] = None,
+    blob_uri: Optional[str] = None,
+) -> ArtifactStorageInfo:
+    _validate_path(path, _type=artifact_type)
+    if not ignore_file:
+        ignore_file = get_ignore_file(path)
+    if not asset_hash:
+        asset_hash = get_object_hash(path, ignore_file)
+    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,
+    )
+    if blob_uri:
+        artifact.storage_account_url = blob_uri
+
+    return artifact
+
+
+def _upload_and_generate_remote_uri(
+    operation_scope: OperationScope,
+    datastore_operation: DatastoreOperations,
+    path: Union[str, Path, os.PathLike],
+    artifact_type: str = ErrorTarget.ARTIFACT,
+    datastore_name: str = WORKSPACE_BLOB_STORE,
+    show_progress: bool = True,
+) -> str:
+    # Asset name is required for uploading to a datastore
+    asset_name = str(uuid.uuid4())
+    artifact_info = _upload_to_datastore(
+        operation_scope=operation_scope,
+        datastore_operation=datastore_operation,
+        path=path,
+        datastore_name=datastore_name,
+        asset_name=asset_name,
+        artifact_type=artifact_type,
+        show_progress=show_progress,
+    )
+
+    path = artifact_info.relative_path
+    datastore = AMLNamedArmId(artifact_info.datastore_arm_id).asset_name
+    return SHORT_URI_FORMAT.format(datastore, path)
+
+
+def _update_metadata(name, version, indicator_file, datastore_info) -> None:
+    storage_client = get_storage_client(**datastore_info)
+
+    if isinstance(storage_client, BlobStorageClient):
+        _update_blob_metadata(name, version, indicator_file, storage_client)
+    elif isinstance(storage_client, Gen2StorageClient):
+        _update_gen2_metadata(name, version, indicator_file, storage_client)
+
+
+def _update_blob_metadata(name, version, indicator_file, storage_client) -> None:
+    container_client = storage_client.container_client
+    if indicator_file.startswith(storage_client.container):
+        indicator_file = indicator_file.split(storage_client.container)[1]
+    blob = container_client.get_blob_client(blob=indicator_file)
+    blob.set_blob_metadata(_build_metadata_dict(name=name, version=version))
+
+
+def _update_gen2_metadata(name, version, indicator_file, storage_client) -> None:
+    artifact_directory_client = storage_client.file_system_client.get_directory_client(indicator_file)
+    artifact_directory_client.set_metadata(_build_metadata_dict(name=name, version=version))
+
+
+T = TypeVar("T", bound=Artifact)
+
+
+def _check_and_upload_path(
+    artifact: T,
+    asset_operations: Union[
+        "DataOperations",
+        "ModelOperations",
+        "EvaluatorOperations",
+        "CodeOperations",
+        "FeatureSetOperations",
+        "IndexOperations",
+    ],
+    artifact_type: str,
+    datastore_name: Optional[str] = None,
+    sas_uri: Optional[str] = None,
+    show_progress: bool = True,
+    blob_uri: Optional[str] = None,
+) -> Tuple[T, Optional[str]]:
+    """Checks whether `artifact` is a path or a uri and uploads it to the datastore if necessary.
+
+    :param artifact: artifact to check and upload param
+    :type artifact: T
+    :param asset_operations: The asset operations to use for uploading
+    :type asset_operations: Union["DataOperations", "ModelOperations", "CodeOperations", "IndexOperations"]
+    :param artifact_type: The artifact type
+    :type artifact_type: str
+    :param datastore_name: the name of the datastore to upload to
+    :type datastore_name: Optional[str]
+    :param sas_uri: the sas uri to use for uploading
+    :type sas_uri: Optional[str]
+    :param show_progress: Whether to show progress on the console. Defaults to True.
+    :type show_progress: bool
+    :param blob_uri: The storage account uri
+    :type blob_uri: Optional[str]
+    :return: A 2-tuple of the uploaded artifact, and the indicator file.
+    :rtype: Tuple[T, Optional[str]]
+    """
+
+    datastore_name = artifact.datastore
+    indicator_file = None
+    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()
+        uploaded_artifact = _upload_to_datastore(
+            asset_operations._operation_scope,
+            asset_operations._datastore_operation,
+            path,
+            datastore_name=datastore_name,
+            asset_name=artifact.name,
+            asset_version=str(artifact.version),
+            asset_hash=getattr(artifact, "_upload_hash", None),
+            sas_uri=sas_uri,
+            artifact_type=artifact_type,
+            show_progress=show_progress,
+            ignore_file=getattr(artifact, "_ignore_file", None),
+            blob_uri=blob_uri,
+        )
+        indicator_file = uploaded_artifact.indicator_file  # reference to storage contents
+        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, indicator_file
+
+
+def _check_and_upload_env_build_context(
+    environment: Environment,
+    operations: "EnvironmentOperations",
+    sas_uri=None,
+    show_progress: bool = True,
+) -> Environment:
+    if environment.path:
+        uploaded_artifact = _upload_to_datastore(
+            operations._operation_scope,
+            operations._datastore_operation,
+            environment.path,
+            asset_name=environment.name,
+            asset_version=str(environment.version),
+            asset_hash=environment._upload_hash,
+            sas_uri=sas_uri,
+            artifact_type=ErrorTarget.ENVIRONMENT,
+            datastore_name=environment.datastore,
+            show_progress=show_progress,
+        )
+        if environment.build is not None:
+            # TODO: Depending on decision trailing "/" needs to stay or not. EMS requires it to be present
+            environment.build.path = str(uploaded_artifact.full_storage_path) + "/"
+    return environment
+
+
+def _get_snapshot_path_info(artifact) -> Optional[Tuple[Path, IgnoreFile, str]]:
+    """
+    Validate an Artifact's local path and get its resolved path, ignore file, and hash. If no local path, return None.
+    :param artifact: Artifact object
+    :type artifact: azure.ai.ml.entities._assets._artifacts.artifact.Artifact
+    :return: Artifact's path, ignorefile, and hash
+    :rtype: Tuple[os.PathLike, IgnoreFile, 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()
+    else:
+        return None
+
+    _validate_path(path, _type=ErrorTarget.CODE)
+
+    # to align with _check_and_upload_path, we need to try getting the ignore file from the artifact first
+    ignore_file = getattr(artifact, "_ignore_file", get_ignore_file(path))
+    # Note that we haven't used getattr(artifact, "_upload_hash", get_content_hash(path, ignore_file)) here, which
+    # is aligned with _check_and_upload_path. Current guess is that content_hash is what we used in blob, so we must
+    # use it to retrieve the artifact.
+    # TODO: Core SDK team to provide more information on this
+    asset_hash = get_content_hash(path, ignore_file)
+
+    return path, ignore_file, asset_hash