Implement tracing in action server

changelog/1016.feature.md

@@ -0,0 +1,2 @@
+Added tracing functionality to the Rasa SDK, bringing enhanced monitoring, execution profiling and debugging capabilities to the Rasa Actions Server.
+See [Rasa Documentation on Tracing](https://rasa.com/docs/rasa/monitoring/tracing/#configuring-a-tracing-backend-or-collector) to know more about configuring a tracing backend or collector.

pyproject.toml

@@ -83,6 +83,10 @@ prompt-toolkit = "^3.0,<3.0.29"
 "ruamel.yaml" = ">=0.16.5,<0.18.0"
 websockets = ">=10.0,<12.0"
 pluggy = "^1.0.0"
+opentelemetry-api = "~1.15.0"
+opentelemetry-sdk = "~1.15.0"
+opentelemetry-exporter-jaeger = "~1.15.0"
+opentelemetry-exporter-otlp = "~1.15.0"
 
 [tool.poetry.dev-dependencies]
 pytest-cov = "^4.1.0"

rasa_sdk/__main__.py

@@ -3,6 +3,7 @@
 from rasa_sdk import utils
 from rasa_sdk.endpoint import create_argument_parser, run
 from rasa_sdk.constants import APPLICATION_ROOT_LOGGER_NAME
+from rasa_sdk.tracing.utils import get_tracer_provider
 
 
 def main_from_args(args):
@@ -17,6 +18,7 @@ def main_from_args(args):
         args.logging_config_file,
     )
     utils.update_sanic_log_level()
+    tracer_provider = get_tracer_provider(args)
 
     run(
         args.actions,
@@ -26,6 +28,7 @@ def main_from_args(args):
         args.ssl_keyfile,
         args.ssl_password,
         args.auto_reload,
+        tracer_provider,
     )
 
 

rasa_sdk/endpoint.py

@@ -4,6 +4,7 @@
 import types
 import zlib
 import json
+from opentelemetry.sdk.trace import TracerProvider
 from typing import List, Text, Union, Optional
 from ssl import SSLContext
 
@@ -18,6 +19,7 @@
 from rasa_sdk.executor import ActionExecutor
 from rasa_sdk.interfaces import ActionExecutionRejection, ActionNotFoundException
 from rasa_sdk.plugin import plugin_manager
+from rasa_sdk.tracing.utils import get_tracer_and_context, set_span_attributes
 
 logger = logging.getLogger(__name__)
 
@@ -65,6 +67,7 @@ def create_app(
     action_package_name: Union[Text, types.ModuleType],
     cors_origins: Union[Text, List[Text], None] = "*",
     auto_reload: bool = False,
+    tracer_provider: Optional[TracerProvider] = None,
 ) -> Sanic:
     """Create a Sanic application and return it.
 
@@ -73,6 +76,7 @@ def create_app(
             from.
         cors_origins: CORS origins to allow.
         auto_reload: When `True`, auto-reloading of actions is enabled.
+        tracer_provider: Tracer provider to use for tracing.
 
     Returns:
         A new Sanic application ready to be run.
@@ -93,34 +97,38 @@ async def health(_) -> HTTPResponse:
     @app.post("/webhook")
     async def webhook(request: Request) -> HTTPResponse:
         """Webhook to retrieve action calls."""
-        if request.headers.get("Content-Encoding") == "deflate":
-            # Decompress the request data using zlib
-            decompressed_data = zlib.decompress(request.body)
-            # Load the JSON data from the decompressed request data
-            action_call = json.loads(decompressed_data)
-        else:
-            action_call = request.json
-        if action_call is None:
-            body = {"error": "Invalid body request"}
-            return response.json(body, status=400)
-
-        utils.check_version_compatibility(action_call.get("version"))
-
-        if auto_reload:
-            executor.reload()
-
-        try:
-            result = await executor.run(action_call)
-        except ActionExecutionRejection as e:
-            logger.debug(e)
-            body = {"error": e.message, "action_name": e.action_name}
-            return response.json(body, status=400)
-        except ActionNotFoundException as e:
-            logger.error(e)
-            body = {"error": e.message, "action_name": e.action_name}
-            return response.json(body, status=404)
-
-        return response.json(result, status=200)
+        tracer, context, span_name = get_tracer_and_context(tracer_provider, request)
+
+        with tracer.start_as_current_span(span_name, context=context) as span:
+            if request.headers.get("Content-Encoding") == "deflate":
+                # Decompress the request data using zlib
+                decompressed_data = zlib.decompress(request.body)
+                # Load the JSON data from the decompressed request data
+                action_call = json.loads(decompressed_data)
+            else:
+                action_call = request.json
+            if action_call is None:
+                body = {"error": "Invalid body request"}
+                return response.json(body, status=400)
+
+            utils.check_version_compatibility(action_call.get("version"))
+
+            if auto_reload:
+                executor.reload()
+            try:
+                result = await executor.run(action_call)
+            except ActionExecutionRejection as e:
+                logger.debug(e)
+                body = {"error": e.message, "action_name": e.action_name}
+                return response.json(body, status=400)
+            except ActionNotFoundException as e:
+                logger.error(e)
+                body = {"error": e.message, "action_name": e.action_name}
+                return response.json(body, status=404)
+
+            set_span_attributes(span, action_call)
+
+            return response.json(result, status=200)
 
     @app.get("/actions")
     async def actions(_) -> HTTPResponse:
@@ -151,11 +159,15 @@ def run(
     ssl_keyfile: Optional[Text] = None,
     ssl_password: Optional[Text] = None,
     auto_reload: bool = False,
+    tracer_provider: Optional[TracerProvider] = None,
 ) -> None:
     """Starts the action endpoint server with given config values."""
     logger.info("Starting action endpoint server...")
     app = create_app(
-        action_package_name, cors_origins=cors_origins, auto_reload=auto_reload
+        action_package_name,
+        cors_origins=cors_origins,
+        auto_reload=auto_reload,
+        tracer_provider=tracer_provider,
     )
     ## Attach additional sanic extensions: listeners, middleware and routing
     logger.info("Starting plugins...")

rasa_sdk/tracing/config.py

@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+import abc
+import logging
+import os
+from typing import Any, Dict, Optional, Text
+
+import grpc
+from opentelemetry.exporter.jaeger.thrift import JaegerExporter
+from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.resources import SERVICE_NAME, Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from rasa_sdk.tracing.endpoints import EndpointConfig, read_endpoint_config
+
+
+TRACING_SERVICE_NAME = os.environ.get("TRACING_SERVICE_NAME", "rasa_sdk")
+
+ENDPOINTS_TRACING_KEY = "tracing"
+
+logger = logging.getLogger(__name__)
+
+
+def get_tracer_provider(endpoints_file: Text) -> Optional[TracerProvider]:
+    """Configure tracing backend.
+
+    When a known tracing backend is defined in the endpoints file, this
+    function will configure the tracing infrastructure. When no or an unknown
+    tracing backend is defined, this function does nothing.
+
+    :param endpoints_file: The configuration file containing information about the
+        tracing backend.
+    :return: The `TracingProvider` to be used for all subsequent tracing.
+    """
+    cfg = read_endpoint_config(endpoints_file, ENDPOINTS_TRACING_KEY)
+
+    if not cfg:
+        logger.info(
+            f"No endpoint for tracing type available in {endpoints_file},"
+            f"tracing will not be configured."
+        )
+        return None
+    if cfg.type == "jaeger":
+        tracer_provider = JaegerTracerConfigurer.configure_from_endpoint_config(cfg)
+    elif cfg.type == "otlp":
+        tracer_provider = OTLPCollectorConfigurer.configure_from_endpoint_config(cfg)
+    else:
+        logger.warning(
+            f"Unknown tracing type {cfg.type} read from {endpoints_file}, ignoring."
+        )
+        return None
+
+    return tracer_provider
+
+
+class TracerConfigurer(abc.ABC):
+    """Abstract superclass for tracing configuration.
+
+    `TracerConfigurer` is the abstract superclass from which all configurers
+    for different supported backends should inherit.
+    """
+
+    @classmethod
+    @abc.abstractmethod
+    def configure_from_endpoint_config(cls, cfg: EndpointConfig) -> TracerProvider:
+        """Configure tracing.
+
+        This abstract method should be implemented by all concrete `TracerConfigurer`s.
+        It shall read the configuration from the supplied argument, configure all
+        necessary infrastructure for tracing, and return the `TracerProvider` to be
+        used for tracing purposes.
+
+        :param cfg: The configuration to be read for configuring tracing.
+        :return: The configured `TracerProvider`.
+        """
+
+
+class JaegerTracerConfigurer(TracerConfigurer):
+    """The `TracerConfigurer` for a Jaeger backend."""
+
+    @classmethod
+    def configure_from_endpoint_config(cls, cfg: EndpointConfig) -> TracerProvider:
+        """Configure tracing for Jaeger.
+
+        This will read the Jaeger-specific configuration from the `EndpointConfig` and
+        create a corresponding `TracerProvider` that exports to the given Jaeger
+        backend.
+
+        :param cfg: The configuration to be read for configuring tracing.
+        :return: The configured `TracerProvider`.
+        """
+        provider = TracerProvider(
+            resource=Resource.create(
+                {SERVICE_NAME: cfg.kwargs.get("service_name", TRACING_SERVICE_NAME)}
+            )
+        )
+
+        jaeger_exporter = JaegerExporter(
+            **cls._extract_config(cfg), udp_split_oversized_batches=True
+        )
+        logger.info(
+            f"Registered {cfg.type} endpoint for tracing. Traces will be exported to"
+            f" {jaeger_exporter.agent_host_name}:{jaeger_exporter.agent_port}"
+        )
+        provider.add_span_processor(BatchSpanProcessor(jaeger_exporter))
+
+        return provider
+
+    @classmethod
+    def _extract_config(cls, cfg: EndpointConfig) -> Dict[str, Any]:
+        return {
+            "agent_host_name": (cfg.kwargs.get("host", "localhost")),
+            "agent_port": (cfg.kwargs.get("port", 6831)),
+            "username": cfg.kwargs.get("username"),
+            "password": cfg.kwargs.get("password"),
+        }
+
+
+class OTLPCollectorConfigurer(TracerConfigurer):
+    """The `TracerConfigurer` for an OTLP collector backend."""
+
+    @classmethod
+    def configure_from_endpoint_config(cls, cfg: EndpointConfig) -> TracerProvider:
+        """Configure tracing for Jaeger.
+
+        This will read the OTLP collector-specific configuration from the
+        `EndpointConfig` and create a corresponding `TracerProvider` that exports to
+        the given OTLP collector.
+        Currently, this only supports insecure connections via gRPC.
+
+        :param cfg: The configuration to be read for configuring tracing.
+        :return: The configured `TracerProvider`.
+        """
+        provider = TracerProvider(
+            resource=Resource.create(
+                {SERVICE_NAME: cfg.kwargs.get("service_name", TRACING_SERVICE_NAME)}
+            )
+        )
+
+        insecure = cfg.kwargs.get("insecure")
+
+        credentials = cls._get_credentials(cfg, insecure)  # type: ignore
+
+        otlp_exporter = OTLPSpanExporter(
+            endpoint=cfg.kwargs["endpoint"],
+            insecure=insecure,
+            credentials=credentials,
+        )
+        logger.info(
+            f"Registered {cfg.type} endpoint for tracing."
+            f"Traces will be exported to {cfg.kwargs['endpoint']}"
+        )
+        provider.add_span_processor(BatchSpanProcessor(otlp_exporter))
+
+        return provider
+
+    @classmethod
+    def _get_credentials(
+        cls, cfg: EndpointConfig, insecure: bool
+    ) -> Optional[grpc.ChannelCredentials]:
+        credentials = None
+        if not insecure and "root_certificates" in cfg.kwargs:
+            with open(cfg.kwargs.get("root_certificates"), "rb") as f:  # type: ignore
+                root_cert = f.read()
+            credentials = grpc.ssl_channel_credentials(root_certificates=root_cert)
+        return credentials

rasa_sdk/tracing/endpoints.py

@@ -0,0 +1,64 @@
+import logging
+
+import os
+from typing import Any, Dict, Optional, Text
+import rasa_sdk.utils
+
+
+logger = logging.getLogger(__name__)
+DEFAULT_ENCODING = "utf-8"
+
+
+def read_endpoint_config(
+    filename: Text, endpoint_type: Text
+) -> Optional["EndpointConfig"]:
+    """Read an endpoint configuration file from disk and extract one
+
+    config."""
+    if not filename:
+        return None
+
+    try:
+        content = rasa_sdk.utils.read_file(filename)
+        content = rasa_sdk.utils.read_yaml(content)
+
+        if content.get(endpoint_type) is None:
+            return None
+
+        return EndpointConfig.from_dict(content[endpoint_type])
+    except FileNotFoundError:
+        logger.error(
+            "Failed to read endpoint configuration "
+            "from {}. No such file.".format(os.path.abspath(filename))
+        )
+        return None
+
+
+class EndpointConfig:
+    """Configuration for an external HTTP endpoint."""
+
+    def __init__(
+        self,
+        url: Optional[Text] = None,
+        params: Optional[Dict[Text, Any]] = None,
+        headers: Optional[Dict[Text, Any]] = None,
+        basic_auth: Optional[Dict[Text, Text]] = None,
+        token: Optional[Text] = None,
+        token_name: Text = "token",
+        cafile: Optional[Text] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Creates an `EndpointConfig` instance."""
+        self.url = url
+        self.params = params or {}
+        self.headers = headers or {}
+        self.basic_auth = basic_auth or {}
+        self.token = token
+        self.token_name = token_name
+        self.type = kwargs.pop("store_type", kwargs.pop("type", None))
+        self.cafile = cafile
+        self.kwargs = kwargs
+
+    @classmethod
+    def from_dict(cls, data: Dict[Text, Any]) -> "EndpointConfig":
+        return EndpointConfig(**data)