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

jina/serve/executors/__init__.py

@@ -15,11 +15,9 @@
 from jina.importer import ImportExtensions
 from jina.jaml import JAML, JAMLCompatible, env_var_regex, internal_var_regex
 from jina.logging.logger import JinaLogger
-from jina.serve.executors.decorators import (
-    avoid_concurrent_lock_cls,
-    store_init_kwargs,
-    wrap_func,
-)
+from jina.serve.executors.decorators import avoid_concurrent_lock_cls, requests
+from jina.serve.executors.metas import get_default_metas, get_executor_taboo
+from jina.serve.helper import store_init_kwargs, wrap_func
 
 if TYPE_CHECKING:
     from prometheus_client import Summary
@@ -190,12 +188,17 @@ def _init_instrumentation(self, _runtime_args: Optional[Dict]):
             else self.__class__.__name__
         )
 
-        self.tracer_provider = _runtime_args.get(
-            'tracer_provider', trace.NoOpTracerProvider()
+        self.tracer_provider = (
+            _runtime_args['tracer_provider']
+            if hasattr(_runtime_args, 'tracer_provider')
+            else trace.NoOpTracerProvider()
         )
+
         self.tracer = self.tracer_provider.get_tracer(instrumentating_module_name)
-        self.meter_provider = _runtime_args.get(
-            'meter_provider', metrics.NoOpMeterProvider()
+        self.meter_provider = (
+            _runtime_args['meter_provider']
+            if hasattr(_runtime_args, 'meter_provider')
+            else metrics.NoOpMeterProvider()
         )
         self.meter = self.meter_provider.get_meter(instrumentating_module_name)
 

jina/serve/gateway.py

@@ -1,8 +1,8 @@
 import abc
 import argparse
-from typing import TYPE_CHECKING, Any, Optional, Sequence
-
-from opentelemetry import metrics, trace
+import functools
+import inspect
+from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence
 
 from jina.helper import convert_tuple_to_list
 from jina.jaml import JAMLCompatible

jina/serve/runtimes/gateway/__init__.py

@@ -13,9 +13,6 @@
 from jina.parsers.helper import _set_gateway_uses
 from jina.serve.gateway import BaseGateway
 from jina.serve.runtimes.asyncio import AsyncNewLoopRuntime
-from jina.serve.runtimes.gateway.grpc import GRPCGateway
-from jina.serve.runtimes.gateway.http import HTTPGateway
-from jina.serve.runtimes.gateway.websocket import WebSocketGateway
 
 if TYPE_CHECKING:
     import multiprocessing
@@ -56,6 +53,7 @@ async def async_setup(self):
             raise PortAlreadyUsed(f'port:{self.args.port}')
 
         uses_with = self.args.uses_with or {}
+        print(f'--->args: {self.args}')
         self.gateway = BaseGateway.load_config(
             self.args.uses,
             uses_with=dict(
@@ -73,6 +71,8 @@ async def async_setup(self):
                 ssl_certfile=self.args.ssl_certfile,
                 uvicorn_kwargs=self.args.uvicorn_kwargs,
                 proxy=self.args.proxy,
+                opentelemetry_tracing=self.opentelemetry_tracing,
+                tracer_provider=self.tracer_provider,
                 **uses_with,
             ),
             uses_metas={},
@@ -88,6 +88,12 @@ async def async_setup(self):
             timeout_send=self.timeout_send,
             metrics_registry=self.metrics_registry,
             runtime_name=self.args.name,
+            aio_tracing_client_interceptors=self.aio_tracing_client_interceptors(
+                self.tracer
+            ),
+            tracing_client_interceptor=self.tracing_client_interceptor(
+                self.tracer_provider
+            ),
         )
         await self.gateway.setup_server()
 

jina/serve/runtimes/gateway/grpc/__init__.py

@@ -1,58 +1 @@
 from jina.serve.runtimes.gateway.grpc.gateway import GRPCGateway
-
-__all__ = ['GRPCGatewayRuntime']
-
-
-class GRPCGatewayRuntime(GatewayRuntime):
-    """Gateway Runtime for gRPC."""
-
-    async def async_setup(self):
-        """
-        The async method to setup.
-
-        Create the gRPC server and expose the port for communication.
-        """
-        if not self.args.proxy and os.name != 'nt':
-            os.unsetenv('http_proxy')
-            os.unsetenv('https_proxy')
-
-        if not (is_port_free(__default_host__, self.args.port)):
-            raise PortAlreadyUsed(f'port:{self.args.port}')
-
-        self.gateway = GRPCGateway(
-            name=self.name,
-            grpc_server_options=self.args.grpc_server_options,
-            grpc_tracing_server_interceptors=self.aio_tracing_server_interceptor(),
-            port=self.args.port,
-            ssl_keyfile=self.args.ssl_keyfile,
-            ssl_certfile=self.args.ssl_certfile,
-        )
-
-        self.gateway.set_streamer(
-            args=self.args,
-            timeout_send=self.timeout_send,
-            metrics_registry=self.metrics_registry,
-            runtime_name=self.name,
-            aio_tracing_client_interceptors=self.aio_tracing_client_interceptors(
-                self.tracer
-            ),
-            tracing_client_interceptor=self.tracing_client_interceptor(
-                self.tracer_provider
-            ),
-        )
-        await self.gateway.setup_server()
-
-    async def async_teardown(self):
-        """Close the connection pool"""
-        # usually async_cancel should already have been called, but then its a noop
-        # if the runtime is stopped without a sigterm (e.g. as a context manager, this can happen)
-        await self.gateway.teardown()
-        await self.async_cancel()
-
-    async def async_cancel(self):
-        """The async method to stop server."""
-        await self.gateway.stop_server()
-
-    async def async_run_forever(self):
-        """The async running of server."""
-        await self.gateway.run_server()

jina/serve/runtimes/gateway/http/__init__.py

@@ -1,71 +1 @@
-import asyncio
-import os
-
-from jina import __default_host__
-from jina.serve.runtimes.gateway import GatewayRuntime
-from jina.serve.runtimes.gateway.http.app import get_fastapi_app
-
-__all__ = ['HTTPGatewayRuntime']
-
-from jina.serve.runtimes.gateway.http.gateway import HTTPGateway
-
-
-class HTTPGatewayRuntime(GatewayRuntime):
-    """Runtime for HTTP interface."""
-
-    async def async_setup(self):
-        """
-        The async method setup the runtime.
-
-        Setup the uvicorn server.
-        """
-        self.gateway = HTTPGateway(
-            name=self.name,
-            port=self.args.port,
-            title=self.args.title,
-            description=self.args.description,
-            no_debug_endpoints=self.args.no_debug_endpoints,
-            no_crud_endpoints=self.args.no_crud_endpoints,
-            expose_endpoints=self.args.expose_endpoints,
-            expose_graphql_endpoint=self.args.expose_graphql_endpoint,
-            cors=self.args.cors,
-            ssl_keyfile=self.args.ssl_keyfile,
-            ssl_certfile=self.args.ssl_certfile,
-            uvicorn_kwargs=self.args.uvicorn_kwargs,
-            opentelemetry_tracing=self.opentelemetry_tracing,
-            tracer_provider=self.tracer_provider,
-        )
-
-        self.gateway.set_streamer(
-            args=self.args,
-            timeout_send=self.timeout_send,
-            metrics_registry=self.metrics_registry,
-            runtime_name=self.args.name,
-            aio_tracing_client_interceptors=self.aio_tracing_client_interceptors(
-                self.tracer
-            ),
-            tracing_client_interceptor=self.tracing_client_interceptor(
-                self.tracer_provider
-            ),
-        )
-        await self.gateway.setup_server()
-
-    async def _wait_for_cancel(self):
-        """Do NOT override this method when inheriting from :class:`GatewayPod`"""
-        # handle terminate signals
-        while not self.is_cancel.is_set() and not self.gateway.should_exit:
-            await asyncio.sleep(0.1)
-
-        await self.async_cancel()
-
-    async def async_teardown(self):
-        """Shutdown the server."""
-        await self.gateway.teardown()
-
-    async def async_cancel(self):
-        """Stop the server."""
-        await self.gateway.stop_server()
-
-    async def async_run_forever(self):
-        """Running method of the server."""
-        await self.gateway.run_server()
+from .gateway import HTTPGateway

jina/serve/runtimes/gateway/http/gateway.py

@@ -26,6 +26,7 @@ def __init__(
         ssl_keyfile: Optional[str] = None,
         ssl_certfile: Optional[str] = None,
         uvicorn_kwargs: Optional[dict] = None,
+        proxy: Optional[bool] = None,
         opentelemetry_tracing: Optional[bool] = None,
         tracer_provider: Optional[trace.TracerProvider] = None,
         **kwargs
@@ -45,6 +46,8 @@ def __init__(
         :param ssl_keyfile: the path to the key file
         :param ssl_certfile: the path to the certificate file
         :param uvicorn_kwargs: Dictionary of kwargs arguments that will be passed to Uvicorn server when starting the server
+        :param proxy: If set, respect the http_proxy and https_proxy environment variables, otherwise, it will unset
+            these proxy variables before start. gRPC seems to prefer no proxy
         :param opentelemetry_tracing: Enables tracing if set to True.
         :param tracer_provider: If tracing is enabled the tracer_provider will be used to instrument the code.
         :param kwargs: keyword args

jina/serve/runtimes/gateway/websocket/__init__.py

@@ -1,57 +1,8 @@
-from jina.serve.runtimes.gateway.websocket.gateway import WebSocketGateway
-
-
-class WebSocketGatewayRuntime(GatewayRuntime):
-    """Runtime for Websocket interface."""
-
-    async def async_setup(self):
-        """
-        The async method setup the runtime.
-
-        Setup the uvicorn server.
-        """
+import asyncio
 
-        self.gateway = WebSocketGateway(
-            name=self.name,
-            port=self.args.port,
-            ssl_keyfile=self.args.ssl_keyfile,
-            ssl_certfile=self.args.ssl_certfile,
-            uvicorn_kwargs=self.args.uvicorn_kwargs,
-            logger=self.logger,
-            opentelemetry_tracing=self.opentelemetry_tracing,
-            tracer_provider=self.tracer_provider,
-        )
+from jina.serve.runtimes.gateway import GatewayRuntime
+from jina.serve.runtimes.gateway.websocket.app import get_fastapi_app
 
-        self.gateway.set_streamer(
-            args=self.args,
-            timeout_send=self.timeout_send,
-            metrics_registry=self.metrics_registry,
-            runtime_name=self.args.name,
-            aio_tracing_client_interceptors=self.aio_tracing_client_interceptors(
-                self.tracer
-            ),
-            tracing_client_interceptor=self.tracing_client_interceptor(
-                self.tracer_provider
-            ),
-        )
-        await self.gateway.setup_server()
+__all__ = ['WebSocketGatewayRuntime']
 
-    async def _wait_for_cancel(self):
-        """Do NOT override this method when inheriting from :class:`GatewayPod`"""
-        # handle terminate signals
-        while not self.is_cancel.is_set() and not self.gateway.should_exit:
-            await asyncio.sleep(0.1)
-
-        await self.async_cancel()
-
-    async def async_teardown(self):
-        """Shutdown the server."""
-        await self.gateway.teardown()
-
-    async def async_cancel(self):
-        """Stop the server."""
-        await self.gateway.stop_server()
-
-    async def async_run_forever(self):
-        """Running method of ther server."""
-        await self.gateway.run_server()
+from jina.serve.runtimes.gateway.websocket.gateway import WebSocketGateway

jina/serve/runtimes/gateway/websocket/gateway.py

@@ -19,6 +19,7 @@ def __init__(
         ssl_keyfile: Optional[str] = None,
         ssl_certfile: Optional[str] = None,
         uvicorn_kwargs: Optional[dict] = None,
+        proxy: Optional[bool] = None,
         opentelemetry_tracing: Optional[bool] = None,
         tracer_provider: Optional[trace.TracerProvider] = None,
         **kwargs
@@ -28,6 +29,8 @@ def __init__(
         :param ssl_keyfile: the path to the key file
         :param ssl_certfile: the path to the certificate file
         :param uvicorn_kwargs: Dictionary of kwargs arguments that will be passed to Uvicorn server when starting the server
+        :param proxy: If set, respect the http_proxy and https_proxy environment variables, otherwise, it will unset
+            these proxy variables before start. gRPC seems to prefer no proxy
         :param opentelemetry_tracing: Enables tracing if set to True.
         :param tracer_provider: If tracing is enabled the tracer_provider will be used to instrument the code.
         :param kwargs: keyword args

tests/integration/clients_extra_kwargs/test_clients_post_extra_kwargs.py

@@ -39,49 +39,14 @@ async def intercept_service(self, continuation, handler_call_details):
             return self._deny
 
     class AlternativeGRPCGateway(GRPCGateway):
-        def __init__(self, namespace_args, *args, **kwargs):
-            super(AlternativeGRPCGateway, self).__init__(
-                args=namespace_args, *args, **kwargs
-            )
+        def __init__(self, *args, **kwargs):
+            super(AlternativeGRPCGateway, self).__init__(*args, **kwargs)
             self.server = grpc.aio.server(
                 interceptors=(AuthInterceptor('access_key'),),
                 options=_get_grpc_server_options(self.grpc_server_options),
             )
 
-    class AlternativeGRPCGatewayRuntime(_GRPCGatewayRuntime):
-        async def async_setup(self):
-            """
-            The async method to setup.
-            Create the gRPC server and expose the port for communication.
-            """
-            if not self.args.proxy and os.name != 'nt':
-                os.unsetenv('http_proxy')
-                os.unsetenv('https_proxy')
-
-            if not (is_port_free(__default_host__, self.args.port)):
-                raise PortAlreadyUsed(f'port:{self.args.port}')
-
-            self.gateway = AlternativeGRPCGateway(
-                name=self.name,
-                namespace_args=self.args,
-                grpc_server_options=self.args.grpc_server_options,
-                port=self.args.port,
-                ssl_keyfile=self.args.ssl_keyfile,
-                ssl_certfile=self.args.ssl_certfile,
-            )
-
-            self.gateway.set_streamer(
-                timeout_send=self.timeout_send,
-                metrics_registry=self.metrics_registry,
-                runtime_name=self.name,
-            )
-            await self.gateway.setup_server()
-
-    monkeypatch.setattr(
-        'jina.serve.runtimes.gateway.grpc.GRPCGatewayRuntime',
-        AlternativeGRPCGatewayRuntime,
-    )
-    return Flow(protocol='grpc').add()
+    return Flow(protocol='grpc', uses=AlternativeGRPCGateway).add()
 
 
 def test_client_grpc_kwargs(flow_with_grpc):