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

rasa/model_manager/model_api.py

@@ -0,0 +1,559 @@
+import asyncio
+from functools import wraps
+import os
+from http import HTTPStatus
+from typing import Any, Callable, Dict, Optional, Union
+import dotenv
+import psutil
+from sanic import Blueprint, Sanic, response
+from sanic.response import json
+from sanic.exceptions import NotFound
+from sanic.request import Request
+import structlog
+from socketio import AsyncServer
+
+from rasa.exceptions import ModelNotFound
+from rasa.model_manager import config
+from rasa.model_manager.config import SERVER_BASE_URL
+from rasa.constants import MODEL_ARCHIVE_EXTENSION
+from rasa.model_manager.runner_service import (
+    BotSession,
+    BotSessionStatus,
+    fetch_remote_model_to_dir,
+    fetch_size_of_remote_model,
+    run_bot,
+    terminate_bot,
+    update_bot_status,
+)
+from rasa.model_manager.socket_bridge import create_bridge_server
+from rasa.model_manager.trainer_service import (
+    TrainingSession,
+    TrainingSessionStatus,
+    run_training,
+    terminate_training,
+    update_training_status,
+)
+from rasa.model_manager.utils import (
+    InvalidPathException,
+    get_logs_content,
+    logs_base_path,
+    models_base_path,
+    subpath,
+)
+from rasa.model_manager.warm_rasa_process import (
+    initialize_warm_rasa_process,
+    shutdown_warm_rasa_processes,
+)
+
+dotenv.load_dotenv()
+
+structlogger = structlog.get_logger()
+
+
+# A simple in-memory store for training sessions and running bots
+trainings: Dict[str, TrainingSession] = {}
+running_bots: Dict[str, BotSession] = {}
+
+
+def prepare_working_directories() -> None:
+    """Make sure all required directories exist."""
+    os.makedirs(logs_base_path(), exist_ok=True)
+    os.makedirs(models_base_path(), exist_ok=True)
+
+
+def cleanup_training_processes() -> None:
+    """Terminate all running training processes."""
+    structlogger.debug("model_trainer.cleanup_processes.started")
+    running = list(trainings.values())
+    for training in running:
+        terminate_training(training)
+
+
+def cleanup_bot_processes() -> None:
+    """Terminate all running bot processes."""
+    structlogger.debug("model_runner.cleanup_processes.started")
+    running = list(running_bots.values())
+    for bot in running:
+        terminate_bot(bot)
+
+
+def update_status_of_all_trainings() -> None:
+    """Update the status of all training processes."""
+    running = list(trainings.values())
+    for training in running:
+        update_training_status(training)
+
+
+async def update_status_of_all_bots() -> None:
+    """Update the status of all bot processes."""
+    # we need to get the values first, because (since we are async and waiting
+    # within the loop) some other job on the asyncio loop could change the dict
+    # (adding or removing). python doesn't like if you change the size of a dict
+    # while iterating over it and will raise a RuntimeError. so we get the values
+    # first and iterate over them to avoid that.
+    running = list(running_bots.values())
+    for bot in running:
+        await update_bot_status(bot)
+
+
+def base_server_url(request: Request) -> str:
+    """Return the base URL of the server."""
+    if SERVER_BASE_URL:
+        return SERVER_BASE_URL.rstrip("/")
+    else:
+        return f"{request.scheme}://{request.host}/{config.DEFAULT_SERVER_PATH_PREFIX}"
+
+
+async def continuously_update_process_status() -> None:
+    """Regularly Update the status of all training and bot processes."""
+    structlogger.debug("model_api.update_process_status.started")
+
+    while True:
+        try:
+            update_status_of_all_trainings()
+            await update_status_of_all_bots()
+        except asyncio.exceptions.CancelledError:
+            structlogger.debug("model_api.update_process_status.cancelled")
+            break
+        except Exception as e:
+            structlogger.error("model_api.update_process_status.error", error=str(e))
+        finally:
+            await asyncio.sleep(0.1)
+
+
+def internal_blueprint() -> Blueprint:
+    """Create a blueprint for the model manager API."""
+    bp = Blueprint("model_api_internal")
+
+    @bp.before_server_stop
+    async def cleanup_processes(app: Sanic, loop: asyncio.AbstractEventLoop) -> None:
+        """Terminate all running processes before the server stops."""
+        structlogger.debug("model_api.cleanup_processes.started")
+        cleanup_training_processes()
+        cleanup_bot_processes()
+        shutdown_warm_rasa_processes()
+
+    @bp.after_server_start
+    async def create_warm_rasa_processes(
+        app: Sanic, loop: asyncio.AbstractEventLoop
+    ) -> None:
+        """Create warm Rasa processes to speed up future training and bot runs."""
+        structlogger.debug("model_api.create_warm_rasa_processes.started")
+        initialize_warm_rasa_process()
+
+    def limit_parallel_training_requests() -> Callable[[Callable], Callable[..., Any]]:
+        """Limit the number of parallel training requests."""
+
+        def decorator(f: Callable) -> Callable:
+            @wraps(f)
+            def decorated(*args: Any, **kwargs: Any) -> Any:
+                running_requests = len(
+                    [
+                        training
+                        for training in trainings.values()
+                        if training.status == TrainingSessionStatus.RUNNING
+                        and training.process.poll() is None
+                    ]
+                )
+
+                if running_requests >= int(config.MAX_PARALLEL_TRAININGS):
+                    return response.json(
+                        {
+                            "message": f"Too many parallel training requests, above "
+                            f"the limit of {config.MAX_PARALLEL_TRAININGS}. "
+                            f"Retry later or increase your server's "
+                            f"memory and CPU resources."
+                        },
+                        status=HTTPStatus.TOO_MANY_REQUESTS,
+                    )
+                return f(*args, **kwargs)
+
+            return decorated
+
+        return decorator
+
+    def limit_parallel_bot_runs() -> Callable[[Callable], Callable[..., Any]]:
+        """Limit the number of parallel training requests."""
+
+        def decorator(f: Callable) -> Callable:
+            @wraps(f)
+            def decorated(*args: Any, **kwargs: Any) -> Any:
+                running_requests = len(
+                    [
+                        bot
+                        for bot in running_bots.values()
+                        if bot.status
+                        in {BotSessionStatus.RUNNING, BotSessionStatus.QUEUED}
+                    ]
+                )
+
+                if running_requests >= int(config.MAX_PARALLEL_BOT_RUNS):
+                    return response.json(
+                        {
+                            "message": f"Too many parallel bot runs, above "
+                            f"the limit of {config.MAX_PARALLEL_BOT_RUNS}. "
+                            f"Retry later or increase your server's "
+                            f"memory and CPU resources."
+                        },
+                        status=HTTPStatus.TOO_MANY_REQUESTS,
+                    )
+
+                return f(*args, **kwargs)
+
+            return decorated
+
+        return decorator
+
+    def ensure_minimum_disk_space() -> Callable[[Callable], Callable[..., Any]]:
+        """Ensure that there is enough disk space before starting a new process."""
+        min_required_disk_space = 1024 * 1024 * config.MIN_REQUIRED_DISCSPACE_MB
+
+        def decorator(f: Callable) -> Callable:
+            @wraps(f)
+            def decorated(*args: Any, **kwargs: Any) -> Any:
+                if os.path.exists(config.SERVER_BASE_WORKING_DIRECTORY):
+                    free_space_bytes = psutil.disk_usage(
+                        config.SERVER_BASE_WORKING_DIRECTORY
+                    ).free
+                    structlogger.debug(
+                        "model_api.storage.available_disk_space",
+                        available_space_mb=free_space_bytes / 1024 / 1024,
+                    )
+
+                    if free_space_bytes < min_required_disk_space:
+                        return response.json(
+                            {
+                                "message": (
+                                    f"Less than {config.MIN_REQUIRED_DISCSPACE_MB} MB "
+                                    f"of free disk space available. "
+                                    f"Please free up some space on the model service."
+                                )
+                            },
+                            status=HTTPStatus.INSUFFICIENT_STORAGE,
+                        )
+
+                return f(*args, **kwargs)
+
+            return decorated
+
+        return decorator
+
+    @bp.get("/")
+    async def health(request: Request) -> response.HTTPResponse:
+        return json(
+            {
+                "status": "ok",
+                "bots": [
+                    {
+                        "deployment_id": bot.deployment_id,
+                        "status": bot.status,
+                        "internal_url": bot.internal_url,
+                        "returncode": bot.returncode,
+                        "url": bot.url,
+                    }
+                    for bot in running_bots.values()
+                ],
+                "trainings": [
+                    {
+                        "training_id": training.training_id,
+                        "assistant_id": training.assistant_id,
+                        "client_id": training.client_id,
+                        "progress": training.progress,
+                        "status": training.status,
+                    }
+                    for training in trainings.values()
+                ],
+            }
+        )
+
+    @bp.get("/training")
+    async def get_training_list(request: Request) -> response.HTTPResponse:
+        """Return a list of all training sessions for an assistant."""
+        assistant_id = request.args.get("assistant_id")
+        sessions = [
+            {
+                "training_id": session.training_id,
+                "assistant_id": session.assistant_id,
+                "client_id": session.client_id,
+                "progress": session.progress,
+                "status": session.status,
+                "model_name": session.model_name,
+                "runtime_metadata": None,
+            }
+            for session in trainings.values()
+            if session.assistant_id == assistant_id
+        ]
+        return json({"training_sessions": sessions, "total_number": len(sessions)})
+
+    @bp.post("/training")
+    @limit_parallel_training_requests()
+    @ensure_minimum_disk_space()
+    async def start_training(request: Request) -> response.HTTPResponse:
+        """Start a new training session."""
+        data = request.json
+        training_id: Optional[str] = data.get("id")
+        assistant_id: Optional[str] = data.get("assistant_id")
+        client_id: Optional[str] = data.get("client_id")
+        encoded_training_data: Dict[str, str] = data.get("bot_config", {}).get(
+            "data", {}
+        )
+
+        if training_id in trainings:
+            # fail, because there apparently is already a training with this id
+            return json({"message": "Training with this id already exists"}, status=409)
+
+        if not assistant_id:
+            return json({"message": "Assistant id is required"}, status=400)
+
+        if not training_id:
+            return json({"message": "Training id is required"}, status=400)
+
+        try:
+            # file deepcode ignore PT: path traversal is prevented
+            # by the `subpath` function found in the `rasa.model_manager.utils` module
+            training_session = run_training(
+                training_id=training_id,
+                assistant_id=assistant_id,
+                client_id=client_id,
+                encoded_training_data=encoded_training_data,
+            )
+            trainings[training_id] = training_session
+            return json(
+                {"training_id": training_id, "model_name": training_session.model_name}
+            )
+        except InvalidPathException as exc:
+            return json({"message": str(exc)}, status=403)
+        except Exception as exc:
+            return json({"message": str(exc)}, status=500)
+
+    @bp.get("/training/<training_id>")
+    async def get_training(request: Request, training_id: str) -> response.HTTPResponse:
+        """Return the status of a training session."""
+        if training := trainings.get(training_id):
+            return json(
+                {
+                    "training_id": training_id,
+                    "assistant_id": training.assistant_id,
+                    "client_id": training.client_id,
+                    "progress": training.progress,
+                    "model_name": training.model_name,
+                    "status": training.status,
+                    "logs": get_logs_content(training.log_id),
+                }
+            )
+        else:
+            return json({"message": "Training not found"}, status=404)
+
+    @bp.delete("/training/<training_id>")
+    async def stop_training(
+        request: Request, training_id: str
+    ) -> response.HTTPResponse:
+        # this is a no-op if the training is already done
+        if not (training := trainings.get(training_id)):
+            return json({"message": "Training session not found"}, status=404)
+
+        terminate_training(training)
+        return json({"training_id": training_id})
+
+    @bp.post("/bot")
+    @limit_parallel_bot_runs()
+    @ensure_minimum_disk_space()
+    async def start_bot(request: Request) -> response.HTTPResponse:
+        data = request.json
+        deployment_id: Optional[str] = data.get("deployment_id")
+        model_name: Optional[str] = data.get("model_name")
+        encoded_configs: Dict[str, str] = data.get("bot_config", {})
+
+        if deployment_id in running_bots:
+            # fail, because there apparently is already a bot running with this id
+            return json(
+                {"message": "Bot with this deployment id already exists"}, status=409
+            )
+
+        if not deployment_id:
+            return json({"message": "Deployment id is required"}, status=400)
+
+        if not model_name:
+            return json({"message": "Model name is required"}, status=400)
+
+        base_url_path = base_server_url(request)
+        try:
+            bot_session = run_bot(
+                deployment_id,
+                model_name,
+                base_url_path,
+                encoded_configs,
+            )
+            running_bots[deployment_id] = bot_session
+            return json(
+                {
+                    "deployment_id": deployment_id,
+                    "status": bot_session.status,
+                    "url": bot_session.url,
+                }
+            )
+        except ModelNotFound:
+            return json(
+                {"message": f"Model with name '{model_name}' could not be found."},
+                status=404,
+            )
+        except Exception as e:
+            return json({"message": str(e)}, status=500)
+
+    @bp.delete("/bot/<deployment_id>")
+    async def stop_bot(request: Request, deployment_id: str) -> response.HTTPResponse:
+        bot = running_bots.get(deployment_id)
+        if bot is None:
+            return json({"message": "Bot not found"}, status=404)
+
+        terminate_bot(bot)
+
+        return json(
+            {"deployment_id": deployment_id, "status": bot.status, "url": bot.url}
+        )
+
+    @bp.get("/bot/<deployment_id>")
+    async def get_bot(request: Request, deployment_id: str) -> response.HTTPResponse:
+        bot = running_bots.get(deployment_id)
+        if bot is None:
+            return json({"message": "Bot not found"}, status=404)
+
+        return json(
+            {
+                "deployment_id": deployment_id,
+                "status": bot.status,
+                "returncode": bot.returncode,
+                "url": bot.url,
+                "logs": get_logs_content(bot.log_id),
+            }
+        )
+
+    @bp.get("/bot")
+    async def list_bots(request: Request) -> response.HTTPResponse:
+        bots = [
+            {
+                "deployment_id": bot.deployment_id,
+                "status": bot.status,
+                "returncode": bot.returncode,
+                "url": bot.url,
+            }
+            for bot in running_bots.values()
+        ]
+        return json({"deployment_sessions": bots, "total_number": len(bots)})
+
+    @bp.route("/models/<model_name>", methods=["GET"])
+    async def send_model(
+        request: Request, model_name: str
+    ) -> Union[response.ResponseStream, response.HTTPResponse]:
+        try:
+            model_path = path_to_model(model_name)
+
+            # get size of model file
+            model_size = os.stat(model_path)
+
+            return await response.file_stream(
+                model_path, headers={"Content-Length": str(model_size.st_size)}
+            )
+        except NotFound:
+            return json({"message": "Model not found"}, status=404)
+        except ModelNotFound:
+            return json({"message": "Model not found"}, status=404)
+
+    @bp.route("/models/<model_name>", methods=["HEAD"])
+    async def head_model(request: Request, model_name: str) -> response.HTTPResponse:
+        try:
+            model_size = size_of_model(model_name)
+
+            structlogger.debug(
+                "model_api.internal.head_model",
+                model_name=model_name,
+                size=model_size,
+            )
+            return response.raw(
+                b"", status=200, headers={"Content-Length": str(model_size)}
+            )
+        except ModelNotFound:
+            return response.raw(b"", status=404)
+
+    return bp
+
+
+def external_blueprint() -> Blueprint:
+    """Create a blueprint for the model manager API."""
+    from rasa.core.channels.socketio import SocketBlueprint
+
+    sio = AsyncServer(async_mode="sanic", cors_allowed_origins=[])
+    bp = SocketBlueprint(sio, "", "model_api_external")
+
+    create_bridge_server(sio, running_bots)
+
+    @bp.get("/health")
+    async def health(request: Request) -> response.HTTPResponse:
+        return json(
+            {
+                "status": "ok",
+                "bots": [
+                    {
+                        "deployment_id": bot.deployment_id,
+                        "status": bot.status,
+                        "internal_url": bot.internal_url,
+                        "url": bot.url,
+                    }
+                    for bot in running_bots.values()
+                ],
+                "trainings": [
+                    {
+                        "training_id": training.training_id,
+                        "assistant_id": training.assistant_id,
+                        "client_id": training.client_id,
+                        "progress": training.progress,
+                        "status": training.status,
+                    }
+                    for training in trainings.values()
+                ],
+            }
+        )
+
+    return bp
+
+
+def size_of_model(model_name: str) -> Optional[int]:
+    """Return the size of a model."""
+    model_file_name = f"{model_name}.{MODEL_ARCHIVE_EXTENSION}"
+    model_path = subpath(models_base_path(), model_file_name)
+
+    if os.path.exists(model_path):
+        return os.path.getsize(model_path)
+
+    if config.SERVER_MODEL_REMOTE_STORAGE:
+        structlogger.debug(
+            "model_api.storage.fetching_remote_model_size",
+            model_name=model_file_name,
+        )
+        return fetch_size_of_remote_model(
+            model_file_name,
+            config.SERVER_MODEL_REMOTE_STORAGE,
+        )
+    raise ModelNotFound("Model not found.")
+
+
+def path_to_model(model_name: str) -> Optional[str]:
+    """Return the path to a local model."""
+    model_file_name = f"{model_name}.{MODEL_ARCHIVE_EXTENSION}"
+    model_path = subpath(models_base_path(), model_file_name)
+
+    if os.path.exists(model_path):
+        return model_path
+
+    if config.SERVER_MODEL_REMOTE_STORAGE:
+        structlogger.info(
+            "model_api.storage.fetching_remote_model",
+            model_name=model_file_name,
+        )
+        return fetch_remote_model_to_dir(
+            model_file_name,
+            models_base_path(),
+            config.SERVER_MODEL_REMOTE_STORAGE,
+        )
+
+    raise ModelNotFound("Model not found.")