Add data/datastore mount in AzureML SDK, depending on azureml-dataprep-rslex (#31223)

* add data.mount in SDK, depending on azureml-dataprep-rslex

* impl datastore.mount()

* Update docstring for datastore.mount

* point to azureml-dataprep-rslex whl in setup.py

* Use build on test index

https://msdata.visualstudio.com/Vienna/_build/results?buildId=100511334&view=results

https://dataprepdownloads.azureedge.net/pypi/test-M3ME5B1GMEM3SW0W/100511334/

* from azureml.dataprep import rslex_fuse_cli

* update rslex version

* pass debug bool argument

* add type annotation; fix datastore docstring

* add TODO comment

* rename rslex_fuse_cli -> rslex_fuse_subprocess_wrapper

* Add azureml-dataprep-rslex to shared_requirements.txt

* # cspell:ignore rslex

* black reformat

* support persistent on CI (TODO: call MLC)

* Add mlc swagger, update autorest config

* overwrite with online version

* Revert "overwrite with online version"

This reverts commit d91db0d.

* Revert "Add mlc swagger, update autorest config"

This reverts commit be24871.

* [stash] manually add updateDataMounts to swagger.json

* generated operation for MLC updateDataMounts

* add code to call compute operation: ?

* [stash]

* add body type in swagger.json

* [stash] update path in command

* generated compute operations with body type, both w/ & wo/ aio

* Add body calling MLC update_data_mounts

* black format

* Fix: [ComputeInstanceDataMount]

* specify api_version; poll; impl datastore op

* pass user specified mount_point in persistent mount

* Fix default mount_point; remove mount_name TODO

* amend: parity in _datastore_operations.py

* credential pass-in for data/datastore mount()

* Add to DatastoreOperations self._compute_operation = serviceclient_2023_04_01_preview.compute

* Revert "add body type in swagger.json"

This reverts commit 5b7d67f.

* Revert "[stash] manually add updateDataMounts to swagger.json"

This reverts commit 9c59d0d.

* Revert "generated compute operations with body type, both w/ & wo/ aio"

This reverts commit a835869.

* Revert "generated operation for MLC updateDataMounts"

This reverts commit 46ef0a7.

* Revert "[stash] update path in command"

This reverts commit f4a682e.

* Revert "Revert "generated operation for MLC updateDataMounts""

This reverts commit 1e686c0.

* Revert "[stash]"

This reverts commit e0586bb.

* Revert "Revert "Revert "generated operation for MLC updateDataMounts"""

This reverts commit 8bbe587.

* Point to 2024-01-01 client

* Fix NameError: name 'experimental' is not defined

* black

* Update mocks for data/datastore operations

* fix mock syntax error

* Add fixture mock_aml_services_2024_01_01_preview

* Fix fixture name mock_aml_services_2024_01_01_preview

* black formatting

* Fix pylint: R1723(no-else-break), W0613(unused-argument), C4739(docstring-missing-param), C4740(docstring-missing-type)

* Remove comments on TBD timeout

* Disable E1123(unexpected-keyword-arg) temporarily

* Revert "Disable E1123(unexpected-keyword-arg) temporarily"

This reverts commit d60f1e4.

* point to azureml-dataprep-rslex>=2.22.0

* add rslex to spell check allowlist

* ", where the 'CI_NAME' environment variable is set."

Co-authored-by: MilesHolland <[email protected]>

* ", where the 'CI_NAME' environment variable is set."

Co-authored-by: MilesHolland <[email protected]>

* Add in _data_operations.py clause ", where the 'CI_NAME' environment variable is set."

* Add clause in both datastore and data ", so that you can access data ..."

* Change azureml-dataprep-rslex to be an extra, operation-specific dependency

* Add unit test (azureml-dataprep-rslex ImportError)

* Update dummy mount path to be under /tmp

* fix errors from pylint

* black reformat

* UT patches azureml.dataprep.rslex_fuse_subprocess_wrapper

---------

Co-authored-by: MilesHolland <[email protected]>

Author: lizy14

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

@@ -39,6 +39,7 @@
 
 # Same object, but was renamed starting in v2023_08_01_preview
 from azure.ai.ml._restclient.v2023_10_01 import AzureMachineLearningServices as ServiceClient102023
+from azure.ai.ml._restclient.v2024_01_01_preview import AzureMachineLearningWorkspaces as ServiceClient012024Preview
 from azure.ai.ml._scope_dependent_operations import OperationConfig, OperationsContainer, OperationScope
 from azure.ai.ml._telemetry.logging_handler import get_appinsights_log_handler
 from azure.ai.ml._user_agent import USER_AGENT
@@ -407,6 +408,15 @@ def __init__(
             **kwargs,
         )
 
+        self._service_client_01_2024_preview = ServiceClient012024Preview(
+            credential=self._credential,
+            subscription_id=self._ws_operation_scope._subscription_id
+            if registry_reference
+            else self._operation_scope._subscription_id,
+            base_url=base_url,
+            **kwargs,
+        )
+
         self._workspaces = WorkspaceOperations(
             self._ws_operation_scope if registry_reference else self._operation_scope,
             self._service_client_08_2023_preview,
@@ -460,6 +470,7 @@ def __init__(
             operation_scope=self._operation_scope,
             operation_config=self._operation_config,
             serviceclient_2023_04_01_preview=self._service_client_04_2023_preview,
+            serviceclient_2024_01_01_preview=self._service_client_01_2024_preview,
             **ops_kwargs,
         )
         self._operation_container.add(AzureMLResourceType.DATASTORE, self._datastores)
@@ -545,6 +556,7 @@ def __init__(
             self._operation_scope,
             self._operation_config,
             self._service_client_10_2021_dataplanepreview if registry_name else self._service_client_04_2023_preview,
+            self._service_client_01_2024_preview,
             self._datastores,
             requests_pipeline=self._requests_pipeline,
             all_operations=self._operation_container,

sdk/ml/azure-ai-ml/azure/ai/ml/operations/_data_operations.py

@@ -5,6 +5,8 @@
 # pylint: disable=protected-access,no-value-for-parameter
 
 import os
+import time
+import uuid
 from contextlib import contextmanager
 from pathlib import Path
 from typing import Any, Dict, Generator, Iterable, List, Optional, Union, cast
@@ -23,6 +25,8 @@
 )
 from azure.ai.ml._restclient.v2023_04_01_preview import AzureMachineLearningWorkspaces as ServiceClient042023_preview
 from azure.ai.ml._restclient.v2023_04_01_preview.models import ListViewType
+from azure.ai.ml._restclient.v2024_01_01_preview import AzureMachineLearningWorkspaces as ServiceClient012024_preview
+from azure.ai.ml._restclient.v2024_01_01_preview.models import ComputeInstanceDataMount
 from azure.ai.ml._scope_dependent_operations import (
     OperationConfig,
     OperationsContainer,
@@ -108,6 +112,7 @@ def __init__(
         operation_scope: OperationScope,
         operation_config: OperationConfig,
         service_client: Union[ServiceClient042023_preview, ServiceClient102021Dataplane],
+        service_client_012024_preview: ServiceClient012024_preview,
         datastore_operations: DatastoreOperations,
         **kwargs: Dict,
     ):
@@ -116,6 +121,7 @@ def __init__(
         self._operation = service_client.data_versions
         self._container_operation = service_client.data_containers
         self._datastore_operation = datastore_operations
+        self._compute_operation = service_client_012024_preview.compute
         self._service_client = service_client
         self._init_kwargs = kwargs
         self._requests_pipeline: HttpPipeline = kwargs.pop("requests_pipeline")
@@ -738,6 +744,95 @@ def share(
         with self._set_registry_client(registry_name):
             return self.create_or_update(data_ref)
 
+    @monitor_with_activity(logger, "data.Mount", ActivityType.PUBLICAPI)
+    @experimental
+    def mount(
+        self,
+        path: str,
+        mount_point: str = None,
+        mode: str = "ro_mount",
+        debug: bool = False,
+        persistent: bool = False,
+        **_kwargs,
+    ) -> None:
+        """Mount a data asset to a local path, so that you can access data inside it
+        under a local path with any tools of your choice.
+
+        :param path: The data asset path to mount, in the form of `azureml:<name>` or `azureml:<name>:<version>`.
+        :type path: str
+        :param mount_point: A local path used as mount point.
+        :type mount_point: str
+        :param mode: Mount mode. Only `ro_mount` (read-only) is supported for data asset mount.
+        :type mode: str
+        :param debug: Whether to enable verbose logging.
+        :type debug: bool
+        :param persistent: Whether to persist the mount after reboot. Applies only when running on Compute Instance,
+                where the 'CI_NAME' environment variable is set."
+        :type persistent: bool
+        :return: None
+        """
+
+        assert mode in ["ro_mount", "rw_mount"], "mode should be either `ro_mount` or `rw_mount`"
+        read_only = mode == "ro_mount"
+        assert read_only, "read-write mount for data asset is not supported yet"
+
+        ci_name = os.environ.get("CI_NAME")
+        assert not persistent or (
+            persistent and ci_name is not None
+        ), "persistent mount is only supported on Compute Instance, where the 'CI_NAME' environment variable is set."
+
+        try:
+            from azureml.dataprep import rslex_fuse_subprocess_wrapper
+        except ImportError as exc:
+            raise Exception(
+                "Mount operations requires package azureml-dataprep-rslex installed. "
+                + "You can install it with Azure ML SDK with `pip install azure-ai-ml[mount]`."
+            ) from exc
+
+        uri = rslex_fuse_subprocess_wrapper.build_data_asset_uri(
+            self._operation_scope._subscription_id, self._resource_group_name, self._workspace_name, path
+        )
+        if persistent and ci_name is not None:
+            mount_name = f"unified_mount_{str(uuid.uuid4()).replace('-', '')}"
+            self._compute_operation.update_data_mounts(
+                self._resource_group_name,
+                self._workspace_name,
+                ci_name,
+                [
+                    ComputeInstanceDataMount(
+                        source=uri,
+                        source_type="URI",
+                        mount_name=mount_name,
+                        mount_action="Mount",
+                        mount_path=mount_point or "",
+                    )
+                ],
+                api_version="2021-01-01",
+            )
+            print(f"Mount requested [name: {mount_name}]. Waiting for completion ...")
+            while True:
+                compute = self._compute_operation.get(self._resource_group_name, self._workspace_name, ci_name)
+                mounts = compute.properties.properties.data_mounts
+                try:
+                    mount = [mount for mount in mounts if mount.mount_name == mount_name][0]
+                    if mount.mount_state == "Mounted":
+                        print(f"Mounted [name: {mount_name}].")
+                        break
+                    if mount.mount_state == "MountRequested":
+                        pass
+                    elif mount.mount_state == "MountFailed":
+                        raise Exception(f"Mount failed [name: {mount_name}]: {mount.error}")
+                    else:
+                        raise Exception(f"Got unexpected mount state [name: {mount_name}]: {mount.mount_state}")
+                except IndexError:
+                    pass
+                time.sleep(5)
+
+        else:
+            rslex_fuse_subprocess_wrapper.start_fuse_mount_subprocess(
+                uri, mount_point, read_only, debug, credential=self._operation._config.credential
+            )
+
     @contextmanager
     # pylint: disable-next=docstring-missing-return,docstring-missing-rtype
     def _set_registry_client(self, registry_name: str) -> Generator:

sdk/ml/azure-ai-ml/azure/ai/ml/operations/_datastore_operations.py

@@ -4,6 +4,8 @@
 
 # pylint: disable=protected-access
 
+import time
+import uuid
 from typing import Dict, Iterable, cast
 
 from marshmallow.exceptions import ValidationError as SchemaValidationError
@@ -12,8 +14,11 @@
 from azure.ai.ml._restclient.v2023_04_01_preview import AzureMachineLearningWorkspaces as ServiceClient042023Preview
 from azure.ai.ml._restclient.v2023_04_01_preview.models import Datastore as DatastoreData
 from azure.ai.ml._restclient.v2023_04_01_preview.models import DatastoreSecrets, NoneDatastoreCredentials
+from azure.ai.ml._restclient.v2024_01_01_preview import AzureMachineLearningWorkspaces as ServiceClient012024Preview
+from azure.ai.ml._restclient.v2024_01_01_preview.models import ComputeInstanceDataMount
 from azure.ai.ml._scope_dependent_operations import OperationConfig, OperationScope, _ScopeDependentOperations
 from azure.ai.ml._telemetry import ActivityType, monitor_with_activity
+from azure.ai.ml._utils._experimental import experimental
 from azure.ai.ml._utils._logger_utils import OpsLogger
 from azure.ai.ml.entities._datastore.datastore import Datastore
 from azure.ai.ml.exceptions import ValidationException
@@ -43,11 +48,13 @@ def __init__(
         operation_scope: OperationScope,
         operation_config: OperationConfig,
         serviceclient_2023_04_01_preview: ServiceClient042023Preview,
-        **kwargs: Dict
+        serviceclient_2024_01_01_preview: ServiceClient012024Preview,
+        **kwargs: Dict,
     ):
         super(DatastoreOperations, self).__init__(operation_scope, operation_config)
         ops_logger.update_info(kwargs)
         self._operation = serviceclient_2023_04_01_preview.datastores
+        self._compute_operation = serviceclient_2024_01_01_preview.compute
         self._credential = serviceclient_2023_04_01_preview._config.credential
         self._init_kwargs = kwargs
 
@@ -81,7 +88,7 @@ def _list_helper(datastore_resource: Datastore, include_secrets: bool) -> Datast
                 resource_group_name=self._operation_scope.resource_group_name,
                 workspace_name=self._workspace_name,
                 cls=lambda objs: [_list_helper(obj, include_secrets) for obj in objs],
-                **self._init_kwargs
+                **self._init_kwargs,
             ),
         )
 
@@ -91,7 +98,7 @@ def _list_secrets(self, name: str) -> DatastoreSecrets:
             name=name,
             resource_group_name=self._operation_scope.resource_group_name,
             workspace_name=self._workspace_name,
-            **self._init_kwargs
+            **self._init_kwargs,
         )
 
     @monitor_with_activity(logger, "Datastore.Delete", ActivityType.PUBLICAPI)
@@ -116,7 +123,7 @@ def delete(self, name: str) -> None:
             name=name,
             resource_group_name=self._operation_scope.resource_group_name,
             workspace_name=self._workspace_name,
-            **self._init_kwargs
+            **self._init_kwargs,
         )
 
     @monitor_with_activity(logger, "Datastore.Get", ActivityType.PUBLICAPI)
@@ -144,7 +151,7 @@ def get(self, name: str, *, include_secrets: bool = False) -> Datastore:
                 name=name,
                 resource_group_name=self._operation_scope.resource_group_name,
                 workspace_name=self._workspace_name,
-                **self._init_kwargs
+                **self._init_kwargs,
             )
             if include_secrets:
                 self._fetch_and_populate_secret(datastore_resource)
@@ -182,7 +189,7 @@ def get_default(self, *, include_secrets: bool = False) -> Datastore:
                 resource_group_name=self._operation_scope.resource_group_name,
                 workspace_name=self._workspace_name,
                 is_default=True,
-                **self._init_kwargs
+                **self._init_kwargs,
             ).next()
             if include_secrets:
                 self._fetch_and_populate_secret(datastore_resource)
@@ -223,3 +230,92 @@ def create_or_update(self, datastore: Datastore) -> Datastore:
                 log_and_raise_error(ex)
             else:
                 raise ex
+
+    @monitor_with_activity(logger, "Datastore.Mount", ActivityType.PUBLICAPI)
+    @experimental
+    def mount(
+        self,
+        path: str,
+        mount_point: str = None,
+        mode: str = "ro_mount",
+        debug: bool = False,
+        persistent: bool = False,
+        **_kwargs,
+    ) -> None:
+        """Mount a datastore to a local path, so that you can access data inside it
+        under a local path with any tools of your choice.
+
+        :param path: The data store path to mount, in the form of `<name>` or `azureml://datastores/<name>`.
+        :type path: str
+        :param mount_point: A local path used as mount point.
+        :type mount_point: str
+        :param mode: Mount mode, either `ro_mount` (read-only) or `rw_mount` (read-write).
+        :type mode: str
+        :param debug: Whether to enable verbose logging.
+        :type debug: bool
+        :param persistent: Whether to persist the mount after reboot. Applies only when running on Compute Instance,
+                where the 'CI_NAME' environment variable is set."
+        :type persistent: bool
+        :return: None
+        """
+
+        assert mode in ["ro_mount", "rw_mount"], "mode should be either `ro_mount` or `rw_mount`"
+        read_only = mode == "ro_mount"
+
+        import os
+
+        ci_name = os.environ.get("CI_NAME")
+        assert not persistent or (
+            persistent and ci_name is not None
+        ), "persistent mount is only supported on Compute Instance, where the 'CI_NAME' environment variable is set."
+
+        try:
+            from azureml.dataprep import rslex_fuse_subprocess_wrapper
+        except ImportError as exc:
+            raise Exception(
+                "Mount operations requires package azureml-dataprep-rslex installed. "
+                + "You can install it with Azure ML SDK with `pip install azure-ai-ml[mount]`."
+            ) from exc
+
+        uri = rslex_fuse_subprocess_wrapper.build_datastore_uri(
+            self._operation_scope._subscription_id, self._resource_group_name, self._workspace_name, path
+        )
+        if persistent and ci_name is not None:
+            mount_name = f"unified_mount_{str(uuid.uuid4()).replace('-', '')}"
+            self._compute_operation.update_data_mounts(
+                self._resource_group_name,
+                self._workspace_name,
+                ci_name,
+                [
+                    ComputeInstanceDataMount(
+                        source=uri,
+                        source_type="URI",
+                        mount_name=mount_name,
+                        mount_action="Mount",
+                        mount_path=mount_point or "",
+                    )
+                ],
+                api_version="2021-01-01",
+            )
+            print(f"Mount requested [name: {mount_name}]. Waiting for completion ...")
+            while True:
+                compute = self._compute_operation.get(self._resource_group_name, self._workspace_name, ci_name)
+                mounts = compute.properties.properties.data_mounts
+                try:
+                    mount = [mount for mount in mounts if mount.mount_name == mount_name][0]
+                    if mount.mount_state == "Mounted":
+                        print(f"Mounted [name: {mount_name}].")
+                        break
+                    if mount.mount_state == "MountRequested":
+                        pass
+                    elif mount.mount_state == "MountFailed":
+                        raise Exception(f"Mount failed [name: {mount_name}]: {mount.error}")
+                    else:
+                        raise Exception(f"Got unexpected mount state [name: {mount_name}]: {mount.mount_state}")
+                except IndexError:
+                    pass
+                time.sleep(5)
+        else:
+            rslex_fuse_subprocess_wrapper.start_fuse_mount_subprocess(
+                uri, mount_point, read_only, debug, credential=self._operation._config.credential
+            )

sdk/ml/azure-ai-ml/cspell.json

@@ -1,6 +1,7 @@
 {
     "version": "0.2",
     "ignoreWords": [
-        "kwoa"
+        "kwoa",
+        "rslex"
     ]
 }

sdk/ml/azure-ai-ml/dev_requirements.txt

@@ -21,4 +21,5 @@ mldesigner
 azure-mgmt-resourcegraph<9.0.0,>=2.0.0
 azure-mgmt-resource<23.0.0,>=3.0.0
 pytest-reportlog
-python-dotenv
+python-dotenv
+azureml-dataprep-rslex>=2.22.0

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

@@ -95,6 +95,11 @@
         "designer": [
             "mldesigner",
         ],
+        # user can run `pip install azure-ai-ml[mount]` to install azureml-dataprep-rslex alone with this package
+        # so user can call data.mount() and datastore.mount() operations supported by it.
+        "mount": [
+            "azureml-dataprep-rslex>=2.22.0",
+        ],
     },
     project_urls={
         "Bug Reports": "https://github.com/Azure/azure-sdk-for-python/issues",

sdk/ml/azure-ai-ml/tests/batch_services/unittests/test_batch_endpoints.py

@@ -53,11 +53,13 @@ def mock_datastore_operations(
     mock_workspace_scope: OperationScope,
     mock_operation_config: OperationConfig,
     mock_aml_services_2023_04_01_preview: Mock,
+    mock_aml_services_2024_01_01_preview: Mock,
 ) -> CodeOperations:
     yield DatastoreOperations(
         operation_scope=mock_workspace_scope,
         operation_config=mock_operation_config,
         serviceclient_2023_04_01_preview=mock_aml_services_2023_04_01_preview,
+        serviceclient_2024_01_01_preview=mock_aml_services_2024_01_01_preview,
     )