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

# ---------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# ---------------------------------------------------------
import contextlib
import json
import os
from collections import defaultdict
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, Generator, List, Optional, Tuple, Union

import yaml  # type: ignore[import]
from marshmallow import EXCLUDE, Schema, ValidationError

from azure.ai.ml.constants._common import (
    BASE_PATH_CONTEXT_KEY,
    COMPONENT_TYPE,
    PROMPTFLOW_AZUREML_OVERRIDE_KEY,
    SOURCE_PATH_CONTEXT_KEY,
    AssetTypes,
    SchemaUrl,
)
from azure.ai.ml.constants._component import ComponentParameterTypes, NodeType

from ..._restclient.v2022_10_01.models import ComponentVersion
from ..._schema import PathAwareSchema
from ..._schema.component.flow import FlowComponentSchema, FlowSchema, RunSchema
from ...exceptions import ErrorCategory, ErrorTarget, ValidationException
from .. import Environment
from .._inputs_outputs import GroupInput, Input, Output
from ._additional_includes import AdditionalIncludesMixin
from .component import Component

# avoid circular import error
if TYPE_CHECKING:
    from azure.ai.ml.entities._builders.parallel import Parallel

# pylint: disable=protected-access


class _FlowPortNames:
    """Common yaml fields.

    Common yaml fields are used to define the common fields in yaml files. It can be one of the following values: type,
    name, $schema.
    """

    DATA = "data"
    RUN_OUTPUTS = "run_outputs"
    CONNECTIONS = "connections"

    FLOW_OUTPUTS = "flow_outputs"
    DEBUG_INFO = "debug_info"


class _FlowComponentPortDict(dict):
    def __init__(self, ports: Dict):
        self._allow_update_item = True
        super().__init__()
        for input_port_name, input_port in ports.items():
            self[input_port_name] = input_port
        self._allow_update_item = False

    def __setitem__(self, key: Any, value: Any) -> None:
        if not self._allow_update_item:
            raise RuntimeError("Ports of flow component are not editable.")
        super().__setitem__(key, value)

    def __delitem__(self, key: Any) -> None:
        if not self._allow_update_item:
            raise RuntimeError("Ports of flow component are not editable.")
        super().__delitem__(key)


class FlowComponentInputDict(_FlowComponentPortDict):
    """Input port dictionary for FlowComponent, with fixed input ports."""

    def __init__(self) -> None:
        super().__init__(
            {
                _FlowPortNames.CONNECTIONS: GroupInput(values={}, _group_class=None),
                _FlowPortNames.DATA: Input(type=AssetTypes.URI_FOLDER, optional=False),
                _FlowPortNames.FLOW_OUTPUTS: Input(type=AssetTypes.URI_FOLDER, optional=True),
            }
        )

    @contextlib.contextmanager
    def _fit_inputs(self, inputs: Optional[Dict]) -> Generator:
        """Add dynamic input ports to the input port dictionary.
        Input ports of a flow component include:
        1. data: required major uri_folder input
        2. run_output: optional uri_folder input
        3. connections.xxx.xxx: group of string parameters, first layer key can be any node name,
           but we won't resolve the exact keys in SDK
        4. xxx: input_mapping parameters, key can be any node name, but we won't resolve the exact keys in SDK

        #3 will be grouped into connections, we make it a fixed group input port.
        #4 are dynamic input ports, we will add them temporarily in this context manager and remove them
        after the context manager is finished.

        :param inputs: The dynamic input to fit.
        :type inputs: Dict[str, Any]
        :return: None
        :rtype: None
        """
        dynamic_columns_mapping_keys = []
        dynamic_connections_inputs = defaultdict(list)
        from azure.ai.ml.entities._job.pipeline._io import _GroupAttrDict
        from azure.ai.ml.entities._job.pipeline._io.mixin import flatten_dict

        flattened_inputs = flatten_dict(inputs, _GroupAttrDict, allow_dict_fields=[_FlowPortNames.CONNECTIONS])

        for flattened_input_key in flattened_inputs:
            if flattened_input_key.startswith(f"{_FlowPortNames.CONNECTIONS}."):
                if flattened_input_key.count(".") != 2:
                    raise ValidationException(
                        message="flattened connection input prot name must be "
                        "in the format of connections.<node_name>.<port_name>, "
                        "but got %s" % flattened_input_key,
                        no_personal_data_message="flattened connection input prot name must be in the format of "
                        "connections.<node_name>.<port_name>",
                        target=ErrorTarget.COMPONENT,
                        error_category=ErrorCategory.USER_ERROR,
                    )
                _, node_name, param_name = flattened_input_key.split(".")
                dynamic_connections_inputs[node_name].append(param_name)
                continue
            if flattened_input_key not in self:
                dynamic_columns_mapping_keys.append(flattened_input_key)

        self._allow_update_item = True
        for flattened_input_key in dynamic_columns_mapping_keys:
            self[flattened_input_key] = Input(type=ComponentParameterTypes.STRING, optional=True)
        if dynamic_connections_inputs:
            self[_FlowPortNames.CONNECTIONS] = GroupInput(
                values={
                    node_name: GroupInput(
                        values={
                            parameter_name: Input(
                                type=ComponentParameterTypes.STRING,
                            )
                            for parameter_name in param_names
                        },
                        _group_class=None,
                    )
                    for node_name, param_names in dynamic_connections_inputs.items()
                },
                _group_class=None,
            )
        self._allow_update_item = False

        yield

        self._allow_update_item = True
        for flattened_input_key in dynamic_columns_mapping_keys:
            del self[flattened_input_key]
        self[_FlowPortNames.CONNECTIONS] = GroupInput(values={}, _group_class=None)
        self._allow_update_item = False


class FlowComponentOutputDict(_FlowComponentPortDict):
    """Output port dictionary for FlowComponent, with fixed output ports."""

    def __init__(self) -> None:
        super().__init__(
            {
                _FlowPortNames.FLOW_OUTPUTS: Output(type=AssetTypes.URI_FOLDER),
                _FlowPortNames.DEBUG_INFO: Output(type=AssetTypes.URI_FOLDER),
            }
        )


class FlowComponent(Component, AdditionalIncludesMixin):
    """Flow component version, used to define a Flow Component or Job.

    :keyword name: The name of the Flow job or component.
    :type name: Optional[str]
    :keyword version: The version of the Flow job or component.
    :type version: Optional[str]
    :keyword description: The description of the component. Defaults to None.
    :type description: Optional[str]
    :keyword tags: Tag dictionary. Tags can be added, removed, and updated. Defaults to None.
    :type tags: Optional[dict]
    :keyword display_name: The display name of the component.
    :type display_name: Optional[str]
    :keyword flow: The path to the flow directory or flow definition file. Defaults to None and base path of this
        component will be used as flow directory.
    :type flow: Optional[Union[str, Path]]
    :keyword column_mappings: The column mapping for the flow. Defaults to None.
    :type column_mapping: Optional[dict[str, str]]
    :keyword variant: The variant of the flow. Defaults to None.
    :type variant: Optional[str]
    :keyword connections: The connections for the flow. Defaults to None.
    :type connections: Optional[dict[str, dict[str, str]]]
    :keyword environment_variables: The environment variables for the flow. Defaults to None.
    :type environment_variables: Optional[dict[str, str]]
    :keyword environment: The environment for the flow component. Defaults to None.
    :type environment: Optional[Union[str, Environment])
    :keyword is_deterministic: Specifies whether the Flow will return the same output given the same input.
        Defaults to True. When True, if a Flow (component) is deterministic and has been run before in the
        current workspace with the same input and settings, it will reuse results from a previous submitted job
        when used as a node or step in a pipeline. In that scenario, no compute resources will be used.
    :type is_deterministic: Optional[bool]
    :keyword additional_includes: A list of shared additional files to be included in the component. Defaults to None.
    :type additional_includes: Optional[list[str]]
    :keyword properties: The job property dictionary. Defaults to None.
    :type properties: Optional[dict[str, str]]
    :raises ~azure.ai.ml.exceptions.ValidationException: Raised if FlowComponent cannot be successfully validated.
        Details will be provided in the error message.
    """

    def __init__(
        self,
        *,
        name: Optional[str] = None,
        version: Optional[str] = None,
        description: Optional[str] = None,
        tags: Optional[Dict] = None,
        display_name: Optional[str] = None,
        flow: Optional[Union[str, Path]] = None,
        column_mapping: Optional[Dict[str, str]] = None,
        variant: Optional[str] = None,
        connections: Optional[Dict[str, Dict[str, str]]] = None,
        environment_variables: Optional[Dict[str, str]] = None,
        environment: Optional[Union[str, Environment]] = None,
        is_deterministic: bool = True,
        additional_includes: Optional[List] = None,
        properties: Optional[Dict] = None,
        **kwargs: Any,
    ) -> None:
        # validate init params are valid type
        kwargs[COMPONENT_TYPE] = NodeType.FLOW_PARALLEL

        # always use flow directory as base path
        # Note: we suppose that there is no relative path in run.yaml other than flow.
        #   If there are any, we will need to rebase them so that they have the same base path as attributes in
        #   flow.dag.yaml
        flow_dir, self._flow = self._get_flow_definition(
            flow=flow,
            base_path=kwargs.pop(BASE_PATH_CONTEXT_KEY, Path.cwd()),
            source_path=kwargs.get(SOURCE_PATH_CONTEXT_KEY, None),
        )
        kwargs[BASE_PATH_CONTEXT_KEY] = flow_dir

        super().__init__(
            name=name or self._normalize_component_name(flow_dir.name),
            version=version or "1",
            description=description,
            tags=tags,
            display_name=display_name,
            inputs={},
            outputs={},
            is_deterministic=is_deterministic,
            properties=properties,
            **kwargs,
        )
        self._environment = environment
        self._column_mapping = column_mapping or {}
        self._variant = variant
        self._connections = connections or {}

        self._inputs = FlowComponentInputDict()
        self._outputs = FlowComponentOutputDict()

        if flow:
            # file existence has been checked in _get_flow_definition
            # we don't need to rebase additional_includes as we have updated base_path
            with open(Path(self.base_path, self._flow), "r", encoding="utf-8") as f:
                flow_content = yaml.safe_load(f.read())
            additional_includes = flow_content.get("additional_includes", None)
            # environment variables in run.yaml have higher priority than those in flow.dag.yaml
            self._environment_variables = flow_content.get("environment_variables", {})
            self._environment_variables.update(environment_variables or {})
        else:
            self._environment_variables = environment_variables or {}

        self._additional_includes = additional_includes or []

        # unlike other Component, code is a private property in FlowComponent and
        # will be used to store the arm id of the created code before constructing rest object
        # we haven't used self.flow directly as self.flow can be a path to the flow dag yaml file instead of a directory
        self._code_arm_id: Optional[str] = None

    # region valid properties
    @property
    def flow(self) -> str:
        """The path to the flow definition file relative to the flow directory.

        :rtype: str
        """
        return self._flow

    @property
    def environment(self) -> Optional[Union[str, Environment]]:
        """The environment for the flow component. Defaults to None.

        :rtype: Union[str, Environment])
        """
        return self._environment

    @environment.setter
    def environment(self, value: Union[str, Environment]) -> None:
        """The environment for the flow component. Defaults to None.

        :param value: The column mapping for the flow.
        :type value: Union[str, Environment])
        """
        self._environment = value

    @property
    def column_mapping(self) -> Dict[str, str]:
        """The column mapping for the flow. Defaults to None.

        :rtype: Dict[str, str]
        """
        return self._column_mapping

    @column_mapping.setter
    def column_mapping(self, value: Optional[Dict[str, str]]) -> None:
        """
        The column mapping for the flow. Defaults to None.

        :param value: The column mapping for the flow.
        :type value: Optional[Dict[str, str]]
        """
        self._column_mapping = value or {}

    @property
    def variant(self) -> Optional[str]:
        """The variant of the flow. Defaults to None.

        :rtype: Optional[str]
        """
        return self._variant

    @variant.setter
    def variant(self, value: Optional[str]) -> None:
        """The variant of the flow. Defaults to None.

        :param value: The variant of the flow.
        :type value: Optional[str]
        """
        self._variant = value

    @property
    def connections(self) -> Dict[str, Dict[str, str]]:
        """The connections for the flow. Defaults to None.

        :rtype: Dict[str, Dict[str, str]]
        """
        return self._connections

    @connections.setter
    def connections(self, value: Optional[Dict[str, Dict[str, str]]]) -> None:
        """
        The connections for the flow. Defaults to None.

        :param value: The connections for the flow.
        :type value: Optional[Dict[str, Dict[str, str]]]
        """
        self._connections = value or {}

    @property
    def environment_variables(self) -> Dict[str, str]:
        """The environment variables for the flow. Defaults to None.

        :rtype: Dict[str, str]
        """
        return self._environment_variables

    @environment_variables.setter
    def environment_variables(self, value: Optional[Dict[str, str]]) -> None:
        """The environment variables for the flow. Defaults to None.

        :param value: The environment variables for the flow.
        :type value: Optional[Dict[str, str]]
        """
        self._environment_variables = value or {}

    @property
    def additional_includes(self) -> List:
        """A list of shared additional files to be included in the component. Defaults to None.

        :rtype: List
        """
        return self._additional_includes

    @additional_includes.setter
    def additional_includes(self, value: Optional[List]) -> None:
        """A list of shared additional files to be included in the component. Defaults to None.
        All local additional includes should be relative to the flow directory.

        :param value: A list of shared additional files to be included in the component.
        :type value: Optional[List]
        """
        self._additional_includes = value or []

    # endregion

    @classmethod
    def _normalize_component_name(cls, value: str) -> str:
        return value.replace("-", "_")

    # region Component
    @classmethod
    def _from_rest_object_to_init_params(cls, obj: ComponentVersion) -> Dict:
        raise RuntimeError("FlowComponent does not support loading from REST object.")

    def _to_rest_object(self) -> ComponentVersion:
        rest_obj = super()._to_rest_object()
        rest_obj.properties.component_spec["code"] = self._code_arm_id
        rest_obj.properties.component_spec["flow_file_name"] = self._flow
        return rest_obj

    def _func(self, **kwargs: Any) -> "Parallel":  # pylint: disable=invalid-overridden-method
        from azure.ai.ml.entities._builders.parallel import Parallel

        with self._inputs._fit_inputs(kwargs):  # type: ignore[attr-defined]
            # pylint: disable=not-callable
            return super()._func(**kwargs)  # type: ignore

    @classmethod
    def _get_flow_definition(
        cls,
        base_path: Path,
        *,
        flow: Optional[Union[str, os.PathLike]] = None,
        source_path: Optional[Union[str, os.PathLike]] = None,
    ) -> Tuple[Path, str]:
        """
        Get the path to the flow directory and the file name of the flow dag yaml file.
        If flow is not specified, we will assume that the source_path is the path to the flow dag yaml file.
        If flow is specified, it can be either a path to the flow dag yaml file or a path to the flow directory.
        If flow is a path to the flow directory, we will assume that the flow dag yaml file is named flow.dag.yaml.

        :param base_path: The base path of the flow component.
        :type base_path: Path
        :keyword flow: The path to the flow directory or flow definition file. Defaults to None and base path of this
            component will be used as flow directory.
        :type flow: Optional[Union[str, Path]]
        :keyword source_path: The source path of the flow component, should be path to the flow dag yaml file
            if specified.
        :type source_path: Optional[Union[str, os.PathLike]]
        :return: The path to the flow directory and the file name of the flow dag yaml file.
        :rtype: Tuple[Path, str]
        """
        flow_file_name = "flow.dag.yaml"

        if flow is None and source_path is None:
            raise cls._create_validation_error(
                message="Either flow or source_path must be specified.",
                no_personal_data_message="Either flow or source_path must be specified.",
            )

        if flow is None:
            # Flow component must be created with a local yaml file, so no need to check if source_path exists
            if isinstance(source_path, (os.PathLike, str)):
                flow_file_name = os.path.basename(source_path)
            return Path(base_path), flow_file_name

        flow_path = Path(flow)
        if not flow_path.is_absolute():
            # if flow_path points to a symlink, we still use the parent of the symlink as origin code
            flow_path = Path(base_path, flow)

        if flow_path.is_dir() and (flow_path / flow_file_name).is_file():
            return flow_path, flow_file_name

        if flow_path.is_file():
            return flow_path.parent, flow_path.name

        raise cls._create_validation_error(
            message="Flow path must be a directory containing flow.dag.yaml or a file, but got %s" % flow_path,
            no_personal_data_message="Flow path must be a directory or a file",
        )

    # endregion

    # region SchemaValidatableMixin
    @classmethod
    def _load_with_schema(
        cls, data: Any, *, context: Optional[Any] = None, raise_original_exception: bool = False, **kwargs: Any
    ) -> Any:
        # FlowComponent should be loaded with FlowSchema or FlowRunSchema instead of FlowComponentSchema
        context = context or {BASE_PATH_CONTEXT_KEY: Path.cwd()}
        _schema = data.get("$schema", None)
        if _schema == SchemaUrl.PROMPTFLOW_RUN:
            schema = RunSchema(context=context)
        elif _schema == SchemaUrl.PROMPTFLOW_FLOW:
            schema = FlowSchema(context=context)
        else:
            raise cls._create_validation_error(
                message="$schema must be specified correctly for loading component from flow, but got %s" % _schema,
                no_personal_data_message="$schema must be specified for loading component from flow",
            )

        # unlike other component, we should ignore unknown fields in flow to keep init_params clean and avoid
        # too much understanding of flow.dag.yaml & run.yaml
        kwargs["unknown"] = EXCLUDE
        try:
            loaded_dict = schema.load(data, **kwargs)
        except ValidationError as e:
            if raise_original_exception:
                raise e
            msg = "Trying to load data with schema failed. Data:\n%s\nError: %s" % (
                json.dumps(data, indent=4) if isinstance(data, dict) else data,
                json.dumps(e.messages, indent=4),
            )
            raise cls._create_validation_error(
                message=msg,
                no_personal_data_message=str(e),
            ) from e
        loaded_dict.update(loaded_dict.pop(PROMPTFLOW_AZUREML_OVERRIDE_KEY, {}))
        return loaded_dict

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

    # endregion

    # region AdditionalIncludesMixin
    def _get_origin_code_value(self) -> Union[str, os.PathLike, None]:
        if self._code_arm_id:
            return self._code_arm_id
        res: Union[str, os.PathLike, None] = self.base_path
        return res

    def _fill_back_code_value(self, value: str) -> None:
        self._code_arm_id = value

    @contextlib.contextmanager
    def _try_build_local_code(self) -> Generator:
        # false-positive by pylint, hence disable it
        # (https://github.com/pylint-dev/pylint/blob/main/doc/data/messages
        # /c/contextmanager-generator-missing-cleanup/details.rst)
        with super()._try_build_local_code() as code:  # pylint:disable=contextmanager-generator-missing-cleanup
            if not code or not code.path:
                yield code
                return

            if not (Path(code.path) / ".promptflow" / "flow.tools.json").is_file():
                raise self._create_validation_error(
                    message="Flow component must be created with a ./promptflow/flow.tools.json, "
                    "please run `pf flow validate` to generate it or skip it in your ignore file.",
                    no_personal_data_message="Flow component must be created with a ./promptflow/flow.tools.json, "
                    "please run `pf flow validate` to generate it or skip it in your ignore file.",
                )
            # TODO: should we remove additional includes from flow.dag.yaml? for now we suppose it will be removed
            #  by mldesigner compile if needed

            yield code

    # endregion