Logging and Observability for ML Systems

import hashlib
import json
import logging
import random
import time

import pandas as pd
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline

train, live = pd.read_csv("train.csv"), pd.read_csv("test.csv")   # Banking77; the test split stands in for traffic
model = make_pipeline(TfidfVectorizer(ngram_range=(1, 2), sublinear_tf=True),
                      LogisticRegression(max_iter=3000, C=10)).fit(train.text, train.category)
VERSION, HANDOFF = "intents-tfidf-v1", 0.30                       # below 0.30 confidence, a person answers

random.seed(7)               # the SDK draws IDs from Python's random module; seeded only to make this page reproducible
spans = InMemorySpanExporter()                                    # in production: an OTLP exporter to a tracing backend
provider = TracerProvider()
provider.add_span_processor(SimpleSpanProcessor(spans))
trace.set_tracer_provider(provider)
tracer = trace.get_tracer("intent-service")


class JsonLines(logging.Formatter):
    def format(self, record):
        return json.dumps(record.fields)


handler = logging.FileHandler("predictions.jsonl", mode="w")
handler.setFormatter(JsonLines())
log = logging.getLogger("predictions")
log.addHandler(handler)
log.setLevel(logging.INFO)


def predict(request_id, text):
    with tracer.start_as_current_span("predict", attributes={"request.id": request_id, "model.version": VERSION}) as root:
        start = time.perf_counter()
        with tracer.start_as_current_span("vectorize"):
            features = model[0].transform([text.strip()])
        with tracer.start_as_current_span("classify") as span:
            probabilities = model[1].predict_proba(features)[0]
            intent, confidence = model.classes_[probabilities.argmax()], round(float(probabilities.max()), 3)
            span.set_attributes({"prediction.intent": intent, "prediction.confidence": confidence})
        if confidence < HANDOFF:
            with tracer.start_as_current_span("handoff_to_agent"):
                pass                                              # a real service would queue the conversation here
        log.info("prediction", extra={"fields": {
            "event": "prediction", "request_id": request_id, "trace_id": format(root.get_span_context().trace_id, "032x"),
            "model_version": VERSION, "text_sha256": hashlib.sha256(text.encode()).hexdigest()[:16], "chars": len(text),
            "intent": intent, "confidence": confidence, "handoff": confidence < HANDOFF,
            "latency_ms": round((time.perf_counter() - start) * 1000, 2)}})


for number, text in enumerate(live.text):
    predict(f"req-{number:05d}", text)
handler.close()

records = [json.loads(line) for line in open("predictions.jsonl", encoding="utf-8")]
logs = pd.DataFrame(records)
finished = spans.get_finished_spans()
print(f"{len(live):,} requests -> {len(logs):,} log lines, {len({s.context.trace_id for s in finished}):,} traces, "
      f"{len(finished):,} spans")
print("spans by name:", {name: int(n) for name, n in pd.Series([s.name for s in finished]).value_counts().items()})

record = next(r for r in records if r["handoff"])
print("\none handed-off request, as logged (latency omitted because it varies):")
print(json.dumps({k: v for k, v in record.items() if k != "latency_ms"}))
print("its trace:")
same_trace = [s for s in finished if format(s.context.trace_id, "032x") == record["trace_id"]]
for span in sorted(same_trace, key=lambda s: s.parent is not None):   # root first, then children in the order they ended
    print(f"  {'  ' if span.parent else ''}{span.name:<18} {dict(span.attributes)}")

truth = pd.DataFrame({"request_id": [f"req-{n:05d}" for n in range(len(live))], "true_intent": live.category})
joined = logs.merge(truth, on="request_id")                       # labels arrive later; join them on the request ID
joined["correct"] = joined.intent == joined.true_intent
band = pd.cut(joined.confidence, [0, 0.3, 0.5, 0.7, 1.001], right=False,
              labels=["below 0.3", "0.3 to 0.5", "0.5 to 0.7", "0.7 and above"])
table = joined.groupby(band, observed=True).correct.agg(requests="size", accuracy="mean", errors=lambda c: (~c).sum())
table["share of errors"] = table.errors / table.errors.sum()
print(f"\naccuracy once labels arrive: {joined.correct.mean():.1%}")
print(table.to_string(formatters={"accuracy": "{:.0%}".format, "share of errors": "{:.0%}".format}))
handed_off = joined[joined.handoff].true_intent.value_counts().head(3)
print("\nintents most often handed off:", {intent: int(n) for intent, n in handed_off.items()})

3,080 requests -> 3,080 log lines, 3,080 traces, 9,518 spans
spans by name: {'vectorize': 3080, 'classify': 3080, 'predict': 3080, 'handoff_to_agent': 278}

one handed-off request, as logged (latency omitted because it varies):
{"event": "prediction", "request_id": "req-00000", "trace_id": "6513270e269e0d37f2a74de452e6b438", "model_version": "intents-tfidf-v1", "text_sha256": "57a5a4533e3e3cf2", "chars": 24, "intent": "get_physical_card", "confidence": 0.106, "handoff": true}
its trace:
  predict            {'request.id': 'req-00000', 'model.version': 'intents-tfidf-v1'}
    vectorize          {}
    classify           {'prediction.intent': 'get_physical_card', 'prediction.confidence': 0.106}
    handoff_to_agent   {}

accuracy once labels arrive: 89.4%
               requests accuracy  errors share of errors
confidence                                              
below 0.3           278      52%     134             41%
0.3 to 0.5          361      73%      99             30%
0.5 to 0.7          453      87%      58             18%
0.7 and above      1988      98%      36             11%

intents most often handed off: {'compromised_card': 9, 'card_acceptance': 8, 'topping_up_by_card': 8}