refactor(tracing): add jaeger thrift exporter (#9)

* refactor(tracing): add jaeger thrift exporter
* chore(ci): dynamic badge

Author: hspedro

README.md

@@ -3,7 +3,7 @@
 [![Tests](https://github.com/hspedro/babeltron/actions/workflows/test.yml/badge.svg?branch=main)](https://github.com/hspedro/babeltron/actions/workflows/test.yml)
 [![Python Version](https://img.shields.io/badge/python-3.9%2B-blue)](https://www.python.org/downloads/)
 [![License](https://img.shields.io/github/license/hspedro/babeltron)](https://github.com/hspedro/babeltron/blob/main/LICENSE)
-[![Version](https://img.shields.io/badge/version-0.3.1-green)](https://github.com/hspedro/babeltron/releases)
+[![Version](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/hspedro/babeltron/main/.github/badges/version.json)](https://github.com/hspedro/babeltron/releases)
 [![PyPI version](https://badge.fury.io/py/babeltron.svg)](https://badge.fury.io/py/babeltron)
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
 [![Coverage](https://img.shields.io/badge/coverage-90%25-brightgreen.svg)](https://github.com/hspedro/babeltron/actions/workflows/test.yml)

babeltron/app/main.py

@@ -1,18 +1,30 @@
+import importlib.metadata
+import logging
 import os
-from importlib.metadata import version
 
-from fastapi import FastAPI, Response
+from fastapi import FastAPI, HTTPException, status
 from fastapi.middleware.cors import CORSMiddleware
+from starlette.responses import JSONResponse, PlainTextResponse
 
 from babeltron.app.middlewares.auth import BasicAuthMiddleware
+from babeltron.app.models.m2m import M2MTranslationModel
 from babeltron.app.monitoring import PrometheusMiddleware, metrics_endpoint
 from babeltron.app.tracing import setup_jaeger
 from babeltron.app.utils import include_routers
 
+# Get version from package metadata
 try:
-    __version__ = version("babeltron")
-except ImportError:
-    __version__ = "dev"
+    __version__ = importlib.metadata.version("babeltron")
+except importlib.metadata.PackageNotFoundError:
+    # If package is not installed, try to get version from pyproject.toml
+    try:
+        import tomli
+
+        with open("pyproject.toml", "rb") as f:
+            pyproject = tomli.load(f)
+            __version__ = pyproject["project"]["version"]
+    except (FileNotFoundError, KeyError, ImportError):
+        __version__ = "dev"
 
 
 def create_app() -> FastAPI:
@@ -68,6 +80,8 @@ def create_app() -> FastAPI:
                 "/healthz",
                 "/readyz",
                 "/metrics",
+                "/version",  # Add version endpoint to excluded paths
+                "/version-badge",  # Add version badge endpoint to excluded paths
             ],
         )
 
@@ -80,17 +94,50 @@ def create_app() -> FastAPI:
     # Add Prometheus middleware
     app.add_middleware(PrometheusMiddleware)
 
+    # Add version endpoint for badge
+    @app.get("/version", response_class=PlainTextResponse, include_in_schema=False)
+    async def get_version():
+        return __version__
+
+    # Add version badge endpoint for Shields.io
+    @app.get("/version-badge", response_class=JSONResponse, include_in_schema=False)
+    async def get_version_badge():
+        return {
+            "schemaVersion": 1,
+            "label": "version",
+            "message": __version__,
+            "color": "green",
+            "cacheSeconds": 3600,  # Cache for 1 hour
+        }
+
+    @app.get("/healthz", include_in_schema=False)
+    async def health():
+        return {"status": "ok"}
+
+    @app.get("/readyz", include_in_schema=False)
+    async def ready():
+        # Check if model is loaded
+        try:
+            # Just initialize the model to check if it loads correctly
+            M2MTranslationModel()
+            return {"status": "ready"}
+        except Exception as e:
+            logging.error(f"Readiness check failed: {e}")
+            raise HTTPException(
+                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+                detail="Service not ready",
+            )
+
+    @app.get("/metrics", include_in_schema=False)
+    async def metrics():
+        return PlainTextResponse(content=metrics_endpoint())
+
     return app
 
 
 app = create_app()
 
 
-@app.get("/metrics", include_in_schema=False)
-async def metrics():
-    return Response(content=metrics_endpoint(), media_type="text/plain")
-
-
 if __name__ == "__main__":
     import uvicorn
 

babeltron/app/routers/translate.py

@@ -29,7 +29,7 @@ class TranslationRequest(BaseModel):
     )
 
     class Config:
-        schema_extra = {
+        json_schema_extra = {
             "example": {
                 "text": "Hello, how are you?",
                 "src_lang": "en",

babeltron/app/tracing.py

@@ -7,6 +7,7 @@
 
 if not IN_TEST:
     from opentelemetry import trace
+    from opentelemetry.exporter.jaeger.thrift import JaegerExporter
     from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
         OTLPSpanExporter as OTLPSpanExporterGRPC,
     )
@@ -15,10 +16,12 @@
     )
     from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
     from opentelemetry.instrumentation.logging import LoggingInstrumentor
+    from opentelemetry.sdk.resources import Resource
     from opentelemetry.sdk.trace import TracerProvider
-    from opentelemetry.sdk.trace.export import BatchSpanProcessor
+    from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
+    from opentelemetry.semconv.resource import ResourceAttributes
 else:
-    # Create dummy classes for test environment
+
     class DummyTracerProvider:
         def add_span_processor(self, processor):
             pass
@@ -43,53 +46,85 @@ def instrument_app(app, **kwargs):
     FastAPIInstrumentor = DummyInstrumentor
     LoggingInstrumentor = DummyInstrumentor
 
-OTLP_MODE = os.environ.get("OTLP_MODE", "otlp-grpc")
-OTLP_GRPC_ENDPOINT = os.environ.get("OTLP_GRPC_ENDPOINT", "otel-collector:4317")
-OTLP_HTTP_ENDPOINT = os.environ.get(
-    "OTLP_HTTP_ENDPOINT", "http://otel-collector:4318/v1/traces"
+# Configuration options
+OTLP_MODE = os.environ.get("OTLP_MODE", "otlp-grpc")  # Default to otlp-grpc
+OTLP_COLLECTOR_HOST = os.environ.get("OTLP_COLLECTOR_HOST", "otel-collector")
+OTLP_COLLECTOR_PORT = os.environ.get("OTLP_COLLECTOR_PORT", "4317")
+JAEGER_AGENT_HOST = os.environ.get(
+    "JAEGER_AGENT_HOST", "localhost" if IN_TEST else "jaeger"
 )
+JAEGER_AGENT_PORT = int(os.environ.get("JAEGER_AGENT_PORT", "6831"))
+SERVICE_NAME = os.environ.get("OTEL_SERVICE_NAME", "babeltron")
 
 
 def setup_jaeger(app: FastAPI, log_correlation: bool = True) -> None:
     if IN_TEST:
-        logging.info("Skipping OpenTelemetry setup in test environment")
+        logging.info("Using minimal OpenTelemetry setup for test environment")
+        # For tests, we'll set up a minimal tracer that doesn't try to connect to external services
+        resource = Resource.create(
+            {ResourceAttributes.SERVICE_NAME: f"{SERVICE_NAME}-test"}
+        )
+        tracer = TracerProvider(resource=resource)
+        trace.set_tracer_provider(tracer)
+
+        # In tests, we'll use the console exporter which doesn't require external connections
+        # This is optional and can be disabled if you don't want any tracing in tests
+        if os.environ.get("OTEL_TEST_EXPORT", "false").lower() == "true":
+            tracer.add_span_processor(BatchSpanProcessor(ConsoleSpanExporter()))
+            logging.info("Console span exporter enabled for tests")
+
+        # Instrument the app for tests if needed
+        FastAPIInstrumentor.instrument_app(
+            app,
+            tracer_provider=tracer,
+            excluded_urls="/metrics,/healthz,/readyz,/docs,/redoc,/openapi.json",
+        )
         return
 
-    if OTLP_GRPC_ENDPOINT.lower() == "disabled":
+    if OTLP_MODE.lower() == "disabled":
         logging.info("OpenTelemetry tracing is disabled")
         return
 
-    tracer = TracerProvider()
+    # Create a resource with service name
+    resource = Resource.create({ResourceAttributes.SERVICE_NAME: SERVICE_NAME})
+
+    # Create tracer provider with resource
+    tracer = TracerProvider(resource=resource)
     trace.set_tracer_provider(tracer)
 
     if OTLP_MODE == "otlp-grpc":
-        if not IN_TEST:
-            tracer.add_span_processor(
-                BatchSpanProcessor(
-                    OTLPSpanExporterGRPC(endpoint=OTLP_GRPC_ENDPOINT, insecure=True)
-                )
-            )
+        # Use OTLP gRPC exporter to send to the OpenTelemetry Collector
+        endpoint = f"{OTLP_COLLECTOR_HOST}:{OTLP_COLLECTOR_PORT}"
+        tracer.add_span_processor(
+            BatchSpanProcessor(OTLPSpanExporterGRPC(endpoint=endpoint, insecure=True))
+        )
+        logging.info(f"OTLP gRPC exporter enabled with endpoint: {endpoint}")
     elif OTLP_MODE == "otlp-http":
-        if not IN_TEST:
-            tracer.add_span_processor(
-                BatchSpanProcessor(OTLPSpanExporterHTTP(endpoint=OTLP_HTTP_ENDPOINT))
-            )
+        # Use OTLP HTTP exporter to send to the OpenTelemetry Collector
+        endpoint = f"{OTLP_COLLECTOR_HOST}:{OTLP_COLLECTOR_PORT}"
+        tracer.add_span_processor(
+            BatchSpanProcessor(OTLPSpanExporterHTTP(endpoint=endpoint))
+        )
+        logging.info(f"OTLP HTTP exporter enabled with endpoint: {endpoint}")
     else:
-        if not IN_TEST:
-            tracer.add_span_processor(
-                BatchSpanProcessor(
-                    OTLPSpanExporterGRPC(endpoint=OTLP_GRPC_ENDPOINT, insecure=True)
-                )
-            )
-
-    if log_correlation and not IN_TEST:
-        LoggingInstrumentor().instrument(set_logging_format=True)
-
-    if not IN_TEST:
-        FastAPIInstrumentor.instrument_app(
-            app,
-            tracer_provider=tracer,
-            excluded_urls="/metrics,/healthz,/readyz,/docs,/redoc,/openapi.json",
+        # Use Jaeger Thrift exporter (deprecated but still functional)
+        jaeger_exporter = JaegerExporter(
+            agent_host_name=JAEGER_AGENT_HOST,
+            agent_port=JAEGER_AGENT_PORT,
+        )
+        tracer.add_span_processor(BatchSpanProcessor(jaeger_exporter))
+        logging.info(
+            f"Jaeger Thrift exporter enabled with agent: {JAEGER_AGENT_HOST}:{JAEGER_AGENT_PORT}"
+        )
+        logging.warning(
+            "Note: The Jaeger Thrift exporter is deprecated. Consider migrating to OTLP."
         )
 
-    logging.info(f"OpenTelemetry tracing enabled with endpoint: {OTLP_GRPC_ENDPOINT}")
+    if log_correlation:
+        LoggingInstrumentor().instrument(set_logging_format=True)
+
+    FastAPIInstrumentor.instrument_app(
+        app,
+        tracer_provider=tracer,
+        excluded_urls="/metrics,/healthz,/readyz,/docs,/redoc,/openapi.json",
+    )

docker-compose.yml

@@ -14,7 +14,8 @@ services:
       - METRICS_ENABLED=true
       - OTLP_MODE=otlp-grpc
       - OTEL_SERVICE_NAME=babeltron
-      - OTLP_GRPC_ENDPOINT=otel-collector:4317
+      - OTLP_COLLECTOR_HOST=otel-collector
+      - OTLP_COLLECTOR_PORT=4317
       - API_PASSWORD=translation2025
       - WORKER_COUNT=2
     restart: unless-stopped
@@ -24,6 +25,8 @@ services:
       timeout: 10s
       retries: 3
       start_period: 40s
+    depends_on:
+      - otel-collector
 
   prometheus:
     image: prom/prometheus:latest
@@ -53,58 +56,41 @@ services:
     depends_on:
       - prometheus
 
-  jaeger-collector:
-    container_name: jaeger-collector
-    image: jaegertracing/jaeger-collector:1.57.0
-    command:
-      - "--cassandra.keyspace=jaeger_v1_dc1"
-      - "--cassandra.servers=cassandra"
-      - "--collector.otlp.enabled=true"
-    environment:
-      - SAMPLING_CONFIG_TYPE=adaptive
-    ports:
-      - "4317" # accept OpenTelemetry Protocol (OTLP) over gRPC
-      - "4318" # accept OpenTelemetry Protocol (OTLP) over HTTP
-    restart: on-failure
-    depends_on:
-      - cassandra-schema
-
-  cassandra:
-    container_name: cassandra
-    image: cassandra:4.1.4
-
-  cassandra-schema:
-    container_name: cassandra-schema
-    image: jaegertracing/jaeger-cassandra-schema:1.57.0
-    depends_on:
-      - cassandra
-
-  jaeger-query:
-    container_name: jaeger-query
-    image: jaegertracing/jaeger-query:1.57.0
-    command:
-      - "--cassandra.keyspace=jaeger_v1_dc1"
-      - "--cassandra.servers=cassandra"
+  jaeger:
+    image: jaegertracing/all-in-one:latest
+    container_name: jaeger
     ports:
-      - "16686:16686"
-      - "16687:16687"
-    restart: on-failure
-    depends_on:
-      - cassandra-schema
+      - "6831:6831/udp"   # Jaeger thrift compact (agent)
+      - "6832:6832/udp"   # Jaeger thrift binary (agent)
+      - "5778:5778"       # Agent configs
+      - "16686:16686"     # Jaeger UI
+      - "14250:14250"     # Model/collector gRPC
+      - "14268:14268"     # Jaeger HTTP (collector)
+      - "14269:14269"     # Admin port
+    environment:
+      - COLLECTOR_OTLP_ENABLED=true
+      - COLLECTOR_ZIPKIN_HOST_PORT=:9411
+      - SAMPLING_STRATEGIES_FILE=/etc/jaeger/sampling.json
+      - LOG_LEVEL=debug
+    volumes:
+      - ./etc/jaeger-sampling.json:/etc/jaeger/sampling.json
+    restart: unless-stopped
 
   otel-collector:
-    image: otel/opentelemetry-collector-contrib:0.100.0
+    image: otel/opentelemetry-collector-contrib:latest
     container_name: otel-collector
-    command:
-      - "--config=/conf/config.yaml"
+    command: ["--config=/etc/otel-collector-config.yaml"]
     volumes:
-      - ./etc/otel-collector-config.yaml:/conf/config.yaml
+      - ./etc/otel-collector-config.yaml:/etc/otel-collector-config.yaml
     ports:
-      - "4317" # OTLP gRPC receiver
-      - "4318" # OTLP http receiver
-    restart: on-failure
+      - "4317:4317"   # OTLP gRPC receiver
+      - "4318:4318"   # OTLP HTTP receiver
+      - "8888:8888"   # Prometheus metrics exposed by the collector
+      - "8889:8889"   # Prometheus exporter metrics
+      - "13133:13133" # Health check extension
+    restart: unless-stopped
     depends_on:
-      - jaeger-collector
+      - jaeger
 
 volumes:
   grafana-storage: