# --------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # --------------------------------------------------------- # pylint: disable=protected-access import copy import json import logging import os import re from enum import Enum from typing import Any, Dict, List, Optional, Tuple, Type, Union, cast from marshmallow import INCLUDE, Schema from azure.ai.ml._schema.core.fields import NestedField, UnionField from azure.ai.ml._schema.job.identity import AMLTokenIdentitySchema, ManagedIdentitySchema, UserIdentitySchema from azure.ai.ml.entities._credentials import ( AmlTokenConfiguration, ManagedIdentityConfiguration, UserIdentityConfiguration, _BaseJobIdentityConfiguration, ) from azure.ai.ml.entities._job.job import Job from azure.ai.ml.entities._job.parallel.run_function import RunFunction from azure.ai.ml.entities._job.pipeline._io import NodeOutput from azure.ai.ml.exceptions import MlException from ..._schema import PathAwareSchema from ..._utils.utils import is_data_binding_expression from ...constants._common import ARM_ID_PREFIX from ...constants._component import NodeType from .._component.component import Component from .._component.flow import FlowComponent from .._component.parallel_component import ParallelComponent from .._inputs_outputs import Input, Output from .._job.job_resource_configuration import JobResourceConfiguration from .._job.parallel.parallel_job import ParallelJob from .._job.parallel.parallel_task import ParallelTask from .._job.parallel.retry_settings import RetrySettings from .._job.pipeline._io import NodeWithGroupInputMixin from .._util import convert_ordered_dict_to_dict, get_rest_dict_for_node_attrs, validate_attribute_type from .base_node import BaseNode module_logger = logging.getLogger(__name__) class Parallel(BaseNode, NodeWithGroupInputMixin): # pylint: disable=too-many-instance-attributes """Base class for parallel node, used for parallel component version consumption. You should not instantiate this class directly. Instead, you should create from builder function: parallel. :param component: Id or instance of the parallel component/job to be run for the step :type component: ~azure.ai.ml.entities._component.parallel_component.parallelComponent :param name: Name of the parallel :type name: str :param description: Description of the commad :type description: str :param tags: Tag dictionary. Tags can be added, removed, and updated :type tags: dict[str, str] :param properties: The job property dictionary :type properties: dict[str, str] :param display_name: Display name of the job :type display_name: str :param retry_settings: Parallel job run failed retry :type retry_settings: BatchRetrySettings :param logging_level: A string of the logging level name :type logging_level: str :param max_concurrency_per_instance: The max parallellism that each compute instance has :type max_concurrency_per_instance: int :param error_threshold: The number of item processing failures should be ignored :type error_threshold: int :param mini_batch_error_threshold: The number of mini batch processing failures should be ignored :type mini_batch_error_threshold: int :param task: The parallel task :type task: ParallelTask :param mini_batch_size: For FileDataset input, this field is the number of files a user script can process in one run() call. For TabularDataset input, this field is the approximate size of data the user script can process in one run() call. Example values are 1024, 1024KB, 10MB, and 1GB. (optional, default value is 10 files for FileDataset and 1MB for TabularDataset.) This value could be set through PipelineParameter :type mini_batch_size: str :param partition_keys: The keys used to partition dataset into mini-batches. If specified, the data with the same key will be partitioned into the same mini-batch. If both partition_keys and mini_batch_size are specified, the partition keys will take effect. The input(s) must be partitioned dataset(s), and the partition_keys must be a subset of the keys of every input dataset for this to work. :keyword identity: The identity that the command job will use while running on compute. :paramtype identity: Optional[Union[ dict[str, str], ~azure.ai.ml.entities.ManagedIdentityConfiguration, ~azure.ai.ml.entities.AmlTokenConfiguration, ~azure.ai.ml.entities.UserIdentityConfiguration]] :type partition_keys: List :param input_data: The input data :type input_data: str :param inputs: Inputs of the component/job :type inputs: dict :param outputs: Outputs of the component/job :type outputs: dict """ # pylint: disable=too-many-statements def __init__( self, *, component: Union[ParallelComponent, str], compute: Optional[str] = None, inputs: Optional[Dict[str, Union[NodeOutput, Input, str, bool, int, float, Enum]]] = None, outputs: Optional[Dict[str, Union[str, Output, "Output"]]] = None, retry_settings: Optional[Union[RetrySettings, Dict[str, str]]] = None, logging_level: Optional[str] = None, max_concurrency_per_instance: Optional[int] = None, error_threshold: Optional[int] = None, mini_batch_error_threshold: Optional[int] = None, input_data: Optional[str] = None, task: Optional[Union[ParallelTask, RunFunction, Dict]] = None, partition_keys: Optional[List] = None, mini_batch_size: Optional[Union[str, int]] = None, resources: Optional[JobResourceConfiguration] = None, environment_variables: Optional[Dict] = None, identity: Optional[ Union[ManagedIdentityConfiguration, AmlTokenConfiguration, UserIdentityConfiguration, Dict] ] = None, **kwargs: Any, ) -> None: # validate init params are valid type validate_attribute_type(attrs_to_check=locals(), attr_type_map=self._attr_type_map()) kwargs.pop("type", None) if isinstance(component, FlowComponent): # make input definition fit actual inputs for flow component with component._inputs._fit_inputs(inputs): # type: ignore[attr-defined] BaseNode.__init__( self, type=NodeType.PARALLEL, component=component, inputs=inputs, outputs=outputs, compute=compute, **kwargs, ) else: BaseNode.__init__( self, type=NodeType.PARALLEL, component=component, inputs=inputs, outputs=outputs, compute=compute, **kwargs, ) # init mark for _AttrDict self._init = True self._task = task if ( mini_batch_size is not None and not isinstance(mini_batch_size, int) and not is_data_binding_expression(mini_batch_size) ): """Convert str to int.""" # pylint: disable=pointless-string-statement pattern = re.compile(r"^\d+([kKmMgG][bB])*$") if not pattern.match(mini_batch_size): raise ValueError(r"Parameter mini_batch_size must follow regex rule ^\d+([kKmMgG][bB])*$") try: mini_batch_size = int(mini_batch_size) except ValueError as e: if not isinstance(mini_batch_size, int): unit = mini_batch_size[-2:].lower() if unit == "kb": mini_batch_size = int(mini_batch_size[0:-2]) * 1024 elif unit == "mb": mini_batch_size = int(mini_batch_size[0:-2]) * 1024 * 1024 elif unit == "gb": mini_batch_size = int(mini_batch_size[0:-2]) * 1024 * 1024 * 1024 else: raise ValueError("mini_batch_size unit must be kb, mb or gb") from e self.mini_batch_size = mini_batch_size self.partition_keys = partition_keys self.input_data = input_data self._retry_settings = retry_settings self.logging_level = logging_level self.max_concurrency_per_instance = max_concurrency_per_instance self.error_threshold = error_threshold self.mini_batch_error_threshold = mini_batch_error_threshold self._resources = resources self.environment_variables = {} if environment_variables is None else environment_variables self._identity = identity if isinstance(self.component, ParallelComponent): self.resources = cast(JobResourceConfiguration, self.resources) or cast( JobResourceConfiguration, copy.deepcopy(self.component.resources) ) # TODO: Bug Item number: 2897665 self.retry_settings = self.retry_settings or copy.deepcopy(self.component.retry_settings) # type: ignore self.input_data = self.input_data or self.component.input_data self.max_concurrency_per_instance = ( self.max_concurrency_per_instance or self.component.max_concurrency_per_instance ) self.mini_batch_error_threshold = ( self.mini_batch_error_threshold or self.component.mini_batch_error_threshold ) self.mini_batch_size = self.mini_batch_size or self.component.mini_batch_size self.partition_keys = self.partition_keys or copy.deepcopy(self.component.partition_keys) if not self.task: self.task = self.component.task # task.code is based on self.component.base_path self._base_path = self.component.base_path self._init = False @classmethod def _get_supported_outputs_types(cls) -> Tuple: return str, Output @property def retry_settings(self) -> RetrySettings: """Get the retry settings for the parallel job. :return: The retry settings for the parallel job. :rtype: ~azure.ai.ml.entities._job.parallel.retry_settings.RetrySettings """ return self._retry_settings # type: ignore @retry_settings.setter def retry_settings(self, value: Union[RetrySettings, Dict]) -> None: """Set the retry settings for the parallel job. :param value: The retry settings for the parallel job. :type value: ~azure.ai.ml.entities._job.parallel.retry_settings.RetrySettings or dict """ if isinstance(value, dict): value = RetrySettings(**value) self._retry_settings = value @property def resources(self) -> Optional[JobResourceConfiguration]: """Get the resource configuration for the parallel job. :return: The resource configuration for the parallel job. :rtype: ~azure.ai.ml.entities._job.job_resource_configuration.JobResourceConfiguration """ return self._resources @resources.setter def resources(self, value: Union[JobResourceConfiguration, Dict]) -> None: """Set the resource configuration for the parallel job. :param value: The resource configuration for the parallel job. :type value: ~azure.ai.ml.entities._job.job_resource_configuration.JobResourceConfiguration or dict """ if isinstance(value, dict): value = JobResourceConfiguration(**value) self._resources = value @property def identity( self, ) -> Optional[Union[ManagedIdentityConfiguration, AmlTokenConfiguration, UserIdentityConfiguration, Dict]]: """The identity that the job will use while running on compute. :return: The identity that the job will use while running on compute. :rtype: Optional[Union[~azure.ai.ml.ManagedIdentityConfiguration, ~azure.ai.ml.AmlTokenConfiguration, ~azure.ai.ml.UserIdentityConfiguration]] """ return self._identity @identity.setter def identity( self, value: Union[Dict, ManagedIdentityConfiguration, AmlTokenConfiguration, UserIdentityConfiguration, None], ) -> None: """Sets the identity that the job will use while running on compute. :param value: The identity that the job will use while running on compute. :type value: Union[dict[str, str], ~azure.ai.ml.ManagedIdentityConfiguration, ~azure.ai.ml.AmlTokenConfiguration, ~azure.ai.ml.UserIdentityConfiguration] """ if isinstance(value, dict): identity_schema = UnionField( [ NestedField(ManagedIdentitySchema, unknown=INCLUDE), NestedField(AMLTokenIdentitySchema, unknown=INCLUDE), NestedField(UserIdentitySchema, unknown=INCLUDE), ] ) value = identity_schema._deserialize(value=value, attr=None, data=None) self._identity = value @property def component(self) -> Union[str, ParallelComponent]: """Get the component of the parallel job. :return: The component of the parallel job. :rtype: str or ~azure.ai.ml.entities._component.parallel_component.ParallelComponent """ res: Union[str, ParallelComponent] = self._component return res @property def task(self) -> Optional[ParallelTask]: """Get the parallel task. :return: The parallel task. :rtype: ~azure.ai.ml.entities._job.parallel.parallel_task.ParallelTask """ return self._task # type: ignore @task.setter def task(self, value: Union[ParallelTask, Dict]) -> None: """Set the parallel task. :param value: The parallel task. :type value: ~azure.ai.ml.entities._job.parallel.parallel_task.ParallelTask or dict """ # base path should be reset if task is set via sdk self._base_path: Optional[Union[str, os.PathLike]] = None if isinstance(value, dict): value = ParallelTask(**value) self._task = value def _set_base_path(self, base_path: Optional[Union[str, os.PathLike]]) -> None: if self._base_path: return super(Parallel, self)._set_base_path(base_path) def set_resources( self, *, instance_type: Optional[Union[str, List[str]]] = None, instance_count: Optional[int] = None, properties: Optional[Dict] = None, docker_args: Optional[str] = None, shm_size: Optional[str] = None, # pylint: disable=unused-argument **kwargs: Any, ) -> None: """Set the resources for the parallel job. :keyword instance_type: The instance type or a list of instance types used as supported by the compute target. :paramtype instance_type: Union[str, List[str]] :keyword instance_count: The number of instances or nodes used by the compute target. :paramtype instance_count: int :keyword properties: The property dictionary for the resources. :paramtype properties: dict :keyword docker_args: Extra arguments to pass to the Docker run command. :paramtype docker_args: str :keyword shm_size: Size of the Docker container's shared memory block. :paramtype shm_size: str """ if self.resources is None: self.resources = JobResourceConfiguration() if instance_type is not None: self.resources.instance_type = instance_type if instance_count is not None: self.resources.instance_count = instance_count if properties is not None: self.resources.properties = properties if docker_args is not None: self.resources.docker_args = docker_args if shm_size is not None: self.resources.shm_size = shm_size # Save the resources to internal component as well, otherwise calling sweep() will loose the settings if isinstance(self.component, Component): self.component.resources = self.resources @classmethod def _attr_type_map(cls) -> dict: return { "component": (str, ParallelComponent, FlowComponent), "retry_settings": (dict, RetrySettings), "resources": (dict, JobResourceConfiguration), "task": (dict, ParallelTask), "logging_level": str, "max_concurrency_per_instance": (str, int), "error_threshold": (str, int), "mini_batch_error_threshold": (str, int), "environment_variables": dict, } def _to_job(self) -> ParallelJob: return ParallelJob( name=self.name, display_name=self.display_name, description=self.description, tags=self.tags, properties=self.properties, compute=self.compute, resources=self.resources, partition_keys=self.partition_keys, mini_batch_size=self.mini_batch_size, task=self.task, retry_settings=self.retry_settings, input_data=self.input_data, logging_level=self.logging_level, identity=self.identity, max_concurrency_per_instance=self.max_concurrency_per_instance, error_threshold=self.error_threshold, mini_batch_error_threshold=self.mini_batch_error_threshold, environment_variables=self.environment_variables, inputs=self._job_inputs, outputs=self._job_outputs, ) def _parallel_attr_to_dict(self, attr: str, base_type: Type) -> dict: # Convert parallel attribute to dict rest_attr = {} parallel_attr = getattr(self, attr) if parallel_attr is not None: if isinstance(parallel_attr, base_type): rest_attr = parallel_attr._to_dict() elif isinstance(parallel_attr, dict): rest_attr = parallel_attr else: msg = f"Expecting {base_type} for {attr}, got {type(parallel_attr)} instead." raise MlException(message=msg, no_personal_data_message=msg) # TODO: Bug Item number: 2897665 res: dict = convert_ordered_dict_to_dict(rest_attr) # type: ignore return res @classmethod def _picked_fields_from_dict_to_rest_object(cls) -> List[str]: return [ "type", "resources", "error_threshold", "mini_batch_error_threshold", "environment_variables", "max_concurrency_per_instance", "task", "input_data", ] def _to_rest_object(self, **kwargs: Any) -> dict: rest_obj: Dict = super(Parallel, self)._to_rest_object(**kwargs) rest_obj.update( convert_ordered_dict_to_dict( { "componentId": self._get_component_id(), "retry_settings": get_rest_dict_for_node_attrs(self.retry_settings), "logging_level": self.logging_level, "mini_batch_size": self.mini_batch_size, "partition_keys": ( json.dumps(self.partition_keys) if self.partition_keys is not None else self.partition_keys ), "identity": get_rest_dict_for_node_attrs(self.identity), "resources": get_rest_dict_for_node_attrs(self.resources), } ) ) return rest_obj @classmethod def _from_rest_object_to_init_params(cls, obj: dict) -> Dict: obj = super()._from_rest_object_to_init_params(obj) # retry_settings if "retry_settings" in obj and obj["retry_settings"]: obj["retry_settings"] = RetrySettings._from_dict(obj["retry_settings"]) if "task" in obj and obj["task"]: obj["task"] = ParallelTask._from_dict(obj["task"]) task_code = obj["task"].code task_env = obj["task"].environment # remove azureml: prefix in code and environment which is added in _to_rest_object if task_code and isinstance(task_code, str) and task_code.startswith(ARM_ID_PREFIX): obj["task"].code = task_code[len(ARM_ID_PREFIX) :] if task_env and isinstance(task_env, str) and task_env.startswith(ARM_ID_PREFIX): obj["task"].environment = task_env[len(ARM_ID_PREFIX) :] if "resources" in obj and obj["resources"]: obj["resources"] = JobResourceConfiguration._from_rest_object(obj["resources"]) if "partition_keys" in obj and obj["partition_keys"]: obj["partition_keys"] = json.dumps(obj["partition_keys"]) if "identity" in obj and obj["identity"]: obj["identity"] = _BaseJobIdentityConfiguration._from_rest_object(obj["identity"]) return obj def _build_inputs(self) -> Dict: inputs = super(Parallel, self)._build_inputs() built_inputs = {} # Validate and remove non-specified inputs for key, value in inputs.items(): if value is not None: built_inputs[key] = value return built_inputs @classmethod def _create_schema_for_validation(cls, context: Any) -> Union[PathAwareSchema, Schema]: from azure.ai.ml._schema.pipeline import ParallelSchema return ParallelSchema(context=context) # pylint: disable-next=docstring-missing-param def __call__(self, *args: Any, **kwargs: Any) -> "Parallel": """Call Parallel as a function will return a new instance each time. :return: A Parallel node :rtype: Parallel """ if isinstance(self._component, Component): # call this to validate inputs node: Parallel = self._component(*args, **kwargs) # merge inputs for name, original_input in self.inputs.items(): if name not in kwargs: # use setattr here to make sure owner of input won't change setattr(node.inputs, name, original_input._data) # get outputs for name, original_output in self.outputs.items(): # use setattr here to make sure owner of input won't change if not isinstance(original_output, str): setattr(node.outputs, name, original_output._data) self._refine_optional_inputs_with_no_value(node, kwargs) # set default values: compute, environment_variables, outputs node._name = self.name node.compute = self.compute node.tags = self.tags node.display_name = self.display_name node.mini_batch_size = self.mini_batch_size node.partition_keys = self.partition_keys node.logging_level = self.logging_level node.max_concurrency_per_instance = self.max_concurrency_per_instance node.error_threshold = self.error_threshold # deep copy for complex object node.retry_settings = copy.deepcopy(self.retry_settings) node.input_data = self.input_data node.task = copy.deepcopy(self.task) node._base_path = self.base_path node.resources = copy.deepcopy(self.resources) node.environment_variables = copy.deepcopy(self.environment_variables) node.identity = copy.deepcopy(self.identity) return node msg = f"Parallel can be called as a function only when referenced component is {type(Component)}, \ currently got {self._component}." raise MlException(message=msg, no_personal_data_message=msg) @classmethod def _load_from_dict(cls, data: Dict, context: Dict, additional_message: str, **kwargs: Any) -> "Job": raise NotImplementedError()