instrumented_engine

instrumented_engine

Transparent telemetry wrapper for inference engines.

Classes

InstrumentedEngine

InstrumentedEngine(engine: InferenceEngine, bus: EventBus, gpu_monitor: Optional[Any] = None, energy_monitor: Optional[Any] = None)

Bases: InferenceEngine

Transparent wrapper that records telemetry around engine calls.

Agents call engine.generate() normally -- they don't know about telemetry. The wrapper publishes INFERENCE_START, INFERENCE_END, and TELEMETRY_RECORD events on the bus.

If an energy_monitor is provided (new multi-vendor :class:~openjarvis.telemetry.energy_monitor.EnergyMonitor), it is preferred over the legacy gpu_monitor for energy measurement.

Source code in src/openjarvis/telemetry/instrumented_engine.py

def __init__(
    self,
    engine: InferenceEngine,
    bus: EventBus,
    gpu_monitor: Optional[Any] = None,
    energy_monitor: Optional[Any] = None,
) -> None:
    self._inner = engine
    self._bus = bus
    self._gpu_monitor = gpu_monitor
    self._energy_monitor = energy_monitor
    self._publishes_events = True

Functions

generate

generate(messages: Sequence[Message], *, model: str = '', temperature: float = 0.7, max_tokens: int = 1024, **kwargs: Any) -> Dict[str, Any]

Generate with telemetry recording.

Source code in src/openjarvis/telemetry/instrumented_engine.py

def generate(
    self,
    messages: Sequence[Message],
    *,
    model: str = "",
    temperature: float = 0.7,
    max_tokens: int = 1024,
    **kwargs: Any,
) -> Dict[str, Any]:
    """Generate with telemetry recording."""
    self._bus.publish(
        EventType.INFERENCE_START,
        {
            "model": model,
            "message_count": len(messages),
        },
    )

    gpu_sample: Optional[GpuSample] = None
    energy_sample: Optional[Any] = None
    t0 = time.time()

    # Prefer EnergyMonitor over legacy GpuMonitor
    if self._energy_monitor is not None:
        with self._energy_monitor.sample() as energy_sample:
            result = self._inner.generate(
                messages,
                model=model,
                temperature=temperature,
                max_tokens=max_tokens,
                **kwargs,
            )
    elif self._gpu_monitor is not None:
        with self._gpu_monitor.sample() as gpu_sample:
            result = self._inner.generate(
                messages,
                model=model,
                temperature=temperature,
                max_tokens=max_tokens,
                **kwargs,
            )
    else:
        result = self._inner.generate(
            messages,
            model=model,
            temperature=temperature,
            max_tokens=max_tokens,
            **kwargs,
        )

    latency = time.time() - t0

    usage = result.get("usage", {})
    completion_tokens = usage.get("completion_tokens", 0)
    prompt_tokens_evaluated = usage.get("prompt_tokens_evaluated", 0)
    ttft = result.get("ttft", 0.0)
    throughput = completion_tokens / latency if latency > 0 else 0.0

    # Energy / GPU metrics from sample
    energy_joules = 0.0
    power_watts = 0.0
    gpu_utilization_pct = 0.0
    gpu_memory_used_gb = 0.0
    gpu_temperature_c = 0.0
    prefill_latency = 0.0
    energy_method = ""
    energy_vendor = ""
    cpu_energy_joules = 0.0
    gpu_energy_joules = 0.0
    dram_energy_joules = 0.0
    ane_energy_joules = 0.0
    soc_energy_joules = 0.0
    energy_basis = ""

    if energy_sample is not None:
        # New multi-vendor EnergyMonitor path
        energy_joules = energy_sample.energy_joules
        power_watts = energy_sample.mean_power_watts
        gpu_utilization_pct = energy_sample.mean_utilization_pct
        gpu_memory_used_gb = energy_sample.peak_memory_used_gb
        gpu_temperature_c = energy_sample.mean_temperature_c
        energy_method = energy_sample.energy_method
        energy_vendor = energy_sample.vendor
        cpu_energy_joules = energy_sample.cpu_energy_joules
        gpu_energy_joules = energy_sample.gpu_energy_joules
        dram_energy_joules = energy_sample.dram_energy_joules
        ane_energy_joules = energy_sample.ane_energy_joules
        soc_energy_joules = energy_sample.soc_energy_joules
        energy_basis = energy_sample.basis
    elif gpu_sample is not None:
        # Legacy GpuMonitor path
        energy_joules = gpu_sample.energy_joules
        power_watts = gpu_sample.mean_power_watts
        gpu_utilization_pct = gpu_sample.mean_utilization_pct
        gpu_memory_used_gb = gpu_sample.peak_memory_used_gb
        gpu_temperature_c = gpu_sample.mean_temperature_c
        energy_method = "polling"
        energy_vendor = "nvidia"
        # The legacy monitor reads one GPU rail and nothing else.
        gpu_energy_joules = gpu_sample.energy_joules
        soc_energy_joules = gpu_sample.energy_joules
        energy_basis = BASIS_GPU

    if ttft > 0:
        prefill_latency = ttft

    # --- Tier 1: Derived metrics ---
    energy_per_output_token = (
        energy_joules / completion_tokens if completion_tokens > 0 else 0.0
    )
    throughput_per_watt = throughput / power_watts if power_watts > 0 else 0.0

    # --- Tier 2.1: Phase energy split ---
    decode_latency = latency - prefill_latency if prefill_latency > 0 else 0.0
    prefill_energy = 0.0
    decode_energy = 0.0
    if energy_joules > 0 and prefill_latency > 0 and latency > 0:
        prefill_frac = prefill_latency / latency
        prefill_energy = energy_joules * prefill_frac
        decode_energy = energy_joules * (1.0 - prefill_frac)

    # --- Tier 3: Non-streaming mean ITL approximation ---
    mean_itl_ms = (
        (decode_latency / completion_tokens) * 1000
        if completion_tokens > 0 and decode_latency > 0
        else 0.0
    )

    # --- Tier 4: Per-inference efficiency ---
    tokens_per_joule = (
        completion_tokens / energy_joules
        if energy_joules > 0 and completion_tokens > 0
        else 0.0
    )

    engine_id = getattr(self._inner, "engine_id", "unknown")

    prompt_tok = usage.get("prompt_tokens", 0)
    record = TelemetryRecord(
        timestamp=t0,
        model_id=model,
        prompt_tokens=prompt_tok,
        prompt_tokens_evaluated=prompt_tokens_evaluated or prompt_tok,
        completion_tokens=completion_tokens,
        total_tokens=prompt_tok + completion_tokens,
        latency_seconds=latency,
        cost_usd=result.get("cost_usd", 0.0),
        ttft=ttft,
        throughput_tok_per_sec=throughput,
        energy_per_output_token_joules=energy_per_output_token,
        throughput_per_watt=throughput_per_watt,
        energy_joules=energy_joules,
        power_watts=power_watts,
        gpu_utilization_pct=gpu_utilization_pct,
        gpu_memory_used_gb=gpu_memory_used_gb,
        gpu_temperature_c=gpu_temperature_c,
        prefill_latency_seconds=prefill_latency,
        decode_latency_seconds=decode_latency,
        prefill_energy_joules=prefill_energy,
        decode_energy_joules=decode_energy,
        mean_itl_ms=mean_itl_ms,
        engine=engine_id,
        energy_method=energy_method,
        energy_vendor=energy_vendor,
        cpu_energy_joules=cpu_energy_joules,
        gpu_energy_joules=gpu_energy_joules,
        dram_energy_joules=dram_energy_joules,
        ane_energy_joules=ane_energy_joules,
        soc_energy_joules=soc_energy_joules,
        energy_basis=energy_basis,
        tokens_per_joule=tokens_per_joule,
        # Stamp every new record with the current methodology version
        # so the leaderboard aggregator can drop legacy rows (those
        # left NULL by pre-fix builds) cleanly. Source of truth for
        # the integer is `server/savings.py`.
        token_counting_version=TOKEN_COUNTING_VERSION,
    )

    event_data = {
        "model": model,
        "latency": latency,
        "usage": usage,
        "ttft": ttft,
        "throughput_tok_per_sec": throughput,
        "energy_per_output_token_joules": energy_per_output_token,
        "throughput_per_watt": throughput_per_watt,
        "energy_joules": energy_joules,
        "power_watts": power_watts,
        "gpu_utilization_pct": gpu_utilization_pct,
        "gpu_memory_used_gb": gpu_memory_used_gb,
        "gpu_temperature_c": gpu_temperature_c,
        "prefill_latency_seconds": prefill_latency,
        "decode_latency_seconds": decode_latency,
        "prefill_energy_joules": prefill_energy,
        "decode_energy_joules": decode_energy,
        "mean_itl_ms": mean_itl_ms,
        "energy_method": energy_method,
        "energy_vendor": energy_vendor,
        # Rich trace data: model response content + structured blocks
        "content": result.get("content", ""),
        "tool_calls": result.get("tool_calls", []),
        "tool_results": result.get("tool_results", []),
        "content_blocks": result.get("content_blocks", []),
        "finish_reason": result.get("finish_reason", ""),
    }

    self._bus.publish(EventType.INFERENCE_END, event_data)
    self._bus.publish(EventType.TELEMETRY_RECORD, {"record": record})

    # Inject telemetry dict into result for downstream consumers (eval backend)
    result["_telemetry"] = {
        "latency": latency,
        "ttft": ttft,
        "throughput_tok_per_sec": throughput,
        "energy_per_output_token_joules": energy_per_output_token,
        "throughput_per_watt": throughput_per_watt,
        "energy_joules": energy_joules,
        "power_watts": power_watts,
        "gpu_utilization_pct": gpu_utilization_pct,
        "gpu_memory_used_gb": gpu_memory_used_gb,
        "gpu_temperature_c": gpu_temperature_c,
        "prefill_latency_seconds": prefill_latency,
        "decode_latency_seconds": decode_latency,
        "prefill_energy_joules": prefill_energy,
        "decode_energy_joules": decode_energy,
        "mean_itl_ms": mean_itl_ms,
        "energy_method": energy_method,
        "energy_vendor": energy_vendor,
        "cpu_energy_joules": cpu_energy_joules,
        "gpu_energy_joules": gpu_energy_joules,
        "dram_energy_joules": dram_energy_joules,
        "ane_energy_joules": ane_energy_joules,
        "soc_energy_joules": soc_energy_joules,
        "energy_basis": energy_basis,
    }

    return result

stream async

stream(messages: Sequence[Message], *, model: str = '', temperature: float = 0.7, max_tokens: int = 1024, **kwargs: Any) -> Any

Stream with per-token timing and full telemetry recording.

Source code in src/openjarvis/telemetry/instrumented_engine.py

async def stream(
    self,
    messages: Sequence[Message],
    *,
    model: str = "",
    temperature: float = 0.7,
    max_tokens: int = 1024,
    **kwargs: Any,
) -> Any:
    """Stream with per-token timing and full telemetry recording."""
    self._bus.publish(
        EventType.INFERENCE_START,
        {
            "model": model,
            "message_count": len(messages),
        },
    )

    t0 = time.time()
    token_timestamps: list[float] = []
    token_count = 0

    energy_sample: Optional[Any] = None
    gpu_sample: Optional[GpuSample] = None

    if self._energy_monitor is not None:
        with self._energy_monitor.sample() as energy_sample:
            async for token in self._inner.stream(
                messages,
                model=model,
                temperature=temperature,
                max_tokens=max_tokens,
                **kwargs,
            ):
                token_timestamps.append(time.time())
                token_count += 1
                yield token
    elif self._gpu_monitor is not None:
        with self._gpu_monitor.sample() as gpu_sample:
            async for token in self._inner.stream(
                messages,
                model=model,
                temperature=temperature,
                max_tokens=max_tokens,
                **kwargs,
            ):
                token_timestamps.append(time.time())
                token_count += 1
                yield token
    else:
        async for token in self._inner.stream(
            messages,
            model=model,
            temperature=temperature,
            max_tokens=max_tokens,
            **kwargs,
        ):
            token_timestamps.append(time.time())
            token_count += 1
            yield token

    self._publish_stream_telemetry(
        model=model,
        t0=t0,
        token_timestamps=token_timestamps,
        fallback_token_count=token_count,
        energy_sample=energy_sample,
        gpu_sample=gpu_sample,
    )

stream_full async

stream_full(messages: Sequence[Message], *, model: str, temperature: float = 0.7, max_tokens: int = 1024, **kwargs: Any) -> AsyncIterator['StreamChunk']

Stream rich chunks with timing, usage, and energy telemetry.

Telemetry is published only after the inner stream completes successfully.

Source code in src/openjarvis/telemetry/instrumented_engine.py

async def stream_full(
    self,
    messages: Sequence[Message],
    *,
    model: str,
    temperature: float = 0.7,
    max_tokens: int = 1024,
    **kwargs: Any,
) -> AsyncIterator["StreamChunk"]:
    """Stream rich chunks with timing, usage, and energy telemetry.

    Telemetry is published only after the inner stream completes
    successfully.
    """
    self._bus.publish(
        EventType.INFERENCE_START,
        {
            "model": model,
            "message_count": len(messages),
        },
    )

    t0 = time.time()
    token_timestamps: list[float] = []
    content_chunk_count = 0
    usage: Optional[Dict[str, Any]] = None
    energy_sample: Optional[Any] = None
    gpu_sample: Optional[GpuSample] = None

    if self._energy_monitor is not None:
        with self._energy_monitor.sample() as energy_sample:
            async for chunk in self._inner.stream_full(
                messages,
                model=model,
                temperature=temperature,
                max_tokens=max_tokens,
                **kwargs,
            ):
                if chunk.content:
                    token_timestamps.append(time.time())
                    content_chunk_count += 1
                if chunk.usage is not None:
                    usage = chunk.usage
                yield chunk
    elif self._gpu_monitor is not None:
        with self._gpu_monitor.sample() as gpu_sample:
            async for chunk in self._inner.stream_full(
                messages,
                model=model,
                temperature=temperature,
                max_tokens=max_tokens,
                **kwargs,
            ):
                if chunk.content:
                    token_timestamps.append(time.time())
                    content_chunk_count += 1
                if chunk.usage is not None:
                    usage = chunk.usage
                yield chunk
    else:
        async for chunk in self._inner.stream_full(
            messages,
            model=model,
            temperature=temperature,
            max_tokens=max_tokens,
            **kwargs,
        ):
            if chunk.content:
                token_timestamps.append(time.time())
                content_chunk_count += 1
            if chunk.usage is not None:
                usage = chunk.usage
            yield chunk

    self._publish_stream_telemetry(
        model=model,
        t0=t0,
        token_timestamps=token_timestamps,
        fallback_token_count=content_chunk_count,
        energy_sample=energy_sample,
        gpu_sample=gpu_sample,
        usage=usage,
    )