genai-otel-instrument 0.1.1.dev0__py3-none-any.whl

genai_otel/gpu_metrics.py

@@ -0,0 +1,240 @@
+"""Module for collecting GPU metrics using nvidia-ml-py and reporting them via OpenTelemetry.
+
+This module provides the `GPUMetricsCollector` class, which periodically collects
+GPU utilization, memory usage, and temperature, and exports these as OpenTelemetry
+metrics. It relies on the `nvidia-ml-py` library for interacting with NVIDIA GPUs.
+"""
+
+import logging
+import threading
+import time
+from typing import Optional
+
+from opentelemetry.metrics import Meter, ObservableCounter, ObservableGauge, Observation
+
+from genai_otel.config import OTelConfig
+
+logger = logging.getLogger(__name__)
+
+# Try to import nvidia-ml-py (official replacement for pynvml)
+try:
+    import pynvml
+
+    NVML_AVAILABLE = True
+except ImportError:
+    NVML_AVAILABLE = False
+    logger.debug("nvidia-ml-py not available, GPU metrics will be disabled")
+
+
+class GPUMetricsCollector:
+    """Collects and reports GPU metrics using nvidia-ml-py."""
+
+    def __init__(self, meter: Meter, config: OTelConfig, interval: int = 10):
+        """Initializes the GPUMetricsCollector.
+
+        Args:
+            meter (Meter): The OpenTelemetry meter to use for recording metrics.
+        """
+        self.meter = meter
+        self.running = False
+        self.thread: Optional[threading.Thread] = None
+        self._thread: Optional[threading.Thread] = None  # Initialize _thread
+        self._stop_event = threading.Event()
+        self.gpu_utilization_counter: Optional[ObservableCounter] = None
+        self.gpu_memory_used_gauge: Optional[ObservableGauge] = None
+        self.gpu_temperature_gauge: Optional[ObservableGauge] = None
+        self.config = config
+        self.interval = interval  # seconds
+        self.gpu_available = False
+
+        self.device_count = 0
+        self.nvml = None
+        if NVML_AVAILABLE:
+            try:
+                pynvml.nvmlInit()
+                self.device_count = pynvml.nvmlDeviceGetCount()
+                if self.device_count > 0:
+                    self.gpu_available = True
+                self.nvml = pynvml
+            except Exception as e:
+                logger.error("Failed to initialize NVML to get device count: %s", e)
+
+        self.cumulative_energy_wh = [0.0] * self.device_count  # Per GPU, in Wh
+        self.last_timestamp = [time.time()] * self.device_count
+        self.co2_counter = meter.create_counter(
+            "gen_ai.co2.emissions",  # Fixed metric name
+            description="Cumulative CO2 equivalent emissions in grams",
+            unit="gCO2e",
+        )
+        if not NVML_AVAILABLE:
+            logger.warning(
+                "GPU metrics collection not available - nvidia-ml-py not installed. "
+                "Install with: pip install genai-otel-instrument[gpu]"
+            )
+            return
+
+        try:
+            # Use ObservableGauge for all GPU metrics (not Counter!)
+            self.gpu_utilization_gauge = self.meter.create_observable_gauge(
+                "gen_ai.gpu.utilization",  # Fixed metric name
+                callbacks=[self._observe_gpu_utilization],
+                description="GPU utilization percentage",
+                unit="%",
+            )
+            self.gpu_memory_used_gauge = self.meter.create_observable_gauge(
+                "gen_ai.gpu.memory.used",  # Fixed metric name
+                callbacks=[self._observe_gpu_memory],
+                description="GPU memory used in MiB",
+                unit="MiB",
+            )
+            self.gpu_temperature_gauge = self.meter.create_observable_gauge(
+                "gen_ai.gpu.temperature",  # Fixed metric name
+                callbacks=[self._observe_gpu_temperature],
+                description="GPU temperature in Celsius",
+                unit="Cel",
+            )
+        except Exception as e:
+            logger.error("Failed to create GPU metrics instruments: %s", e, exc_info=True)
+
+    def _get_device_name(self, handle, index):
+        """Get GPU device name safely."""
+        try:
+            device_name = pynvml.nvmlDeviceGetName(handle)
+            if isinstance(device_name, bytes):
+                device_name = device_name.decode("utf-8")
+            return device_name
+        except Exception as e:
+            logger.debug("Failed to get GPU name: %s", e)
+            return f"GPU_{index}"
+
+    def _observe_gpu_utilization(self, options):
+        """Observable callback for GPU utilization."""
+        if not NVML_AVAILABLE or not self.gpu_available:
+            return
+
+        try:
+            pynvml.nvmlInit()
+            device_count = pynvml.nvmlDeviceGetCount()
+
+            for i in range(device_count):
+                handle = pynvml.nvmlDeviceGetHandleByIndex(i)
+                device_name = self._get_device_name(handle, i)
+
+                try:
+                    utilization = pynvml.nvmlDeviceGetUtilizationRates(handle)
+                    yield Observation(
+                        value=utilization.gpu,
+                        attributes={"gpu_id": str(i), "gpu_name": device_name},
+                    )
+                except Exception as e:
+                    logger.debug("Failed to get GPU utilization for GPU %d: %s", i, e)
+
+            pynvml.nvmlShutdown()
+        except Exception as e:
+            logger.error("Error observing GPU utilization: %s", e)
+
+    def _observe_gpu_memory(self, options):
+        """Observable callback for GPU memory usage."""
+        if not NVML_AVAILABLE or not self.gpu_available:
+            return
+
+        try:
+            pynvml.nvmlInit()
+            device_count = pynvml.nvmlDeviceGetCount()
+
+            for i in range(device_count):
+                handle = pynvml.nvmlDeviceGetHandleByIndex(i)
+                device_name = self._get_device_name(handle, i)
+
+                try:
+                    memory_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+                    gpu_memory_used = memory_info.used / (1024**2)  # Convert to MiB
+                    yield Observation(
+                        value=gpu_memory_used,
+                        attributes={"gpu_id": str(i), "gpu_name": device_name},
+                    )
+                except Exception as e:
+                    logger.debug("Failed to get GPU memory for GPU %d: %s", i, e)
+
+            pynvml.nvmlShutdown()
+        except Exception as e:
+            logger.error("Error observing GPU memory: %s", e)
+
+    def _observe_gpu_temperature(self, options):
+        """Observable callback for GPU temperature."""
+        if not NVML_AVAILABLE or not self.gpu_available:
+            return
+
+        try:
+            pynvml.nvmlInit()
+            device_count = pynvml.nvmlDeviceGetCount()
+
+            for i in range(device_count):
+                handle = pynvml.nvmlDeviceGetHandleByIndex(i)
+                device_name = self._get_device_name(handle, i)
+
+                try:
+                    gpu_temp = pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU)
+                    yield Observation(
+                        value=gpu_temp, attributes={"gpu_id": str(i), "gpu_name": device_name}
+                    )
+                except Exception as e:
+                    logger.debug("Failed to get GPU temperature for GPU %d: %s", i, e)
+
+            pynvml.nvmlShutdown()
+        except Exception as e:
+            logger.error("Error observing GPU temperature: %s", e)
+
+    def start(self):
+        """Starts the GPU metrics collection.
+
+        ObservableGauges are automatically collected by the MeterProvider,
+        so we only need to start the CO2 collection thread.
+        """
+        if not NVML_AVAILABLE:
+            logger.warning("Cannot start GPU metrics collection - nvidia-ml-py not available")
+            return
+
+        if not self.gpu_available:
+            return
+
+        logger.info("Starting GPU metrics collection (CO2 tracking)")
+        # Only start CO2 collection thread - ObservableGauges are auto-collected
+        self._thread = threading.Thread(target=self._collect_loop, daemon=True)
+        self._thread.start()
+
+    def _collect_loop(self):
+        while not self._stop_event.wait(self.interval):
+            current_time = time.time()
+            for i in range(self.device_count):
+                try:
+                    handle = self.nvml.nvmlDeviceGetHandleByIndex(i)
+                    power_w = self.nvml.nvmlDeviceGetPowerUsage(handle) / 1000.0  # Watts
+                    delta_time_hours = (current_time - self.last_timestamp[i]) / 3600.0
+                    delta_energy_wh = (power_w / 1000.0) * (
+                        delta_time_hours * 3600.0
+                    )  # Wh (power in kW * hours = kWh, but track in Wh for precision)
+                    self.cumulative_energy_wh[i] += delta_energy_wh
+                    if self.config.enable_co2_tracking:
+                        delta_co2_g = (
+                            delta_energy_wh / 1000.0
+                        ) * self.config.carbon_intensity  # gCO2e
+                        self.co2_counter.add(delta_co2_g, {"gpu_id": str(i)})
+                    self.last_timestamp[i] = current_time
+                except Exception as e:
+                    logger.error(f"Error collecting GPU {i} metrics: {e}")
+
+    def stop(self):
+        """Stops the GPU metrics collection thread."""
+        # Stop CO2 collection thread
+        self._stop_event.set()
+        if self._thread and self._thread.is_alive():
+            self._thread.join(timeout=5)
+            logger.info("GPU CO2 metrics collection thread stopped.")
+
+        # ObservableGauges will automatically stop when MeterProvider is shutdown
+        if self.gpu_available and NVML_AVAILABLE:
+            try:
+                pynvml.nvmlShutdown()
+            except Exception as e:
+                logger.debug("Error shutting down NVML: %s", e)

genai_otel/instrumentors/__init__.py

@@ -0,0 +1,47 @@
+"""Module for OpenTelemetry instrumentors for various LLM providers and frameworks.
+
+This package contains individual instrumentor classes for different Generative AI
+libraries and frameworks, allowing for automatic tracing and metric collection
+of their operations.
+
+All imports are done lazily to avoid ImportError when optional dependencies
+are not installed.
+"""
+
+from .anthropic_instrumentor import AnthropicInstrumentor
+from .anyscale_instrumentor import AnyscaleInstrumentor
+from .aws_bedrock_instrumentor import AWSBedrockInstrumentor
+from .azure_openai_instrumentor import AzureOpenAIInstrumentor
+from .cohere_instrumentor import CohereInstrumentor
+from .google_ai_instrumentor import GoogleAIInstrumentor
+from .groq_instrumentor import GroqInstrumentor
+from .huggingface_instrumentor import HuggingFaceInstrumentor
+from .langchain_instrumentor import LangChainInstrumentor
+from .llamaindex_instrumentor import LlamaIndexInstrumentor
+from .mistralai_instrumentor import MistralAIInstrumentor
+from .ollama_instrumentor import OllamaInstrumentor
+
+# Import instrumentors only - they handle their own dependency checking
+from .openai_instrumentor import OpenAIInstrumentor
+from .replicate_instrumentor import ReplicateInstrumentor
+from .togetherai_instrumentor import TogetherAIInstrumentor
+from .vertexai_instrumentor import VertexAIInstrumentor
+
+__all__ = [
+    "OpenAIInstrumentor",
+    "AnthropicInstrumentor",
+    "GoogleAIInstrumentor",
+    "AWSBedrockInstrumentor",
+    "AzureOpenAIInstrumentor",
+    "CohereInstrumentor",
+    "MistralAIInstrumentor",
+    "TogetherAIInstrumentor",
+    "GroqInstrumentor",
+    "OllamaInstrumentor",
+    "VertexAIInstrumentor",
+    "ReplicateInstrumentor",
+    "AnyscaleInstrumentor",
+    "LangChainInstrumentor",
+    "LlamaIndexInstrumentor",
+    "HuggingFaceInstrumentor",
+]

genai_otel/instrumentors/anthropic_instrumentor.py

@@ -0,0 +1,134 @@
+"""OpenTelemetry instrumentor for the Anthropic Claude SDK.
+
+This instrumentor automatically traces calls to the Anthropic API, capturing
+relevant attributes such as model name, message count, and token usage.
+"""
+
+import logging
+from typing import Any, Dict, Optional
+
+from ..config import OTelConfig
+from .base import BaseInstrumentor
+
+logger = logging.getLogger(__name__)
+
+
+class AnthropicInstrumentor(BaseInstrumentor):
+    """Instrumentor for Anthropic Claude SDK"""
+
+    def __init__(self):
+        """Initialize the instrumentor."""
+        super().__init__()
+        self._anthropic_available = False
+        self._check_availability()
+
+    def _check_availability(self):
+        """Check if Anthropic library is available."""
+        try:
+            import anthropic
+
+            self._anthropic_available = True
+            logger.debug("Anthropic library detected and available for instrumentation")
+        except ImportError:
+            logger.debug("Anthropic library not installed, instrumentation will be skipped")
+            self._anthropic_available = False
+
+    def instrument(self, config: OTelConfig):
+        """Instrument Anthropic SDK if available.
+
+        Args:
+            config (OTelConfig): The OpenTelemetry configuration object.
+        """
+        if not self._anthropic_available:
+            logger.debug("Skipping Anthropic instrumentation - library not available")
+            return
+
+        self.config = config
+
+        try:
+            import anthropic
+            import wrapt
+
+            if hasattr(anthropic, "Anthropic"):
+                original_init = anthropic.Anthropic.__init__
+
+                def wrapped_init(wrapped, instance, args, kwargs):
+                    result = wrapped(*args, **kwargs)
+                    self._instrument_client(instance)
+                    return result
+
+                anthropic.Anthropic.__init__ = wrapt.FunctionWrapper(original_init, wrapped_init)
+                self._instrumented = True
+                logger.info("Anthropic instrumentation enabled")
+
+        except Exception as e:
+            logger.error("Failed to instrument Anthropic: %s", e, exc_info=True)
+            if config.fail_on_error:
+                raise
+
+    def _instrument_client(self, client):
+        """Instrument Anthropic client methods.
+
+        Args:
+            client: The Anthropic client instance to instrument.
+        """
+        if hasattr(client, "messages") and hasattr(client.messages, "create"):
+            original_create = client.messages.create
+            instrumented_create_method = self.create_span_wrapper(
+                span_name="anthropic.messages.create",
+                extract_attributes=self._extract_anthropic_attributes,
+            )(original_create)
+            client.messages.create = instrumented_create_method
+
+    def _extract_anthropic_attributes(
+        self, instance: Any, args: Any, kwargs: Any
+    ) -> Dict[str, Any]:
+        """Extract attributes from Anthropic API call.
+
+        Args:
+            instance: The client instance.
+            args: Positional arguments.
+            kwargs: Keyword arguments.
+
+        Returns:
+            Dict[str, Any]: Dictionary of attributes to set on the span.
+        """
+        attrs = {}
+        model = kwargs.get("model", "unknown")
+        messages = kwargs.get("messages", [])
+
+        attrs["gen_ai.system"] = "anthropic"
+        attrs["gen_ai.request.model"] = model
+        attrs["gen_ai.request.message_count"] = len(messages)
+        return attrs
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        """Extract token usage from Anthropic response.
+
+        Args:
+            result: The API response object.
+
+        Returns:
+            Optional[Dict[str, int]]: Dictionary with token counts or None.
+        """
+        if hasattr(result, "usage") and result.usage:
+            usage = result.usage
+            usage_dict = {
+                "prompt_tokens": getattr(usage, "input_tokens", 0),
+                "completion_tokens": getattr(usage, "output_tokens", 0),
+                "total_tokens": getattr(usage, "input_tokens", 0)
+                + getattr(usage, "output_tokens", 0),
+            }
+
+            # Extract cache tokens for Anthropic models (Phase 3.2)
+            # cache_read_input_tokens: Tokens that were read from cache
+            # cache_creation_input_tokens: Tokens that were written to cache
+            if hasattr(usage, "cache_read_input_tokens"):
+                usage_dict["cache_read_input_tokens"] = getattr(usage, "cache_read_input_tokens", 0)
+            if hasattr(usage, "cache_creation_input_tokens"):
+                usage_dict["cache_creation_input_tokens"] = getattr(
+                    usage, "cache_creation_input_tokens", 0
+                )
+
+            return usage_dict
+        return None

genai_otel/instrumentors/anyscale_instrumentor.py

@@ -0,0 +1,27 @@
+"""OpenTelemetry instrumentor for Anyscale Endpoints.
+
+This instrumentor integrates with Anyscale Endpoints, which often leverage
+OpenAI-compatible APIs. It ensures that calls made to Anyscale services are
+properly traced and attributed within the OpenTelemetry ecosystem.
+"""
+
+from typing import Dict, Optional
+
+from ..config import OTelConfig
+from .base import BaseInstrumentor
+
+
+class AnyscaleInstrumentor(BaseInstrumentor):
+    """Instrumentor for Anyscale Endpoints"""
+
+    def instrument(self, config: OTelConfig):
+        self.config = config
+        try:
+            # Anyscale uses OpenAI SDK, already instrumented
+            pass
+
+        except ImportError:
+            pass
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        return None

genai_otel/instrumentors/aws_bedrock_instrumentor.py

@@ -0,0 +1,94 @@
+import json
+import logging
+from typing import Any, Dict, Optional
+
+from ..config import OTelConfig
+from .base import BaseInstrumentor
+
+logger = logging.getLogger(__name__)
+
+
+class AWSBedrockInstrumentor(BaseInstrumentor):
+    """Instrumentor for AWS Bedrock"""
+
+    def __init__(self):
+        """Initialize the instrumentor."""
+        super().__init__()
+        self._boto3_available = False
+        self._check_availability()
+
+    def _check_availability(self):
+        """Check if boto3 library is available."""
+        try:
+            import boto3  # Moved to top
+
+            self._boto3_available = True
+            logger.debug("boto3 library detected and available for instrumentation")
+        except ImportError:
+            logger.debug("boto3 library not installed, instrumentation will be skipped")
+            self._boto3_available = False
+
+    def instrument(self, config: OTelConfig):
+        self.config = config
+        try:
+            import boto3  # Moved to top
+
+            original_client = boto3.client
+
+            def wrapped_client(*args, **kwargs):
+                client = original_client(*args, **kwargs)
+                if args and args[0] == "bedrock-runtime":
+                    self._instrument_bedrock_client(client)
+                return client
+
+            boto3.client = wrapped_client
+
+        except ImportError:
+            pass
+
+    def _instrument_bedrock_client(self, client):
+        if hasattr(client, "invoke_model"):
+            instrumented_invoke_method = self.create_span_wrapper(
+                span_name="aws.bedrock.invoke_model",
+                extract_attributes=self._extract_aws_bedrock_attributes,
+            )
+            client.invoke_model = instrumented_invoke_method
+
+    def _extract_aws_bedrock_attributes(
+        self, instance: Any, args: Any, kwargs: Any
+    ) -> Dict[str, Any]:  # pylint: disable=W0613
+        attrs = {}
+        model_id = kwargs.get("modelId", "unknown")
+
+        attrs["gen_ai.system"] = "aws_bedrock"
+        attrs["gen_ai.request.model"] = model_id
+        return attrs
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:  # pylint: disable=R1705
+        if hasattr(result, "get"):
+            content_type = result.get("contentType", "").lower()
+            body_str = result.get("body", "")
+
+            if "application/json" in content_type and body_str:
+                try:
+                    body = json.loads(body_str)
+                    if "usage" in body and isinstance(body["usage"], dict):
+                        usage = body["usage"]
+                        return {
+                            "prompt_tokens": getattr(usage, "inputTokens", 0),
+                            "completion_tokens": getattr(usage, "outputTokens", 0),
+                            "total_tokens": getattr(usage, "inputTokens", 0)
+                            + getattr(usage, "outputTokens", 0),
+                        }
+                    elif "usageMetadata" in body and isinstance(body["usageMetadata"], dict):
+                        usage = body["usageMetadata"]
+                        return {
+                            "prompt_tokens": getattr(usage, "promptTokenCount", 0),
+                            "completion_tokens": getattr(usage, "candidatesTokenCount", 0),
+                            "total_tokens": getattr(usage, "totalTokenCount", 0),
+                        }
+                except json.JSONDecodeError:
+                    logger.debug("Failed to parse Bedrock response body as JSON.")
+                except Exception as e:
+                    logger.debug("Error extracting usage from Bedrock response: %s", e)
+        return None

genai_otel/instrumentors/azure_openai_instrumentor.py

@@ -0,0 +1,69 @@
+"""OpenTelemetry instrumentor for Azure OpenAI SDK.
+
+This instrumentor automatically traces calls to Azure OpenAI models, capturing
+relevant attributes such as model name and token usage.
+"""
+
+import logging
+from typing import Dict, Optional
+
+from ..config import OTelConfig
+from .base import BaseInstrumentor
+
+logger = logging.getLogger(__name__)
+
+
+class AzureOpenAIInstrumentor(BaseInstrumentor):
+    """Instrumentor for Azure OpenAI"""
+
+    def __init__(self):
+        """Initialize the instrumentor."""
+        super().__init__()
+        self._azure_openai_available = False
+        self._check_availability()
+
+    def _check_availability(self):
+        """Check if Azure AI OpenAI library is available."""
+        try:
+            import azure.ai.openai  # Moved to top
+
+            self._azure_openai_available = True
+            logger.debug("Azure AI OpenAI library detected and available for instrumentation")
+        except ImportError:
+            logger.debug("Azure AI OpenAI library not installed, instrumentation will be skipped")
+            self._azure_openai_available = False
+
+    def instrument(self, config: OTelConfig):
+        self.config = config
+        try:
+            from azure.ai.openai import OpenAIClient
+
+            original_complete = OpenAIClient.complete
+
+            def wrapped_complete(instance, *args, **kwargs):
+                with self.tracer.start_as_current_span("azure.openai.complete") as span:
+                    model = kwargs.get("model", "unknown")
+
+                    span.set_attribute("gen_ai.system", "azure_openai")
+                    span.set_attribute("gen_ai.request.model", model)
+
+                    if self.request_counter:
+                        self.request_counter.add(1, {"model": model, "provider": "azure_openai"})
+
+                    result = original_complete(instance, *args, **kwargs)
+                    self._record_result_metrics(span, result, 0)
+                    return result
+
+            OpenAIClient.complete = wrapped_complete
+
+        except ImportError:
+            pass
+
+    def _extract_usage(self, result) -> Optional[Dict[str, int]]:
+        if hasattr(result, "usage") and result.usage:
+            return {
+                "prompt_tokens": result.usage.prompt_tokens,
+                "completion_tokens": result.usage.completion_tokens,
+                "total_tokens": result.usage.total_tokens,
+            }
+        return None