[ML][Pipelines] feat: support concurrent artifact downloading (Azure#29597)

* feat: support concurrent artifact downloading

* refactor: unify cache folder

* feat: enable concurrent artifact uploading within a component

* fix: typo

* feat: resolve all leaf nodes concurrently

* fix: handle loop node

* revert: revert leaf node first optimization

* fix: pylint

Author: elliotzh

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/_additional_includes.py

@@ -6,12 +6,16 @@
 import shutil
 import tempfile
 import zipfile
+from concurrent.futures import ThreadPoolExecutor
+from multiprocessing import cpu_count
 from pathlib import Path
 from typing import Union
 
 import yaml
 
+from ..._artifacts._constants import PROCESSES_PER_CORE
 from ..._utils._asset_utils import IgnoreFile, traverse_directory
+from ..._utils.utils import is_concurrent_component_registration_enabled, is_private_preview_enabled
 from ...entities._util import _general_copy
 from ...entities._validation import MutableValidationResult, _ValidationResultBuilder
 from ._artifact_cache import ArtifactCache
@@ -269,6 +273,50 @@ def _artifact_validate_result(self):
             self._load_artifact_additional_includes()
         return self.__artifact_validate_result
 
+    @classmethod
+    def merge_local_path_to_additional_includes(cls, local_path, config_info, conflict_files):
+        file_name = Path(local_path).name
+        conflicts = conflict_files.get(file_name, set())
+        conflicts.add(config_info)
+        conflict_files[file_name] = conflicts
+
+    @classmethod
+    def _get_artifacts_by_config(cls, artifact_config):
+        artifact_cache = ArtifactCache()
+        if any(item not in artifact_config for item in ["feed", "name", "version"]):
+            raise RuntimeError("Feed, name and version are required for artifacts config.")
+        return artifact_cache.get(
+            organization=artifact_config.get("organization", None),
+            project=artifact_config.get("project", None),
+            feed=artifact_config["feed"],
+            name=artifact_config["name"],
+            version=artifact_config["version"],
+            scope=artifact_config.get("scope", "organization"),
+            resolve=True,
+        )
+
+    def _resolve_additional_include_config(self, additional_include_config):
+        result = []
+        if isinstance(additional_include_config, dict) and additional_include_config.get("type") == ARTIFACT_KEY:
+            try:
+                # Get the artifacts package from devops to the local
+                artifact_path = self._get_artifacts_by_config(additional_include_config)
+                for item in os.listdir(artifact_path):
+                    config_info = (
+                        f"{additional_include_config['name']}:{additional_include_config['version']} in "
+                        f"{additional_include_config['feed']}"
+                    )
+                    result.append((os.path.join(artifact_path, item), config_info))
+            except Exception as e:  # pylint: disable=broad-except
+                self._artifact_validate_result.append_error(message=e.args[0])
+        elif isinstance(additional_include_config, str):
+            result.append((additional_include_config, additional_include_config))
+        else:
+            self._artifact_validate_result.append_error(
+                message=f"Unexpected format in additional_includes, {additional_include_config}"
+            )
+        return result
+
     def _load_artifact_additional_includes(self):
         """
         Load the additional includes by yaml format.
@@ -290,57 +338,34 @@ def _load_artifact_additional_includes(self):
         :return additional_includes: Path list of additional_includes
         :rtype additional_includes: List[str]
         """
-        additional_includes, conflict_files = [], {}
         self.__artifact_validate_result = _ValidationResultBuilder.success()
 
-        def merge_local_path_to_additional_includes(local_path, config_info):
-            additional_includes.append(local_path)
-            file_name = Path(local_path).name
-            conflicts = conflict_files.get(file_name, set())
-            conflicts.add(config_info)
-            conflict_files[file_name] = conflicts
-
-        def get_artifacts_by_config(artifact_config):
-            artifact_cache = ArtifactCache()
-            if any(item not in artifact_config for item in ["feed", "name", "version"]):
-                raise RuntimeError("Feed, name and version are required for artifacts config.")
-            artifact_path = artifact_cache.get(
-                organization=artifact_config.get("organization", None),
-                project=artifact_config.get("project", None),
-                feed=artifact_config["feed"],
-                name=artifact_config["name"],
-                version=artifact_config["version"],
-                scope=artifact_config.get("scope", "organization"),
-                resolve=True,
-            )
-            return artifact_path
-
         # Load the artifacts config from additional_includes
         with open(self._additional_includes_file_path) as f:
             additional_includes_configs = yaml.safe_load(f)
             additional_includes_configs = additional_includes_configs.get(ADDITIONAL_INCLUDES_KEY, [])
 
-        for additional_include in additional_includes_configs:
-            if isinstance(additional_include, dict) and additional_include.get("type") == ARTIFACT_KEY:
-                try:
-                    # Get the artifacts package from devops to the local
-                    artifact_path = get_artifacts_by_config(additional_include)
-                    for item in os.listdir(artifact_path):
-                        config_info = (
-                            f"{additional_include['name']}:{additional_include['version']} in "
-                            f"{additional_include['feed']}"
-                        )
-                        merge_local_path_to_additional_includes(
-                            local_path=os.path.join(artifact_path, item), config_info=config_info
+        additional_includes, conflict_files = [], {}
+        num_threads = int(cpu_count()) * PROCESSES_PER_CORE
+        if (
+            len(additional_includes_configs) > 1
+            and is_concurrent_component_registration_enabled()
+            and is_private_preview_enabled()
+        ):
+            with ThreadPoolExecutor(max_workers=num_threads) as executor:
+                for result in executor.map(self._resolve_additional_include_config, additional_includes_configs):
+                    for local_path, config_info in result:
+                        additional_includes.append(local_path)
+                        self.merge_local_path_to_additional_includes(
+                            local_path=local_path, config_info=config_info, conflict_files=conflict_files
                         )
-                except Exception as e:  # pylint: disable=broad-except
-                    self._artifact_validate_result.append_error(message=e.args[0])
-            elif isinstance(additional_include, str):
-                merge_local_path_to_additional_includes(local_path=additional_include, config_info=additional_include)
-            else:
-                self._artifact_validate_result.append_error(
-                    message=f"Unexpected format in additional_includes, {additional_include}"
-                )
+        else:
+            for result in map(self._resolve_additional_include_config, additional_includes_configs):
+                for local_path, config_info in result:
+                    additional_includes.append(local_path)
+                    self.merge_local_path_to_additional_includes(
+                        local_path=local_path, config_info=config_info, conflict_files=conflict_files
+                    )
 
         # Check the file conflict in local path and artifact package.
         conflict_files = {k: v for k, v in conflict_files.items() if len(v) > 1}

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

@@ -10,22 +10,20 @@
 from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
-from functools import partial
 from pathlib import Path
-from typing import List, Dict, Optional, Union, Callable
+from typing import Callable, Dict, List, Optional, Union
 
 from azure.ai.ml._utils._asset_utils import get_object_hash
 from azure.ai.ml._utils.utils import (
-    is_on_disk_cache_enabled,
     is_concurrent_component_registration_enabled,
+    is_on_disk_cache_enabled,
     is_private_preview_enabled,
     write_to_shared_file,
 )
-from azure.ai.ml.constants._common import AzureMLResourceType, AZUREML_COMPONENT_REGISTRATION_MAX_WORKERS
+from azure.ai.ml.constants._common import AZUREML_COMPONENT_REGISTRATION_MAX_WORKERS, AzureMLResourceType
 from azure.ai.ml.entities import Component
 from azure.ai.ml.entities._builders import BaseNode
 
-
 logger = logging.getLogger(__name__)
 
 _ANONYMOUS_HASH_PREFIX = "anonymous-component-"
@@ -47,6 +45,9 @@ class _CacheContent:
     on_disk_hash: Optional[str] = None
     arm_id: Optional[str] = None
 
+    def update_on_disk_hash(self):
+        self.on_disk_hash = CachedNodeResolver.calc_on_disk_hash_for_component(self.component_ref, self.in_memory_hash)
+
 
 class CachedNodeResolver(object):
     """Class to resolve component in nodes with cached component resolution results.
@@ -176,7 +177,7 @@ def _get_in_memory_hash_for_component(component: Component) -> str:
         return _ANONYMOUS_HASH_PREFIX + component._get_anonymous_hash()  # pylint: disable=protected-access
 
     @staticmethod
-    def _get_on_disk_hash_for_component(component: Component, in_memory_hash: str) -> str:
+    def calc_on_disk_hash_for_component(component: Component, in_memory_hash: str) -> str:
         """Get a hash for a component.
 
         This function will calculate the hash based on the component's code folder if the component has code, so it's
@@ -261,21 +262,23 @@ def _save_to_on_disk_cache(self, on_disk_hash: str, arm_id: str) -> None:
 
     def _resolve_cache_contents(self, cache_contents_to_resolve: List[_CacheContent], resolver):
         """Resolve all components to resolve and save the results in cache."""
-        _components = list(map(lambda x: x.component_ref, cache_contents_to_resolve))
-        _map_func = partial(resolver, azureml_type=AzureMLResourceType.COMPONENT)
 
-        if len(_components) > 1 and is_concurrent_component_registration_enabled() and is_private_preview_enabled():
+        def _map_func(_cache_content: _CacheContent):
+            _cache_content.arm_id = resolver(_cache_content.component_ref, azureml_type=AzureMLResourceType.COMPONENT)
+            if is_on_disk_cache_enabled() and is_private_preview_enabled():
+                self._save_to_on_disk_cache(_cache_content.on_disk_hash, _cache_content.arm_id)
+
+        if (
+            len(cache_contents_to_resolve) > 1
+            and is_concurrent_component_registration_enabled()
+            and is_private_preview_enabled()
+        ):
             # given deduplication has already been done, we can safely assume that there is no
             # conflict in concurrent local cache access
             with ThreadPoolExecutor(max_workers=self._get_component_registration_max_workers()) as executor:
-                resolution_results = executor.map(_map_func, _components)
+                list(executor.map(_map_func, cache_contents_to_resolve))
         else:
-            resolution_results = map(_map_func, _components)
-
-        for cache_content, resolution_results in zip(cache_contents_to_resolve, resolution_results):
-            cache_content.arm_id = resolution_results
-            if is_on_disk_cache_enabled() and is_private_preview_enabled():
-                self._save_to_on_disk_cache(cache_content.on_disk_hash, cache_content.arm_id)
+            list(map(_map_func, cache_contents_to_resolve))
 
     def _prepare_items_to_resolve(self):
         """Pop all nodes in self._nodes_to_resolve to prepare cache contents to resolve and nodes to resolve. Nodes in
@@ -306,11 +309,21 @@ def _resolve_cache_contents_from_disk(self, cache_contents_to_resolve: List[_Cac
         """Check on-disk cache to resolve cache contents in cache_contents_to_resolve and return unresolved cache
         contents."""
         # Note that we should recalculate the hash based on code for local cache, as
-        # we can't assume that the code folder won't change among dependency resolution
-        for cache_content in cache_contents_to_resolve:
-            cache_content.on_disk_hash = self._get_on_disk_hash_for_component(
-                cache_content.component_ref, cache_content.in_memory_hash
-            )
+        # we can't assume that the code folder won't change among dependency
+        # On-disk hash calculation can be slow as it involved data copying and artifact downloading.
+        # It is thread-safe given:
+        # 1. artifact downloading is thread-safe as we have a lock in ArtifactCache
+        # 2. data copying is thread-safe as there is only read operation on source folder
+        #    and target folder is unique for each thread
+        if (
+            len(cache_contents_to_resolve) > 1
+            and is_concurrent_component_registration_enabled()
+            and is_private_preview_enabled()
+        ):
+            with ThreadPoolExecutor(max_workers=self._get_component_registration_max_workers()) as executor:
+                executor.map(_CacheContent.update_on_disk_hash, cache_contents_to_resolve)
+        else:
+            list(map(_CacheContent.update_on_disk_hash, cache_contents_to_resolve))
 
         left_cache_contents_to_resolve = []
         # need to deduplicate disk hash first if concurrent resolution is enabled

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

@@ -6,20 +6,19 @@
 
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import TYPE_CHECKING, Dict, Optional, Union
+from typing import TYPE_CHECKING, Dict, Union
 from urllib.parse import urlparse
 
 import yaml
 
-from azure.ai.ml._restclient.v2023_02_01_preview.operations import (  # pylint: disable = unused-import
+from .._artifacts._artifact_utilities import get_datastore_info, get_storage_client
+from .._restclient.v2023_02_01_preview.operations import (  # pylint: disable = unused-import
     FeaturesetContainersOperations,
     FeaturesetVersionsOperations,
     FeaturestoreEntityContainersOperations,
     FeaturestoreEntityVersionsOperations,
 )
-from azure.ai.ml._artifacts._artifact_utilities import get_datastore_info, get_storage_client
-from azure.ai.ml.operations._datastore_operations import DatastoreOperations
-
+from ..operations._datastore_operations import DatastoreOperations
 from ._storage_utils import AzureMLDatastorePathUri
 from .utils import load_yaml
 

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

@@ -8,7 +8,7 @@
 import types
 from functools import partial
 from inspect import Parameter, signature
-from typing import Callable, Dict, Iterable, Optional, Union, List, Any
+from typing import Any, Callable, Dict, Iterable, List, Optional, Union
 
 from azure.ai.ml._restclient.v2021_10_01_dataplanepreview import (
     AzureMachineLearningWorkspaces as ServiceClient102021Dataplane,
@@ -21,12 +21,7 @@
     OperationScope,
     _ScopeDependentOperations,
 )
-
-from azure.ai.ml._telemetry import (
-    ActivityType,
-    monitor_with_activity,
-    monitor_with_telemetry_mixin,
-)
+from azure.ai.ml._telemetry import ActivityType, monitor_with_activity, monitor_with_telemetry_mixin
 from azure.ai.ml._utils._asset_utils import (
     _archive_or_restore,
     _create_or_update_autoincrement,
@@ -45,8 +40,8 @@
 )
 from azure.ai.ml.entities import Component, ValidationResult
 from azure.ai.ml.exceptions import ComponentException, ErrorCategory, ErrorTarget, ValidationException
-from .._utils._cache_utils import CachedNodeResolver
 
+from .._utils._cache_utils import CachedNodeResolver
 from .._utils._experimental import experimental
 from .._utils.utils import is_data_binding_expression
 from ..entities._builders import BaseNode

sdk/ml/azure-ai-ml/tests/internal_utils/unittests/test_cache_utils.py

@@ -5,7 +5,6 @@
 
 import mock
 import pytest
-
 from azure.ai.ml import MLClient, load_job
 from azure.ai.ml._utils._cache_utils import CachedNodeResolver
 from azure.ai.ml.entities import Component, PipelineJob
@@ -24,7 +23,7 @@ def _mock_resolver(component: Union[str, Component], azureml_type: str) -> str:
     @staticmethod
     def _get_cache_path(component: Component, resolver: CachedNodeResolver) -> Path:
         in_memory_hash = resolver._get_in_memory_hash_for_component(component)
-        on_disk_hash = resolver._get_on_disk_hash_for_component(component=component, in_memory_hash=in_memory_hash)
+        on_disk_hash = resolver.calc_on_disk_hash_for_component(component=component, in_memory_hash=in_memory_hash)
         return resolver._get_on_disk_cache_path(on_disk_hash)
 
     @staticmethod