Skip to content

EmbeddingModelOperationalMetric metric

Bases: Metric

Base class for EmbeddingModel runtime-counter metrics.

Subclasses set _phase to one of "inference", "reward", or "optimizer" to read the corresponding counter set on each bound embedding model. Counters are populated by the EM based on the active op_scope (contextvar) the trainer sets for each phase.

Binds itself automatically to every EmbeddingModel reachable from the program (and their .fallback chains) on program.compile().

Example:

program.compile(
    metrics=[
        synalinks.metrics.EmbeddingModelOperationalMetric(),
    ],
)
Source code in synalinks/src/metrics/em_metrics.py 
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
@synalinks_export(
    [
        "synalinks.metrics.EmbeddingModelOperationalMetric",
        "synalinks.EmbeddingModelOperationalMetric",
    ]
)
class EmbeddingModelOperationalMetric(Metric):
    """Base class for `EmbeddingModel` runtime-counter metrics.

    Subclasses set `_phase` to one of ``"inference"``, ``"reward"``, or
    ``"optimizer"`` to read the corresponding counter set on each bound
    embedding model. Counters are populated by the EM based on the active
    ``op_scope`` (contextvar) the trainer sets for each phase.

    Binds itself automatically to every `EmbeddingModel` reachable from
    the program (and their `.fallback` chains) on `program.compile()`.

    Example:

    ```python
    program.compile(
        metrics=[
            synalinks.metrics.EmbeddingModelOperationalMetric(),
        ],
    )
    ```
    """

    _phase = "inference"

    def __init__(self, name=None):
        super().__init__(name=name)
        self._embedding_models = []
        self._baselines = {suffix: 0 for suffix in _TRACKED_SUFFIXES}
        self._wall_baseline = 0.0

    @property
    def embedding_models(self):
        return list(self._embedding_models)

    def bind_program(self, program):
        self._embedding_models = _collect_embedding_models(program)
        self._snapshot()

    def _attr(self, suffix):
        return f"{self._phase}_cumulated_{suffix}"

    def _read(self, suffix):
        attr = self._attr(suffix)
        return sum(getattr(em, attr, 0) for em in self._embedding_models)

    def _snapshot(self):
        for suffix in _TRACKED_SUFFIXES:
            self._baselines[suffix] = self._read(suffix)
        self._wall_baseline = read_phase_wall_clock_s(self._phase)

    def _delta(self, suffix):
        return self._read(suffix) - self._baselines.get(suffix, 0)

    def _wall_clock_delta(self):
        """Wall-clock seconds the trainer spent in this metric's phase since
        the last snapshot: the throughput denominator (concurrency-safe,
        unlike summed `elapsed_s`).
        """
        return read_phase_wall_clock_s(self._phase) - self._wall_baseline

    def reset_state(self):
        self._snapshot()

    async def update_state(self, *args, **kwargs):
        return

    def result(self):
        raise NotImplementedError

    def get_config(self):
        return {"name": self.name}