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

rasa/model_manager/trainer_service.py

@@ -0,0 +1,325 @@
+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_manager.warm_rasa_process import (
+    start_rasa_process,
+)
+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
+    log_id: str
+
+    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 + "/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 + "/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) + "/models",
+        f"{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():
+        parent_path, file_name = os.path.split(file)
+
+        write_encoded_data_to_file(
+            encoded_training_data.get(key, ""),
+            subpath(training_base_path + "/" + parent_path, file_name),
+        )
+
+
+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:
+    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
+    arguments = [
+        "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:
+        arguments.extend(
+            [
+                "--keep-local-model-copy",
+                "--remote-storage",
+                config.SERVER_MODEL_REMOTE_STORAGE,
+            ]
+        )
+
+    structlogger.debug(
+        "model_trainer.training_arguments", arguments=" ".join(arguments)
+    )
+
+    warm_process = start_rasa_process(cwd=training_base_path, arguments=arguments)
+
+    structlogger.info(
+        "model_trainer.training_started",
+        training_id=training_id,
+        assistant_id=assistant_id,
+        model_name=model_name,
+        client_id=client_id,
+        log=logs_path(warm_process.log_id),
+        pid=warm_process.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=warm_process.process,  # Store the process handle
+        log_id=warm_process.log_id,
+    )
+
+
+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,87 @@
+import os
+import base64
+from typing import Optional
+
+import structlog
+
+from rasa.model_manager import config
+from rasa.shared.exceptions import RasaException
+
+structlogger = structlog.get_logger()
+
+
+class InvalidPathException(RasaException):
+    """Raised if a path is invalid - e.g. path traversal is detected."""
+
+
+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. Raises an InvalidPathException if the path is invalid.
+
+    Based on Snyk's directory traversal mitigation:
+    https://learn.snyk.io/lesson/directory-traversal/
+    """
+    safe_path = os.path.abspath(os.path.join(parent, child))
+    parent = os.path.abspath(parent)
+
+    common_base = os.path.commonpath([parent, safe_path])
+    if common_base != parent:
+        raise InvalidPathException(f"Invalid path: {safe_path}")
+
+    if os.path.basename(safe_path) != child:
+        raise InvalidPathException(
+            f"Invalid path - path traversal detected: {safe_path}"
+        )
+
+    return safe_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:
+        structlogger.debug("model_service.logs.not_found", action_id=action_id)
+        return None

rasa/model_manager/warm_rasa_process.py

@@ -0,0 +1,187 @@
+import shlex
+import subprocess
+from rasa.__main__ import main
+import os
+from typing import List
+import structlog
+from dataclasses import dataclass
+import uuid
+
+from rasa.model_manager import config
+from rasa.model_manager.utils import ensure_base_directory_exists, logs_path
+
+structlogger = structlog.get_logger(__name__)
+
+warm_rasa_processes: List["WarmRasaProcess"] = []
+
+NUMBER_OF_INITIAL_PROCESSES = 3
+
+
+@dataclass
+class WarmRasaProcess:
+    """Data class to store a warm Rasa process.
+
+    A "warm" Rasa process is one where we've done the heavy lifting of
+    importing key modules ahead of time (e.g. litellm). This is to avoid
+    long import times when we actually want to run a command.
+
+    This is a started process waiting for a Rasa CLI command. It's
+    output is stored in a log file identified by `log_id`.
+    """
+
+    process: subprocess.Popen
+    log_id: str
+
+
+def _create_warm_rasa_process() -> WarmRasaProcess:
+    """Create a new warm Rasa process."""
+    command = [
+        config.RASA_PYTHON_PATH,
+        "-m",
+        "rasa.model_manager.warm_rasa_process",
+    ]
+
+    envs = os.environ.copy()
+    envs["RASA_TELEMETRY_ENABLED"] = "false"
+
+    log_id = uuid.uuid4().hex
+    log_path = logs_path(log_id)
+
+    ensure_base_directory_exists(log_path)
+
+    process = subprocess.Popen(
+        command,
+        stdout=open(log_path, "w"),
+        stderr=subprocess.STDOUT,
+        stdin=subprocess.PIPE,
+        env=envs,
+    )
+
+    structlogger.debug(
+        "model_trainer.created_warm_rasa_process",
+        pid=process.pid,
+        command=command,
+        log_path=log_path,
+    )
+
+    return WarmRasaProcess(process=process, log_id=log_id)
+
+
+def initialize_warm_rasa_process() -> None:
+    """Initialize the warm Rasa processes."""
+    global warm_rasa_processes
+    for _ in range(NUMBER_OF_INITIAL_PROCESSES):
+        warm_rasa_processes.append(_create_warm_rasa_process())
+
+
+def shutdown_warm_rasa_processes() -> None:
+    """Shutdown all warm Rasa processes."""
+    global warm_rasa_processes
+    for warm_rasa_process in warm_rasa_processes:
+        warm_rasa_process.process.terminate()
+    warm_rasa_processes = []
+
+
+def start_rasa_process(cwd: str, arguments: List[str]) -> WarmRasaProcess:
+    """Start a Rasa process.
+
+    This will start a Rasa process with the given current working directory
+    and arguments. The process will be a warm one, meaning that it has already
+    imported all necessary modules.
+    """
+    warm_rasa_process = _get_warm_rasa_process()
+    _pass_arguments_to_process(warm_rasa_process.process, cwd, arguments)
+    return warm_rasa_process
+
+
+def _get_warm_rasa_process() -> WarmRasaProcess:
+    """Get a warm Rasa process.
+
+    This will return a warm Rasa process from the pool and create a
+    new one to replace it.
+    """
+    global warm_rasa_processes
+
+    if not warm_rasa_processes:
+        warm_rasa_processes = [_create_warm_rasa_process()]
+
+    previous_warm_rasa_process = warm_rasa_processes.pop(0)
+
+    if previous_warm_rasa_process.process.poll() is not None:
+        # process has finished (for some reason...)
+        # back up plan is to create a new one on the spot.
+        # this should not happen, but let's be safe
+        structlogger.warning(
+            "model_trainer.warm_rasa_process_finished_unexpectedly",
+            pid=previous_warm_rasa_process.process.pid,
+        )
+        previous_warm_rasa_process = _create_warm_rasa_process()
+
+    warm_rasa_processes.append(_create_warm_rasa_process())
+    return previous_warm_rasa_process
+
+
+def _pass_arguments_to_process(
+    process: subprocess.Popen, cwd: str, arguments: List[str]
+) -> None:
+    """Pass arguments to a warm Rasa process.
+
+    The process is waiting for input on stdin. We pass the current working
+    directory and the arguments to run a Rasa CLI command.
+    """
+    arguments_string = " ".join(arguments)
+    # send arguments to stdin
+    process.stdin.write(cwd.encode())  # type: ignore[union-attr]
+    process.stdin.write("\n".encode())  # type: ignore[union-attr]
+    process.stdin.write(arguments_string.encode())  # type: ignore[union-attr]
+    process.stdin.write("\n".encode())  # type: ignore[union-attr]
+    process.stdin.flush()  # type: ignore[union-attr]
+
+
+def warmup() -> None:
+    """Import all necessary modules to warm up the process.
+
+    This should include all the modules that take a long time to import.
+    We import them now, so that the training / deployment can later
+    directly start.
+    """
+    try:
+        import presidio_analyzer  # noqa: F401
+        import litellm  # noqa: F401
+        import langchain  # noqa: F401
+        import tensorflow  # noqa: F401
+        import matplotlib  # noqa: F401
+        import pandas  # noqa: F401
+        import numpy  # noqa: F401
+        import spacy  # noqa: F401
+        import rasa.validator  # noqa: F401
+    except ImportError:
+        pass
+
+
+def warm_rasa_main() -> None:
+    """Entry point for processes waiting for their command to run.
+
+    The process will wait for the current working directory and the command
+    to run. These will be send on stdin by the parent process. After receiving
+    the input, we will kick things of starting or running a bot.
+
+    Uses the normal Rasa CLI entry point (e.g. `rasa train --data ...`).
+    """
+    warmup()
+
+    cwd = input()
+
+    # this should be `train --data ...` or similar
+    cli_arguments_str = input()
+    # splits the arguments string into a list of arguments as expected by `argparse`
+    arguments = shlex.split(cli_arguments_str)
+
+    # needed to make sure the passed arguments are relative to the working directory
+    os.chdir(cwd)
+
+    main(arguments)
+
+
+if __name__ == "__main__":
+    warm_rasa_main()