rasa-pro 3.12.0.dev1__py3-none-any.whl

rasa/engine/training/components.py

@@ -0,0 +1,176 @@
+from __future__ import annotations
+from typing import Any, Dict, Optional, Text, Type
+import dataclasses
+import uuid
+
+from rasa.engine.caching import Cacheable, TrainingCache
+from rasa.engine.graph import ExecutionContext, GraphComponent, SchemaNode
+from rasa.engine.storage.resource import Resource
+from rasa.engine.storage.storage import ModelStorage
+from rasa.engine.training import fingerprinting
+
+
+class PrecomputedValueProvider(GraphComponent):
+    """Holds the precomputed values of a `GraphNode` from a previous training.
+
+    Pre-computed values can either be
+    - values loaded from cache
+    - values which were provided during the fingerprint run by input nodes
+    """
+
+    def __init__(self, output: Cacheable):
+        """Initializes a `PrecomputedValueProvider`.
+
+        Args:
+            output: The precomputed output to return.
+        """
+        self._output = output
+
+    @classmethod
+    def create(
+        cls,
+        config: Dict[Text, Any],
+        model_storage: ModelStorage,
+        resource: Resource,
+        execution_context: ExecutionContext,
+    ) -> PrecomputedValueProvider:
+        """Creates instance (see parent class for full docstring)."""
+        return cls(output=config["output"])
+
+    def get_value(self) -> Cacheable:
+        """Returns the precomputed output."""
+        return self._output
+
+    @classmethod
+    def replace_schema_node(cls, node: SchemaNode, output: Any) -> None:
+        """Updates a `SchemaNode` to use a `PrecomputedValueProvider`.
+
+        This is for when we want to use the precomputed output value of a node from a
+        previous training in a subsequent training. We replace the class in the `uses`
+        of the node to a be a `PrecomputedValueProvider` configured to return the
+        precomputed value.
+
+        Args:
+            node: The node to update.
+            output: precomputed cached output that the `PrecomputedValueProvider` will
+            return.
+        """
+        node.uses = cls
+        node.config = {"output": output}
+        node.fn = cls.get_value.__name__
+        node.constructor_name = cls.create.__name__
+
+
+@dataclasses.dataclass
+class FingerprintStatus:
+    """Holds the output of a `FingerprintComponent` and is used to prune the graph.
+
+    Attributes:
+        output_fingerprint: A fingerprint of the node's output value.
+        is_hit: `True` if node's fingerprint key exists in the cache, `False` otherwise.
+    """
+
+    output_fingerprint: Optional[Text]
+    is_hit: bool
+
+    def fingerprint(self) -> Text:
+        """Returns the internal fingerprint.
+
+        If there is no fingerprint returns a random string that will never match.
+        """
+        return self.output_fingerprint or uuid.uuid4().hex
+
+
+class FingerprintComponent(GraphComponent):
+    """Replaces non-input nodes during a fingerprint run."""
+
+    def __init__(
+        self,
+        cache: TrainingCache,
+        config_of_replaced_component: Dict[Text, Any],
+        class_of_replaced_component: Type,
+    ) -> None:
+        """Initializes a `FingerprintComponent`.
+
+        Args:
+            cache: Training cache used to determine if the run is a hit or not.
+            config_of_replaced_component: Needed to generate the fingerprint key.
+            class_of_replaced_component: Needed to generate the fingerprint key.
+        """
+        self._cache = cache
+        self._config_of_replaced_component = config_of_replaced_component
+        self._class_of_replaced_component = class_of_replaced_component
+
+    @classmethod
+    def create(
+        cls,
+        config: Dict[Text, Any],
+        model_storage: ModelStorage,
+        resource: Resource,
+        execution_context: ExecutionContext,
+    ) -> FingerprintComponent:
+        """Creates a `FingerprintComponent` (see parent class for full docstring)."""
+        return cls(
+            cache=config["cache"],
+            config_of_replaced_component=config["config_of_replaced_component"],
+            class_of_replaced_component=config["graph_component_class"],
+        )
+
+    def run(self, **kwargs: Any) -> FingerprintStatus:
+        """Calculates the fingerprint key to determine if cached output can be used.
+
+        If the fingerprint key matches an entry in the cache it means that there has
+        been a previous node execution which matches the same component class, component
+        config and input values. This means that we can potentially prune this node
+        from the schema, or replace it with a cached value before the next graph run.
+
+        Args:
+            **kwargs: Inputs from all parent nodes.
+
+        Returns:
+            A `FingerprintStatus` determining if the run was a hit, and if it was a hit
+            also the output fingerprint from the cache.
+        """
+        fingerprint_key = fingerprinting.calculate_fingerprint_key(
+            graph_component_class=self._class_of_replaced_component,
+            config={
+                **self._class_of_replaced_component.get_default_config(),
+                **self._config_of_replaced_component,
+            },
+            inputs=kwargs,
+        )
+
+        output_fingerprint = self._cache.get_cached_output_fingerprint(fingerprint_key)
+
+        return FingerprintStatus(
+            is_hit=output_fingerprint is not None, output_fingerprint=output_fingerprint
+        )
+
+    @classmethod
+    def replace_schema_node(cls, node: SchemaNode, cache: TrainingCache) -> None:
+        """Updates a `SchemaNode` to use a `FingerprintComponent`.
+
+        This is for when we want to do a fingerprint run. During the fingerprint run we
+        replace all non-input nodes with `FingerprintComponent`s so we can determine
+        whether they are able to be pruned or cached before the next graph run without
+        running the actual components.
+
+
+        Args:
+            node: The node to update.
+            cache: The cache is needed to determine of there is cache hit for the
+                fingerprint key.
+        """
+        graph_component_class = node.uses
+        node.uses = cls
+        # We update the node to be "eager" so that `FingerprintComponent.run` sees
+        # ALL the inputs to the node. If it was not eager, we would miss any args used
+        # by the constructor.
+        node.eager = True
+        node.constructor_name = cls.create.__name__
+        node.fn = cls.run.__name__
+        node.config = {
+            "config_of_replaced_component": node.config,
+            "cache": cache,
+            "graph_component_class": graph_component_class,
+        }

rasa/engine/training/fingerprinting.py

@@ -0,0 +1,64 @@
+import inspect
+import logging
+from typing import Any, Dict, Text, Type
+from typing_extensions import Protocol, runtime_checkable
+import importlib_metadata
+import rasa.utils.common
+import rasa.shared.utils.io
+from rasa.engine.graph import GraphComponent
+
+logger = logging.getLogger(__name__)
+
+import_name_to_package_map = {"sklearn": "scikit_learn"}
+
+
+@runtime_checkable
+class Fingerprintable(Protocol):
+    """Interface that enforces training data can be fingerprinted."""
+
+    def fingerprint(self) -> Text:
+        """Returns a unique stable fingerprint of the data."""
+        ...
+
+
+def calculate_fingerprint_key(
+    graph_component_class: Type[GraphComponent],
+    config: Dict[Text, Any],
+    inputs: Dict[Text, Fingerprintable],
+) -> Text:
+    """Calculates a fingerprint key that uniquely represents a single node's execution.
+
+    Args:
+        graph_component_class: The graph component class.
+        config: The component config.
+        inputs: The inputs as a mapping of parent node name to input value.
+
+    Returns:
+        The fingerprint key.
+    """
+    dependency_versions = {
+        package: importlib_metadata.version(
+            import_name_to_package_map.get(package, package)
+        )
+        for package in graph_component_class.required_packages()
+    }
+    fingerprint_data = {
+        "node_name": rasa.utils.common.module_path_from_class(graph_component_class),
+        "component_implementation": inspect.getsource(graph_component_class),
+        "config": config,
+        "inputs": inputs,
+        "dependency_versions": dependency_versions,
+    }
+
+    fingerprint_addon = graph_component_class.fingerprint_addon(config)
+    if fingerprint_addon is not None:
+        fingerprint_data["addon"] = fingerprint_addon
+
+    fingerprint_key = rasa.shared.utils.io.deep_container_fingerprint(fingerprint_data)
+
+    logger.debug(
+        f"Calculated fingerprint_key '{fingerprint_key}' for class "
+        f"'{graph_component_class.__name__}'."
+    )
+
+    return fingerprint_key

rasa/engine/training/graph_trainer.py

@@ -0,0 +1,256 @@
+import copy
+import logging
+from pathlib import Path
+from typing import Any, Dict, Text, Type, Union
+
+from rasa.engine.caching import TrainingCache
+from rasa.engine.graph import ExecutionContext, GraphSchema, GraphModelConfiguration
+from rasa.engine.constants import PLACEHOLDER_IMPORTER
+from rasa.engine.runner.interface import GraphRunner
+from rasa.engine.storage.storage import ModelStorage, ModelMetadata
+from rasa.engine.training.components import (
+    PrecomputedValueProvider,
+    FingerprintComponent,
+    FingerprintStatus,
+)
+from rasa.engine.training.hooks import TrainingHook, LoggingHook
+from rasa.shared.importers.importer import TrainingDataImporter
+
+logger = logging.getLogger(__name__)
+
+
+class GraphTrainer:
+    """Trains a model using a graph schema."""
+
+    def __init__(
+        self,
+        model_storage: ModelStorage,
+        cache: TrainingCache,
+        graph_runner_class: Type[GraphRunner],
+    ) -> None:
+        """Initializes a `GraphTrainer`.
+
+        Args:
+            model_storage: Storage which graph components can use to persist and load.
+                Also used for packaging the trained model.
+            cache: Cache used to store fingerprints and outputs.
+            graph_runner_class: The class to instantiate the runner from.
+        """
+        self._model_storage = model_storage
+        self._cache = cache
+        self._graph_runner_class = graph_runner_class
+
+    async def train(
+        self,
+        model_configuration: GraphModelConfiguration,
+        importer: TrainingDataImporter,
+        output_filename: Path,
+        force_retraining: bool = False,
+        is_finetuning: bool = False,
+    ) -> ModelMetadata:
+        """Trains and packages a model and returns the prediction graph runner.
+
+        Args:
+            model_configuration: The model configuration (schemas, language, etc.)
+            importer: The importer which provides the training data for the training.
+            output_filename: The location to save the packaged model.
+            force_retraining: If `True` then the cache is skipped and all components
+                are retrained.
+            is_finetuning: `True` if we want to finetune the model.
+
+        Returns:
+            The metadata describing the trained model.
+        """
+        logger.debug("Starting training.")
+
+        # Retrieve the domain for the model metadata right at the start.
+        # This avoids that something during the graph runs mutates it.
+        domain = copy.deepcopy(importer.get_domain())
+
+        if force_retraining:
+            logger.debug(
+                "Skip fingerprint run as a full training of the model was enforced."
+            )
+            pruned_training_schema = model_configuration.train_schema
+        else:
+            fingerprint_run_outputs = await self.fingerprint(
+                model_configuration.train_schema,
+                importer=importer,
+                is_finetuning=is_finetuning,
+            )
+            pruned_training_schema = self._prune_schema(
+                model_configuration.train_schema, fingerprint_run_outputs
+            )
+
+        hooks = [
+            LoggingHook(pruned_schema=pruned_training_schema),
+            TrainingHook(
+                cache=self._cache,
+                model_storage=self._model_storage,
+                pruned_schema=pruned_training_schema,
+            ),
+        ]
+
+        graph_runner = self._graph_runner_class.create(
+            graph_schema=pruned_training_schema,
+            model_storage=self._model_storage,
+            execution_context=ExecutionContext(
+                graph_schema=model_configuration.train_schema,
+                is_finetuning=is_finetuning,
+            ),
+            hooks=hooks,
+        )
+
+        logger.debug("Running the pruned train graph with real node execution.")
+
+        await graph_runner.run(inputs={PLACEHOLDER_IMPORTER: importer})
+
+        return self._model_storage.create_model_package(
+            output_filename, model_configuration, domain
+        )
+
+    async def fingerprint(
+        self,
+        train_schema: GraphSchema,
+        importer: TrainingDataImporter,
+        is_finetuning: bool = False,
+    ) -> Dict[Text, Union[FingerprintStatus, Any]]:
+        """Runs the graph using fingerprints to determine which nodes need to re-run.
+
+        Nodes which have a matching fingerprint key in the cache can either be removed
+        entirely from the graph, or replaced with a cached value if their output is
+        needed by descendent nodes.
+
+        Args:
+            train_schema: The train graph schema that will be run in fingerprint mode.
+            importer: The importer which provides the training data for the training.
+            is_finetuning: `True` if we want to finetune the model.
+
+        Returns:
+            Mapping of node names to fingerprint results.
+        """
+        fingerprint_schema = self._create_fingerprint_schema(train_schema)
+
+        fingerprint_graph_runner = self._graph_runner_class.create(
+            graph_schema=fingerprint_schema,
+            model_storage=self._model_storage,
+            execution_context=ExecutionContext(
+                graph_schema=train_schema, is_finetuning=is_finetuning
+            ),
+        )
+
+        logger.debug("Running the train graph in fingerprint mode.")
+        return await fingerprint_graph_runner.run(
+            inputs={PLACEHOLDER_IMPORTER: importer}
+        )
+
+    def _create_fingerprint_schema(self, train_schema: GraphSchema) -> GraphSchema:
+        fingerprint_schema = copy.deepcopy(train_schema)
+        for node_name, schema_node in fingerprint_schema.nodes.items():
+            # We make every node a target so that `graph_runner.run(...)` returns
+            # the output for each node. We need the output of each node
+            # to decide which nodes we can prune.
+            schema_node.is_target = True
+
+            # We do not replace the input nodes as we need an up-to-date fingerprint of
+            # any input data to the graph. This means we can prune according to what
+            # has actually changed.
+            if not schema_node.is_input:
+                FingerprintComponent.replace_schema_node(schema_node, self._cache)
+        return fingerprint_schema
+
+    def _prune_schema(
+        self,
+        schema: GraphSchema,
+        fingerprint_run_outputs: Dict[Text, Union[FingerprintStatus, Any]],
+    ) -> GraphSchema:
+        """Uses the fingerprint statuses to prune the graph schema.
+
+        Walks the graph starting at each target node. If a node has a cache hit we
+        replace it with a `PrecomputedValueProvider` and remove its input dependencies.
+        At the end, any node that is not an ancestor of a target node will be pruned
+        when we call `minimal_graph_schema()`.
+
+        Args:
+            schema: The graph to prune.
+            fingerprint_run_outputs: Node outputs from the fingerprint run as a mapping
+                from node name to output.
+
+        Returns:
+            The pruned schema.
+        """
+        pruned_schema = copy.deepcopy(schema)
+        target_node_names = pruned_schema.target_names
+
+        for target_node_name in target_node_names:
+            self._walk_and_prune(
+                pruned_schema, target_node_name, fingerprint_run_outputs
+            )
+
+        return pruned_schema.minimal_graph_schema()
+
+    def _walk_and_prune(
+        self,
+        schema: GraphSchema,
+        current_node_name: Text,
+        fingerprint_run_outputs: Dict[Text, Union[FingerprintStatus, Any]],
+    ) -> None:
+        """Recursively walks backwards though a graph checking the status of each node.
+
+        If node has a fingerprint key hit then we check if there is a cached output.
+        If there is a cached output we will replace the node with a
+        `PrecomputedValueProvider` and remove all its dependencies (`.needs`). If
+        there is not a fingerprint key hit, or there is no cached output, the node is
+        left untouched and will be executed again next run unless it is no longer the
+        ancestor of a target node.
+
+        Args:
+            schema: The graph we are currently walking.
+            current_node_name: The current node on the walk.
+            fingerprint_run_outputs: The fingerprint statuses of every node as a mapping
+                from node name to status.
+        """
+        fingerprint_run_output = fingerprint_run_outputs[current_node_name]
+        node = schema.nodes[current_node_name]
+
+        # If we have replaced this node with a `PrecomputedValueProvider` we have
+        # already visited this node. A `PrecomputedValueProvider` is updated to have
+        # no parent nodes, so
+        # we can end the walk here.
+        if node.uses == PrecomputedValueProvider:
+            return
+
+        # If the output was a `FingerprintStatus` we must check the cache and status.
+        if isinstance(fingerprint_run_output, FingerprintStatus):
+            # If there is a fingerprint key hit we can potentially use a cached output.
+            if fingerprint_run_output.is_hit:
+                output_result = self._cache.get_cached_result(
+                    output_fingerprint_key=fingerprint_run_output.output_fingerprint,
+                    node_name=current_node_name,
+                    model_storage=self._model_storage,
+                )
+                if output_result:
+                    logger.debug(
+                        f"Updating '{current_node_name}' to use a "
+                        f"'{PrecomputedValueProvider.__name__}'."
+                    )
+                    PrecomputedValueProvider.replace_schema_node(node, output_result)
+                    # We remove all parent dependencies as the cached output value will
+                    # be used.
+                    node.needs = {}
+                else:
+                    # If there is no cached output the node must be re-run if it ends
+                    # up as an ancestor of a target node.
+                    fingerprint_run_output.is_hit = False
+
+        # Else the node was an input node and the output is the actual node's output.
+        else:
+            # As fingerprint_run_output is just the node's output there is no need to
+            # execute the node again. We can just return it from a
+            # `PrecomputedValueProvider`.
+            PrecomputedValueProvider.replace_schema_node(node, fingerprint_run_output)
+            node.needs = {}
+
+        # Continue walking for every parent node.
+        for parent_node_name in node.needs.values():
+            self._walk_and_prune(schema, parent_node_name, fingerprint_run_outputs)

rasa/engine/training/hooks.py

@@ -0,0 +1,164 @@
+import logging
+from typing import Any, Dict, Text, Type
+
+from rasa.engine.caching import TrainingCache
+from rasa.engine.graph import ExecutionContext, GraphNodeHook, GraphSchema, SchemaNode
+from rasa.engine.storage.storage import ModelStorage
+from rasa.engine.training.components import PrecomputedValueProvider
+import rasa.shared.utils.io
+from rasa.engine.training import fingerprinting
+
+logger = logging.getLogger(__name__)
+
+
+class TrainingHook(GraphNodeHook):
+    """Caches fingerprints and outputs of nodes during model training."""
+
+    def __init__(
+        self,
+        cache: TrainingCache,
+        model_storage: ModelStorage,
+        pruned_schema: GraphSchema,
+    ) -> None:
+        """Initializes a `TrainingHook`.
+
+        Args:
+            cache: Cache used to store fingerprints and outputs.
+            model_storage: Used to cache `Resource`s.
+            pruned_schema: The pruned training schema.
+        """
+        self._cache = cache
+        self._model_storage = model_storage
+        self._pruned_schema = pruned_schema
+
+    def on_before_node(
+        self,
+        node_name: Text,
+        execution_context: ExecutionContext,
+        config: Dict[Text, Any],
+        received_inputs: Dict[Text, Any],
+    ) -> Dict:
+        """Calculates the run fingerprint for use in `on_after_node`."""
+        graph_component_class = self._get_graph_component_class(
+            execution_context, node_name
+        )
+        graph_component_config = self._get_graph_component_config(
+            execution_context, node_name
+        )
+
+        merged_config = {
+            **config,
+            **graph_component_config,
+        }
+
+        fingerprint_key = fingerprinting.calculate_fingerprint_key(
+            graph_component_class=graph_component_class,
+            config=merged_config,
+            inputs=received_inputs,
+        )
+
+        return {"fingerprint_key": fingerprint_key}
+
+    def on_after_node(
+        self,
+        node_name: Text,
+        execution_context: ExecutionContext,
+        config: Dict[Text, Any],
+        output: Any,
+        input_hook_data: Dict,
+    ) -> None:
+        """Stores the fingerprints and caches the output of the node."""
+        # We should not re-cache the output of a PrecomputedValueProvider.
+        graph_component_class = self._pruned_schema.nodes[node_name].uses
+
+        if graph_component_class == PrecomputedValueProvider:
+            return None
+
+        output_fingerprint = rasa.shared.utils.io.deep_container_fingerprint(output)
+        fingerprint_key = input_hook_data["fingerprint_key"]
+
+        logger.debug(
+            f"Caching '{output.__class__.__name__}' with fingerprint_key: "
+            f"'{fingerprint_key}' and output_fingerprint '{output_fingerprint}'."
+        )
+
+        self._cache.cache_output(
+            fingerprint_key=fingerprint_key,
+            output=output,
+            output_fingerprint=output_fingerprint,
+            model_storage=self._model_storage,
+        )
+
+    @staticmethod
+    def _get_graph_component_class(
+        execution_context: ExecutionContext, node_name: Text
+    ) -> Type:
+        graph_component_class = execution_context.graph_schema.nodes[node_name].uses
+        return graph_component_class
+
+    @staticmethod
+    def _get_graph_component_config(
+        execution_context: ExecutionContext, node_name: str
+    ) -> Dict[str, Any]:
+        return execution_context.graph_schema.nodes[node_name].config
+
+
+class LoggingHook(GraphNodeHook):
+    """Logs the training of components."""
+
+    def __init__(self, pruned_schema: GraphSchema) -> None:
+        """Creates hook.
+
+        Args:
+            pruned_schema: The pruned schema provides us with the information whether
+                a component is cached or not.
+        """
+        self._pruned_schema = pruned_schema
+
+    def on_before_node(
+        self,
+        node_name: Text,
+        execution_context: ExecutionContext,
+        config: Dict[Text, Any],
+        received_inputs: Dict[Text, Any],
+    ) -> Dict:
+        """Logs the training start of a graph node."""
+        node = self._pruned_schema.nodes[node_name]
+
+        if not self._is_cached_node(node) and self._does_node_train(node):
+            logger.info(f"Starting to train component '{node.uses.__name__}'.")
+
+        return {}
+
+    @staticmethod
+    def _does_node_train(node: SchemaNode) -> bool:
+        # Nodes which train are always targets so that they store their output in the
+        # model storage. `is_input` filters out nodes which don't really train but e.g.
+        # persist some training data.
+        return node.is_target and not node.is_input
+
+    @staticmethod
+    def _is_cached_node(node: SchemaNode) -> bool:
+        return node.uses == PrecomputedValueProvider
+
+    def on_after_node(
+        self,
+        node_name: Text,
+        execution_context: ExecutionContext,
+        config: Dict[Text, Any],
+        output: Any,
+        input_hook_data: Dict,
+    ) -> None:
+        """Logs when a component finished its training."""
+        node = self._pruned_schema.nodes[node_name]
+
+        if not self._does_node_train(node):
+            return
+
+        if self._is_cached_node(node):
+            actual_component = execution_context.graph_schema.nodes[node_name]
+            logger.info(
+                f"Restored component '{actual_component.uses.__name__}' from cache."
+            )
+        else:
+            logger.info(f"Finished training component '{node.uses.__name__}'.")