Real-Time ML Data Pipelines with Python Apache Kafka and Scikit-Learn

By Khaled Md Tuhidul Hossain · 2024-11-11 · Technology

Batch processing has long been the default approach for machine learning pipelines: collect data, train a model, run predictions on a static dataset, repeat. But many real-world problems demand something faster. Fraud detection, sensor monitoring, recommendation engines, and anomaly detection all benefit from processing data the moment it arrives rather than hours or days later.

I will walk through building a real-time ML pipeline using Apache Kafka for message streaming, Scikit-Learn for model training and inference, and kafka-python as the client library. The architecture is straightforward: data flows into Kafka, a consumer picks it up, applies feature engineering, runs it through a trained model, and publishes predictions back to another Kafka topic.

Architecture overview

The pipeline has three main components:

1. Producer: Simulates or ingests live data and publishes it to a Kafka topic.
2. ML Consumer: Reads from the input topic, transforms the data, runs model inference, and writes predictions to an output topic.
3. Prediction Consumer: Reads predictions for downstream use (dashboards, alerts, storage).

This separation of concerns means each component can scale independently, and the Kafka topics act as durable buffers between them.

Prerequisites

You will need Kafka running locally. The simplest approach is Docker:

docker run -d --name kafka \
  -p 9092:9092 \
  -e KAFKA_CFG_NODE_ID=0 \
  -e KAFKA_CFG_PROCESS_ROLES=controller,broker \
  -e KAFKA_CFG_LISTENERS=PLAINTEXT://:9092,CONTROLLER://:9093 \
  -e KAFKA_CFG_ADVERTISED_LISTENERS=PLAINTEXT://localhost:9092 \
  -e KAFKA_CFG_CONTROLLER_QUORUM_VOTERS=0@localhost:9093 \
  -e KAFKA_CFG_CONTROLLER_LISTENER_NAMES=CONTROLLER \
  -e KAFKA_CFG_LISTENER_SECURITY_PROTOCOL_MAP=CONTROLLER:PLAINTEXT,PLAINTEXT:PLAINTEXT \
  bitnami/kafka:latest

Install the Python dependencies:

pip install kafka-python scikit-learn pandas numpy joblib

Training the model

Before building the streaming pipeline, we need a trained model. Let us create a simple anomaly detection model for network traffic data:

import numpy as np
import pandas as pd
from sklearn.ensemble import IsolationForest
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
import joblib

# Generate synthetic training data representing normal network traffic
np.random.seed(42)
n_samples = 5000

normal_data = pd.DataFrame({
    "bytes_sent": np.random.exponential(scale=500, size=n_samples),
    "bytes_received": np.random.exponential(scale=800, size=n_samples),
    "packets_per_second": np.random.normal(loc=150, scale=30, size=n_samples),
    "connection_duration": np.random.exponential(scale=60, size=n_samples),
    "unique_ports": np.random.poisson(lam=3, size=n_samples),
})

# Build and train the pipeline
model_pipeline = Pipeline([
    ("scaler", StandardScaler()),
    ("detector", IsolationForest(
        n_estimators=200,
        contamination=0.05,
        random_state=42,
    )),
])

model_pipeline.fit(normal_data)

# Save the trained model
joblib.dump(model_pipeline, "anomaly_detector.joblib")
print("Model trained and saved successfully.")

The IsolationForest is well-suited for streaming anomaly detection because it is fast at inference time and does not require labeled data for training. The contamination parameter sets the expected proportion of anomalies, which influences the decision threshold.

Building the Kafka producer

The producer simulates incoming network traffic events and publishes them to a Kafka topic:

import json
import time
import random
import numpy as np
from kafka import KafkaProducer
from datetime import datetime

producer = KafkaProducer(
    bootstrap_servers=["localhost:9092"],
    value_serializer=lambda v: json.dumps(v).encode("utf-8"),
    key_serializer=lambda k: k.encode("utf-8") if k else None,
)

TOPIC = "network-traffic"

def generate_normal_event():
    """Generate a normal network traffic event."""
    return {
        "timestamp": datetime.utcnow().isoformat(),
        "source_ip": f"192.168.1.{random.randint(1, 254)}",
        "bytes_sent": max(0, np.random.exponential(scale=500)),
        "bytes_received": max(0, np.random.exponential(scale=800)),
        "packets_per_second": max(0, np.random.normal(loc=150, scale=30)),
        "connection_duration": max(0, np.random.exponential(scale=60)),
        "unique_ports": max(0, int(np.random.poisson(lam=3))),
    }

def generate_anomalous_event():
    """Generate a suspicious network traffic event."""
    return {
        "timestamp": datetime.utcnow().isoformat(),
        "source_ip": f"10.0.0.{random.randint(1, 254)}",
        "bytes_sent": np.random.exponential(scale=5000),
        "bytes_received": np.random.exponential(scale=50),
        "packets_per_second": np.random.normal(loc=800, scale=100),
        "connection_duration": np.random.exponential(scale=300),
        "unique_ports": int(np.random.poisson(lam=25)),
    }

def run_producer(num_events=1000, anomaly_rate=0.05):
    """Produce a stream of network traffic events."""
    print(f"Producing {num_events} events to topic '{TOPIC}'...")

    for i in range(num_events):
        if random.random() < anomaly_rate:
            event = generate_anomalous_event()
        else:
            event = generate_normal_event()

        producer.send(TOPIC, key=event["source_ip"], value=event)

        if (i + 1) % 100 == 0:
            print(f"Produced {i + 1} events")

        time.sleep(0.01)  # Simulate real-time arrival

    producer.flush()
    print("Producer finished.")

if __name__ == "__main__":
    run_producer()

Using the source IP as the Kafka key ensures that events from the same IP are always routed to the same partition, which can be useful for stateful processing.

Building the ML consumer

This is the core of the pipeline. The consumer reads events, extracts features, runs inference, and publishes predictions:

import json
import pandas as pd
import joblib
from kafka import KafkaConsumer, KafkaProducer
from datetime import datetime

# Load the trained model
model = joblib.load("anomaly_detector.joblib")

# Feature columns used during training
FEATURE_COLUMNS = [
    "bytes_sent",
    "bytes_received",
    "packets_per_second",
    "connection_duration",
    "unique_ports",
]

INPUT_TOPIC = "network-traffic"
OUTPUT_TOPIC = "anomaly-predictions"

consumer = KafkaConsumer(
    INPUT_TOPIC,
    bootstrap_servers=["localhost:9092"],
    value_deserializer=lambda v: json.loads(v.decode("utf-8")),
    group_id="ml-inference-group",
    auto_offset_reset="latest",
    enable_auto_commit=True,
)

producer = KafkaProducer(
    bootstrap_servers=["localhost:9092"],
    value_serializer=lambda v: json.dumps(v).encode("utf-8"),
)

def extract_features(event: dict) -> pd.DataFrame:
    """Extract model features from a raw event."""
    features = {col: [event.get(col, 0)] for col in FEATURE_COLUMNS}
    return pd.DataFrame(features)

def run_inference_consumer():
    """Consume events, run inference, and produce predictions."""
    print(f"ML Consumer started. Listening on '{INPUT_TOPIC}'...")
    event_count = 0
    anomaly_count = 0

    for message in consumer:
        event = message.value
        features = extract_features(event)

        # Run inference: IsolationForest returns -1 for anomalies, 1 for normal
        prediction = model.predict(features)[0]
        anomaly_score = model.decision_function(features)[0]

        is_anomaly = prediction == -1

        result = {
            "timestamp": event["timestamp"],
            "processed_at": datetime.utcnow().isoformat(),
            "source_ip": event.get("source_ip", "unknown"),
            "is_anomaly": is_anomaly,
            "anomaly_score": round(float(anomaly_score), 4),
            "features": {col: event.get(col) for col in FEATURE_COLUMNS},
        }

        producer.send(OUTPUT_TOPIC, value=result)

        event_count += 1
        if is_anomaly:
            anomaly_count += 1
            print(
                f"ANOMALY DETECTED | IP: {result['source_ip']} | "
                f"Score: {result['anomaly_score']:.4f}"
            )

        if event_count % 100 == 0:
            rate = (anomaly_count / event_count) * 100
            print(
                f"Processed {event_count} events | "
                f"Anomalies: {anomaly_count} ({rate:.1f}%)"
            )

if __name__ == "__main__":
    run_inference_consumer()

A few important design decisions here:

• auto_offset_reset="latest": The consumer starts reading from the most recent messages, which is appropriate for real-time inference where historical data has already been processed.
• group_id: Enables consumer group coordination. Multiple instances of this consumer can run in parallel, with Kafka automatically partitioning the work.
• Decision function scores: The decision_function provides a continuous anomaly score, which is more useful than the binary prediction for setting alert thresholds.

Building the prediction consumer

A simple consumer to read and act on predictions:

import json
from kafka import KafkaConsumer

consumer = KafkaConsumer(
    "anomaly-predictions",
    bootstrap_servers=["localhost:9092"],
    value_deserializer=lambda v: json.loads(v.decode("utf-8")),
    group_id="alert-consumer-group",
    auto_offset_reset="latest",
)

def run_alert_consumer():
    """Consume predictions and trigger alerts for anomalies."""
    print("Alert Consumer started. Watching for anomalies...")

    for message in consumer:
        prediction = message.value

        if prediction["is_anomaly"]:
            print(
                f"[ALERT] Anomalous traffic from {prediction['source_ip']} "
                f"at {prediction['timestamp']} "
                f"(score: {prediction['anomaly_score']})"
            )
            # In production: send to PagerDuty, Slack, email, etc.

if __name__ == "__main__":
    run_alert_consumer()

Real-time feature engineering

For more sophisticated pipelines, you often need to compute features over sliding windows. Here is a pattern for maintaining a rolling window of events per IP address:

from collections import defaultdict, deque
from datetime import datetime, timedelta

class StreamingFeatureStore:
    """Maintains sliding window features for streaming data."""

    def __init__(self, window_seconds=300):
        self.window_seconds = window_seconds
        self.event_windows = defaultdict(deque)

    def add_event(self, source_ip: str, event: dict):
        """Add an event and clean up expired entries."""
        now = datetime.utcnow()
        cutoff = now - timedelta(seconds=self.window_seconds)

        window = self.event_windows[source_ip]
        window.append({"timestamp": now, **event})

        # Remove expired events
        while window and window[0]["timestamp"] < cutoff:
            window.popleft()

    def get_window_features(self, source_ip: str) -> dict:
        """Compute aggregate features over the current window."""
        window = self.event_windows[source_ip]

        if not window:
            return {"event_count": 0, "avg_bytes_sent": 0, "port_diversity": 0}

        bytes_sent = [e.get("bytes_sent", 0) for e in window]
        ports = [e.get("unique_ports", 0) for e in window]

        return {
            "event_count": len(window),
            "avg_bytes_sent": sum(bytes_sent) / len(bytes_sent),
            "max_bytes_sent": max(bytes_sent),
            "port_diversity": len(set(ports)),
        }

These windowed features capture temporal patterns that single-event features miss, such as a gradual increase in traffic volume or a sudden burst of connections from a single IP.

Production considerations

When moving this pipeline to production, keep these points in mind:

• Schema validation: Use a schema registry (like Confluent Schema Registry with Avro or JSON Schema) to enforce data contracts between producers and consumers.
• Model versioning: Tag each prediction with the model version used, so you can trace results back to a specific model.
• Monitoring: Track consumer lag, inference latency, and prediction distributions. A sudden shift in the anomaly rate could indicate data drift.
• Batched inference: For higher throughput, accumulate a small batch of events before running inference. Scikit-Learn models are significantly faster on batches than on individual samples.

Conclusion

Combining Kafka's streaming capabilities with Scikit-Learn's inference speed gives you a practical real-time ML pipeline. The architecture decouples data ingestion from inference from downstream consumption, making each component independently scalable and replaceable. The same pattern applies to fraud detection, infrastructure monitoring, and recommendation scoring.