Untitled

from prometheus_client import Counter, Histogram


class MetricsUtils:
    DEFAULT_BUCKETS = (100, 200, 400, 800, 1600, 3000, 6000, 12000, 25000, 50000)

    tetris_step_latency = Histogram(
        "tetris_step_latency",
        "Latency of major steps in the tetris execution",
        labelnames=["step", "algorithm", "centerId", "groupId"],
        buckets=DEFAULT_BUCKETS,
    )

    model_exception_count = Counter(
        name="tetris_model_exception_count",
        documentation="model exception count",
        labelnames=["algorithm", "centerId", "groupId", "errorType"],
    )

    @staticmethod
    def record_step_latency(step_name: str, algorithm: str, center_id: int, group_id: str, latency: int):
        """Record the step latency and emit the metric"""
        MetricsUtils.tetris_step_latency.labels(
            step=step_name, algorithm=algorithm, centerId=center_id, groupId=group_id
        ).observe(latency)

    @staticmethod
    def record_model_exception(algorithm: str, center_id: int, group_id: str, error_type: str):
        """Record the step latency and emit the metric"""
        MetricsUtils.model_exception_count.labels(
            algorithm=algorithm, centerId=center_id, groupId=group_id, errorType=error_type
        ).inc()