-
Notifications
You must be signed in to change notification settings - Fork 233
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #1016 from RasaHQ/implement-tracing-in-action-server
Implement tracing in action server
- Loading branch information
Showing
17 changed files
with
1,112 additions
and
38 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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. |
Large diffs are not rendered by default.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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) |
Oops, something went wrong.