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

rasa/engine/storage/local_model_storage.py

@@ -0,0 +1,244 @@
+from __future__ import annotations
+
+import logging
+import shutil
+import sys
+import tempfile
+import uuid
+from contextlib import contextmanager
+from datetime import datetime
+from pathlib import Path
+from tarsafe import TarSafe
+from typing import Generator, Optional, Text, Tuple, Union
+
+import rasa.utils.common
+import rasa.shared.utils.io
+from rasa.engine.storage.storage import ModelMetadata, ModelStorage
+from rasa.engine.graph import GraphModelConfiguration
+from rasa.engine.storage.resource import Resource
+from rasa.exceptions import UnsupportedModelVersionError
+from rasa.shared.core.domain import Domain
+import rasa.model
+
+logger = logging.getLogger(__name__)
+
+# Paths within model archive
+MODEL_ARCHIVE_COMPONENTS_DIR = "components"
+MODEL_ARCHIVE_METADATA_FILE = "metadata.json"
+
+
+@contextmanager
+def windows_safe_temporary_directory(
+    suffix: Optional[Text] = None,
+    prefix: Optional[Text] = None,
+    dir: Optional[Text] = None,
+) -> Generator[Text, None, None]:
+    """Like `tempfile.TemporaryDirectory`, but works with Windows and long file names.
+
+    On Windows by default there is a restriction on long path names.
+    Using the prefix below allows to bypass this restriction in environments
+    where it's not possible to override this behavior, mostly for internal
+    policy reasons.
+
+    Reference: https://stackoverflow.com/a/49102229
+    """
+    if sys.platform == "win32":
+        directory = tempfile.mkdtemp(suffix, prefix, dir)
+        directory = rasa.utils.common.decode_bytes(directory)
+
+        try:
+            yield directory
+        finally:
+            shutil.rmtree(f"\\\\?\\{directory}")
+    else:
+        with tempfile.TemporaryDirectory() as temporary_directory:
+            temporary_directory = rasa.utils.common.decode_bytes(temporary_directory)
+            yield temporary_directory
+
+
+class LocalModelStorage(ModelStorage):
+    """Stores and provides output of `GraphComponents` on local disk."""
+
+    def __init__(self, storage_path: Path) -> None:
+        """Creates storage (see parent class for full docstring)."""
+        self._storage_path = storage_path
+
+    @classmethod
+    def create(cls, storage_path: Path) -> ModelStorage:
+        """Creates a new instance (see parent class for full docstring)."""
+        return cls(storage_path)
+
+    @classmethod
+    def from_model_archive(
+        cls, storage_path: Path, model_archive_path: Union[Text, Path]
+    ) -> Tuple[LocalModelStorage, ModelMetadata]:
+        """Initializes storage from archive (see parent class for full docstring)."""
+        if next(storage_path.glob("*"), None):
+            raise ValueError(
+                f"The model storage with path '{storage_path}' is "
+                f"not empty. You can only unpack model archives into an "
+                f"empty model storage."
+            )
+
+        with windows_safe_temporary_directory() as temporary_directory:
+            temporary_directory_path = Path(temporary_directory)
+
+            cls._extract_archive_to_directory(
+                model_archive_path, temporary_directory_path
+            )
+            logger.debug(f"Extracted model to '{temporary_directory_path}'.")
+            cls._initialize_model_storage_from_model_archive(
+                temporary_directory_path, storage_path
+            )
+
+            metadata = cls._load_metadata(temporary_directory_path)
+
+            return (cls(storage_path), metadata)
+
+    @classmethod
+    def metadata_from_archive(
+        cls, model_archive_path: Union[Text, Path]
+    ) -> ModelMetadata:
+        """Retrieves metadata from archive (see parent class for full docstring)."""
+        with windows_safe_temporary_directory() as temporary_directory:
+            temporary_directory_path = Path(temporary_directory)
+
+            cls._extract_archive_to_directory(
+                model_archive_path, temporary_directory_path
+            )
+            metadata = cls._load_metadata(temporary_directory_path)
+
+            return metadata
+
+    @staticmethod
+    def _extract_archive_to_directory(
+        model_archive_path: Union[Text, Path], temporary_directory: Path
+    ) -> None:
+        with TarSafe.open(model_archive_path, mode="r:gz") as tar:
+            if sys.platform == "win32":
+                # on Windows by default there is a restriction on long
+                # path names; using the prefix below allows to bypass
+                # this restriction in environments where it's not possible
+                # to override this behavior, mostly for internal policy reasons
+                # reference: https://stackoverflow.com/a/49102229
+                tar.extractall(f"\\\\?\\{temporary_directory}")
+            else:
+                tar.extractall(temporary_directory)
+        LocalModelStorage._assert_not_rasa2_archive(temporary_directory)
+
+    @staticmethod
+    def _assert_not_rasa2_archive(temporary_directory: Union[Text, Path]) -> None:
+        fingerprint_file = Path(temporary_directory) / "fingerprint.json"
+        if fingerprint_file.is_file():
+            serialized_fingerprint = rasa.shared.utils.io.read_json_file(
+                fingerprint_file
+            )
+            raise UnsupportedModelVersionError(
+                model_version=serialized_fingerprint["version"]
+            )
+
+    @staticmethod
+    def _initialize_model_storage_from_model_archive(
+        temporary_directory: Path, storage_path: Path
+    ) -> None:
+        for path in (temporary_directory / MODEL_ARCHIVE_COMPONENTS_DIR).glob("*"):
+            shutil.move(str(path), str(storage_path))
+
+    @staticmethod
+    def _load_metadata(directory: Path) -> ModelMetadata:
+        serialized_metadata = rasa.shared.utils.io.read_json_file(
+            directory / MODEL_ARCHIVE_METADATA_FILE
+        )
+
+        return ModelMetadata.from_dict(serialized_metadata)
+
+    @contextmanager
+    def write_to(self, resource: Resource) -> Generator[Path, None, None]:
+        """Persists data for a resource (see parent class for full docstring)."""
+        logger.debug(f"Resource '{resource.name}' was requested for writing.")
+        directory = self._directory_for_resource(resource)
+
+        if not directory.exists():
+            directory.mkdir()
+
+        yield directory
+
+        logger.debug(f"Resource '{resource.name}' was persisted.")
+
+    def _directory_for_resource(self, resource: Resource) -> Path:
+        return self._storage_path / resource.name
+
+    @contextmanager
+    def read_from(self, resource: Resource) -> Generator[Path, None, None]:
+        """Provides the data of a `Resource` (see parent class for full docstring)."""
+        logger.debug(f"Resource '{resource.name}' was requested for reading.")
+        directory = self._directory_for_resource(resource)
+
+        if not directory.exists():
+            raise ValueError(
+                f"Resource '{resource.name}' does not exist. Please make "
+                f"sure that the graph component providing the resource "
+                f"is a parent node of the current graph node "
+                f"(in case this happens during training) or that the "
+                f"resource was actually persisted during training "
+                f"(in case this happens during inference)."
+            )
+
+        yield directory
+
+    def create_model_package(
+        self,
+        model_archive_path: Union[Text, Path],
+        model_configuration: GraphModelConfiguration,
+        domain: Domain,
+    ) -> ModelMetadata:
+        """Creates model package (see parent class for full docstring)."""
+        logger.debug(f"Start to created model package for path '{model_archive_path}'.")
+
+        with windows_safe_temporary_directory() as temp_dir:
+            temporary_directory = Path(temp_dir)
+
+            shutil.copytree(
+                self._storage_path, temporary_directory / MODEL_ARCHIVE_COMPONENTS_DIR
+            )
+
+            model_metadata = self._create_model_metadata(domain, model_configuration)
+            self._persist_metadata(model_metadata, temporary_directory)
+
+            if isinstance(model_archive_path, str):
+                model_archive_path = Path(model_archive_path)
+
+            if not model_archive_path.parent.exists():
+                model_archive_path.parent.mkdir(parents=True)
+
+            with TarSafe.open(model_archive_path, "w:gz") as tar:
+                tar.add(temporary_directory, arcname="")
+
+        logger.debug(f"Model package created in path '{model_archive_path}'.")
+
+        return model_metadata
+
+    @staticmethod
+    def _persist_metadata(metadata: ModelMetadata, temporary_directory: Path) -> None:
+        rasa.shared.utils.io.dump_obj_as_json_to_file(
+            temporary_directory / MODEL_ARCHIVE_METADATA_FILE, metadata.as_dict()
+        )
+
+    @staticmethod
+    def _create_model_metadata(
+        domain: Domain, model_configuration: GraphModelConfiguration
+    ) -> ModelMetadata:
+        return ModelMetadata(
+            trained_at=datetime.utcnow(),
+            rasa_open_source_version=rasa.__version__,
+            model_id=uuid.uuid4().hex,
+            assistant_id=model_configuration.assistant_id,
+            domain=domain,
+            train_schema=model_configuration.train_schema,
+            predict_schema=model_configuration.predict_schema,
+            training_type=model_configuration.training_type,
+            project_fingerprint=rasa.model.project_fingerprint(),
+            language=model_configuration.language,
+            core_target=model_configuration.core_target,
+            nlu_target=model_configuration.nlu_target,
+        )

rasa/engine/storage/resource.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+import logging
+import typing
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Text
+import uuid
+
+import rasa.utils.common
+import rasa.utils.io
+
+if typing.TYPE_CHECKING:
+    from rasa.engine.storage.storage import ModelStorage
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Resource:
+    """Represents a persisted graph component in the graph.
+
+    Attributes:
+        name: The unique identifier for the `Resource`. Used to locate the associated
+            data from a `ModelStorage`. Normally matches the name of the node which
+            created it.
+        output_fingerprint: An unique identifier for a specific instantiation of a
+            `Resource`. Used to distinguish a specific persistence for the same
+            `Resource` when saving to the cache.
+
+    """
+
+    name: Text
+    output_fingerprint: Text = field(
+        default_factory=lambda: uuid.uuid4().hex,
+        # We do not use this for comparison as it is not consistent after serialization.
+        compare=False,
+    )
+
+    @classmethod
+    def from_cache(
+        cls,
+        node_name: Text,
+        directory: Path,
+        model_storage: ModelStorage,
+        output_fingerprint: Text,
+    ) -> Resource:
+        """Loads a `Resource` from the cache.
+
+        This automatically loads the persisted resource into the given `ModelStorage`.
+
+        Args:
+            node_name: The node name of the `Resource`.
+            directory: The directory with the cached `Resource`.
+            model_storage: The `ModelStorage` which the cached `Resource` will be added
+                to so that the `Resource` is accessible for other graph nodes.
+            output_fingerprint: The fingerprint of the cached `Resource`.
+
+        Returns:
+            The ready-to-use and accessible `Resource`.
+        """
+        logger.debug(f"Loading resource '{node_name}' from cache.")
+
+        resource = Resource(node_name, output_fingerprint=output_fingerprint)
+        if not any(directory.glob("*")):
+            logger.debug(f"Cached resource for '{node_name}' was empty.")
+            return resource
+
+        try:
+            with model_storage.write_to(resource) as resource_directory:
+                rasa.utils.common.copy_directory(directory, resource_directory)
+        except ValueError:
+            # This might happen during finetuning as in this case the model storage
+            # is already filled
+            if not rasa.utils.io.are_directories_equal(directory, resource_directory):
+                # We skip caching in case we see the cached output and output
+                # from the model which we want to finetune are not the same
+                raise
+
+        logger.debug(f"Successfully initialized resource '{node_name}' from cache.")
+
+        return resource
+
+    def to_cache(self, directory: Path, model_storage: ModelStorage) -> None:
+        """Persists the `Resource` to the cache.
+
+        Args:
+            directory: The directory which receives the persisted `Resource`.
+            model_storage: The model storage which currently contains the persisted
+                `Resource`.
+        """
+        try:
+            with model_storage.read_from(self) as resource_directory:
+                rasa.utils.common.copy_directory(resource_directory, directory)
+        except ValueError:
+            logger.debug(
+                f"Skipped caching resource '{self.name}' as no persisted "
+                f"data was found."
+            )
+
+    def fingerprint(self) -> Text:
+        """Provides fingerprint for `Resource`.
+
+        A unique fingerprint is created on initialization of a `Resource` however we
+        also allow a value to be provided for when we retrieve a `Resource` from the
+        cache (see `Resource.from_cache`).
+
+        Returns:
+            Fingerprint for `Resource`.
+        """
+        return self.output_fingerprint

rasa/engine/storage/storage.py

@@ -0,0 +1,199 @@
+from __future__ import annotations
+import abc
+import logging
+import typing
+from contextlib import contextmanager
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import Tuple, Union, Text, Generator, Dict, Any, Optional
+from packaging import version
+
+from rasa.constants import MINIMUM_COMPATIBLE_VERSION
+from rasa.exceptions import UnsupportedModelVersionError
+from rasa.engine.storage.resource import Resource
+from rasa.shared.core.domain import Domain
+from rasa.shared.data import TrainingType
+
+if typing.TYPE_CHECKING:
+    from rasa.engine.graph import GraphSchema, GraphModelConfiguration
+
+logger = logging.getLogger(__name__)
+
+
+class ModelStorage(abc.ABC):
+    """Serves as storage backend for `GraphComponents` which need persistence."""
+
+    @classmethod
+    @abc.abstractmethod
+    def create(cls, storage_path: Path) -> ModelStorage:
+        """Creates the storage.
+
+        Args:
+            storage_path: Directory which will contain the persisted graph components.
+        """
+        ...
+
+    @classmethod
+    @abc.abstractmethod
+    def from_model_archive(
+        cls, storage_path: Path, model_archive_path: Union[Text, Path]
+    ) -> Tuple[ModelStorage, ModelMetadata]:
+        """Unpacks a model archive and initializes a `ModelStorage`.
+
+        Args:
+            storage_path: Directory which will contain the persisted graph components.
+            model_archive_path: The path to the model archive.
+
+        Returns:
+            Initialized model storage, and metadata about the model.
+
+        Raises:
+            `UnsupportedModelError` if the loaded meta data indicates that the model
+            has been created with an outdated Rasa version.
+        """
+        ...
+
+    @classmethod
+    @abc.abstractmethod
+    def metadata_from_archive(
+        cls, model_archive_path: Union[Text, Path]
+    ) -> ModelMetadata:
+        """Retrieves metadata from archive.
+
+        Args:
+            model_archive_path: The path to the model archive.
+
+        Returns:
+            Metadata about the model.
+
+        Raises:
+            `UnsupportedModelError` if the loaded meta data indicates that the model
+            has been created with an outdated Rasa version.
+        """
+        ...
+
+    @contextmanager
+    @abc.abstractmethod
+    def write_to(self, resource: Resource) -> Generator[Path, None, None]:
+        """Persists data for a given resource.
+
+        This `Resource` can then be accessed in dependent graph nodes via
+        `model_storage.read_from`.
+
+        Args:
+            resource: The resource which should be persisted.
+
+        Returns:
+            A directory which can be used to persist data for the given `Resource`.
+        """
+        ...
+
+    @contextmanager
+    @abc.abstractmethod
+    def read_from(self, resource: Resource) -> Generator[Path, None, None]:
+        """Provides the data of a persisted `Resource`.
+
+        Args:
+            resource: The `Resource` whose persisted should be accessed.
+
+        Returns:
+            A directory containing the data of the persisted `Resource`.
+
+        Raises:
+            ValueError: In case no persisted data for the given `Resource` exists.
+        """
+        ...
+
+    @abc.abstractmethod
+    def create_model_package(
+        self,
+        model_archive_path: Union[Text, Path],
+        model_configuration: GraphModelConfiguration,
+        domain: Domain,
+    ) -> ModelMetadata:
+        """Creates a model archive containing all data to load and run the model.
+
+        Args:
+            model_archive_path: The path to the archive which should be created.
+            model_configuration: The model configuration (schemas, language, etc.)
+            domain: The `Domain` which was used to train the model.
+
+        Returns:
+            The model metadata.
+        """
+        ...
+
+
+@dataclass()
+class ModelMetadata:
+    """Describes a trained model."""
+
+    trained_at: datetime
+    rasa_open_source_version: Text
+    model_id: Text
+    assistant_id: Optional[Text]
+    domain: Domain
+    train_schema: GraphSchema
+    predict_schema: GraphSchema
+    project_fingerprint: Text
+    core_target: Optional[Text]
+    nlu_target: Text
+    language: Optional[Text]
+    training_type: TrainingType = TrainingType.BOTH
+
+    def __post_init__(self) -> None:
+        """Raises an exception when the metadata indicates an unsupported version.
+
+        Raises:
+            `UnsupportedModelException` if the `rasa_open_source_version` is lower
+            than the minimum compatible version
+        """
+        minimum_version = version.parse(MINIMUM_COMPATIBLE_VERSION)
+        model_version = version.parse(self.rasa_open_source_version)
+        if model_version < minimum_version:
+            raise UnsupportedModelVersionError(model_version=model_version)
+
+    def as_dict(self) -> Dict[Text, Any]:
+        """Returns serializable version of the `ModelMetadata`."""
+        return {
+            "domain": self.domain.as_dict(),
+            "trained_at": self.trained_at.isoformat(),
+            "model_id": self.model_id,
+            "assistant_id": self.assistant_id,
+            "rasa_open_source_version": self.rasa_open_source_version,
+            "train_schema": self.train_schema.as_dict(),
+            "predict_schema": self.predict_schema.as_dict(),
+            "training_type": self.training_type.value,
+            "project_fingerprint": self.project_fingerprint,
+            "core_target": self.core_target,
+            "nlu_target": self.nlu_target,
+            "language": self.language,
+        }
+
+    @classmethod
+    def from_dict(cls, serialized: Dict[Text, Any]) -> ModelMetadata:
+        """Loads `ModelMetadata` which has been serialized using `metadata.as_dict()`.
+
+        Args:
+            serialized: Serialized `ModelMetadata` (e.g. read from disk).
+
+        Returns:
+            Instantiated `ModelMetadata`.
+        """
+        from rasa.engine.graph import GraphSchema
+
+        return ModelMetadata(
+            trained_at=datetime.fromisoformat(serialized["trained_at"]),
+            rasa_open_source_version=serialized["rasa_open_source_version"],
+            model_id=serialized["model_id"],
+            assistant_id=serialized.get("assistant_id"),
+            domain=Domain.from_dict(serialized["domain"]),
+            train_schema=GraphSchema.from_dict(serialized["train_schema"]),
+            predict_schema=GraphSchema.from_dict(serialized["predict_schema"]),
+            training_type=TrainingType(serialized["training_type"]),
+            project_fingerprint=serialized["project_fingerprint"],
+            core_target=serialized["core_target"],
+            nlu_target=serialized["nlu_target"],
+            language=serialized["language"],
+        )