modaltrace 0.1.0__py3-none-any.whl

modaltrace/__init__.py

@@ -0,0 +1,244 @@
+"""modaltrace: OpenTelemetry observability for real-time AI avatar and video pipelines.
+
+Everything a user needs is importable from `modaltrace` directly:
+
+    import modaltrace
+
+    sdk = modaltrace.init(service_name="artalk-avatar")
+
+    @modaltrace.pipeline_stage("flame_inference")
+    async def run_model(...): ...
+
+    async with modaltrace.stage("render") as s:
+        s.record("vertex_count", 12345)
+
+    modaltrace.info("Pipeline started", target_fps=30)
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from modaltrace._version import __version__
+from modaltrace.logging.api import (
+    debug,
+    error,
+    exception,
+    info,
+    notice,
+    warning,
+)
+from modaltrace.logging.api import (
+    trace_log as trace,
+)
+from modaltrace.tracing.pipeline import async_stage, pipeline_stage, stage
+
+__all__ = [
+    "__version__",
+    "init",
+    "pipeline_stage",
+    "stage",
+    "async_stage",
+    "trace",
+    "debug",
+    "info",
+    "notice",
+    "warning",
+    "error",
+    "exception",
+]
+
+
+class ModalTraceSDK:
+    """SDK handle returned by init(). Context-manager compatible."""
+
+    def __init__(
+        self,
+        config,
+        tracer_provider,
+        meter_provider,
+        logger_provider,
+        frame_aggregator,
+        av_tracker,
+        gpu_monitor=None,
+        pending_processor=None,
+    ):
+        self._config = config
+        self._tracer_provider = tracer_provider
+        self._meter_provider = meter_provider
+        self._logger_provider = logger_provider
+        self.frame_aggregator = frame_aggregator
+        self.av_tracker = av_tracker
+        self._gpu_monitor = gpu_monitor
+        self._pending_processor = pending_processor
+        self._stopped = False
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc):
+        self.stop()
+
+    def flush(self) -> None:
+        """Force-flush all exporters."""
+        self._tracer_provider.force_flush()
+        self._meter_provider.force_flush()
+
+    def stop(self) -> None:
+        """Flush and shut down all components."""
+        if self._stopped:
+            return
+        self._stopped = True
+
+        self.frame_aggregator.stop()
+
+        if self._pending_processor is not None:
+            self._pending_processor.stop()
+
+        if self._gpu_monitor is not None:
+            self._gpu_monitor.stop()
+
+        from modaltrace.instrumentation.eventloop import uninstall_eventloop_monitor
+        from modaltrace.instrumentation.pytorch import uninstrument_pytorch
+        from modaltrace.tracing.propagation import unpatch_all
+
+        uninstrument_pytorch()
+        unpatch_all()
+        uninstall_eventloop_monitor()
+
+        self._tracer_provider.shutdown()
+        self._meter_provider.shutdown()
+        self._logger_provider.shutdown()
+
+
+def init(**kwargs: Any) -> ModalTraceSDK:
+    """Initialize modaltrace. One-liner to get full observability.
+
+    All kwargs are passed to ModalTraceConfig (Pydantic Settings), which also
+    reads from MODALTRACE_* environment variables and .env files.
+
+    Returns an ModalTraceSDK handle with .frame_aggregator, .av_tracker,
+    .flush(), and .stop() methods. Also works as a context manager.
+    """
+    from modaltrace import _registry
+    from modaltrace.config import ModalTraceConfig
+    from modaltrace.exporters.setup import (
+        create_resource,
+        setup_logger_provider,
+        setup_meter_provider,
+        setup_tracer_provider,
+    )
+    from modaltrace.instrumentation.eventloop import install_eventloop_monitor
+    from modaltrace.instrumentation.gpu import GPUMonitor
+    from modaltrace.instrumentation.pytorch import instrument_pytorch
+    from modaltrace.logging.api import _init_logging
+    from modaltrace.logging.scrubber import ScrubbingSpanProcessor
+    from modaltrace.metrics.aggregator import FrameMetricsAggregator
+    from modaltrace.metrics.av_sync import AVSyncTracker
+    from modaltrace.metrics.instruments import MetricInstruments
+    from modaltrace.tracing.pending import PendingSpanProcessor
+    from modaltrace.tracing.propagation import patch_all
+    from modaltrace.tracing.sampler import AdaptiveSampler
+
+    config = ModalTraceConfig(**kwargs)
+    _registry._config = config
+
+    resource = create_resource(config)
+    tracer_provider = setup_tracer_provider(config, resource)
+    meter_provider = setup_meter_provider(config, resource)
+    logger_provider = setup_logger_provider(config, resource)
+
+    tracer = tracer_provider.get_tracer("modaltrace", __version__)
+    meter = meter_provider.get_meter("modaltrace", __version__)
+    _registry._tracer = tracer
+    _registry._meter = meter
+
+    instruments = MetricInstruments(meter)
+
+    if config.scrubbing_enabled:
+        scrubber = ScrubbingSpanProcessor(
+            extra_patterns=config.scrubbing_patterns,
+            callback=config.scrubbing_callback,
+        )
+        tracer_provider.add_span_processor(scrubber)
+
+    from modaltrace.exporters.setup import _create_span_exporter
+
+    pending_exporter = _create_span_exporter(config)
+    pending_processor = PendingSpanProcessor(
+        exporter=pending_exporter,
+        flush_interval_ms=config.pending_span_flush_interval_ms,
+    )
+    tracer_provider.add_span_processor(pending_processor)
+    pending_processor.start()
+
+    aggregator = FrameMetricsAggregator(
+        instruments=instruments,
+        buffer_size=config.ring_buffer_size,
+        flush_interval_ms=config.metrics_flush_interval_ms,
+    )
+    aggregator.start()
+
+    av_tracker = AVSyncTracker(
+        instruments=instruments,
+        drift_warning_ms=config.av_drift_warning_ms,
+        chunk_ttl_s=config.av_chunk_ttl_s,
+        jitter_window=config.av_jitter_window,
+        warning_callback=warning,
+    )
+
+    sampler = AdaptiveSampler(
+        window_s=config.span_window_s,
+        anomaly_threshold_ms=config.anomaly_threshold_ms,
+    )
+    _registry._sampler = sampler
+
+    gpu_monitor = None
+    if config.gpu_monitoring:
+        gpu_monitor = GPUMonitor(
+            poll_interval_s=config.gpu_poll_interval_s,
+            device_indices=config.gpu_device_indices,
+        )
+        if gpu_monitor.start():
+            gpu_monitor.register_gauges(meter)
+        else:
+            gpu_monitor = None
+
+    if config.pytorch_instrumentation:
+        instrument_pytorch(
+            tracer=tracer,
+            sample_rate=config.pytorch_sample_rate,
+            anomaly_threshold_ms=config.anomaly_threshold_ms,
+            track_memory=config.pytorch_track_memory,
+            track_shapes=config.pytorch_track_shapes,
+            aggregator=aggregator,
+        )
+
+    if config.threadpool_propagation:
+        patch_all()
+
+    if config.eventloop_monitoring:
+        install_eventloop_monitor(
+            threshold_ms=config.eventloop_lag_threshold_ms,
+            warning_callback=warning,
+        )
+
+    _init_logging(
+        logger_provider=logger_provider,
+        log_level=config.log_level,
+        log_console=config.log_console,
+    )
+
+    sdk = ModalTraceSDK(
+        config=config,
+        tracer_provider=tracer_provider,
+        meter_provider=meter_provider,
+        logger_provider=logger_provider,
+        frame_aggregator=aggregator,
+        av_tracker=av_tracker,
+        gpu_monitor=gpu_monitor,
+        pending_processor=pending_processor,
+    )
+    _registry._sdk = sdk
+
+    return sdk

modaltrace/_registry.py

@@ -0,0 +1,25 @@
+"""Module-level singletons for the modaltrace SDK.
+
+Holds references to the active config, providers, and components
+so that module-level functions (pipeline_stage, info, etc.) can
+access them without requiring the user to pass the SDK instance.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from modaltrace.config import ModalTraceConfig
+
+_config: ModalTraceConfig | None = None
+_tracer = None
+_meter = None
+_logger_provider = None
+_sdk = None
+
+
+def get_config() -> ModalTraceConfig:
+    if _config is None:
+        raise RuntimeError("modaltrace.init() has not been called yet")
+    return _config

modaltrace/_version.py

@@ -0,0 +1,2 @@
+# This file is managed by hatch-vcs. Fallback for editable installs.
+__version__ = "0.0.0"

modaltrace/config.py

@@ -0,0 +1,92 @@
+"""Pydantic Settings model — single source of truth for all configuration.
+
+Users can configure via Python kwargs, environment variables (MODALTRACE_*), or .env files.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Literal
+
+from pydantic import AnyHttpUrl, Field, field_validator
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+class ModalTraceConfig(BaseSettings):
+    """Configuration for modaltrace SDK."""
+
+    # ── Identity ──────────────────────────────────────────────────────────
+    service_name: str = "modaltrace-pipeline"
+    service_version: str = "0.0.0"
+    deployment_environment: str = "development"
+
+    # ── OTLP Export ───────────────────────────────────────────────────────
+    otlp_endpoint: AnyHttpUrl = "http://localhost:4318"  # type: ignore[assignment]
+    otlp_protocol: Literal["http", "grpc"] = "http"
+    otlp_headers: dict[str, str] = Field(default_factory=dict)
+    otlp_timeout_ms: int = 10_000
+
+    # ── Feature Flags ─────────────────────────────────────────────────────
+    pytorch_instrumentation: bool = True
+    gpu_monitoring: bool = True
+    webrtc_monitoring: bool = False
+    eventloop_monitoring: bool = True
+    threadpool_propagation: bool = True
+
+    # ── Frame Metrics Aggregator ──────────────────────────────────────────
+    metrics_flush_interval_ms: int = 1_000
+    ring_buffer_size: int = 512
+
+    # ── Adaptive Sampler ──────────────────────────────────────────────────
+    span_window_s: float = 1.0
+    anomaly_threshold_ms: float = 50.0
+    pytorch_sample_rate: float = 0.01
+
+    # ── Pending Spans ─────────────────────────────────────────────────────
+    pending_span_flush_interval_ms: int = 5_000
+
+    # ── A/V Sync ──────────────────────────────────────────────────────────
+    av_drift_warning_ms: float = 40.0
+    av_chunk_ttl_s: float = 5.0
+    av_jitter_window: int = 30
+
+    # ── GPU Monitor ───────────────────────────────────────────────────────
+    gpu_poll_interval_s: float = 1.0
+    gpu_device_indices: list[int] | None = None
+
+    # ── PyTorch Instrumentation ───────────────────────────────────────────
+    pytorch_track_memory: bool = True
+    pytorch_track_shapes: bool = False
+
+    # ── PII Scrubbing ─────────────────────────────────────────────────────
+    scrubbing_enabled: bool = True
+    scrubbing_patterns: list[str] = Field(default_factory=list)
+    scrubbing_callback: Callable | None = None
+
+    # ── Structured Logging ────────────────────────────────────────────────
+    log_level: str = "info"
+    log_console: bool = True
+
+    # ── Event Loop Monitor ────────────────────────────────────────────────
+    eventloop_lag_threshold_ms: float = 100.0
+
+    model_config = SettingsConfigDict(
+        env_prefix="MODALTRACE_",
+        env_file=".env",
+        env_file_encoding="utf-8",
+        arbitrary_types_allowed=True,
+    )
+
+    @field_validator("ring_buffer_size")
+    @classmethod
+    def must_be_power_of_two(cls, v: int) -> int:
+        if v & (v - 1) != 0:
+            raise ValueError(f"ring_buffer_size must be a power of 2, got {v}")
+        return v
+
+    @field_validator("pytorch_sample_rate")
+    @classmethod
+    def must_be_fraction(cls, v: float) -> float:
+        if not 0.0 <= v <= 1.0:
+            raise ValueError(f"pytorch_sample_rate must be between 0 and 1, got {v}")
+        return v

modaltrace/conventions/__init__.py

@@ -0,0 +1,19 @@
+from modaltrace.conventions.attributes import (
+    AVSyncAttributes,
+    EventLoopAttributes,
+    GPUAttributes,
+    InferenceAttributes,
+    ModalAttributes,
+    PipelineAttributes,
+    TransportAttributes,
+)
+
+__all__ = [
+    "PipelineAttributes",
+    "InferenceAttributes",
+    "ModalAttributes",
+    "AVSyncAttributes",
+    "GPUAttributes",
+    "TransportAttributes",
+    "EventLoopAttributes",
+]

modaltrace/conventions/attributes.py

@@ -0,0 +1,66 @@
+"""Semantic convention string constants for modaltrace.
+
+All attribute keys live here — zero magic strings in the rest of the codebase.
+"""
+
+
+class PipelineAttributes:
+    ID = "modaltrace.pipeline.id"
+    SESSION_ID = "modaltrace.pipeline.session_id"
+    STAGE_NAME = "modaltrace.pipeline.stage.name"
+    STAGE_DURATION_MS = "modaltrace.pipeline.stage.duration_ms"
+    FRAME_SEQ = "modaltrace.pipeline.frame.sequence_number"
+    TARGET_FPS = "modaltrace.pipeline.target_fps"
+    SPAN_PENDING = "modaltrace.span.pending"
+
+
+class InferenceAttributes:
+    MODEL_NAME = "modaltrace.inference.model_name"
+    FORWARD_PASS_MS = "modaltrace.inference.forward_pass_ms"
+    BATCH_SIZE = "modaltrace.inference.batch_size"
+    GPU_MEMORY_MB = "modaltrace.inference.gpu.memory_allocated_mb"
+    GPU_MEMORY_DELTA_MB = "modaltrace.inference.gpu.memory_delta_mb"
+    INPUT_SHAPES = "modaltrace.inference.input_shapes"
+    DEVICE = "modaltrace.inference.device"
+
+
+class ModalAttributes:
+    FLAME_INFERENCE_MS = "modaltrace.flame.inference_ms"
+    FLAME_PARAM_COUNT = "modaltrace.flame.parameter_count"
+    RENDER_FRAME_MS = "modaltrace.render.frame_ms"
+    RENDER_RESOLUTION = "modaltrace.render.resolution"
+    MESH_VERTEX_COUNT = "modaltrace.mesh.vertex_count"
+    FRAME_SEQ = "modaltrace.frame.sequence_number"
+
+
+class AVSyncAttributes:
+    DRIFT_MS = "modaltrace.av_sync.drift_ms"
+    JITTER_MS = "modaltrace.av_sync.jitter_ms"
+    THRESHOLD_MS = "modaltrace.av_sync.threshold_ms"
+    UNMATCHED_CHUNKS = "modaltrace.av_sync.unmatched_chunks"
+    CHUNK_ID = "modaltrace.av_sync.chunk_id"
+
+
+class GPUAttributes:
+    DEVICE_INDEX = "modaltrace.gpu.device_index"
+    UTILIZATION_PCT = "modaltrace.gpu.utilization"
+    MEMORY_USED_MB = "modaltrace.gpu.memory.used"
+    MEMORY_FREE_MB = "modaltrace.gpu.memory.free"
+    TEMPERATURE_C = "modaltrace.gpu.temperature"
+    POWER_W = "modaltrace.gpu.power.draw"
+
+
+class TransportAttributes:
+    PROTOCOL = "modaltrace.transport.protocol"
+    RTT_MS = "modaltrace.transport.rtt_ms"
+    JITTER_MS = "modaltrace.transport.jitter_ms"
+    PACKET_LOSS_PCT = "modaltrace.transport.packet_loss_percent"
+    BITRATE_KBPS = "modaltrace.transport.bitrate_kbps"
+    FRAME_RATE_ACTUAL = "modaltrace.transport.frame_rate_actual"
+    STREAM = "modaltrace.transport.stream"
+
+
+class EventLoopAttributes:
+    ELAPSED_MS = "modaltrace.eventloop.blocked_ms"
+    THRESHOLD_MS = "modaltrace.eventloop.threshold_ms"
+    HANDLE_CALLBACK = "modaltrace.eventloop.handle_callback"

modaltrace/exporters/setup.py

@@ -0,0 +1,126 @@
+"""OTLP exporter and provider setup.
+
+Configures TracerProvider, MeterProvider, and LoggerProvider with OTLP
+HTTP or gRPC exporters based on ModalTraceConfig.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from opentelemetry import trace
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+
+if TYPE_CHECKING:
+    from modaltrace.config import ModalTraceConfig
+
+
+def create_resource(config: ModalTraceConfig) -> Resource:
+    return Resource.create(
+        {
+            "service.name": config.service_name,
+            "service.version": config.service_version,
+            "deployment.environment": config.deployment_environment,
+        }
+    )
+
+
+def setup_tracer_provider(config: ModalTraceConfig, resource: Resource) -> TracerProvider:
+    provider = TracerProvider(resource=resource)
+    exporter = _create_span_exporter(config)
+    provider.add_span_processor(BatchSpanProcessor(exporter))
+    trace.set_tracer_provider(provider)
+    return provider
+
+
+def setup_meter_provider(config: ModalTraceConfig, resource: Resource):
+    from opentelemetry import metrics
+    from opentelemetry.sdk.metrics import MeterProvider
+    from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
+
+    exporter = _create_metric_exporter(config)
+    reader = PeriodicExportingMetricReader(
+        exporter,
+        export_interval_millis=config.metrics_flush_interval_ms,
+    )
+    provider = MeterProvider(resource=resource, metric_readers=[reader])
+    metrics.set_meter_provider(provider)
+    return provider
+
+
+def setup_logger_provider(config: ModalTraceConfig, resource: Resource):
+    from opentelemetry.sdk._logs import LoggerProvider
+    from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
+
+    exporter = _create_log_exporter(config)
+    provider = LoggerProvider(resource=resource)
+    provider.add_log_record_processor(BatchLogRecordProcessor(exporter))
+    return provider
+
+
+def _create_span_exporter(config: ModalTraceConfig):
+    endpoint = str(config.otlp_endpoint)
+    headers = config.otlp_headers
+    timeout = config.otlp_timeout_ms // 1000
+
+    if config.otlp_protocol == "grpc":
+        from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
+
+        return OTLPSpanExporter(
+            endpoint=endpoint,
+            headers=tuple(headers.items()) if headers else None,
+            timeout=timeout,
+        )
+    from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+
+    return OTLPSpanExporter(
+        endpoint=f"{endpoint}/v1/traces",
+        headers=headers,
+        timeout=timeout,
+    )
+
+
+def _create_metric_exporter(config: ModalTraceConfig):
+    endpoint = str(config.otlp_endpoint)
+    headers = config.otlp_headers
+    timeout = config.otlp_timeout_ms // 1000
+
+    if config.otlp_protocol == "grpc":
+        from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import OTLPMetricExporter
+
+        return OTLPMetricExporter(
+            endpoint=endpoint,
+            headers=tuple(headers.items()) if headers else None,
+            timeout=timeout,
+        )
+    from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
+
+    return OTLPMetricExporter(
+        endpoint=f"{endpoint}/v1/metrics",
+        headers=headers,
+        timeout=timeout,
+    )
+
+
+def _create_log_exporter(config: ModalTraceConfig):
+    endpoint = str(config.otlp_endpoint)
+    headers = config.otlp_headers
+    timeout = config.otlp_timeout_ms // 1000
+
+    if config.otlp_protocol == "grpc":
+        from opentelemetry.exporter.otlp.proto.grpc._log_exporter import OTLPLogExporter
+
+        return OTLPLogExporter(
+            endpoint=endpoint,
+            headers=tuple(headers.items()) if headers else None,
+            timeout=timeout,
+        )
+    from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter
+
+    return OTLPLogExporter(
+        endpoint=f"{endpoint}/v1/logs",
+        headers=headers,
+        timeout=timeout,
+    )

modaltrace/instrumentation/eventloop.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+
+logger = logging.getLogger("modaltrace.eventloop")
+
+_original_handle_run = None
+_warning_callback = None
+
+
+def install_eventloop_monitor(
+    threshold_ms: float = 100.0,
+    warning_callback=None,
+):
+    global _original_handle_run, _warning_callback
+    _original_handle_run = asyncio.events.Handle._run
+    _warning_callback = warning_callback
+
+    def patched_run(self):
+        start = time.perf_counter()
+        try:
+            _original_handle_run(self)
+        finally:
+            elapsed_ms = (time.perf_counter() - start) * 1000
+            if elapsed_ms > threshold_ms:
+                cb_name = getattr(self._callback, "__qualname__", str(self._callback))
+                if _warning_callback:
+                    _warning_callback(
+                        f"Event loop blocked for {elapsed_ms:.1f}ms",
+                        elapsed_ms=elapsed_ms,
+                        threshold_ms=threshold_ms,
+                        handle_callback=cb_name,
+                    )
+                else:
+                    logger.warning("Event loop blocked for %.1fms by %s", elapsed_ms, cb_name)
+
+    asyncio.events.Handle._run = patched_run
+
+
+def uninstall_eventloop_monitor():
+    global _original_handle_run, _warning_callback
+    if _original_handle_run is not None:
+        asyncio.events.Handle._run = _original_handle_run
+        _original_handle_run = None
+        _warning_callback = None