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

# ---------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# ---------------------------------------------------------
from pathlib import Path
from typing import Any, Dict, NoReturn, Optional, Union, cast

from marshmallow import Schema

from azure.ai.ml._schema.component.data_transfer_component import (
    DataTransferCopyComponentSchema,
    DataTransferExportComponentSchema,
    DataTransferImportComponentSchema,
)
from azure.ai.ml._utils._experimental import experimental
from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, COMPONENT_TYPE, AssetTypes
from azure.ai.ml.constants._component import DataTransferTaskType, ExternalDataType, NodeType
from azure.ai.ml.entities._inputs_outputs.external_data import Database, FileSystem
from azure.ai.ml.entities._inputs_outputs.output import Output
from azure.ai.ml.entities._validation.core import MutableValidationResult
from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationErrorType, ValidationException

from ..._schema import PathAwareSchema
from .._util import convert_ordered_dict_to_dict, validate_attribute_type
from .component import Component


class DataTransferComponent(Component):
    """DataTransfer component version, used to define a data transfer component.

    :param task: Task type in the data transfer component. Possible values are "copy_data",
                 "import_data", and "export_data".
    :type task: str
    :param inputs: Mapping of input data bindings used in the job.
    :type inputs: dict
    :param outputs: Mapping of output data bindings used in the job.
    :type outputs: dict
    :param kwargs: Additional parameters for the data transfer component.
    :raises ~azure.ai.ml.exceptions.ValidationException: Raised if the component cannot be successfully validated.
        Details will be provided in the error message.
    """

    def __init__(
        self,
        *,
        task: Optional[str] = None,
        inputs: Optional[Dict] = None,
        outputs: Optional[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[COMPONENT_TYPE] = NodeType.DATA_TRANSFER
        # Set default base path
        if BASE_PATH_CONTEXT_KEY not in kwargs:
            kwargs[BASE_PATH_CONTEXT_KEY] = Path(".")

        super().__init__(
            inputs=inputs,
            outputs=outputs,
            **kwargs,
        )
        self._task = task

    @classmethod
    def _attr_type_map(cls) -> dict:
        return {}

    @property
    def task(self) -> Optional[str]:
        """Task type of the component.

        :return: Task type of the component.
        :rtype: str
        """
        return self._task

    def _to_dict(self) -> Dict:
        return cast(
            dict,
            convert_ordered_dict_to_dict({**self._other_parameter, **super(DataTransferComponent, self)._to_dict()}),
        )

    def __str__(self) -> str:
        try:
            _toYaml: str = self._to_yaml()
            return _toYaml
        except BaseException:  # pylint: disable=W0718
            _toStr: str = super(DataTransferComponent, self).__str__()
            return _toStr

    @classmethod
    def _build_source_sink(cls, io_dict: Union[Dict, Database, FileSystem]) -> Union[Database, FileSystem]:
        component_io: Union[Database, FileSystem] = Database()

        if isinstance(io_dict, Database):
            component_io = Database()
        elif isinstance(io_dict, FileSystem):
            component_io = FileSystem()
        else:
            if isinstance(io_dict, dict):
                data_type = io_dict.pop("type", None)
                if data_type == ExternalDataType.DATABASE:
                    component_io = Database()
                elif data_type == ExternalDataType.FILE_SYSTEM:
                    component_io = FileSystem()
                else:
                    msg = "Type in source or sink only support {} and {}, currently got {}."
                    raise ValidationException(
                        message=msg.format(
                            ExternalDataType.DATABASE,
                            ExternalDataType.FILE_SYSTEM,
                            data_type,
                        ),
                        no_personal_data_message=msg.format(
                            ExternalDataType.DATABASE,
                            ExternalDataType.FILE_SYSTEM,
                            "data_type",
                        ),
                        target=ErrorTarget.COMPONENT,
                        error_category=ErrorCategory.USER_ERROR,
                        error_type=ValidationErrorType.INVALID_VALUE,
                    )
            else:
                msg = "Source or sink only support dict, Database and FileSystem"
                raise ValidationException(
                    message=msg,
                    no_personal_data_message=msg,
                    target=ErrorTarget.COMPONENT,
                    error_category=ErrorCategory.USER_ERROR,
                    error_type=ValidationErrorType.INVALID_VALUE,
                )

        return component_io


@experimental
class DataTransferCopyComponent(DataTransferComponent):
    """DataTransfer copy component version, used to define a data transfer copy component.

    :param data_copy_mode: Data copy mode in the copy task.
                           Possible values are "merge_with_overwrite" and "fail_if_conflict".
    :type data_copy_mode: str
    :param inputs: Mapping of input data bindings used in the job.
    :type inputs: dict
    :param outputs: Mapping of output data bindings used in the job.
    :type outputs: dict
    :param kwargs: Additional parameters for the data transfer copy component.
    :raises ~azure.ai.ml.exceptions.ValidationException: Raised if the component cannot be successfully validated.
        Details will be provided in the error message.
    """

    def __init__(
        self,
        *,
        data_copy_mode: Optional[str] = None,
        inputs: Optional[Dict] = None,
        outputs: Optional[Dict] = None,
        **kwargs: Any,
    ) -> None:
        kwargs["task"] = DataTransferTaskType.COPY_DATA
        super().__init__(
            inputs=inputs,
            outputs=outputs,
            **kwargs,
        )

        self._data_copy_mode = data_copy_mode

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

    @property
    def data_copy_mode(self) -> Optional[str]:
        """Data copy mode of the component.

        :return: Data copy mode of the component.
        :rtype: str
        """
        return self._data_copy_mode

    def _customized_validate(self) -> MutableValidationResult:
        validation_result = super(DataTransferCopyComponent, self)._customized_validate()
        validation_result.merge_with(self._validate_input_output_mapping())
        return validation_result

    def _validate_input_output_mapping(self) -> MutableValidationResult:
        validation_result = self._create_empty_validation_result()
        inputs_count = len(self.inputs)
        outputs_count = len(self.outputs)
        if outputs_count != 1:
            msg = "Only support single output in {}, but there're {} outputs."
            validation_result.append_error(
                message=msg.format(DataTransferTaskType.COPY_DATA, outputs_count),
                yaml_path="outputs",
            )
        else:
            input_type = None
            output_type = None
            if inputs_count == 1:
                for _, input_data in self.inputs.items():
                    input_type = input_data.type
                for _, output_data in self.outputs.items():
                    output_type = output_data.type
                if input_type is None or output_type is None or input_type != output_type:
                    msg = "Input type {} doesn't exactly match with output type {} in task {}"
                    validation_result.append_error(
                        message=msg.format(input_type, output_type, DataTransferTaskType.COPY_DATA),
                        yaml_path="outputs",
                    )
            elif inputs_count > 1:
                for _, output_data in self.outputs.items():
                    output_type = output_data.type
                if output_type is None or output_type != AssetTypes.URI_FOLDER:
                    msg = "output type {} need to be {} in task {}"
                    validation_result.append_error(
                        message=msg.format(
                            output_type,
                            AssetTypes.URI_FOLDER,
                            DataTransferTaskType.COPY_DATA,
                        ),
                        yaml_path="outputs",
                    )
            else:
                msg = "Inputs must be set in task {}."
                validation_result.append_error(
                    message=msg.format(DataTransferTaskType.COPY_DATA),
                    yaml_path="inputs",
                )
        return validation_result


@experimental
class DataTransferImportComponent(DataTransferComponent):
    """DataTransfer import component version, used to define a data transfer import component.

    :param source: The data source of the file system or database.
    :type source: dict
    :param outputs: Mapping of output data bindings used in the job.
                    Default value is an output port with the key "sink" and the type "mltable".
    :type outputs: dict
    :param kwargs: Additional parameters for the data transfer import component.
    :raises ~azure.ai.ml.exceptions.ValidationException: Raised if the component cannot be successfully validated.
        Details will be provided in the error message.
    """

    def __init__(
        self,
        *,
        source: Optional[Dict] = None,
        outputs: Optional[Dict] = None,
        **kwargs: Any,
    ) -> None:
        outputs = outputs or {"sink": Output(type=AssetTypes.MLTABLE)}
        kwargs["task"] = DataTransferTaskType.IMPORT_DATA
        super().__init__(
            outputs=outputs,
            **kwargs,
        )

        source = source if source else {}
        self.source = self._build_source_sink(source)

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

    # pylint: disable-next=docstring-missing-param
    def __call__(self, *args: Any, **kwargs: Any) -> NoReturn:
        """Call ComponentVersion as a function and get a Component object."""

        msg = "DataTransfer component is not callable for import task."
        raise ValidationException(
            message=msg,
            no_personal_data_message=msg,
            target=ErrorTarget.COMPONENT,
            error_category=ErrorCategory.USER_ERROR,
        )


@experimental
class DataTransferExportComponent(DataTransferComponent):
    """DataTransfer export component version, used to define a data transfer export component.

    :param sink: The sink of external data and databases.
    :type sink: Union[Dict, Database, FileSystem]
    :param inputs: Mapping of input data bindings used in the job.
    :type inputs: dict
    :param kwargs: Additional parameters for the data transfer export component.
    :raises ~azure.ai.ml.exceptions.ValidationException: Raised if the component cannot be successfully validated.
        Details will be provided in the error message.
    """

    def __init__(
        self,
        *,
        inputs: Optional[Dict] = None,
        sink: Optional[Dict] = None,
        **kwargs: Any,
    ) -> None:
        kwargs["task"] = DataTransferTaskType.EXPORT_DATA
        super().__init__(
            inputs=inputs,
            **kwargs,
        )

        sink = sink if sink else {}
        self.sink = self._build_source_sink(sink)

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

    # pylint: disable-next=docstring-missing-param
    def __call__(self, *args: Any, **kwargs: Any) -> NoReturn:
        """Call ComponentVersion as a function and get a Component object."""

        msg = "DataTransfer component is not callable for export task."
        raise ValidationException(
            message=msg,
            no_personal_data_message=msg,
            target=ErrorTarget.COMPONENT,
            error_category=ErrorCategory.USER_ERROR,
        )