feat(instrumentation): add OpenTelemetry tracing and metrics with basic configurations

extra-requirements.txt

@@ -35,8 +35,14 @@ packaging>=20.0:            core
 docarray>=0.16.4:           core
 jina-hubble-sdk>=0.15.1:    core
 jcloud>=0.0.35:             core
+opentelemetry-api>=1.12.0:  core
 uvloop:                     perf,standard,devel
 prometheus_client:          perf,standard,devel
+opentelemetry-sdk>=1.12.0:   perf,standard,devel
+opentelemetry-exporter-otlp>=1.12.0:  perf,standard,devel
+opentelemetry-exporter-prometheus>=1.12.0rc1:  perf,standard,devel
+opentelemetry-semantic-conventions>=0.33b0:    perf,standard,devel
+opentelemetry-instrumentation-grpc>=0.33b0:    perf,standard,devel
 fastapi>=0.76.0:            standard,devel
 uvicorn[standard]:          standard,devel
 docarray[common]>=0.16.3:   standard,devel

jina/__init__.py

@@ -102,6 +102,7 @@ def _warning_on_one_line(message, category, filename, lineno, *args, **kwargs):
     'JINA_OPTOUT_TELEMETRY',
     'JINA_RANDOM_PORT_MAX',
     'JINA_RANDOM_PORT_MIN',
+    'JINA_ENABLE_OTEL_TRACING',
 )
 
 __default_host__ = _os.environ.get(

jina/clients/__init__.py

@@ -24,6 +24,8 @@ def Client(
     protocol: Optional[str] = 'GRPC',
     proxy: Optional[bool] = False,
     tls: Optional[bool] = False,
+    opentelemetry_tracing: Optional[bool] = False,
+    opentelemetry_metrics: Optional[bool] = False,
     **kwargs
 ) -> Union[
     'AsyncWebSocketClient',

jina/clients/base/__init__.py

@@ -17,9 +17,10 @@
 
     InputType = Union[GeneratorSourceType, Callable[..., GeneratorSourceType]]
     CallbackFnType = Optional[Callable[[Response], None]]
+from jina.serve.instrumentation import InstrumentationMixin
 
 
-class BaseClient(ABC):
+class BaseClient(InstrumentationMixin, ABC):
     """A base client for connecting to the Flow Gateway.
 
     :param args: the Namespace from argparse
@@ -47,6 +48,7 @@ def __init__(
             os.unsetenv('https_proxy')
         self._inputs = None
         send_telemetry_event(event='start', obj=self)
+        self._setup_instrumentation()
 
     @staticmethod
     def check_input(inputs: Optional['InputType'] = None, **kwargs) -> None:

jina/parsers/client.py

@@ -30,3 +30,19 @@ def mixin_client_features_parser(parser):
         default=False,
         help='If set, then the input and output of this Client work in an asynchronous manner. ',
     )
+
+    parser.add_argument(
+        '--opentelemetry-tracing',
+        action='store_true',
+        default=False,
+        help='If set, real implementation of the tracer will be available and will be enabled for automatic tracing of requests and customer span creation. '
+        'Otherwise a no-op implementation will be provided.',
+    )
+
+    parser.add_argument(
+        '--opentelemetry-metrics',
+        action='store_true',
+        default=False,
+        help='If set, real implementation of the metrics will be available for default monitoring and custom measurements. '
+        'Otherwise a no-op implementation will be provided.',
+    )

jina/parsers/orchestrate/pod.py

@@ -119,3 +119,19 @@ def mixin_pod_parser(parser):
         help='If set, the current Pod/Deployment can not be further chained, '
         'and the next `.add()` will chain after the last Pod/Deployment not this current one.',
     )
+
+    gp.add_argument(
+        '--opentelemetry-tracing',
+        action='store_true',
+        default=False,
+        help='If set, real implementation of the tracer will be available and will be enabled for automatic tracing of requests and customer span creation. '
+        'Otherwise a no-op implementation will be provided.',
+    )
+
+    gp.add_argument(
+        '--opentelemetry-metrics',
+        action='store_true',
+        default=False,
+        help='If set, real implementation of the metrics will be available for default monitoring and custom measurements. '
+        'Otherwise a no-op implementation will be provided.',
+    )

jina/resources/extra-requirements.txt

@@ -35,8 +35,14 @@ packaging>=20.0:            core
 docarray>=0.16.4:           core
 jina-hubble-sdk>=0.15.1:    core
 jcloud>=0.0.35:             core
+opentelemetry-api>=1.12.0:  core
 uvloop:                     perf,standard,devel
 prometheus_client:          perf,standard,devel
+opentelemetry-sdk>=1.12.0:   perf,standard,devel
+opentelemetry-exporter-otlp>=1.12.0:  perf,standard,devel
+opentelemetry-exporter-prometheus>=1.12.0rc1:  perf,standard,devel
+opentelemetry-semantic-conventions>=0.33b0:    perf,standard,devel
+opentelemetry-instrumentation-grpc>=0.33b0:    perf,standard,devel
 fastapi>=0.76.0:            standard,devel
 uvicorn[standard]:          standard,devel
 docarray[common]>=0.16.3:   standard,devel

jina/serve/executors/__init__.py

@@ -7,6 +7,8 @@
 from types import SimpleNamespace
 from typing import TYPE_CHECKING, Any, Dict, Optional, Type, Union
 
+from opentelemetry import metrics, trace
+
 from jina import __args_executor_init__, __cache_path__, __default_endpoint__
 from jina.enums import BetterEnum
 from jina.helper import (
@@ -140,6 +142,7 @@ def __init__(
         self._add_requests(requests)
         self._add_runtime_args(runtime_args)
         self._init_monitoring()
+        self._init_instrumentation()
         self._init_workspace = workspace
         self.logger = JinaLogger(self.__class__.__name__)
         if __dry_run_endpoint__ not in self.requests:
@@ -186,6 +189,10 @@ def _init_monitoring(self):
             self._summary_method = None
             self._metrics_buffer = None
 
+    def _init_instrumentation(self):
+        self.tracer = trace.get_tracer(self.runtime_args.name)
+        self.meter = metrics.get_meter(self.runtime_args.name)
+
     def _add_requests(self, _requests: Optional[Dict]):
         if not hasattr(self, 'requests'):
             self.requests = {}

jina/serve/instrumentation/__init__.py

@@ -0,0 +1,77 @@
+from opentelemetry import metrics, trace
+from opentelemetry.instrumentation.grpc import (
+    client_interceptor as grpc_client_interceptor,
+)
+from opentelemetry.sdk.metrics import MeterProvider
+from opentelemetry.sdk.metrics.export import (
+    ConsoleMetricExporter,
+    PeriodicExportingMetricReader,
+)
+from opentelemetry.sdk.resources import SERVICE_NAME, Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
+
+from jina.serve.instrumentation._aio_client import (
+    StreamStreamAioClientInterceptor,
+    StreamUnaryAioClientInterceptor,
+    UnaryStreamAioClientInterceptor,
+    UnaryUnaryAioClientInterceptor,
+)
+
+
+class InstrumentationMixin:
+    '''Instrumentation mixin for OpenTelemetery Tracing and Metrics handling'''
+
+    def __init__(self) -> None:
+        self.tracer = trace.NoOpTracer()
+        self.meter = metrics.NoOpMeter(name='no-op')
+
+    def _setup_instrumentation(self) -> None:
+        name = self.__class__.__name__
+        if hasattr(self, 'name') and self.name:
+            name = self.name
+        resource = Resource(attributes={SERVICE_NAME: name})
+
+        if self.args.opentelemetry_tracing:
+            provider = TracerProvider(resource=resource)
+            processor = BatchSpanProcessor(ConsoleSpanExporter())
+            provider.add_span_processor(processor)
+            trace.set_tracer_provider(provider)
+            self.tracer = trace.get_tracer(name)
+
+        if self.args.opentelemetry_metrics:
+            metric_reader = PeriodicExportingMetricReader(ConsoleMetricExporter())
+            meter_provider = MeterProvider(
+                metric_readers=[metric_reader], resource=resource
+            )
+            metrics.set_meter_provider(meter_provider)
+            self.meter = metrics.get_meter(name)
+
+    def aio_tracing_server_interceptor(self):
+        '''Create a gRPC aio server interceptor.
+        :returns: A service-side aio interceptor object.
+        '''
+        from . import _aio_server
+
+        return _aio_server.OpenTelemetryAioServerInterceptor(self.tracer)
+
+    @staticmethod
+    def aio_tracing_client_interceptors():
+        '''Create a gRPC client aio channel interceptor.
+        :returns: An invocation-side list of aio interceptor objects.
+        '''
+        tracer = trace.get_tracer(__name__)
+
+        return [
+            UnaryUnaryAioClientInterceptor(tracer),
+            UnaryStreamAioClientInterceptor(tracer),
+            StreamUnaryAioClientInterceptor(tracer),
+            StreamStreamAioClientInterceptor(tracer),
+        ]
+
+    @staticmethod
+    def tracing_client_interceptor():
+        '''
+        :returns: a gRPC client interceptor with the global tracing provider.
+        '''
+        return grpc_client_interceptor(trace.get_tracer_provider())