sample_prometheus_equipped_predictor.py

import json
import random
import time
import os

from spell.serving import BasePredictor
import spell.serving.metrics as m


class Predictor(BasePredictor):
    """More complex metrics example. """
    def __init__(self):
        self.pid = str(os.getpid())
        self.inference_time_hist = m.prometheus.Histogram(
            'inference_time', 'Model inference time',
            labelnames=['pid'], labelvalues=[self.pid],
        )
        self.prediction_confidence_hist = m.prometheus.Histogram(
            'prediction_confidence',
            'Model (self-reported) prediction confidence',
            buckets=[round(0.1 * x, 1) for x in range(0, 11)],
            labelnames=['pid'], labelvalues=[self.pid],
        )

    def predict(self, payload):
        prediction_time = 0.5 + ((random.random() - 0.5) * 0.2)

        with self.inference_time_hist.time():
            # simulate an inference event, randomly takes 0.4 to 0.6 seconds
            time.sleep(prediction_time)

        # simulate model confidence, randomly in the range 0.7 and 0.9
        prediction_confidence = 0.8 + ((random.random() - 0.5) * 0.2)
        self.prediction_confidence_hist.observe(prediction_confidence)

        return json.dumps({"response": payload})