path: root/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_component/parallel_component.py

# ---------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# ---------------------------------------------------------

import json
import os
import re
from typing import Any, Dict, List, Optional, Union, cast

from marshmallow import Schema

from azure.ai.ml._restclient.v2022_10_01.models import ComponentVersion
from azure.ai.ml._schema.component.parallel_component import ParallelComponentSchema
from azure.ai.ml.constants._common import COMPONENT_TYPE
from azure.ai.ml.constants._component import NodeType
from azure.ai.ml.entities._job.job_resource_configuration import JobResourceConfiguration
from azure.ai.ml.entities._job.parallel.parallel_task import ParallelTask
from azure.ai.ml.entities._job.parallel.parameterized_parallel import ParameterizedParallel
from azure.ai.ml.entities._job.parallel.retry_settings import RetrySettings
from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationException

from ..._schema import PathAwareSchema
from .._util import validate_attribute_type
from .._validation import MutableValidationResult
from .code import ComponentCodeMixin
from .component import Component


class ParallelComponent(
    Component, ParameterizedParallel, ComponentCodeMixin
):  # pylint: disable=too-many-instance-attributes
    """Parallel component version, used to define a parallel component.

    :param name: Name of the component. Defaults to None
    :type name: str
    :param version: Version of the component. Defaults to None
    :type version: str
    :param description: Description of the component. Defaults to None
    :type description: str
    :param tags: Tag dictionary. Tags can be added, removed, and updated. Defaults to None
    :type tags: dict
    :param display_name: Display name of the component. Defaults to None
    :type display_name: str
    :param retry_settings: parallel component run failed retry. Defaults to None
    :type retry_settings: BatchRetrySettings
    :param logging_level: A string of the logging level name. Defaults to None
    :type logging_level: str
    :param max_concurrency_per_instance: The max parallellism that each compute instance has. Defaults to None
    :type max_concurrency_per_instance: int
    :param error_threshold: The number of item processing failures should be ignored. Defaults to None
    :type error_threshold: int
    :param mini_batch_error_threshold: The number of mini batch processing failures should be ignored. Defaults to None
    :type mini_batch_error_threshold: int
    :param task: The parallel task. Defaults to None
    :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. Defaults to None
        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, 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.
    :type partition_keys: list
    :param input_data: The input data. Defaults to None
    :type input_data: str
    :param resources: Compute Resource configuration for the component. Defaults to None
    :type resources: Union[dict, ~azure.ai.ml.entities.JobResourceConfiguration]
    :param inputs: Inputs of the component. Defaults to None
    :type inputs: dict
    :param outputs: Outputs of the component. Defaults to None
    :type outputs: dict
    :param code: promoted property from task.code
    :type code: str
    :param instance_count: promoted property from resources.instance_count. Defaults to None
    :type instance_count: int
    :param is_deterministic: Whether the parallel component is deterministic. Defaults to True
    :type is_deterministic: bool
    :raises ~azure.ai.ml.exceptions.ValidationException: Raised if ParallelComponent cannot be successfully validated.
        Details will be provided in the error message.
    """

    def __init__(  # pylint: disable=too-many-locals
        self,
        *,
        name: Optional[str] = None,
        version: Optional[str] = None,
        description: Optional[str] = None,
        tags: Optional[Dict[str, Any]] = None,
        display_name: Optional[str] = None,
        retry_settings: Optional[RetrySettings] = 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,
        task: Optional[ParallelTask] = None,
        mini_batch_size: Optional[str] = None,
        partition_keys: Optional[List] = None,
        input_data: Optional[str] = None,
        resources: Optional[JobResourceConfiguration] = None,
        inputs: Optional[Dict] = None,
        outputs: Optional[Dict] = None,
        code: Optional[str] = None,  # promoted property from task.code
        instance_count: Optional[int] = None,  # promoted property from resources.instance_count
        is_deterministic: bool = True,
        **kwargs: Any,
    ):
        # validate init params are valid type
        validate_attribute_type(attrs_to_check=locals(), attr_type_map=self._attr_type_map())

        kwargs[COMPONENT_TYPE] = NodeType.PARALLEL

        super().__init__(
            name=name,
            version=version,
            description=description,
            tags=tags,
            display_name=display_name,
            inputs=inputs,
            outputs=outputs,
            is_deterministic=is_deterministic,
            **kwargs,
        )

        # No validation on value passed here because in pipeline job, required code&environment maybe absent
        # and fill in later with job defaults.
        self.task = task
        self.mini_batch_size: int = 0
        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

        # check mutual exclusivity of promoted properties
        if self.resources is not None and instance_count is not None:
            msg = "instance_count and resources are mutually exclusive"
            raise ValidationException(
                message=msg,
                target=ErrorTarget.COMPONENT,
                no_personal_data_message=msg,
                error_category=ErrorCategory.USER_ERROR,
            )
        self.instance_count = instance_count
        self.code = code

        if mini_batch_size is not None:
            # Convert str to int.
            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:
                self.mini_batch_size = int(mini_batch_size)
            except ValueError as e:
                unit = mini_batch_size[-2:].lower()
                if unit == "kb":
                    self.mini_batch_size = int(mini_batch_size[0:-2]) * 1024
                elif unit == "mb":
                    self.mini_batch_size = int(mini_batch_size[0:-2]) * 1024 * 1024
                elif unit == "gb":
                    self.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

    @property
    def instance_count(self) -> Optional[int]:
        """Return value of promoted property resources.instance_count.

        :return: Value of resources.instance_count.
        :rtype: Optional[int]
        """
        return self.resources.instance_count if self.resources and not isinstance(self.resources, dict) else None

    @instance_count.setter
    def instance_count(self, value: int) -> None:
        """Set the value of the promoted property resources.instance_count.

        :param value: The value to set for resources.instance_count.
        :type value: int
        """
        if not value:
            return
        if not self.resources:
            self.resources = JobResourceConfiguration(instance_count=value)
        else:
            if not isinstance(self.resources, dict):
                self.resources.instance_count = value

    @property
    def code(self) -> Optional[str]:
        """Return value of promoted property task.code, which is a local or
        remote path pointing at source code.

        :return: Value of task.code.
        :rtype: Optional[str]
        """
        return self.task.code if self.task else None

    @code.setter
    def code(self, value: str) -> None:
        """Set the value of the promoted property task.code.

        :param value: The value to set for task.code.
        :type value: str
        """
        if not value:
            return
        if not self.task:
            self.task = ParallelTask(code=value)
        else:
            self.task.code = value

    def _to_ordered_dict_for_yaml_dump(self) -> Dict:
        """Dump the component content into a sorted yaml string.

        :return: The ordered dict
        :rtype: Dict
        """

        obj: dict = super()._to_ordered_dict_for_yaml_dump()
        # dict dumped base on schema will transfer code to an absolute path, while we want to keep its original value
        if self.code and isinstance(self.code, str):
            obj["task"]["code"] = self.code
        return obj

    @property
    def environment(self) -> Optional[str]:
        """Return value of promoted property task.environment, indicate the
        environment that training job will run in.

        :return: Value of task.environment.
        :rtype: Optional[Environment, str]
        """
        if self.task:
            return cast(Optional[str], self.task.environment)
        return None

    @environment.setter
    def environment(self, value: str) -> None:
        """Set the value of the promoted property task.environment.

        :param value: The value to set for task.environment.
        :type value: str
        """
        if not value:
            return
        if not self.task:
            self.task = ParallelTask(environment=value)
        else:
            self.task.environment = value

    def _customized_validate(self) -> MutableValidationResult:
        validation_result = super()._customized_validate()
        self._append_diagnostics_and_check_if_origin_code_reliable_for_local_path_validation(validation_result)
        return validation_result

    @classmethod
    def _attr_type_map(cls) -> dict:
        return {
            "retry_settings": (dict, RetrySettings),
            "task": (dict, ParallelTask),
            "logging_level": str,
            "max_concurrency_per_instance": int,
            "input_data": str,
            "error_threshold": int,
            "mini_batch_error_threshold": int,
            "code": (str, os.PathLike),
            "resources": (dict, JobResourceConfiguration),
        }

    def _to_rest_object(self) -> ComponentVersion:
        rest_object = super()._to_rest_object()
        # schema required list while backend accept json string
        if self.partition_keys:
            rest_object.properties.component_spec["partition_keys"] = json.dumps(self.partition_keys)
        return rest_object

    @classmethod
    def _from_rest_object_to_init_params(cls, obj: ComponentVersion) -> Dict:
        # schema required list while backend accept json string
        # update rest obj as it will be
        partition_keys = obj.properties.component_spec.get("partition_keys", None)
        if partition_keys:
            obj.properties.component_spec["partition_keys"] = json.loads(partition_keys)
        res: dict = super()._from_rest_object_to_init_params(obj)
        return res

    @classmethod
    def _create_schema_for_validation(cls, context: Any) -> Union[PathAwareSchema, Schema]:
        return ParallelComponentSchema(context=context)

    def __str__(self) -> str:
        try:
            toYaml: str = self._to_yaml()
            return toYaml
        except BaseException:  # pylint: disable=W0718
            toStr: str = super(ParallelComponent, self).__str__()
            return toStr