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

rasa/core/persistor.py

@@ -4,19 +4,16 @@ 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,
@@ -31,7 +28,6 @@ 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()
 
@@ -90,15 +86,14 @@ def get_persistor(storage: StorageType) -> Optional[Persistor]:
 
     if storage == RemoteStorageType.AWS.value:
         return AWSPersistor(
-            os.environ.get(BUCKET_NAME_ENV, DEFAULT_BUCKET_NAME),
-            os.environ.get(AWS_ENDPOINT_URL_ENV),
+            os.environ.get(BUCKET_NAME_ENV), os.environ.get(AWS_ENDPOINT_URL_ENV)
         )
     if storage == RemoteStorageType.GCS.value:
-        return GCSPersistor(os.environ.get(BUCKET_NAME_ENV, DEFAULT_BUCKET_NAME))
+        return GCSPersistor(os.environ.get(BUCKET_NAME_ENV))
 
     if storage == RemoteStorageType.AZURE.value:
         return AzurePersistor(
-            os.environ.get(AZURE_CONTAINER_ENV, DEFAULT_BUCKET_NAME),
+            os.environ.get(AZURE_CONTAINER_ENV),
             os.environ.get(AZURE_ACCOUNT_NAME_ENV),
             os.environ.get(AZURE_ACCOUNT_KEY_ENV),
         )
@@ -123,8 +118,7 @@ class Persistor(abc.ABC):
 
     def persist(self, trained_model: str) -> None:
         """Uploads a trained model persisted in the `target_dir` to cloud storage."""
-        absolute_file_key = self._create_file_key(trained_model)
-        file_key = Path(absolute_file_key).name
+        file_key = self._create_file_key(trained_model)
         self._persist_tar(file_key, trained_model)
 
     def retrieve(self, model_name: Text, target_path: Text) -> Text:
@@ -143,8 +137,7 @@ class Persistor(abc.ABC):
             # 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)
-        self._retrieve_tar(target_filename)
+        self._retrieve_tar(tar_name)
         self._copy(os.path.basename(tar_name), target_path)
 
         if os.path.isdir(target_path):
@@ -152,36 +145,6 @@ class Persistor(abc.ABC):
 
         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.
-
-        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."
-            ),
-        )
-        self._retrieve_tar(filename)
-        return os.path.getsize(os.path.basename(filename))
-
     @abc.abstractmethod
     def _retrieve_tar(self, filename: Text) -> None:
         """Downloads a model previously persisted to cloud storage."""
@@ -218,7 +181,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,
@@ -230,7 +193,10 @@ 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.",
+                "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.",
             )
 
         file_key = os.path.basename(model_path)
@@ -263,13 +229,14 @@ class AWSPersistor(Persistor):
     def _ensure_bucket_exists(
         self, bucket_name: Text, region_name: Optional[Text] = None
     ) -> None:
-        from botocore import exceptions
+        import botocore
 
         # noinspection PyUnresolvedReferences
         try:
             self.s3.meta.client.head_bucket(Bucket=bucket_name)
-        except exceptions.ClientError as exc:
-            if self._error_code(exc) == HTTP_STATUS_FORBIDDEN:
+        except botocore.exceptions.ClientError as e:
+            error_code = int(e.response["Error"]["Code"])
+            if error_code == HTTP_STATUS_FORBIDDEN:
                 log = (
                     f"Access to the specified bucket '{bucket_name}' is forbidden. "
                     "Please make sure you have the necessary "
@@ -281,7 +248,7 @@ class AWSPersistor(Persistor):
                     event_info=log,
                 )
                 raise RasaException(log)
-            elif self._error_code(exc) == HTTP_STATUS_NOT_FOUND:
+            elif error_code == HTTP_STATUS_NOT_FOUND:
                 log = (
                     f"The specified bucket '{bucket_name}' does not exist. "
                     "Please make sure to create the bucket first."
@@ -293,57 +260,16 @@ 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."""
         with open(tar_path, "rb") as f:
             self.s3.Object(self.bucket_name, file_key).put(Body=f)
 
-    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:
+    def _retrieve_tar(self, model_path: Text) -> None:
         """Downloads a model that has previously been persisted to s3."""
-        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
+        tar_name = os.path.basename(model_path)
+        with open(tar_name, "wb") as f:
+            self.bucket.download_fileobj(model_path, f)
 
 
 class GCSPersistor(Persistor):
@@ -368,95 +294,42 @@ 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 auth_exceptions.GoogleAuthError as exc:
+        except exceptions.NotFound:
             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 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."
+                f"The specified bucket '{bucket_name}' does not exist. "
+                "Please make sure to create the bucket first."
             )
             structlogger.error(
                 "gcp_persistor.ensure_bucket_exists.bucket_not_found",
                 bucket_name=bucket_name,
                 event_info=log,
             )
-            raise RasaException(log) from exc
-        except exceptions.Forbidden as exc:
+            raise RasaException(log)
+        except exceptions.Forbidden:
             log = (
-                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. "
+                f"Access to the specified bucket '{bucket_name}' is forbidden. "
+                "Please make sure you have the necessary "
+                "permission to access the bucket. "
             )
             structlogger.error(
                 "gcp_persistor.ensure_bucket_exists.bucket_access_forbidden",
                 bucket_name=bucket_name,
                 event_info=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
+            raise RasaException(log)
 
     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)
-        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
+        blob.download_to_filename(target_filename)
 
 
 class AzurePersistor(Persistor):
@@ -482,8 +355,7 @@ class AzurePersistor(Persistor):
         else:
             log = (
                 f"The specified container '{self.container_name}' does not exist."
-                "Please make sure to create the bucket first or "
-                f"provide an alternative valid bucket name."
+                "Please make sure to create the container first."
             )
             structlogger.error(
                 "azure_persistor.ensure_container_exists.container_not_found",
@@ -500,39 +372,10 @@ class AzurePersistor(Persistor):
         with open(tar_path, "rb") as data:
             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."""
-        from azure.core.exceptions import AzureError
+        blob_client = self._container_client().get_blob_client(target_filename)
 
-        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
+        with open(target_filename, "wb") as blob:
+            download_stream = blob_client.download_blob()
+            blob.write(download_stream.readall())