rasa-pro 3.10.16__py3-none-any.whl → 3.11.0__py3-none-any.whl

rasa/{nlu → core}/persistor.py

@@ -4,16 +4,19 @@ import abc
 import os
 import shutil
 from enum import Enum
+from pathlib import Path
 from typing import TYPE_CHECKING, List, Optional, Text, Tuple, Union
 
 import structlog
 
+from rasa.exceptions import ModelNotFound
 import rasa.shared.utils.common
 import rasa.utils.common
 from rasa.constants import (
     HTTP_STATUS_FORBIDDEN,
     HTTP_STATUS_NOT_FOUND,
     MODEL_ARCHIVE_EXTENSION,
+    DEFAULT_BUCKET_NAME,
 )
 from rasa.env import (
     AWS_ENDPOINT_URL_ENV,
@@ -28,6 +31,7 @@ from rasa.shared.utils.io import raise_warning
 
 if TYPE_CHECKING:
     from azure.storage.blob import ContainerClient
+    from botocore.exceptions import ClientError
 
 structlogger = structlog.get_logger()
 
@@ -82,16 +86,19 @@ def get_persistor(storage: StorageType) -> Optional[Persistor]:
     Currently, `aws`, `gcs`, `azure` and providing module paths are supported remote
     storages.
     """
-    if storage == RemoteStorageType.AWS:
+    storage = storage.value if isinstance(storage, RemoteStorageType) else storage
+
+    if storage == RemoteStorageType.AWS.value:
         return AWSPersistor(
-            os.environ.get(BUCKET_NAME_ENV), os.environ.get(AWS_ENDPOINT_URL_ENV)
+            os.environ.get(BUCKET_NAME_ENV, DEFAULT_BUCKET_NAME),
+            os.environ.get(AWS_ENDPOINT_URL_ENV),
         )
-    if storage == RemoteStorageType.GCS:
-        return GCSPersistor(os.environ.get(BUCKET_NAME_ENV))
+    if storage == RemoteStorageType.GCS.value:
+        return GCSPersistor(os.environ.get(BUCKET_NAME_ENV, DEFAULT_BUCKET_NAME))
 
-    if storage == RemoteStorageType.AZURE:
+    if storage == RemoteStorageType.AZURE.value:
         return AzurePersistor(
-            os.environ.get(AZURE_CONTAINER_ENV),
+            os.environ.get(AZURE_CONTAINER_ENV, DEFAULT_BUCKET_NAME),
             os.environ.get(AZURE_ACCOUNT_NAME_ENV),
             os.environ.get(AZURE_ACCOUNT_KEY_ENV),
         )
@@ -116,44 +123,64 @@ class Persistor(abc.ABC):
 
     def persist(self, trained_model: str) -> None:
         """Uploads a trained model persisted in the `target_dir` to cloud storage."""
-        file_key = self._create_file_key(trained_model)
+        absolute_file_key = self._create_file_key(trained_model)
+        file_key = Path(absolute_file_key).name
         self._persist_tar(file_key, trained_model)
 
     def retrieve(self, model_name: Text, target_path: Text) -> Text:
         """Downloads a model that has been persisted to cloud storage.
 
         Downloaded model will be saved to the `target_path`.
-        If `target_path` is a directory, the model will be downloaded to that directory.
-        If `target_path` is a file, the model will be downloaded to that file.
+        If `target_path` is a directory, the model will be saved to that directory.
+        If `target_path` is a file, the model will be saved to that file.
 
         Args:
             model_name: The name of the model to retrieve.
-            target_path: The path to which the model should be downloaded.
+            target_path: The path to which the model should be saved.
         """
         tar_name = model_name
         if not model_name.endswith(MODEL_ARCHIVE_EXTENSION):
             # ensure backward compatibility
             tar_name = self._tar_name(model_name)
-        remote_object_path = self._create_file_key(tar_name)
-        self._retrieve_tar(remote_object_path)
+        tar_name = self._create_file_key(tar_name)
+        target_filename = os.path.basename(tar_name)
+        self._retrieve_tar(target_filename)
+        self._copy(os.path.basename(tar_name), target_path)
 
-        target_tar_file_name = os.path.basename(tar_name)
         if os.path.isdir(target_path):
-            target_path = os.path.join(target_path, target_tar_file_name)
+            return os.path.join(target_path, model_name)
 
-        if not os.path.exists(target_path):
-            structlogger.debug(
-                "persistor.retrieve.copy_model",
-                event_info=f"Copying model '{target_tar_file_name}' to "
-                f"'{target_path}'.",
-            )
-            self._copy(target_tar_file_name, target_path)
+        return target_path
+
+    def size_of_persisted_model(self, model_name: Text) -> int:
+        """Returns the size of the model that has been persisted to cloud storage.
 
-        structlogger.debug(
-            "persistor.retrieve.model_retrieved",
-            event_info=f"Model retrieved and saved to '{target_path}'.",
+        Args:
+            model_name: The name of the model to retrieve.
+        """
+        tar_name = model_name
+        if not model_name.endswith(MODEL_ARCHIVE_EXTENSION):
+            # ensure backward compatibility
+            tar_name = self._tar_name(model_name)
+        tar_name = self._create_file_key(tar_name)
+        target_filename = os.path.basename(tar_name)
+        return self._retrieve_tar_size(target_filename)
+
+    def _retrieve_tar_size(self, filename: Text) -> int:
+        """Returns the size of the model that has been persisted to cloud storage."""
+        structlogger.warning(
+            "persistor.retrieve_tar_size.not_implemented",
+            filename=filename,
+            event_info=(
+                "This method should be implemented in the persistor. "
+                "The default implementation will download the model "
+                "to calculate the size. Most persistors should override "
+                "this method to avoid downloading the model and get the "
+                "size directly from the cloud storage."
+            ),
         )
-        return target_path
+        self._retrieve_tar(filename)
+        return os.path.getsize(os.path.basename(filename))
 
     @abc.abstractmethod
     def _retrieve_tar(self, filename: Text) -> None:
@@ -175,7 +202,7 @@ class Persistor(abc.ABC):
             os.path.join(dirpath, base_name),
             "gztar",
             root_dir=model_directory,
-            base_dir=".",
+            base_dir="../nlu",
         )
         file_key = os.path.basename(tar_name)
         return file_key, tar_name
@@ -191,7 +218,7 @@ class Persistor(abc.ABC):
 
     @staticmethod
     def _create_file_key(model_path: str) -> Text:
-        """Appends remote storage folders when provided to upload or retrieve file"""
+        """Appends remote storage folders when provided to upload or retrieve file."""
         bucket_object_path = os.environ.get(REMOTE_STORAGE_PATH_ENV)
 
         # To keep the backward compatibility, if REMOTE_STORAGE_PATH is not provided,
@@ -203,10 +230,7 @@ class Persistor(abc.ABC):
                 f"{REMOTE_STORAGE_PATH_ENV} is deprecated and will be "
                 "removed in future versions. "
                 "Please use the -m path/to/model.tar.gz option to "
-                "specify the model path when loading a model."
-                "Or use --output and --fixed-model-name to specify the "
-                "output directory and the model name when saving a "
-                "trained model to remote storage.",
+                "specify the model path when loading a model.",
             )
 
         file_key = os.path.basename(model_path)
@@ -239,14 +263,13 @@ class AWSPersistor(Persistor):
     def _ensure_bucket_exists(
         self, bucket_name: Text, region_name: Optional[Text] = None
     ) -> None:
-        import botocore
+        from botocore import exceptions
 
         # noinspection PyUnresolvedReferences
         try:
             self.s3.meta.client.head_bucket(Bucket=bucket_name)
-        except botocore.exceptions.ClientError as e:
-            error_code = int(e.response["Error"]["Code"])
-            if error_code == HTTP_STATUS_FORBIDDEN:
+        except exceptions.ClientError as exc:
+            if self._error_code(exc) == HTTP_STATUS_FORBIDDEN:
                 log = (
                     f"Access to the specified bucket '{bucket_name}' is forbidden. "
                     "Please make sure you have the necessary "
@@ -258,7 +281,7 @@ class AWSPersistor(Persistor):
                     event_info=log,
                 )
                 raise RasaException(log)
-            elif error_code == HTTP_STATUS_NOT_FOUND:
+            elif self._error_code(exc) == HTTP_STATUS_NOT_FOUND:
                 log = (
                     f"The specified bucket '{bucket_name}' does not exist. "
                     "Please make sure to create the bucket first."
@@ -270,21 +293,57 @@ class AWSPersistor(Persistor):
                 )
                 raise RasaException(log)
 
+    @staticmethod
+    def _error_code(e: "ClientError") -> int:
+        return int(e.response["Error"]["Code"])
+
     def _persist_tar(self, file_key: Text, tar_path: Text) -> None:
         """Uploads a model persisted in the `target_dir` to s3."""
-        structlogger.debug(
-            "aws_persistor.persist_tar.uploading_model",
-            event_info=f"Uploading tar archive {file_key} to "
-            f"s3 bucket '{self.bucket_name}'.",
-        )
         with open(tar_path, "rb") as f:
             self.s3.Object(self.bucket_name, file_key).put(Body=f)
 
-    def _retrieve_tar(self, model_path: Text) -> None:
+    def _retrieve_tar_size(self, model_path: Text) -> int:
+        """Returns the size of the model that has been persisted to s3."""
+        try:
+            obj = self.s3.Object(self.bucket_name, model_path)
+            return obj.content_length
+        except Exception:
+            raise ModelNotFound()
+
+    def _retrieve_tar(self, target_filename: str) -> None:
         """Downloads a model that has previously been persisted to s3."""
-        tar_name = os.path.basename(model_path)
-        with open(tar_name, "wb") as f:
-            self.bucket.download_fileobj(model_path, f)
+        from botocore import exceptions
+
+        log = (
+            f"Model '{target_filename}' not found in the specified bucket "
+            f"'{self.bucket_name}'. Please make sure the model exists "
+            f"in the bucket."
+        )
+
+        try:
+            with open(target_filename, "wb") as f:
+                self.bucket.download_fileobj(target_filename, f)
+
+            structlogger.debug(
+                "aws_persistor.retrieve_tar.object_found", object_key=target_filename
+            )
+        except exceptions.ClientError as exc:
+            if self._error_code(exc) == HTTP_STATUS_NOT_FOUND:
+                structlogger.error(
+                    "aws_persistor.retrieve_tar.model_not_found",
+                    bucket_name=self.bucket_name,
+                    target_filename=target_filename,
+                    event_info=log,
+                )
+                raise ModelNotFound() from exc
+        except exceptions.BotoCoreError as exc:
+            structlogger.error(
+                "aws_persistor.retrieve_tar.model_download_error",
+                bucket_name=self.bucket_name,
+                target_filename=target_filename,
+                event_info=log,
+            )
+            raise ModelNotFound() from exc
 
 
 class GCSPersistor(Persistor):
@@ -309,42 +368,95 @@ class GCSPersistor(Persistor):
 
     def _ensure_bucket_exists(self, bucket_name: Text) -> None:
         from google.cloud import exceptions
+        from google.auth import exceptions as auth_exceptions
 
         try:
             self.storage_client.get_bucket(bucket_name)
-        except exceptions.NotFound:
+        except auth_exceptions.GoogleAuthError as exc:
+            log = (
+                f"An error occurred while authenticating with Google Cloud "
+                f"Storage. Please make sure you have the necessary credentials "
+                f"to access the bucket '{bucket_name}'."
+            )
+            structlogger.error(
+                "gcp_persistor.ensure_bucket_exists.authentication_error",
+                bucket_name=bucket_name,
+                event_info=log,
+            )
+            raise RasaException(log) from exc
+        except exceptions.NotFound as exc:
             log = (
-                f"The specified bucket '{bucket_name}' does not exist. "
-                "Please make sure to create the bucket first."
+                f"The specified Google Cloud Storage bucket '{bucket_name}' "
+                f"does not exist. Please make sure to create the bucket first or "
+                f"provide an alternative valid bucket name."
             )
             structlogger.error(
                 "gcp_persistor.ensure_bucket_exists.bucket_not_found",
                 bucket_name=bucket_name,
                 event_info=log,
             )
-            raise RasaException(log)
-        except exceptions.Forbidden:
+            raise RasaException(log) from exc
+        except exceptions.Forbidden as exc:
             log = (
-                f"Access to the specified bucket '{bucket_name}' is forbidden. "
-                "Please make sure you have the necessary "
-                "permission to access the bucket. "
+                f"Access to the specified Google Cloud storage bucket '{bucket_name}' "
+                f"is forbidden. Please make sure you have the necessary "
+                f"permissions to access the bucket. "
             )
             structlogger.error(
                 "gcp_persistor.ensure_bucket_exists.bucket_access_forbidden",
                 bucket_name=bucket_name,
                 event_info=log,
             )
-            raise RasaException(log)
+            raise RasaException(log) from exc
+        except ValueError as exc:
+            # bucket_name is None
+            log = (
+                "The specified Google Cloud Storage bucket name is None. Please "
+                "make sure to provide a valid bucket name."
+            )
+            structlogger.error(
+                "gcp_persistor.ensure_bucket_exists.bucket_name_none",
+                event_info=log,
+            )
+            raise RasaException(log) from exc
 
     def _persist_tar(self, file_key: Text, tar_path: Text) -> None:
         """Uploads a model persisted in the `target_dir` to GCS."""
         blob = self.bucket.blob(file_key)
         blob.upload_from_filename(tar_path)
 
+    def _retrieve_tar_size(self, target_filename: Text) -> int:
+        """Returns the size of the model that has been persisted to GCS."""
+        try:
+            blob = self.bucket.blob(target_filename)
+            return blob.size
+        except Exception:
+            raise ModelNotFound()
+
     def _retrieve_tar(self, target_filename: Text) -> None:
         """Downloads a model that has previously been persisted to GCS."""
+        from google.api_core import exceptions
+
         blob = self.bucket.blob(target_filename)
-        blob.download_to_filename(os.path.basename(target_filename))
+        try:
+            blob.download_to_filename(target_filename)
+
+            structlogger.debug(
+                "gcs_persistor.retrieve_tar.object_found", object_key=target_filename
+            )
+        except exceptions.NotFound as exc:
+            log = (
+                f"Model '{target_filename}' not found in the specified bucket "
+                f"'{self.bucket_name}'. Please make sure the model exists "
+                f"in the bucket."
+            )
+            structlogger.error(
+                "gcp_persistor.retrieve_tar.model_not_found",
+                bucket_name=self.bucket_name,
+                target_filename=target_filename,
+                event_info=log,
+            )
+            raise ModelNotFound() from exc
 
 
 class AzurePersistor(Persistor):
@@ -370,7 +482,8 @@ class AzurePersistor(Persistor):
         else:
             log = (
                 f"The specified container '{self.container_name}' does not exist."
-                "Please make sure to create the container first."
+                "Please make sure to create the bucket first or "
+                f"provide an alternative valid bucket name."
             )
             structlogger.error(
                 "azure_persistor.ensure_container_exists.container_not_found",
@@ -385,19 +498,41 @@ class AzurePersistor(Persistor):
     def _persist_tar(self, file_key: Text, tar_path: Text) -> None:
         """Uploads a model persisted in the `target_dir` to Azure."""
         with open(tar_path, "rb") as data:
-            self._container_client().upload_blob(
-                name=file_key,
-                data=data,
-                # overwrite is set to True to keep in line with
-                # how GCS and AWS APIs work this enables easy
-                # updating of models in the cloud
-                overwrite=True,
-            )
+            self._container_client().upload_blob(name=file_key, data=data)
+
+    def _retrieve_tar_size(self, target_filename: Text) -> int:
+        """Returns the size of the model that has been persisted to Azure."""
+        try:
+            blob_client = self._container_client().get_blob_client(target_filename)
+            properties = blob_client.get_blob_properties()
+            return properties.size
+        except Exception:
+            raise ModelNotFound()
 
     def _retrieve_tar(self, target_filename: Text) -> None:
         """Downloads a model that has previously been persisted to Azure."""
-        blob_client = self._container_client().get_blob_client(target_filename)
+        from azure.core.exceptions import AzureError
 
-        with open(os.path.basename(target_filename), "wb") as blob:
-            download_stream = blob_client.download_blob()
-            blob.write(download_stream.readall())
+        try:
+            with open(target_filename, "wb") as model_file:
+                blob_client = self._container_client().get_blob_client(target_filename)
+                download_stream = blob_client.download_blob()
+                model_file.write(download_stream.readall())
+            structlogger.debug(
+                "azure_persistor.retrieve_tar.blob_found", blob_name=target_filename
+            )
+        except AzureError as exc:
+            log = (
+                f"An exception occurred while trying to download "
+                f"the model '{target_filename}' in the specified container "
+                f"'{self.container_name}'. Please make sure the model exists "
+                f"in the container."
+            )
+            structlogger.error(
+                "azure_persistor.retrieve_tar.model_download_error",
+                container_name=self.container_name,
+                target_filename=target_filename,
+                event_info=log,
+                exception=exc,
+            )
+            raise ModelNotFound() from exc