chore(ibis): enable OpenTelemetry and multiple workers for the ibis-server image (#1187)

Author: goldmedal

ibis-server/Dockerfile

@@ -64,6 +64,13 @@ COPY --from=builder ${VIRTUAL_ENV} ${VIRTUAL_ENV}
 COPY app app
 COPY resources resources
 
+# Install Opentelemetry zero-instrumentation python
+RUN pip install opentelemetry-distro opentelemetry-exporter-otlp \
+    && opentelemetry-bootstrap -a install
+
+COPY entrypoint.sh ./
+RUN chmod +x ./entrypoint.sh
+
 EXPOSE 8000
 
-CMD ["fastapi", "run"]
+CMD ["./entrypoint.sh"]

ibis-server/README.md

@@ -7,7 +7,7 @@ This module is the API server of Wren Engine. It's built on top of [FastAPI](htt
 You can follow the steps below to run the Java engine and ibis.
 > Wren Engine is migrating to [wren-core](../wren-core/). However, we still recommend starting [the Java engine](../wren-core-legacy/) to enable the query fallback mechanism.
 
-Create `compose.yaml` file and add the following content, edit environment variables if needed (see [Environment Variables](docs/development#environment-variables))
+Create `compose.yaml` file and add the following content, edit environment variables if needed (see [Environment Variables](docs/development#environment-variables)).
 ```yaml
 services:
   ibis:
@@ -38,6 +38,9 @@ Run the docker compose
 docker compose up
 ```
 
+Set up [OpenTelemetry Envrionment Variables](docs/development#environment-variable) to enable tracing log.
+See [Tracing with Jaeger](#tracing-with-jaeger) to start up a Jaeger Server.
+
 ### Running on Local
 Requirements:
 - Python 3.11
@@ -92,6 +95,11 @@ docker run --rm --name jaeger \
   -p 9411:9411 \
   jaegertracing/jaeger:2.5.0
 ```
+- Install [OpenTelemetry Python zero-code instrumentation](https://opentelemetry.io/docs/zero-code/python/#setup)
+```
+pip install opentelemetry-distro opentelemetry-exporter-otlp
+opentelemetry-bootstrap -a install
+```
 - Use the following `just` command to start the `ibis-server` and export tracing logs to Jaeger:
 ```
 just run-trace-otlp 

ibis-server/app/dependencies.py

@@ -29,6 +29,8 @@ def _filter_headers(header_string: str) -> bool:
         return True
     elif header_string.startswith("x-user-"):
         return True
+    elif header_string.startswith("x-correlation-id"):
+        return True
     elif header_string == "traceparent":
         return True
     elif header_string == "tracestate":

ibis-server/app/routers/v2/connector.py

@@ -22,7 +22,13 @@
 from app.model.metadata.factory import MetadataFactory
 from app.model.validator import Validator
 from app.query_cache import QueryCacheManager
-from app.util import build_context, get_fallback_message, pushdown_limit, to_json
+from app.util import (
+    build_context,
+    get_fallback_message,
+    pushdown_limit,
+    set_attribute,
+    to_json,
+)
 
 router = APIRouter(prefix="/connector", tags=["connector"])
 tracer = trace.get_tracer(__name__)
@@ -70,6 +76,7 @@ async def query(
     with tracer.start_as_current_span(
         name=span_name, kind=trace.SpanKind.SERVER, context=build_context(headers)
     ) as span:
+        set_attribute(headers, span)
         try:
             sql = pushdown_limit(dto.sql, limit)
         except Exception as e:
@@ -176,7 +183,8 @@ async def validate(
     span_name = f"v2_validate_{data_source}"
     with tracer.start_as_current_span(
         name=span_name, kind=trace.SpanKind.SERVER, context=build_context(headers)
-    ):
+    ) as span:
+        set_attribute(headers, span)
         validator = Validator(
             Connector(data_source, dto.connection_info),
             Rewriter(
@@ -208,7 +216,8 @@ def get_table_list(
     span_name = f"v2_metadata_tables_{data_source}"
     with tracer.start_as_current_span(
         name=span_name, kind=trace.SpanKind.SERVER, context=build_context(headers)
-    ):
+    ) as span:
+        set_attribute(headers, span)
         return MetadataFactory.get_metadata(
             data_source, dto.connection_info
         ).get_table_list()
@@ -228,7 +237,8 @@ def get_constraints(
     span_name = f"v2_metadata_constraints_{data_source}"
     with tracer.start_as_current_span(
         name=span_name, kind=trace.SpanKind.SERVER, context=build_context(headers)
-    ):
+    ) as span:
+        set_attribute(headers, span)
         return MetadataFactory.get_metadata(
             data_source, dto.connection_info
         ).get_constraints()
@@ -252,7 +262,8 @@ async def dry_plan(
 ) -> str:
     with tracer.start_as_current_span(
         name="dry_plan", kind=trace.SpanKind.SERVER, context=build_context(headers)
-    ):
+    ) as span:
+        set_attribute(headers, span)
         sql = await Rewriter(
             dto.manifest_str, java_engine_connector=java_engine_connector
         ).rewrite(dto.sql)
@@ -278,7 +289,8 @@ async def dry_plan_for_data_source(
     span_name = f"v2_dry_plan_{data_source}"
     with tracer.start_as_current_span(
         name=span_name, kind=trace.SpanKind.SERVER, context=build_context(headers)
-    ):
+    ) as span:
+        set_attribute(headers, span)
         sql = await Rewriter(
             dto.manifest_str,
             data_source=data_source,
@@ -306,7 +318,8 @@ async def model_substitute(
     span_name = f"v2_model_substitute_{data_source}"
     with tracer.start_as_current_span(
         name=span_name, kind=trace.SpanKind.SERVER, context=build_context(headers)
-    ):
+    ) as span:
+        set_attribute(headers, span)
         sql = ModelSubstitute(data_source, dto.manifest_str, headers).substitute(
             dto.sql, write="trino"
         )

ibis-server/app/routers/v3/connector.py

@@ -34,6 +34,7 @@
     build_context,
     pushdown_limit,
     safe_strtobool,
+    set_attribute,
     to_json,
 )
 
@@ -73,6 +74,7 @@ async def query(
     with tracer.start_as_current_span(
         name=span_name, kind=trace.SpanKind.SERVER, context=build_context(headers)
     ) as span:
+        set_attribute(headers, span)
         try:
             if dry_run:
                 sql = pushdown_limit(dto.sql, limit)
@@ -188,6 +190,7 @@ async def dry_plan(
     with tracer.start_as_current_span(
         name="dry_plan", kind=trace.SpanKind.SERVER, context=build_context(headers)
     ) as span:
+        set_attribute(headers, span)
         try:
             return await Rewriter(
                 dto.manifest_str, experiment=True, properties=dict(headers)
@@ -228,6 +231,7 @@ async def dry_plan_for_data_source(
     with tracer.start_as_current_span(
         name=span_name, kind=trace.SpanKind.SERVER, context=build_context(headers)
     ) as span:
+        set_attribute(headers, span)
         try:
             return await Rewriter(
                 dto.manifest_str,
@@ -273,6 +277,7 @@ async def validate(
     with tracer.start_as_current_span(
         name=span_name, kind=trace.SpanKind.SERVER, context=build_context(headers)
     ) as span:
+        set_attribute(headers, span)
         try:
             validator = Validator(
                 Connector(data_source, dto.connection_info),
@@ -342,6 +347,7 @@ async def model_substitute(
     with tracer.start_as_current_span(
         name=span_name, kind=trace.SpanKind.SERVER, context=build_context(headers)
     ) as span:
+        set_attribute(headers, span)
         try:
             sql = ModelSubstitute(data_source, dto.manifest_str, headers).substitute(
                 dto.sql

ibis-server/app/util.py

@@ -114,6 +114,16 @@ def build_context(headers: Header) -> Context:
     return extract(headers)
 
 
+def set_attribute(
+    header: Header,
+    span: trace.Span,
+) -> None:
+    if header is None:
+        return
+    if "X-Correlation-ID" in header:
+        span.set_attribute("correlation_id", header["X-Correlation-ID"])
+
+
 def append_fallback_context(headers: Header, span: trace.Span) -> Headers:
     if headers is None:
         headers = {}

ibis-server/app/worker.py

@@ -0,0 +1,5 @@
+from uvicorn.workers import UvicornWorker
+
+
+class WrenUvicornWorker(UvicornWorker):
+    CONFIG_KWARGS = {"loop": "uvloop", "http": "httptools"}

ibis-server/docs/development.md

@@ -63,6 +63,20 @@ To start the server:
 ### Environment Variables
 
 - `WREN_ENGINE_ENDPOINT`: The endpoint of the Wren Java engine
+- `WREN_NUM_WORKERS`: The number of gunicoron workers
+
+### OpenTelemetry Envrionment Variables
+- `OTLP_ENABLED`: Enable the tracing for Ibis Server.
+- See more [OpenTelemetry environment variables](https://opentelemetry.io/docs/specs/otel/configuration/sdk-environment-variables/).
+- The minimize setting:
+  ```
+  OTLP_ENABLED=true
+  OTEL_EXPORTER_OTLP_TRACES_ENDPOINT=http://host.docker.internal:4317
+  OTEL_SERVICE_NAME=wren-engine
+  OTEL_TRACES_EXPORTER=otlp
+  OTEL_METRICS_EXPORTER=none
+  OTEL_LOGS_EXPORTER=none
+  ```
 
 ## How to add new data source
 

ibis-server/entrypoint.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+
+if [[ -z "${WREN_NUM_WORKERS}" ]]; then
+    echo "WREN_NUM_WORKERS is not set. Using default value of 2."
+    WREN_NUM_WORKERS=2
+else
+    WREN_NUM_WORKERS=${WREN_NUM_WORKERS}
+fi
+
+echo "Number of workers: ${WREN_NUM_WORKERS}"
+
+# Determine the command prefix based on OTLP_ENABLED
+if [[ "${OTLP_ENABLED}" == "true" ]]; then
+    CMD_PREFIX="opentelemetry-instrument"
+else
+    CMD_PREFIX=""
+fi
+
+# Start the WrenUvicornWorker with the specified configuration
+${CMD_PREFIX} gunicorn app.main:app --bind 0.0.0.0:8000 \
+    -k app.worker.WrenUvicornWorker \
+    --workers ${WREN_NUM_WORKERS} \
+    --max-requests 1000 \
+    --max-requests-jitter 100 \
+    --timeout 300 \
+    --graceful-timeout 60

ibis-server/justfile

@@ -23,6 +23,17 @@ port := "8000"
 run:
     poetry run fastapi run --port {{ port }}
 
+workers := "2"
+
+run-gunicorn:
+    gunicorn app.main:app --bind 0.0.0.0:{{ port }} \
+        -k app.worker.WrenUvicornWorker \
+        --workers {{ workers }} \
+        --max-requests 1000 \
+        --max-requests-jitter 100 \
+        --timeout 300 \
+        --graceful-timeout 60
+
 run-trace-console:
     opentelemetry-instrument --traces_exporter console --metrics_exporter none fastapi run --port {{ port }}