modelexpress-client 0.7.0

# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Structured, normalized timing for ModelExpress refit cycles.

The recorder is deliberately independent of torch and vLLM.  A cycle owner
activates it with :func:`use_refit_timing`; lower layers then contribute spans
without adding timing parameters to their public APIs.
"""

from __future__ import annotations

import contextlib
import contextvars
import json
import logging
import os
import sys
import time
from collections.abc import Callable, Iterator
from dataclasses import dataclass, field
from typing import Any

MX_REFIT_TIMING_PREFIX = "MX_REFIT_TIMING"
REFIT_TIMING_STAGES = (
    "control_discovery",
    "source_preparation",
    "setup_registration",
    "transfer_planning",
    "wire_transfer",
    "receive_sync",
    "transformation",
    "installation",
    "post_install",
    "rollout_readiness",
)
_STAGE_SET = frozenset(REFIT_TIMING_STAGES)
_current_recorder: contextvars.ContextVar[RefitTimingRecorder | None] = (
    contextvars.ContextVar("mx_refit_timing_recorder", default=None)
)


@dataclass
class _Stage:
    duration_s: float = 0.0
    count: int = 0
    statuses: list[str] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)


class RefitTimingRecorder:
    """Collect repeated normalized spans and emit one stable JSON record."""

    def __init__(
        self,
        *,
        backend: str,
        version: int | str,
        rank: int | None = None,
        tp_rank: int | None = None,
        tp_size: int | None = None,
        ep_rank: int | None = None,
        ep_size: int | None = None,
        cold: bool | None = None,
        clock: Callable[[], float] = time.perf_counter,
    ) -> None:
        self.backend = backend
        self.version = version
        self.rank = rank
        self.tp_rank = tp_rank
        self.tp_size = tp_size
        self.ep_rank = ep_rank
        self.ep_size = ep_size
        self.cold = cold
        self.bytes = 0
        self._clock = clock
        self._started_at = clock()
        self._finished_at: float | None = None
        self._stages = {name: _Stage() for name in REFIT_TIMING_STAGES}
        self._emitted_payload: dict[str, Any] | None = None

    @contextlib.contextmanager
    def span(
        self,
        stage: str,
        *,
        status: str = "ok",
        metadata: dict[str, Any] | None = None,
        accumulate_metadata: bool = False,
        duration_key: str | None = None,
    ) -> Iterator[dict[str, Any]]:
        """Measure a stage; failed spans are retained and re-raised.

        Yields a dict for counters only knowable once the work has run -- a
        shard count, a cache outcome, a byte total. Without it a caller with
        something to report on exit has to time the region by hand instead,
        which is how a span turns into a pair of clock reads.

        ``duration_key`` also files the measured duration under that metadata
        name. Several spans usually share one normalized stage, and the stage
        total cannot say which of them was expensive; naming them is what
        separates a slow fetch from a slow hash inside the same stage.
        """
        self._validate_stage(stage)
        discovered: dict[str, Any] = {}
        started = self._clock()

        def record(final_status: str) -> None:
            elapsed = self._clock() - started
            extra = {**(metadata or {}), **discovered}
            if duration_key is not None:
                extra[duration_key] = elapsed
            self.add_duration(
                stage,
                elapsed,
                status=final_status,
                metadata=extra,
                accumulate_metadata=accumulate_metadata,
            )

        try:
            yield discovered
        except BaseException:
            record("error")
            raise
        else:
            record(status)

    def add_duration(
        self,
        stage: str,
        duration_s: float,
        *,
        status: str = "ok",
        metadata: dict[str, Any] | None = None,
        accumulate_metadata: bool = False,
    ) -> None:
        """Add an externally measured duration to a normalized stage."""
        self._validate_stage(stage)
        if duration_s < 0:
            raise ValueError("duration_s must be non-negative")
        item = self._stages[stage]
        item.duration_s += float(duration_s)
        item.count += 1
        if status not in item.statuses:
            item.statuses.append(status)
        if metadata:
            for name, value in metadata.items():
                if (
                    accumulate_metadata
                    and isinstance(value, (int, float))
                    and isinstance(item.metadata.get(name, 0), (int, float))
                ):
                    item.metadata[name] = item.metadata.get(name, 0) + value
                else:
                    item.metadata[name] = value

    def add_metadata(
        self,
        stage: str,
        metadata: dict[str, Any],
        *,
        accumulate: bool = False,
    ) -> None:
        """Attach facts to a stage without claiming to have measured anything.

        For sizes, counts and cache outcomes that describe a stage but have no
        duration of their own. Passing them through :meth:`add_duration` with a
        zero would work and would also add a span, leaving the stage looking
        like it ran more times than it did.
        """
        self._validate_stage(stage)
        item = self._stages[stage]
        for name, value in metadata.items():
            if (
                accumulate
                and isinstance(value, (int, float))
                and isinstance(item.metadata.get(name, 0), (int, float))
            ):
                item.metadata[name] = item.metadata.get(name, 0) + value
            else:
                item.metadata[name] = value

    def mark_not_applicable(
        self,
        stage: str,
        *,
        combined_with: str | None = None,
        reason: str | None = None,
    ) -> None:
        """Explicitly mark a stage absent from this backend's live path."""
        self._validate_stage(stage)
        metadata = {}
        if combined_with is not None:
            self._validate_stage(combined_with)
            metadata["combined_with"] = combined_with
        if reason is not None:
            metadata["reason"] = reason
        item = self._stages[stage]
        if item.count:
            raise ValueError(f"stage {stage!r} already has measured spans")
        item.statuses = ["combined" if combined_with else "not_applicable"]
        item.metadata.update(metadata)

    def add_bytes(self, count: int) -> None:
        if count < 0:
            raise ValueError("byte count must be non-negative")
        self.bytes += int(count)

    def set_cold(self, cold: bool) -> None:
        self.cold = bool(cold)

    def has_measurements(self, stage: str) -> bool:
        """Return whether ``stage`` contains at least one measured span."""
        self._validate_stage(stage)
        return self._stages[stage].count > 0

    def finish(self) -> None:
        if self._finished_at is None:
            self._finished_at = self._clock()

    def as_dict(self) -> dict[str, Any]:
        """Return a JSON-serializable record in normalized stage order."""
        end = self._finished_at if self._finished_at is not None else self._clock()
        e2e_s = max(0.0, end - self._started_at)
        measured_s = sum(item.duration_s for item in self._stages.values())
        stages: dict[str, dict[str, Any]] = {}
        for name in REFIT_TIMING_STAGES:
            item = self._stages[name]
            if not item.statuses:
                status = "not_recorded"
            elif len(item.statuses) == 1:
                status = item.statuses[0]
            else:
                status = "mixed"
            value: dict[str, Any] = {
                "status": status,
                "duration_ms": round(item.duration_s * 1000.0, 3),
                "count": item.count,
            }
            if len(item.statuses) > 1:
                value["statuses"] = list(item.statuses)
            if item.metadata:
                value["metadata"] = dict(item.metadata)
            stages[name] = value
        return {
            "backend": self.backend,
            "version": self.version,
            "rank": self.rank,
            "tp_rank": self.tp_rank,
            "tp_size": self.tp_size,
            "ep_rank": self.ep_rank,
            "ep_size": self.ep_size,
            "bytes": self.bytes,
            "cold_warm": (
                "unknown" if self.cold is None else ("cold" if self.cold else "warm")
            ),
            "stages": stages,
            "e2e_ms": round(e2e_s * 1000.0, 3),
            "unattributed_ms": round(max(0.0, e2e_s - measured_s) * 1000.0, 3),
        }

    def emit(self, logger: logging.Logger) -> dict[str, Any]:
        """Log the completed record once and return its dictionary form.

        The record is logged at INFO. Because some deployments (e.g. the Dynamo
        rollout) filter INFO from worker log capture, the marker line is also
        written to stdout with an explicit flush when ``MX_REFIT_TIMING_STDOUT``
        is set to a non-``0`` value, so benchmark harnesses can reliably grep the
        per-stage record regardless of logger configuration.
        """
        if self._emitted_payload is not None:
            return self._emitted_payload
        self.finish()
        payload = self.as_dict()
        encoded = json.dumps(
            payload,
            separators=(",", ":"),
            sort_keys=False,
            default=str,
        )
        line = f"{MX_REFIT_TIMING_PREFIX} {encoded}"
        logger.info("%s", line)
        if os.environ.get("MX_REFIT_TIMING_STDOUT", "0") != "0":
            print(line, flush=True, file=sys.stdout)
        self._emitted_payload = payload
        return self._emitted_payload

    @staticmethod
    def _validate_stage(stage: str) -> None:
        if stage not in _STAGE_SET:
            raise ValueError(
                f"unknown refit timing stage {stage!r}; expected one of "
                f"{REFIT_TIMING_STAGES}"
            )


def current_refit_timing() -> RefitTimingRecorder | None:
    """Return the recorder active in this context, if any."""
    return _current_recorder.get()


@contextlib.contextmanager
def use_refit_timing(recorder: RefitTimingRecorder) -> Iterator[RefitTimingRecorder]:
    """Make ``recorder`` visible to nested ModelExpress layers."""
    token = _current_recorder.set(recorder)
    try:
        yield recorder
    finally:
        _current_recorder.reset(token)


@contextlib.contextmanager
def refit_span(
    stage: str,
    *,
    status: str = "ok",
    metadata: dict[str, Any] | None = None,
    accumulate_metadata: bool = False,
    duration_key: str | None = None,
) -> Iterator[dict[str, Any]]:
    """Record a span when a cycle is active; otherwise act as a no-op.

    The first choice for work the calling function owns and performs itself.
    Reach for :func:`add_refit_duration` only when the duration was produced
    somewhere else.
    """
    recorder = current_refit_timing()
    if recorder is None:
        yield {}
        return
    with recorder.span(
        stage,
        status=status,
        metadata=metadata,
        accumulate_metadata=accumulate_metadata,
        duration_key=duration_key,
    ) as discovered:
        yield discovered


def add_refit_duration(
    stage: str,
    seconds: float,
    *,
    status: str = "ok",
    metadata: dict[str, Any] | None = None,
    accumulate_metadata: bool = False,
) -> None:
    """Attribute an already-measured duration to a normalized stage.

    For durations another layer produced, where the work cannot be wrapped in
    a span from here: a transport that times its own wire read and hands the
    number back, or a figure a legacy metric already reports. Work the calling
    function performs itself belongs in :func:`refit_span`, which cannot
    disagree with the clock about what it covered.

    Negative values are clamped rather than rejected, since these arrive from
    a clock this function did not read.
    """
    recorder = current_refit_timing()
    if recorder is not None:
        recorder.add_duration(
            stage,
            max(0.0, float(seconds)),
            status=status,
            metadata=metadata,
            accumulate_metadata=accumulate_metadata,
        )


def add_refit_metadata(
    stage: str,
    metadata: dict[str, Any],
    *,
    accumulate: bool = False,
) -> None:
    """Attach facts to a stage without claiming to have measured anything."""
    recorder = current_refit_timing()
    if recorder is not None:
        recorder.add_metadata(stage, metadata, accumulate=accumulate)


def add_refit_bytes(count: int) -> None:
    recorder = current_refit_timing()
    if recorder is not None:
        recorder.add_bytes(count)


def set_refit_cold(cold: bool) -> None:
    """Mark whether this cycle built its transfer plan or reused one.

    The distinction dominates ``transfer_planning``: a reused plan makes it
    almost free, so a mean over both is a number that describes neither.
    """
    recorder = current_refit_timing()
    if recorder is not None:
        recorder.set_cold(cold)


__all__ = [
    "MX_REFIT_TIMING_PREFIX",
    "REFIT_TIMING_STAGES",
    "RefitTimingRecorder",
    "add_refit_bytes",
    "add_refit_duration",
    "add_refit_metadata",
    "current_refit_timing",
    "refit_span",
    "set_refit_cold",
    "use_refit_timing",
]