Introduce tracing with OpenTelemetry

mlserver/grpc/server.py

@@ -4,6 +4,7 @@
 
 from ..handlers import DataPlane, ModelRepositoryHandlers
 from ..settings import Settings
+from ..tracing import get_tracer_provider
 
 from .servicers import InferenceServicer
 from .model_repository import ModelRepositoryServicer
@@ -12,6 +13,8 @@
 from .interceptors import LoggingInterceptor, PromServerInterceptor
 from .logging import logger
 
+from opentelemetry.instrumentation.grpc import aio_server_interceptor, filters
+
 # Workers used for non-AsyncIO workloads (which aren't any in our case)
 DefaultGrpcWorkers = 5
 
@@ -46,6 +49,25 @@ def _create_server(self):
                 PromServerInterceptor(enable_handling_time_histogram=True)
             )
 
+        if self._settings.tracing_server:
+            tracer_provider = get_tracer_provider(self._settings)
+            excluded_urls = filters.negate(
+                filters.any_of(
+                    filters.full_method_name(
+                        "/inference.GRPCInferenceService/ServerLive"
+                    ),
+                    filters.full_method_name(
+                        "/inference.GRPCInferenceService/ServerReady"
+                    ),
+                )
+            )
+
+            interceptors.append(
+                aio_server_interceptor(
+                    tracer_provider=tracer_provider, filter_=excluded_urls
+                )
+            )
+
         self._server = aio.server(
             ThreadPoolExecutor(max_workers=DefaultGrpcWorkers),
             interceptors=tuple(interceptors),

mlserver/rest/app.py

@@ -4,6 +4,8 @@
 from fastapi.routing import APIRoute as FastAPIRoute
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.middleware.gzip import GZipMiddleware
+from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
+
 from starlette_exporter import PrometheusMiddleware
 
 from .endpoints import Endpoints, ModelRepositoryEndpoints
@@ -13,6 +15,7 @@
 
 from ..settings import Settings
 from ..handlers import DataPlane, ModelRepositoryHandlers
+from ..tracing import get_tracer_provider
 
 
 class APIRoute(FastAPIRoute):
@@ -145,6 +148,22 @@ def create_app(
         docs_url=None,
         redoc_url=None,
     )
+
+    if settings.tracing_server:
+        tracer_provider = get_tracer_provider(settings)
+        excluded_urls = ",".join(
+            [
+                "/v2/health/live",
+                "/v2/health/ready",
+            ]
+        )
+
+        FastAPIInstrumentor.instrument_app(
+            app,
+            tracer_provider=tracer_provider,
+            excluded_urls=excluded_urls,
+        )
+
     app.router.route_class = APIRoute
     app.add_middleware(GZipMiddleware)
     if settings.cors_settings is not None:

mlserver/settings.py

@@ -224,12 +224,16 @@ class Config:
     logging_settings: Optional[Union[str, Dict]] = None
     """Path to logging config file or dictionary configuration."""
 
-    # Kakfa Server settings
+    # Kafka Server settings
     kafka_enabled: bool = False
     kafka_servers: str = "localhost:9092"
     kafka_topic_input: str = "mlserver-input"
     kafka_topic_output: str = "mlserver-output"
 
+    # OpenTelemetry Tracing settings
+    tracing_server: Optional[str] = None
+    """Server name used to export OpenTelemetry tracing to collector service."""
+
     # Custom server settings
     _custom_rest_server_settings: Optional[dict] = None
     _custom_metrics_server_settings: Optional[dict] = None

mlserver/tracing.py

@@ -0,0 +1,39 @@
+from typing import Optional
+
+from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.resources import Resource, SERVICE_NAME, SERVICE_VERSION
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter
+
+from mlserver.settings import Settings
+
+
+_TRACER_PROVIDER: Optional[TracerProvider] = None
+
+
+def _create_resource(settings: Settings) -> Resource:
+    return Resource(
+        attributes={
+            SERVICE_NAME: settings.server_name,
+            SERVICE_VERSION: settings.server_version,
+        }
+    )
+
+
+def _create_span_exporter(settings: Settings) -> SpanExporter:
+    return OTLPSpanExporter(insecure=True, endpoint=settings.tracing_server)
+
+
+def get_tracer_provider(settings: Settings) -> TracerProvider:
+    global _TRACER_PROVIDER
+    if _TRACER_PROVIDER is not None:
+        return _TRACER_PROVIDER
+
+    resource = _create_resource(settings)
+    tracer_provider = TracerProvider(resource=resource)
+    span_exporter = _create_span_exporter(settings)
+    span_processor = BatchSpanProcessor(span_exporter)
+    tracer_provider.add_span_processor(span_processor)
+
+    _TRACER_PROVIDER = tracer_provider
+    return _TRACER_PROVIDER