rasa-pro 3.11.0a2__py3-none-any.whl → 3.11.0a4.dev1__py3-none-any.whl

rasa/model_manager/trainer_service.py

@@ -0,0 +1,332 @@
+import os
+from typing import Any, Dict, Optional
+import shutil
+import structlog
+import subprocess
+from rasa.constants import MODEL_ARCHIVE_EXTENSION
+from rasa.model_manager.utils import (
+    models_base_path,
+    subpath,
+    write_encoded_data_to_file,
+)
+from pydantic import BaseModel, ConfigDict
+from enum import Enum
+
+from rasa.model_manager import config
+from rasa.model_training import generate_random_model_name
+from rasa.model_manager.utils import ensure_base_directory_exists, logs_path
+
+structlogger = structlog.get_logger()
+
+
+class TrainingSessionStatus(str, Enum):
+    """Enum for the training status."""
+
+    RUNNING = "running"
+    STOPPED = "stopped"
+    DONE = "done"
+    ERROR = "error"
+
+
+class TrainingSession(BaseModel):
+    """Store information about a training session."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    training_id: str
+    assistant_id: str
+    client_id: Optional[str]
+    progress: int
+    model_name: str
+    status: TrainingSessionStatus
+    process: subprocess.Popen
+
+    def is_status_indicating_alive(self) -> bool:
+        """Check if the training is running."""
+        return self.status == TrainingSessionStatus.RUNNING
+
+    def model_path(self) -> str:
+        """Return the path to the model."""
+        return subpath(models_base_path(), f"{self.model_name}.tar.gz")
+
+
+def train_path(training_id: str) -> str:
+    """Return the path to the training directory for a given training id."""
+    return subpath(config.SERVER_BASE_WORKING_DIRECTORY, f"trainings/{training_id}")
+
+
+def cache_for_assistant_path(assistant_id: str) -> str:
+    """Return the path to the cache directory for a given assistant id."""
+    return subpath(config.SERVER_BASE_WORKING_DIRECTORY, f"caches/{assistant_id}")
+
+
+def terminate_training(training: TrainingSession) -> None:
+    if not training.is_status_indicating_alive():
+        # no-op if the training is not running
+        return
+
+    structlogger.info(
+        "model_trainer.user_stopping_training", training_id=training.training_id
+    )
+    try:
+        training.process.terminate()
+        training.status = TrainingSessionStatus.STOPPED
+    except ProcessLookupError:
+        structlogger.debug(
+            "model_trainer.training_process_not_found",
+            training_id=training.training_id,
+        )
+    finally:
+        clean_up_after_training(training)
+
+
+def update_training_status(training: TrainingSession) -> None:
+    if not training.is_status_indicating_alive():
+        # skip if the training is not running
+        return
+    if training.process.poll() is None:
+        # process is still running
+        return
+
+    complete_training(training)
+
+
+def complete_training(training: TrainingSession) -> None:
+    """Complete a training session.
+
+    Transitions the status of a training process to "done" if the process has
+    finished successfully, and to "error" if the process has finished with an
+    error.
+    """
+    if training.process.returncode == 0:
+        training.status = TrainingSessionStatus.DONE
+    else:
+        training.status = TrainingSessionStatus.ERROR
+
+    training.progress = 100
+
+    structlogger.info(
+        "model_trainer.training_finished",
+        training_id=training.training_id,
+        status=training.status,
+    )
+
+    # persist the assistant cache to speed up future training runs for this
+    # assistant
+    persist_rasa_cache(training.assistant_id, train_path(training.training_id))
+    move_model_to_local_storage(training)
+
+    clean_up_after_training(training)
+
+
+def clean_up_after_training(training: TrainingSession) -> None:
+    """Clean up the training directory."""
+    structlogger.debug(
+        "model_trainer.cleaning_up_training", training_id=training.training_id
+    )
+    shutil.rmtree(train_path(training.training_id), ignore_errors=True)
+
+
+def move_model_to_local_storage(training: TrainingSession) -> None:
+    """Persist the model to the remote storage."""
+    ensure_base_directory_exists(models_base_path())
+
+    model_path = subpath(
+        train_path(training.training_id),
+        f"models/{training.model_name}.{MODEL_ARCHIVE_EXTENSION}",
+    )
+
+    if os.path.exists(model_path):
+        structlogger.debug(
+            "model_trainer.persisting_model_to_models_dir",
+            training_model_path=model_path,
+            storage_model_path=models_base_path(),
+        )
+        shutil.move(model_path, models_base_path())
+    else:
+        structlogger.warning(
+            "model_trainer.model_not_found_after_training",
+            training_id=training.training_id,
+            model_path=model_path,
+        )
+
+
+def seed_training_directory_with_rasa_cache(
+    training_base_path: str, assistant_id: str
+) -> None:
+    """Populate the training directory with the cache of a previous training."""
+    # check if there is a cache for this assistant
+    cache_path = cache_for_assistant_path(assistant_id)
+
+    if os.path.exists(cache_path):
+        structlogger.debug(
+            "model_trainer.populating_training_dir_with_cache",
+            assistant_id=assistant_id,
+            training_base_path=training_base_path,
+        )
+        # copy the cache to the training directory
+        shutil.copytree(src=cache_path, dst=subpath(training_base_path, ".rasa"))
+
+
+def persist_rasa_cache(assistant_id: str, training_base_path: str) -> None:
+    """Persist the cache of a training session to speed up future trainings."""
+    # copy the cache from the training directory to the cache directory
+    # cache files are stored inside of `/.rasa/` of the training folder
+    structlogger.debug(
+        "model_trainer.persisting_assistant_cache", assistant_id=assistant_id
+    )
+    cache_path = cache_for_assistant_path(assistant_id)
+
+    # if the training failed and didn't create a cache, skip this step
+    if not os.path.exists(subpath(training_base_path, ".rasa")):
+        return
+
+    # clean up the cache directory first
+    shutil.rmtree(cache_path, ignore_errors=True)
+    shutil.copytree(src=subpath(training_base_path, ".rasa"), dst=cache_path)
+
+
+def write_training_data_to_files(
+    encoded_training_data: Dict[str, Any], training_base_path: str
+) -> None:
+    """Write the training data to files in the training directory.
+
+    Incoming data format, all keys being optional:
+    ````
+    {
+        "domain": "base64 encoded domain.yml",
+        "credentials": "base64 encoded credentials.yml",
+        "endpoints": "base64 encoded endpoints.yml",
+        "flows": "base64 encoded flows.yml",
+        "config": "base64 encoded config.yml",
+        "stories": "base64 encoded stories.yml",
+        "rules": "base64 encoded rules.yml",
+        "nlu": "base64 encoded nlu.yml"
+    }
+    ```
+    """
+    data_to_be_written_to_files = {
+        "domain": "domain.yml",
+        "credentials": "credentials.yml",
+        "endpoints": "endpoints.yml",
+        "flows": "data/flows.yml",
+        "config": "config.yml",
+        "stories": "data/stories.yml",
+        "rules": "data/rules.yml",
+        "nlu": "data/nlu.yml",
+    }
+
+    for key, file in data_to_be_written_to_files.items():
+        write_encoded_data_to_file(
+            encoded_training_data.get(key, ""),
+            subpath(training_base_path, file),
+        )
+
+
+def prepare_training_directory(
+    training_base_path: str, assistant_id: str, encoded_training_data: Dict[str, Any]
+) -> None:
+    """Prepare the training directory for a new training session."""
+    # create a new working directory and store the training data from the
+    # request there. the training data in the request is base64 encoded
+    os.makedirs(training_base_path, exist_ok=True)
+
+    seed_training_directory_with_rasa_cache(training_base_path, assistant_id)
+    write_training_data_to_files(encoded_training_data, training_base_path)
+    structlogger.debug("model_trainer.prepared_training", path=training_base_path)
+
+
+def start_training_process(
+    training_id: str,
+    assistant_id: str,
+    client_id: str,
+    training_base_path: str,
+) -> TrainingSession:
+    log_path = logs_path(training_id)
+
+    ensure_base_directory_exists(log_path)
+
+    model_name = generate_random_model_name()
+    # Start the training in a subprocess
+    # set the working directory to the training directory
+    # run the rasa train command as a subprocess, activating poetry before running
+    # pipe the stdout and stderr to the same file
+    full_command = [
+        config.RASA_PYTHON_PATH,
+        "-m",
+        "rasa.__main__",
+        "train",
+        "--debug",
+        "--data",
+        "data",
+        "--config",
+        "config.yml",
+        "--domain",
+        "domain.yml",
+        "--endpoints",
+        "endpoints.yml",
+        "--fixed-model-name",
+        f"{model_name}.{MODEL_ARCHIVE_EXTENSION}",
+        "--out",
+        "models",
+    ]
+
+    if config.SERVER_MODEL_REMOTE_STORAGE:
+        full_command.extend(
+            [
+                "--keep-local-model-copy",
+                "--remote-storage",
+                config.SERVER_MODEL_REMOTE_STORAGE,
+            ]
+        )
+
+    structlogger.debug("model_trainer.training_command", command=" ".join(full_command))
+
+    envs = os.environ.copy()
+    envs["RASA_TELEMETRY_ENABLED"] = "false"
+
+    process = subprocess.Popen(
+        full_command,
+        cwd=training_base_path,
+        stdout=open(log_path, "w"),
+        stderr=subprocess.STDOUT,
+        env=envs,
+    )
+
+    structlogger.info(
+        "model_trainer.training_started",
+        training_id=training_id,
+        assistant_id=assistant_id,
+        model_name=model_name,
+        client_id=client_id,
+        log=log_path,
+        pid=process.pid,
+    )
+
+    return TrainingSession(
+        training_id=training_id,
+        assistant_id=assistant_id,
+        client_id=client_id,
+        model_name=model_name,
+        progress=0,
+        status=TrainingSessionStatus.RUNNING,
+        process=process,  # Store the process handle
+    )
+
+
+def run_training(
+    training_id: str,
+    assistant_id: str,
+    client_id: str,
+    encoded_training_data: Dict,
+) -> TrainingSession:
+    """Run a training session."""
+    training_base_path = train_path(training_id)
+
+    prepare_training_directory(training_base_path, assistant_id, encoded_training_data)
+    return start_training_process(
+        training_id=training_id,
+        assistant_id=assistant_id,
+        client_id=client_id,
+        training_base_path=training_base_path,
+    )

rasa/model_manager/utils.py

@@ -0,0 +1,66 @@
+import os
+import base64
+from typing import Optional
+
+from rasa.model_manager import config
+
+
+def write_encoded_data_to_file(encoded_data: bytes, file: str) -> None:
+    """Write base64 encoded data to a file."""
+    # create the directory if it does not exist of the parent directory
+    os.makedirs(os.path.dirname(file), exist_ok=True)
+
+    with open(file, "w") as f:
+        decoded = base64.b64decode(encoded_data)
+        text = decoded.decode("utf-8")
+        f.write(text)
+
+
+def logs_base_path() -> str:
+    """Return the path to the logs directory."""
+    return subpath(config.SERVER_BASE_WORKING_DIRECTORY, "logs")
+
+
+def ensure_base_directory_exists(directory: str) -> None:
+    """Ensure that a files parent directory exists.
+
+    Args:
+        directory: The directory to check.
+    """
+    os.makedirs(os.path.dirname(directory), exist_ok=True)
+
+
+def models_base_path() -> str:
+    """Return the path to the models directory."""
+    return subpath(config.SERVER_BASE_WORKING_DIRECTORY, "models")
+
+
+def logs_path(action_id: str) -> str:
+    """Return the path to the log file for a given action id.
+
+    Args:
+        action_id: can either be a training_id or a deployment_id
+    """
+    return subpath(logs_base_path(), f"{action_id}.txt")
+
+
+def subpath(parent: str, child: str) -> str:
+    """Return the path to the child directory of the parent directory.
+
+    Ensures, that child doesn't navigate to parent directories. Prevents
+    path traversal.
+    """
+    path = os.path.abspath(os.path.join(parent, child))
+    if not path.startswith(os.path.abspath(parent)):
+        raise ValueError(f"Invalid path: {path}")
+
+    return path
+
+
+def get_logs_content(action_id: str) -> Optional[str]:
+    """Return the content of the log file for a given action id."""
+    try:
+        with open(logs_path(action_id), "r") as file:
+            return file.read()
+    except FileNotFoundError:
+        return None

rasa/model_service.py

@@ -0,0 +1,109 @@
+import logging
+
+from sanic import Sanic
+import structlog
+
+from rasa.cli.scaffold import print_error_and_exit
+from rasa.core.persistor import RemoteStorageType, get_persistor
+from rasa.core.utils import list_routes
+from rasa.model_manager import model_api
+from rasa.model_manager import config
+from rasa.model_manager.config import SERVER_BASE_URL
+from rasa.utils.common import configure_logging_and_warnings
+import rasa.utils.licensing
+from urllib.parse import urlparse
+
+structlogger = structlog.get_logger()
+
+MODEL_SERVICE_PORT = 8000
+
+DEFAULT_SERVER_PATH_PREFIX = "bot"
+
+
+def url_prefix_from_base_url() -> str:
+    """Return the path prefix from the base URL."""
+    # return path without any trailing slashes
+    prefix = urlparse(SERVER_BASE_URL).path.rstrip("/") if SERVER_BASE_URL else ""
+
+    # can't be empty
+    return prefix or DEFAULT_SERVER_PATH_PREFIX
+
+
+def validate_model_storage_type() -> None:
+    """Validate the storage type if remote storage is used for models."""
+    if config.SERVER_MODEL_REMOTE_STORAGE:
+        if config.SERVER_MODEL_REMOTE_STORAGE not in RemoteStorageType.list():
+            print_error_and_exit(
+                f"Invalid storage type '{config.SERVER_MODEL_REMOTE_STORAGE}'. "
+                f"Supported storage types: {', '.join(RemoteStorageType.list())}."
+                f"Alternatively, unset the remote storage type to store models locally."
+            )
+        else:
+            structlogger.info(
+                "model_api.storage.remote_storage_enabled",
+                remote_storage=config.SERVER_MODEL_REMOTE_STORAGE,
+            )
+        # try to create a client to validate the configuration
+        get_persistor(config.SERVER_MODEL_REMOTE_STORAGE)
+    else:
+        structlogger.info(
+            "model_api.storage.local_storage_enabled",
+            base_path=config.SERVER_BASE_WORKING_DIRECTORY,
+        )
+
+
+def _register_update_task(app: Sanic) -> None:
+    app.add_task(
+        model_api.continuously_update_process_status,
+        name="continuously_update_process_status",
+    )
+
+
+def main() -> None:
+    """Start the Rasa Model Manager server.
+
+    The API server can receive requests to train models, run bots, and manage
+    the lifecycle of models and bots.
+    """
+    configure_logging_and_warnings(
+        log_level=logging.DEBUG,
+        logging_config_file=None,
+        warn_only_once=True,
+        filter_repeated_logs=True,
+    )
+
+    rasa.utils.licensing.validate_license_from_env()
+
+    try:
+        model_api.prepare_working_directories()
+    except Exception as e:
+        structlogger.error(
+            "model_api.prepare_directories.failed",
+            error=str(e),
+            base_directory=config.SERVER_BASE_WORKING_DIRECTORY,
+        )
+        print_error_and_exit(
+            f"Failed to create working directories. Please make sure the "
+            f"server base directory at '{config.SERVER_BASE_WORKING_DIRECTORY}' "
+            f"is writable by the current user."
+        )
+
+    validate_model_storage_type()
+
+    structlogger.debug("model_api.starting_server", port=MODEL_SERVICE_PORT)
+
+    url_prefix = url_prefix_from_base_url()
+    # configure the sanic application
+    app = Sanic("RasaModelService")
+    app.after_server_start(_register_update_task)
+    app.blueprint(model_api.external_blueprint(), url_prefix=url_prefix)
+    app.blueprint(model_api.internal_blueprint())
+
+    # list all routes
+    list_routes(app)
+
+    app.run(host="0.0.0.0", port=MODEL_SERVICE_PORT, legacy=True)
+
+
+if __name__ == "__main__":
+    main()

rasa/model_training.py

@@ -157,6 +157,7 @@ async def train(
     finetuning_epoch_fraction: float = 1.0,
     remote_storage: Optional[StorageType] = None,
     file_importer: Optional[TrainingDataImporter] = None,
+    keep_local_model_copy: bool = False,
 ) -> TrainingResult:
     """Trains a Rasa model (Core and NLU).
 
@@ -182,6 +183,8 @@ async def train(
             use for storing the model.
         file_importer: Instance of `TrainingDataImporter` to use for training.
             If it is not provided, a new instance will be created.
+        keep_local_model_copy: If `True` the model will be stored locally even if
+            remote storage is configured.
 
     Returns:
         An instance of `TrainingResult`.
@@ -263,6 +266,7 @@ async def train(
             finetuning_epoch_fraction=finetuning_epoch_fraction,
             dry_run=dry_run,
             remote_storage=remote_storage,
+            keep_local_model_copy=keep_local_model_copy,
             **(core_additional_arguments or {}),
             **(nlu_additional_arguments or {}),
         )
@@ -277,6 +281,7 @@ async def _train_graph(
     force_full_training: bool = False,
     dry_run: bool = False,
     remote_storage: Optional[StorageType] = None,
+    keep_local_model_copy: bool = False,
     **kwargs: Any,
 ) -> TrainingResult:
     if model_to_finetune:
@@ -339,7 +344,7 @@ async def _train_graph(
             )
             return _dry_run_result(fingerprint_status, force_full_training)
 
-        model_name = _determine_model_name(fixed_model_name, training_type)
+        model_name = determine_model_name(fixed_model_name, training_type)
         full_model_path = Path(output_path, model_name)
 
         with telemetry.track_model_training(
@@ -354,7 +359,8 @@ async def _train_graph(
             )
             if remote_storage:
                 push_model_to_remote_storage(full_model_path, remote_storage)
-                full_model_path.unlink()
+                if not keep_local_model_copy:
+                    full_model_path.unlink()
                 structlogger.info(
                     "model_training.train.finished_training",
                     event_info=(
@@ -386,9 +392,14 @@ def _create_model_storage(
     return model_storage
 
 
-def _determine_model_name(
+def generate_random_model_name() -> str:
+    time_format = "%Y%m%d-%H%M%S"
+    return f"{time.strftime(time_format)}-{randomname.get_name()}"
+
+
+def determine_model_name(
     fixed_model_name: Optional[Text], training_type: TrainingType
-) -> Text:
+) -> str:
     if fixed_model_name:
         if not fixed_model_name.endswith(".tar.gz"):
             return f"{fixed_model_name}.tar.gz"
@@ -398,8 +409,7 @@ def _determine_model_name(
     if training_type in [TrainingType.CORE, TrainingType.NLU]:
         prefix = f"{training_type.model_type}-"
 
-    time_format = "%Y%m%d-%H%M%S"
-    return f"{prefix}{time.strftime(time_format)}-{randomname.get_name()}.tar.gz"
+    return f"{prefix}{generate_random_model_name()}.tar.gz"
 
 
 async def train_core(
@@ -411,6 +421,7 @@ async def train_core(
     additional_arguments: Optional[Dict] = None,
     model_to_finetune: Optional[Text] = None,
     finetuning_epoch_fraction: float = 1.0,
+    keep_local_model_copy: bool = False,
 ) -> Optional[Text]:
     """Trains a Core model.
 
@@ -425,6 +436,8 @@ async def train_core(
             a directory in case the latest trained model should be used.
         finetuning_epoch_fraction: The fraction currently specified training epochs
             in the model configuration which should be used for finetuning.
+        keep_local_model_copy: If `True` the model will be stored locally even if
+            remote storage is configured.
 
     Returns:
         Path to the model archive.
@@ -480,6 +493,7 @@ async def train_core(
             model_to_finetune=model_to_finetune,
             fixed_model_name=fixed_model_name,
             finetuning_epoch_fraction=finetuning_epoch_fraction,
+            keep_local_model_copy=keep_local_model_copy,
             **(additional_arguments or {}),
         )
     ).model
@@ -495,6 +509,7 @@ async def train_nlu(
     domain: Optional[Union[Domain, Text]] = None,
     model_to_finetune: Optional[Text] = None,
     finetuning_epoch_fraction: float = 1.0,
+    keep_local_model_copy: bool = False,
 ) -> Optional[Text]:
     """Trains an NLU model.
 
@@ -512,6 +527,8 @@ async def train_nlu(
             a directory in case the latest trained model should be used.
         finetuning_epoch_fraction: The fraction currently specified training epochs
             in the model configuration which should be used for finetuning.
+        keep_local_model_copy: If `True` the model will be stored locally even if
+            remote storage is configured.
 
     Returns:
         Path to the model archive.
@@ -553,13 +570,14 @@ async def train_nlu(
             fixed_model_name=fixed_model_name,
             finetuning_epoch_fraction=finetuning_epoch_fraction,
             persist_nlu_training_data=persist_nlu_training_data,
+            keep_local_model_copy=keep_local_model_copy,
             **(additional_arguments or {}),
         )
     ).model
 
 
 def push_model_to_remote_storage(model_path: Path, remote_storage: StorageType) -> None:
-    """push model to remote storage"""
+    """Push model to remote storage."""
     from rasa.core.persistor import get_persistor
 
     persistor = get_persistor(remote_storage)

rasa/shared/constants.py

@@ -212,6 +212,9 @@ AZURE_OPENAI_PROVIDER = "azure"
 SELF_HOSTED_PROVIDER = "self-hosted"
 HUGGINGFACE_LOCAL_EMBEDDING_PROVIDER = "huggingface_local"
 
+SELF_HOSTED_VLLM_PREFIX = "hosted_vllm"
+SELF_HOSTED_VLLM_API_KEY_ENV_VAR = "HOSTED_VLLM_API_KEY"
+
 AZURE_API_TYPE = "azure"
 OPENAI_API_TYPE = "openai"
 
@@ -241,3 +244,6 @@ RASA_PATTERN_CANNOT_HANDLE_INVALID_INTENT = (
 )
 
 ROUTE_TO_CALM_SLOT = "route_session_to_calm"
+
+ORIGINAL_VALUE = "original_value"
+RESOLVED_VALUE = "resolved_value"

rasa/shared/core/constants.py

@@ -48,6 +48,7 @@ ACTION_TRIGGER_SEARCH = "action_trigger_search"
 ACTION_TRIGGER_CHITCHAT = "action_trigger_chitchat"
 ACTION_RESET_ROUTING = "action_reset_routing"
 ACTION_HANGUP = "action_hangup"
+ACTION_REPEAT_BOT_MESSAGES = "action_repeat_bot_messages"
 
 
 DEFAULT_ACTION_NAMES = [
@@ -74,6 +75,7 @@ DEFAULT_ACTION_NAMES = [
     ACTION_TRIGGER_CHITCHAT,
     ACTION_RESET_ROUTING,
     ACTION_HANGUP,
+    ACTION_REPEAT_BOT_MESSAGES,
 ]
 
 ACTION_SHOULD_SEND_DOMAIN = "send_domain"