feat: add default sampling rate to mlobs logger

caraml-dev · Apr 4, 2024 · 252d6ae · 252d6ae
1 parent 1381a0b
commit 252d6ae
Showing 1 changed file with 18 additions and 11 deletions.
diff --git a/api/pkg/inference-logger/logger/mlobs_sink.go b/api/pkg/inference-logger/logger/mlobs_sink.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"math/rand"
 
 	"google.golang.org/protobuf/types/known/timestamppb"
 
@@ -19,6 +20,10 @@ var (
 	ErrMalformedLogEntry = errors.New("malformed log entry")
 )
 
+const (
+	SamplingRate = 0.02
+)
+
 type MLObsSink struct {
 	logger   *zap.SugaredLogger
 	producer KafkaProducer
@@ -163,17 +168,19 @@ func (m *MLObsSink) buildNewKafkaMessage(predictionLog *upiv1.PredictionLog) (*k
 
 func (m *MLObsSink) Sink(rawLogEntries []*LogEntry) error {
 	for _, rawLogEntry := range rawLogEntries {
-		predictionLog, err := m.newPredictionLog(rawLogEntry)
-		if err != nil {
-			m.logger.Errorf("unable to convert log entry: %v", err)
-		}
-		kafkaMessage, err := m.buildNewKafkaMessage(predictionLog)
-		if err != nil {
-			m.logger.Errorf("unable to build kafka message: %v", err)
-		}
-		err = m.producer.Produce(kafkaMessage, m.producer.Events())
-		if err != nil {
-			m.logger.Errorf("unable to produce kafka message: %v", err)
+		if rand.Float64() <= SamplingRate {
+			predictionLog, err := m.newPredictionLog(rawLogEntry)
+			if err != nil {
+				m.logger.Errorf("unable to convert log entry: %v", err)
+			}
+			kafkaMessage, err := m.buildNewKafkaMessage(predictionLog)
+			if err != nil {
+				m.logger.Errorf("unable to build kafka message: %v", err)
+			}
+			err = m.producer.Produce(kafkaMessage, m.producer.Events())
+			if err != nil {
+				m.logger.Errorf("unable to produce kafka message: %v", err)
+			}
 		}
 	}
 	return nil