From ddf9a7001a0050cc9dac9a623282c5764f417117 Mon Sep 17 00:00:00 2001 From: pavsoss Date: Wed, 29 Jul 2026 13:45:10 -0400 Subject: [PATCH] Add opt-in NDJSON streaming for bulk prediction --- backend/bulk_predict.py | 409 +++++++++++++++++----- backend/errors.py | 10 + backend/tests/test_bulk_predict_robust.py | 192 ++++++++++ 3 files changed, 529 insertions(+), 82 deletions(-) create mode 100644 backend/tests/test_bulk_predict_robust.py diff --git a/backend/bulk_predict.py b/backend/bulk_predict.py index 52fa6a5c..373f1b17 100644 --- a/backend/bulk_predict.py +++ b/backend/bulk_predict.py @@ -1,85 +1,124 @@ +"""Bulk file prediction endpoints (issue #1021). + +Scoring runs against a single :class:`serving_state.ServingSnapshot` taken at the +start of each request, so an in-flight ``/reload-model`` hot-swap can never pair +a new model with an old vectorizer part-way through a file. Inference is +row-isolated: an empty, over-length or un-transformable row is recorded in a +structured ``skipped`` list with a typed reason instead of aborting the whole +upload, and responses (JSON and exported CSV) carry the serving model +``version`` for provenance. + +Caps are configurable via the environment, mirroring ``BULK_PREDICT_BATCH_SIZE``: +``BULK_PREDICT_MAX_ROWS`` bounds the total data rows accepted (exceeding it is a +fatal, typed error) and ``BULK_PREDICT_MAX_ROW_LEN`` bounds a single row's length +(over-length rows are skipped, not fatal). + +Clients that prefer not to buffer a large result set can opt into an NDJSON +stream (``?stream=ndjson`` or ``Accept: application/x-ndjson``): the same rows +are scored and emitted one newline-delimited JSON object at a time. The default +(buffered JSON) response is unchanged. +""" + import csv import io +import json import os - import numpy as np -from flask import (Blueprint, current_app, jsonify, - request, send_file) +from flask import (Blueprint, Response, jsonify, request, + send_file, stream_with_context) +from errors import ApiError, ErrorCode from rate_limiting import RateLimitPolicy, rate_limit +import serving_state + +__all__ = ["bulk_predict_bp"] bulk_predict_bp = Blueprint("bulk_predict", __name__) +NDJSON_MIMETYPE = "application/x-ndjson" + +# Longest row message echoed back in a skip record; full over-length payloads +# would bloat the response, so only a preview is returned for context. +_SKIP_PREVIEW_LEN = 200 + + +def _int_env(name, default): + """Read a positive int cap from the environment, falling back on garbage.""" + try: + value = int(os.getenv(name, str(default))) + except (TypeError, ValueError): + return default + return value if value > 0 else default + -def parse_and_predict_file(file): - # Check file extension +def _resolve_snapshot(): + """Return a coherent serving snapshot or raise a typed, fatal error. + + A missing snapshot means the process never finished loading (or a reload is + mid-flight with no objects installed); scoring against ``None`` would 500, + so surface it as a typed 503 the client can branch on. + """ + state = serving_state.STATE + snapshot = state.snapshot() if state is not None else None + if ( + snapshot is None + or snapshot.model is None + or snapshot.vectorizer is None + or snapshot.label_encoder is None + ): + raise ApiError( + ErrorCode.BULK_MODEL_UNAVAILABLE, + "Model dependencies are not loaded.", + 503, + ) + return snapshot + + +def _wants_ndjson(): + """True when the caller opted into the NDJSON streaming response.""" + if request.args.get("stream", "").lower() == "ndjson": + return True + accept = request.headers.get("Accept", "") + return NDJSON_MIMETYPE in accept.lower() + + +def _skip_record(index, message, code, detail): + return { + "row": index, + "message": (message or "")[:_SKIP_PREVIEW_LEN], + "reason": code.value, + "detail": detail, + } + + +def _extract_rows(file): + """Return ``(rows, error)`` where rows is a list of ``(index, raw_value)``. + + ``index`` is the 1-based data-row position (useful for the client to locate + a skipped row); ``raw_value`` may be ``None``/blank and is validated later. + Fatal, file-level problems (bad type, size, structure) come back as the + ``error`` string so the caller can map them to the legacy status codes. + """ filename = file.filename.lower() if file.filename else "" if not (filename.endswith(".csv") or filename.endswith(".txt")): return None, "Unsupported file type. Only CSV and TXT files are supported." - # Check file size file.seek(0, os.SEEK_END) file_size = file.tell() file.seek(0) - if file_size > 2 * 1024 * 1024: # 2MB limit return None, "File size exceeds the limit of 2MB." if file_size == 0: return None, "Empty file uploaded." try: - # Use a text wrapper around the uploaded file stream for proper decoding. text_wrapper = io.TextIOWrapper(file.stream, encoding="utf-8", errors="replace") except Exception: return None, "Failed to read uploaded file." - # Helper for batch inference - def _batch_predict(batch_messages): - vectorizer = getattr(current_app, "vectorizer", None) - model = getattr(current_app, "model", None) - label_encoder = getattr(current_app, "label_encoder", None) - if not vectorizer or not model or not label_encoder: - raise RuntimeError("ML Model dependencies are not loaded.") - text_vectors = vectorizer.transform(batch_messages) - predictions = model.predict(text_vectors) - final_outputs = label_encoder.inverse_transform(predictions) - - # Compute decisions - decisions = model.decision_function(text_vectors) - - batch_results = [] - for i, (msg, pred) in enumerate(zip(batch_messages, final_outputs)): - pred_str = str(pred) - dec_score = float(np.max(np.abs(decisions[i]))) - prob = 1.0 / (1.0 + np.exp(-dec_score)) - conf_score = round(prob * 100, 2) - - if conf_score >= 80: - conf_level = "high" - elif conf_score >= 60: - conf_level = "medium" - else: - conf_level = "low" - - batch_results.append( - { - "message": msg, - "prediction": pred_str, - "result": pred_str, - "confidence": round(conf_score / 100.0, 4), - "confidence_score": conf_score, - "decision_score": dec_score, - "confidence_level": conf_level, - } - ) - return batch_results - - BATCH_SIZE = int(os.getenv("BULK_PREDICT_BATCH_SIZE", "256")) - results = [] - batch = [] - + rows = [] if filename.endswith(".csv"): try: reader = csv.DictReader(text_wrapper) @@ -96,29 +135,214 @@ def _batch_predict(batch_messages): None, "CSV file must contain either a 'text' or 'message' column.", ) - for row in reader: - val = row.get(col_name) - if val is not None and val.strip(): - batch.append(val.strip()) - if len(batch) >= BATCH_SIZE: - results.extend(_batch_predict(batch)) - batch = [] + for index, row in enumerate(reader, start=1): + rows.append((index, row.get(col_name))) except Exception as e: return None, f"Failed to parse CSV: {str(e)}" - else: # TXT file + else: # TXT file: one message per non-blank line. + index = 0 for line in text_wrapper: - line = line.strip() - if line: - batch.append(line) - if len(batch) >= BATCH_SIZE: - results.extend(_batch_predict(batch)) - batch = [] - # Process any remaining messages - if batch: - results.extend(_batch_predict(batch)) - if not results: - return None, "No valid messages found in the file." - return results, None + stripped = line.strip() + if not stripped: + continue + index += 1 + rows.append((index, stripped)) + + return rows, None + + +def _partition_rows(rows, max_row_len): + """Split extracted rows into scorable rows and typed per-row skips.""" + valid_rows = [] + skipped = [] + for index, value in rows: + if value is None or not value.strip(): + skipped.append( + _skip_record( + index, value, ErrorCode.BULK_ROW_EMPTY, "Empty row skipped." + ) + ) + continue + msg = value.strip() + if len(msg) > max_row_len: + skipped.append( + _skip_record( + index, + msg, + ErrorCode.BULK_ROW_TOO_LONG, + f"Row exceeds maximum length of {max_row_len} characters.", + ) + ) + continue + valid_rows.append((index, msg)) + return valid_rows, skipped + + +def _predict_batch(messages, snapshot): + """Score a batch of messages against ``snapshot`` and shape the results. + + Raises on any transform/predict failure; callers isolate failures by + retrying the offending batch one row at a time. + """ + text_vectors = snapshot.vectorizer.transform(messages) + predictions = snapshot.model.predict(text_vectors) + final_outputs = snapshot.label_encoder.inverse_transform(predictions) + decisions = snapshot.model.decision_function(text_vectors) + + batch_results = [] + for i, (msg, pred) in enumerate(zip(messages, final_outputs)): + pred_str = str(pred) + dec_score = float(np.max(np.abs(decisions[i]))) + prob = 1.0 / (1.0 + np.exp(-dec_score)) + conf_score = round(prob * 100, 2) + + if conf_score >= 80: + conf_level = "high" + elif conf_score >= 60: + conf_level = "medium" + else: + conf_level = "low" + + batch_results.append( + { + "message": msg, + "prediction": pred_str, + "result": pred_str, + "confidence": round(conf_score / 100.0, 4), + "confidence_score": conf_score, + "decision_score": dec_score, + "confidence_level": conf_level, + } + ) + return batch_results + + +def _score_batch_isolated(chunk, snapshot): + """Score one chunk, degrading to per-row scoring if the batch call fails. + + Returns ``(results, skipped)`` for the chunk so healthy rows survive a + single poison row that would otherwise raise for the whole batch. + """ + messages = [msg for _, msg in chunk] + try: + return _predict_batch(messages, snapshot), [] + except Exception: + results = [] + skipped = [] + for index, msg in chunk: + try: + results.extend(_predict_batch([msg], snapshot)) + except Exception: + skipped.append( + _skip_record( + index, + msg, + ErrorCode.BULK_ROW_UNPROCESSABLE, + "Row could not be transformed or scored.", + ) + ) + return results, skipped + + +def _score_rows(valid_rows, snapshot, batch_size): + """Score pre-validated rows in batches, isolating per-row failures.""" + results = [] + skipped = [] + for start in range(0, len(valid_rows), batch_size): + chunk = valid_rows[start : start + batch_size] + chunk_results, chunk_skipped = _score_batch_isolated(chunk, snapshot) + results.extend(chunk_results) + skipped.extend(chunk_skipped) + return results, skipped + + +def parse_and_predict_file(file, snapshot): + """Parse an uploaded file and score it against ``snapshot``. + + Returns ``(results, skipped, error)``. ``error`` is a non-empty string only + for fatal, file-level problems the legacy handlers map to 400/413; row-level + problems never populate ``error`` -- they go into ``skipped``. + """ + rows, error = _extract_rows(file) + if error: + return None, None, error + if not rows: + return None, None, "No valid messages found in the file." + + _enforce_row_cap(rows) + + max_row_len = _int_env("BULK_PREDICT_MAX_ROW_LEN", 10000) + valid_rows, skipped = _partition_rows(rows, max_row_len) + + batch_size = _int_env("BULK_PREDICT_BATCH_SIZE", 256) + results, score_skipped = _score_rows(valid_rows, snapshot, batch_size) + skipped.extend(score_skipped) + return results, skipped, None + + +def _enforce_row_cap(rows): + max_rows = _int_env("BULK_PREDICT_MAX_ROWS", 10000) + if len(rows) > max_rows: + raise ApiError( + ErrorCode.BULK_TOO_MANY_ROWS, + f"File contains {len(rows)} rows, exceeding the limit of {max_rows}.", + 413, + ) + + +def _summarize(results): + total = len(results) + spam_count = sum( + 1 for r in results if r["prediction"].lower() not in ("ham", "safe") + ) + non_spam_count = total - spam_count + spam_pct = round((spam_count / total) * 100, 2) if total > 0 else 0.0 + return total, spam_count, non_spam_count, spam_pct + + +def _iter_ndjson(rows, snapshot): + """Yield the bulk result set as newline-delimited JSON, row by row. + + The stream opens with a ``meta`` line (model version + row count), then one + ``result`` line per scored row, then any ``skipped`` lines, and closes with + a ``summary`` line. Scoring is done in batches for speed but flushed a row + at a time so the client can start consuming before the whole file is done. + """ + max_row_len = _int_env("BULK_PREDICT_MAX_ROW_LEN", 10000) + batch_size = _int_env("BULK_PREDICT_BATCH_SIZE", 256) + valid_rows, skipped = _partition_rows(rows, max_row_len) + + yield json.dumps( + {"type": "meta", "model_version": snapshot.version, "total_rows": len(rows)} + ) + "\n" + + total = 0 + spam_count = 0 + for start in range(0, len(valid_rows), batch_size): + chunk = valid_rows[start : start + batch_size] + chunk_results, chunk_skipped = _score_batch_isolated(chunk, snapshot) + skipped.extend(chunk_skipped) + for r in chunk_results: + total += 1 + if r["prediction"].lower() not in ("ham", "safe"): + spam_count += 1 + yield json.dumps({"type": "result", **r}) + "\n" + + for s in skipped: + yield json.dumps({"type": "skipped", **s}) + "\n" + + spam_pct = round((spam_count / total) * 100, 2) if total > 0 else 0.0 + yield json.dumps( + { + "type": "summary", + "total_messages": total, + "spam_count": spam_count, + "non_spam_count": total - spam_count, + "spam_percentage": spam_pct, + "skipped_count": len(skipped), + "model_version": snapshot.version, + } + ) + "\n" @bulk_predict_bp.route("/bulk-predict", methods=["POST"]) @@ -131,17 +355,30 @@ def bulk_predict(): if not file or file.filename == "": return jsonify({"error": "No file uploaded"}), 400 - results, error = parse_and_predict_file(file) + snapshot = _resolve_snapshot() + + if _wants_ndjson(): + rows, error = _extract_rows(file) + if error: + status_code = 413 if "exceeds the limit" in error.lower() else 400 + return jsonify({"error": error}), status_code + if not rows: + return jsonify({"error": "No valid messages found in the file."}), 400 + # Enforce the row cap before opening the stream so it still surfaces as + # a normal typed error rather than mid-stream. + _enforce_row_cap(rows) + return Response( + stream_with_context(_iter_ndjson(rows, snapshot)), + mimetype=NDJSON_MIMETYPE, + headers={"X-Model-Version": str(snapshot.version)}, + ) + + results, skipped, error = parse_and_predict_file(file, snapshot) if error: status_code = 413 if "exceeds the limit" in error.lower() else 400 return jsonify({"error": error}), status_code - total = len(results) - spam_count = sum( - 1 for r in results if r["prediction"].lower() not in ("ham", "safe") - ) - non_spam_count = total - spam_count - spam_pct = round((spam_count / total) * 100, 2) if total > 0 else 0.0 + total, spam_count, non_spam_count, spam_pct = _summarize(results) return jsonify( { @@ -149,7 +386,10 @@ def bulk_predict(): "spam_count": spam_count, "non_spam_count": non_spam_count, "spam_percentage": spam_pct, + "model_version": snapshot.version, "results": results, + "skipped": skipped, + "skipped_count": len(skipped), } ) @@ -164,7 +404,8 @@ def bulk_predict_export(): if not file or file.filename == "": return jsonify({"error": "No file uploaded"}), 400 - results, error = parse_and_predict_file(file) + snapshot = _resolve_snapshot() + results, _skipped, error = parse_and_predict_file(file, snapshot) if error: status_code = 413 if "exceeds the limit" in error.lower() else 400 return jsonify({"error": error}), status_code @@ -180,6 +421,7 @@ def bulk_predict_export(): "confidence_score", "decision_score", "confidence_level", + "model_version", ] ) for r in results: @@ -191,17 +433,20 @@ def bulk_predict_export(): r["confidence_score"], r["decision_score"], r["confidence_level"], + snapshot.version, ] ) output_io.seek(0) mem = io.BytesIO(output_io.getvalue().encode("utf-8")) - return send_file( + response = send_file( mem, mimetype="text/csv", as_attachment=True, download_name="bulk_spam_predictions.csv", ) + response.headers["X-Model-Version"] = str(snapshot.version) + return response except Exception as e: return jsonify({"error": f"Failed to generate CSV report: {str(e)}"}), 500 diff --git a/backend/errors.py b/backend/errors.py index 5f96cb87..fe93f41c 100644 --- a/backend/errors.py +++ b/backend/errors.py @@ -75,6 +75,16 @@ class ErrorCode(StrEnum): # ── Readiness / graceful shutdown (issue #1009) ────────────────────── NOT_READY = "NOT_READY" + # ── Bulk prediction (issue #1021) ──────────────────────────────────── + # BULK_MODEL_UNAVAILABLE / BULK_TOO_MANY_ROWS are fatal (whole request); + # the BULK_ROW_* codes are per-row skip reasons returned in the "skipped" + # list so one bad row never fails the batch. + BULK_MODEL_UNAVAILABLE = "BULK_MODEL_UNAVAILABLE" + BULK_TOO_MANY_ROWS = "BULK_TOO_MANY_ROWS" + BULK_ROW_EMPTY = "BULK_ROW_EMPTY" + BULK_ROW_TOO_LONG = "BULK_ROW_TOO_LONG" + BULK_ROW_UNPROCESSABLE = "BULK_ROW_UNPROCESSABLE" + # ── Email / OAuth / IMAP provider paths (PR 2/2) ───────────────────── MISSING_USERNAME = "MISSING_USERNAME" MISSING_AUTH_CODE = "MISSING_AUTH_CODE" diff --git a/backend/tests/test_bulk_predict_robust.py b/backend/tests/test_bulk_predict_robust.py new file mode 100644 index 00000000..68530abb --- /dev/null +++ b/backend/tests/test_bulk_predict_robust.py @@ -0,0 +1,192 @@ +"""Robustness + streaming coverage for the bulk-prediction endpoints (#1021). + +Covers the resilient buffered path (valid rows returned alongside a typed +``skipped`` list, row-cap enforcement, version stamping) and the opt-in NDJSON +streaming path (per-row lines plus meta/summary framing). A fake serving +snapshot is installed so the tests depend only on the endpoint logic, not on +the real ``.pkl`` artifacts. +""" + +import io +import json +from pathlib import Path +import sys + +import numpy as np +import pytest + +BASE_DIR = Path(__file__).resolve().parents[2] +BACKEND_DIR = BASE_DIR / "backend" +sys.path.insert(0, str(BACKEND_DIR)) + +import api as api_module # noqa: E402 +import serving_state # noqa: E402 + +# A vectorizer that refuses to transform any row carrying this token, so a +# single un-transformable row inside an otherwise-valid batch can be simulated. +_POISON = "POISON" + + +class _FakeVectorizer: + def transform(self, messages): + if any(_POISON in m for m in messages): + raise ValueError("cannot transform poisoned row") + return np.array([[float(len(m))] for m in messages]) + + +class _FakeModel: + def predict(self, X): + return np.array([1 if row[0] > 10 else 0 for row in X]) + + def decision_function(self, X): + return np.array([float(row[0]) for row in X]) + + +class _FakeLabelEncoder: + def inverse_transform(self, predictions): + return np.array(["spam" if p == 1 else "ham" for p in predictions]) + + +@pytest.fixture +def client(): + api_module.app.config["TESTING"] = True + with api_module.app.test_client() as c: + yield c + + +@pytest.fixture +def fake_state(): + return serving_state.init_state( + model=_FakeModel(), + vectorizer=_FakeVectorizer(), + label_encoder=_FakeLabelEncoder(), + xai_service=None, + loader=lambda: {}, + metadata=None, + ) + + +def _upload(client, csv_text, path="/bulk-predict", **kwargs): + data = {"file": (io.BytesIO(csv_text.encode("utf-8")), "test.csv")} + return client.post(path, data=data, content_type="multipart/form-data", **kwargs) + + +class TestRowIsolation: + def test_mixed_valid_and_invalid_rows(self, client, fake_state, monkeypatch): + monkeypatch.setenv("BULK_PREDICT_MAX_ROW_LEN", "50") + csv_text = ( + "text\n" + "this is a long spam message\n" # valid -> spam + "hi\n" # valid -> ham + "\n" # empty -> skipped + + ("x" * 80 + "\n") # over-length -> skipped + + f"{_POISON} row\n" # un-transformable -> skipped + ) + res = _upload(client, csv_text) + assert res.status_code == 200 + body = res.get_json() + + messages = {r["message"] for r in body["results"]} + assert "this is a long spam message" in messages + assert "hi" in messages + + reasons = {s["reason"] for s in body["skipped"]} + assert "BULK_ROW_EMPTY" in reasons + assert "BULK_ROW_TOO_LONG" in reasons + assert "BULK_ROW_UNPROCESSABLE" in reasons + assert body["skipped_count"] == len(body["skipped"]) == 3 + + def test_single_bad_row_does_not_500(self, client, fake_state): + csv_text = "text\ngood message here now\n" + f"{_POISON}\n" + res = _upload(client, csv_text) + assert res.status_code == 200 + body = res.get_json() + assert len(body["results"]) == 1 + assert body["skipped_count"] == 1 + + +class TestRowCap: + def test_row_cap_returns_typed_error(self, client, fake_state, monkeypatch): + monkeypatch.setenv("BULK_PREDICT_MAX_ROWS", "2") + csv_text = "text\na message\nb message\nc message\n" + res = _upload(client, csv_text) + assert res.status_code == 413 + body = res.get_json() + assert body["error_detail"]["code"] == "BULK_TOO_MANY_ROWS" + + def test_under_cap_is_accepted(self, client, fake_state, monkeypatch): + monkeypatch.setenv("BULK_PREDICT_MAX_ROWS", "10") + csv_text = "text\na message\nb message\n" + res = _upload(client, csv_text) + assert res.status_code == 200 + + +class TestVersionStamping: + def test_json_response_carries_model_version(self, client, fake_state): + res = _upload(client, "text\nsome message content\n") + assert res.status_code == 200 + assert res.get_json()["model_version"] == fake_state.version + + def test_export_carries_model_version(self, client, fake_state): + res = _upload( + client, "text\nsome message content\n", path="/bulk-predict/export" + ) + assert res.status_code == 200 + assert res.headers["X-Model-Version"] == str(fake_state.version) + body = res.data.decode("utf-8") + assert "model_version" in body.splitlines()[0] + + +def _parse_ndjson(raw): + return [ + json.loads(line) for line in raw.decode("utf-8").splitlines() if line.strip() + ] + + +class TestNdjsonStreaming: + def test_stream_query_param_shape(self, client, fake_state, monkeypatch): + monkeypatch.setenv("BULK_PREDICT_MAX_ROW_LEN", "50") + csv_text = "text\n" "this is a long spam message\n" "hi\n" + f"{_POISON} row\n" + res = _upload(client, csv_text, path="/bulk-predict?stream=ndjson") + assert res.status_code == 200 + assert res.headers["Content-Type"].startswith("application/x-ndjson") + assert res.headers["X-Model-Version"] == str(fake_state.version) + + records = _parse_ndjson(res.data) + assert records[0]["type"] == "meta" + assert records[0]["model_version"] == fake_state.version + + results = [r for r in records if r["type"] == "result"] + skipped = [r for r in records if r["type"] == "skipped"] + summary = records[-1] + assert {r["message"] for r in results} == { + "this is a long spam message", + "hi", + } + assert any(s["reason"] == "BULK_ROW_UNPROCESSABLE" for s in skipped) + assert summary["type"] == "summary" + assert summary["total_messages"] == len(results) + assert summary["skipped_count"] == len(skipped) + + def test_stream_via_accept_header(self, client, fake_state): + res = _upload( + client, + "text\nsome message content\n", + headers={"Accept": "application/x-ndjson"}, + ) + assert res.status_code == 200 + assert res.headers["Content-Type"].startswith("application/x-ndjson") + records = _parse_ndjson(res.data) + assert records[0]["type"] == "meta" + + def test_row_cap_enforced_before_stream(self, client, fake_state, monkeypatch): + monkeypatch.setenv("BULK_PREDICT_MAX_ROWS", "1") + csv_text = "text\na message\nb message\n" + res = _upload(client, csv_text, path="/bulk-predict?stream=ndjson") + assert res.status_code == 413 + assert res.get_json()["error_detail"]["code"] == "BULK_TOO_MANY_ROWS" + + def test_non_streaming_is_default(self, client, fake_state): + res = _upload(client, "text\nsome message content\n") + assert res.headers["Content-Type"].startswith("application/json") + assert "results" in res.get_json()