tencent-wedata-feature-engineering-dev 0.1.0__py3-none-any.whl

wedata/feature_store/utils/signature_utils.py

@@ -0,0 +1,202 @@
+import logging
+from typing import Any, Dict, Optional
+
+import mlflow
+from mlflow.models import ModelSignature
+from mlflow.types import ColSpec
+from mlflow.types import DataType as MlflowDataType
+from mlflow.types import ParamSchema, Schema
+
+from wedata.feature_store.entities.feature_column_info import FeatureColumnInfo
+from wedata.feature_store.entities.feature_spec import FeatureSpec
+from wedata.feature_store.entities.on_demand_column_info import OnDemandColumnInfo
+from wedata.feature_store.entities.source_data_column_info import SourceDataColumnInfo
+
+_logger = logging.getLogger(__name__)
+
+# Some types (array, map, decimal, timestamp_ntz) are unsupported due to MLflow signatures
+# lacking any equivalent types. We thus cannot construct a ColSpec for any column
+# that uses these types.
+SUPPORTED_TYPE_MAP = {
+    "smallint": MlflowDataType.integer,  # Upcast to integer
+    "int": MlflowDataType.integer,
+    "bigint": MlflowDataType.long,
+    "float": MlflowDataType.float,
+    "double": MlflowDataType.double,
+    "boolean": MlflowDataType.boolean,
+    "date": MlflowDataType.datetime,
+    "timestamp": MlflowDataType.datetime,
+    "string": MlflowDataType.string,
+    "binary": MlflowDataType.binary,
+}
+
+
+def is_unsupported_type(type_str: str):
+    return type_str not in SUPPORTED_TYPE_MAP
+
+
+def convert_spark_data_type_to_mlflow_signature_type(spark_type):
+    return SUPPORTED_TYPE_MAP.get(spark_type)
+
+
+def get_input_schema_from_feature_spec(feature_spec: FeatureSpec) -> Schema:
+    """
+    Produces an MLflow signature schema from a feature spec.
+    Source data columns are marked as required inputs and feature columns
+    (both lookups and on-demand features) are marked as optional inputs.
+
+    :param feature_spec: FeatureSpec object with datatypes for each column.
+    """
+    # If we're missing any data types for any column, we are likely dealing with a
+    # malformed feature spec and should halt signature construction.
+    if any([ci.data_type is None for ci in feature_spec.column_infos]):
+        raise Exception("Training set does not contain column data types.")
+
+    source_data_cols = [
+        ci
+        for ci in feature_spec.column_infos
+        if isinstance(ci.info, SourceDataColumnInfo)
+    ]
+    # Don't create signature if any source data columns (required) are of complex types.
+    if any(
+        [
+            ci.data_type is None or is_unsupported_type(ci.data_type)
+            for ci in source_data_cols
+        ]
+    ):
+        raise Exception(
+            "Input DataFrame contains column data types not supported by "
+            "MLflow model signatures."
+        )
+    required_input_colspecs = [
+        ColSpec(
+            convert_spark_data_type_to_mlflow_signature_type(ci.data_type),
+            ci.info.output_name,
+            required=True,
+        )
+        for ci in source_data_cols
+    ]
+    feature_cols = [
+        ci
+        for ci in feature_spec.column_infos
+        if isinstance(ci.info, (FeatureColumnInfo, OnDemandColumnInfo))
+    ]
+    unsupported_feature_cols = [
+        ci for ci in feature_cols if is_unsupported_type(ci.data_type)
+    ]
+    optional_input_colspecs = [
+        ColSpec(
+            convert_spark_data_type_to_mlflow_signature_type(ci.data_type),
+            ci.output_name,
+            required=True,
+        )
+        for ci in feature_cols
+        if not is_unsupported_type(ci.data_type)
+    ]
+    if unsupported_feature_cols:
+        feat_string = ", ".join(
+            [f"{ci.output_name} ({ci.data_type})" for ci in unsupported_feature_cols]
+        )
+        _logger.warning(
+            f"The following features will not be included in the input schema because their"
+            f" data types are not supported by MLflow model signatures: {feat_string}. "
+            f"These features cannot be overridden during model serving."
+        )
+
+    return Schema(optional_input_colspecs)
+
+
+def get_output_schema_from_labels(label_type_map: Optional[Dict[str, str]]) -> Schema:
+    """
+    Produces an MLflow signature schema from the provided label type map.
+    :param label_type_map: Map label column name -> data type
+    """
+    if not label_type_map:
+        raise Exception("Training set does not contain a label.")
+    if any([is_unsupported_type(dtype) for dtype in label_type_map.values()]):
+        raise Exception(
+            "Labels are of data types not supported by MLflow model signatures."
+        )
+    else:
+        output_colspecs = [
+            ColSpec(
+                convert_spark_data_type_to_mlflow_signature_type(spark_type),
+                col_name,
+                required=True,
+            )
+            for col_name, spark_type in label_type_map.items()
+        ]
+        return Schema(output_colspecs)
+
+
+def get_mlflow_signature_from_feature_spec(
+    feature_spec: FeatureSpec,
+    label_type_map: Optional[Dict[str, str]],
+    override_output_schema: Optional[Schema],
+    params: Optional[Dict[str, Any]] = None,
+) -> Optional[ModelSignature]:
+    """
+    Produce an MLflow signature from a feature spec and label type map.
+    Source data columns are marked as required inputs and feature columns
+    (both lookups and on-demand features) are marked as optional inputs.
+
+    Reads output types from the cached label -> datatype map in the training set.
+    If override_output_schema is provided, it will always be used as the output schema.
+
+    :param feature_spec: FeatureSpec object with datatypes for each column.
+    :param label_type_map: Map of label column name -> datatype
+    :param override_output_schema: User-provided output schema to use if provided.
+    """
+    kwargs = {}
+    kwargs["inputs"] = get_input_schema_from_feature_spec(feature_spec)
+    try:
+        output_schema = override_output_schema or get_output_schema_from_labels(
+            label_type_map
+        )
+        kwargs["outputs"] = output_schema
+    except Exception as e:
+        _logger.warning(f"Could not infer an output schema: {e}")
+
+    if params:
+        try:
+            from mlflow.types.utils import _infer_param_schema
+
+            kwargs["params"] = _infer_param_schema(params)
+        except Exception as e:
+            _logger.warning(f"Could not infer params schema: {e}")
+
+    return mlflow.models.ModelSignature(**kwargs)
+
+
+def drop_signature_inputs_and_invalid_params(signature):
+    """
+    Drop ModelSignature inputs field and invalid params from params field.
+    This is useful for feature store model's raw_model.
+    Feature store model's input schema does not apply to raw_model's input,
+    so we drop the inputs field of raw_model's signature.
+    Feature store model's result_type param enables setting and overriding
+    a default result_type for predictions, but this interferes with params
+    passed to MLflow's predict function, so we drop result_type from
+    the params field of raw_model's signature.
+
+    :param signature: ModelSignature object.
+    """
+    if signature:
+        outputs_schema = signature.outputs
+        params_schema = signature.params if hasattr(signature, "params") else None
+        try:
+            # Only for mlflow>=2.6.0 ModelSignature contains params attribute
+            if params_schema:
+                updated_params_schema = ParamSchema(
+                    [param for param in params_schema if param.name != "result_type"]
+                )
+                return ModelSignature(
+                    outputs=outputs_schema, params=updated_params_schema
+                )
+            if outputs_schema:
+                return ModelSignature(outputs=outputs_schema)
+        except TypeError:
+            _logger.warning(
+                "ModelSignature without inputs is not supported, please upgrade "
+                "mlflow >= 2.7.0 to use the feature."
+            )

wedata/feature_store/utils/topological_sort.py

@@ -0,0 +1,158 @@
+from collections import defaultdict, deque
+from queue import PriorityQueue
+from typing import Callable, Dict, Hashable, List, Optional
+
+__all__ = ["find_cycle", "topological_sort"]
+
+
+class _NodeInfo:
+    def __init__(self, node):
+        self.node = node
+        # number of non-processed predecessors.
+        self.n_blockers = 0
+        # list of nodes that depend on this node.
+        self.successors = []
+
+
+def find_cycle(
+    node_dependencies: Dict[Hashable, List[Hashable]]
+) -> Optional[List[Hashable]]:
+    """
+    Finds a cycle in the node_dependencies graph. Returns a list of node(s) that forms a cycle or
+    None if no cycle can be found.
+    :param node_dependencies: A dict with hashable objects as keys and their list of dependency
+                    nodes as values.
+    """
+    # A stack used to perform DFS on the graph.
+    stack = deque()
+    # Another stack storing the path from root to current node in DFS. Used to detect cycle.
+    backtrack_stack = []
+    # A set of nodes that no cycle can be found starting from these nodes.
+    resolved = set()
+    # Create a copy of the dependency graph with defaultdict for convenience.
+    default_dependency = defaultdict(list)
+    default_dependency.update(node_dependencies)
+
+    # Perform DFS on every node in the graph.
+    for node in node_dependencies.keys():
+        if node in resolved:
+            # Skip a node if it's already resolved.
+            continue
+        # DFS from the node
+        stack.append(node)
+        while stack:
+            top = stack[-1]
+            if top not in backtrack_stack:
+                # First time visiting this node. There will be a second visit after the dependencies
+                # are resolved if it has dependencies.
+                backtrack_stack.append(top)
+            # If not expended after traversing the dependencies, meaning there is no dependency or
+            # all dependencies are resolved.
+            expanded = False
+            for depend in default_dependency[top]:
+                if depend in backtrack_stack:
+                    # found a cycle
+                    index = backtrack_stack.index(depend)
+                    return backtrack_stack[index:]
+                if depend in resolved:
+                    continue
+                # Only adding node to stack. backtrack_stack only contains nodes in the current DFS
+                # path.
+                stack.append(depend)
+                expanded = True
+            if not expanded:
+                stack.pop()
+                resolved.add(top)
+                backtrack_stack.pop()
+    return None
+
+
+def _all_items_in_queue_should_be_grouped(
+    queue: PriorityQueue, should_be_grouped: Callable
+) -> bool:
+    temp = []
+    should_group = True
+    # note: avoid using queue.qsize() because it's not guaranteed to be accurate.
+    while not queue.empty():
+        k, node = queue.get()
+        temp.append((k, node))
+        if not should_be_grouped(node):
+            should_group = False
+    for item in temp:
+        queue.put(item)
+    return should_group
+
+
+def topological_sort(
+    node_dependencies: Dict[Hashable, List[Hashable]],
+    key: Callable = None,
+    should_be_grouped: Callable = None,
+) -> List[Hashable]:
+    """
+    Topological sort the given node_dependencies graph. Returns a sorted list of nodes.
+    :param node_dependencies: A dict with hashable objects as keys and their list of dependency
+                    nodes as values.
+    :param key: a Callable that returns a sort key when called with a hashable object. The key is
+                    used to break ties in topological sorting. An object with smaller key is added
+                    to the result list first.
+    :raises ValueError if a cycle is found in the graph.
+    """
+    # Calling a dedicated find_cycle function to be able to give a detailed error message.
+    cycle = find_cycle(node_dependencies)
+    if cycle is not None:
+        raise ValueError(
+            "Following nodes form a cycle: ",
+            cycle,
+            ". Please resolve any circular dependencies before calling Feature Store.",
+        )
+
+    # A priority-queue storing the nodes whose dependency has been resolved.
+    # priority is determined by the given key function.
+    ready_queue = PriorityQueue()
+    # Map from node to _NodeInfo.
+    nodes = {}
+    if key is None:
+        key = hash  # use the built-in hash function by default
+    if should_be_grouped is None:
+        should_be_grouped = lambda _: False
+    # Perform Kahn's algorithm by traversing the graph starting from nodes without dependency.
+    # Node is removed from its successors' dependency once resolved. And node whose dependency gets
+    # all resolved is added to the priority queue.
+    for node, dependencies in node_dependencies.items():
+        # Initialize the graph to topologically sort based on the input node_dependencies.
+        # All nodes, its successors and number of predecessors should be populated.
+        if node not in nodes:
+            nodes[node] = _NodeInfo(node)
+        for dependency in dependencies:
+            if dependency not in nodes:
+                nodes[dependency] = _NodeInfo(dependency)
+            nodes[dependency].successors.append(node)
+        if len(dependencies):
+            nodes[node].n_blockers = len(dependencies)
+    # Initialize the ready_queue to start traversing the graph from nodes without any dependencies.
+    for node, node_info in nodes.items():
+        if node_info.n_blockers == 0:
+            ready_queue.put((key(node), node))
+    # At the end of the algorithm, result_list will have a topologically sorted listed of nodes.
+    result_list = []
+
+    def process_nodes(node_buffer, queue):
+        for node in node_buffer:
+            result_list.append(node)
+            for successor in nodes[node].successors:
+                s_info = nodes[successor]
+                s_info.n_blockers -= 1
+                if s_info.n_blockers == 0:
+                    queue.put((key(successor), successor))
+
+    while not ready_queue.empty():
+        if _all_items_in_queue_should_be_grouped(ready_queue, should_be_grouped):
+            batch_buffer = []
+            while not ready_queue.empty():
+                _, node = ready_queue.get()
+                batch_buffer.append(node)
+            process_nodes(batch_buffer, ready_queue)
+        else:
+            _, node = ready_queue.get()
+            process_nodes([node], ready_queue)
+    return result_list