arima_predictor.py

"""
arima_predictor.py
ARIMA-based predictor server for edge-faas-cpp.

Algorithm:
  - Maintains a sliding window of the last WINDOW_SIZE RPS observations.
  - If history < MIN_SAMPLES: falls back to reactive (ceil(rps * T) + margin).
  - Otherwise: fits ARIMA(2,1,2), forecasts next-window RPS, applies Little's Law.
  - Reports inference latency per call so C++ can log it.

CSCI 599: Network Systems for Cloud Computing
University of Southern California
"""

import os
import sys
import math
import time
import socket
import threading
import warnings
from collections import deque

warnings.filterwarnings("ignore")

ARIMA_SOCK   = "/tmp/faas_arima.sock"
WINDOW_SIZE  = 30   # sliding history length
MIN_SAMPLES  = 10   # minimum before ARIMA is reliable
SAFETY_MARGIN = 1   # matches Predictor._safety_margin

try:
    from statsmodels.tsa.arima.model import ARIMA
    ARIMA_AVAILABLE = True
except ImportError:
    ARIMA_AVAILABLE = False

_history      = deque(maxlen=WINDOW_SIZE)
_history_lock = threading.Lock()


def predict(rps: int, avg_service_time: float) -> tuple[int, float]:
    """
    Record rps, fit ARIMA if enough data, return (target_workers, latency_ms).
    Little's Law: target = ceil(predicted_lambda * T) + margin
    """
    t0 = time.perf_counter()

    with _history_lock:
        _history.append(rps)
        data = list(_history)

    predicted_lambda = float(rps)  # fallback: reactive

    if ARIMA_AVAILABLE and len(data) >= MIN_SAMPLES:
        try:
            fit = ARIMA(data, order=(2, 1, 2)).fit()
            fc  = fit.forecast(steps=1)[0]
            if fc > 0:
                predicted_lambda = fc
        except Exception:
            pass  # stay with reactive fallback

    target      = int(math.ceil(predicted_lambda * avg_service_time)) + SAFETY_MARGIN
    latency_ms  = (time.perf_counter() - t0) * 1000.0
    return target, latency_ms


def handle_client(conn: socket.socket) -> None:
    try:
        buf = b""
        while b"\n" not in buf:
            chunk = conn.recv(64)
            if not chunk:
                return
            buf += chunk
        line = buf.decode().strip()
        parts = line.split(",")
        rps = int(parts[0])
        avg_t = float(parts[1]) if len(parts) > 1 else 0.1

        target, latency_ms = predict(rps, avg_t)
        conn.sendall(f"{target},{latency_ms:.3f}\n".encode())
    except Exception:
        pass
    finally:
        conn.close()


def main() -> None:
    ready_fd = int(sys.argv[1])

    try:
        os.unlink(ARIMA_SOCK)
    except FileNotFoundError:
        pass

    srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    srv.bind(ARIMA_SOCK)
    srv.listen(16)

    # Signal C++ that the socket is up and ready
    os.write(ready_fd, b"READY\n")
    os.close(ready_fd)

    while True:
        conn, _ = srv.accept()
        threading.Thread(target=handle_client, args=(conn,), daemon=True).start()


if __name__ == "__main__":
    main()