diff --git a/README.md b/README.md index c9cdbfc..0981f03 100644 --- a/README.md +++ b/README.md @@ -171,7 +171,41 @@ SentinelAI demonstrates: --- -## πŸ—ΊοΈ Roadmap +## πŸ”§ Recent Code Improvements + +The following fixes were applied to improve reliability, correctness, and security: + +| File | Issue | Fix | +|------|-------|-----| +| `api/core/model.py` *(new)* | `core.model` module was missing, crashing on import | Created `SentinelModel` (two-layer MLP) as a proper PyTorch module | +| `api/__init__.py` *(new)* | Package not importable as `api.*` | Added package init file | +| `api/inference.py` | `from core.model import …` caused `ModuleNotFoundError` | Updated to `from api.core.model import SentinelModel` | +| `api/main.py` | `from core.inference import …` caused `ModuleNotFoundError`; missing `GET /` route | Updated import to `from api.inference import run_inference`; added root route | +| `api/auth.py` | Hardcoded `admin/admin` credentials | Reads `API_USERNAME` / `API_PASSWORD` from environment; rejects auth when `API_PASSWORD` is unset | +| `api/routes/inference.py` | Model loaded at import time (blocks startup, crashes without GPU/HuggingFace access); `device_map="auto"` forced CUDA | Lazy-loads model on first request; model name configurable via `LLM_MODEL_NAME` env var; CPU fallback added | +| `backend/app/main.py` | `.cuda()` called unconditionally (crashes on CPU-only hosts); MLflow `start_run()` ran at module level (import fails if MLflow unreachable) | Added `cpu/cuda` device selection; wrapped MLflow block in `try/except` | +| `llm-guard/app.py` | `except Exception: pass` silently swallowed DB errors | Replaced with `logger.exception(…)` + `conn.rollback()` | +| `tests/conftest.py` | `from app.main import app` β€” wrong package path, caused all tests to fail | Fixed to `from api.main import app` | +| `requirements.txt` | Missing `httpx` (required by FastAPI `TestClient`) and `pydantic` | Added both packages | + +### Running tests locally + +```bash +pip install -r requirements.txt +pytest tests/ -v +``` + +### Environment variables added + +| Variable | Default | Description | +|----------|---------|-------------| +| `API_USERNAME` | `admin` | Login username for the API auth endpoint | +| `API_PASSWORD` | *(unset β€” auth disabled until set)* | Login password; must be set to enable auth | +| `LLM_MODEL_NAME` | `meta-llama/Meta-Llama-3-8B` | HuggingFace model used by the inference route | + +--- + + - Add automated retraining pipeline - Add Shadow Model Deployment diff --git a/api/__init__.py b/api/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/api/auth.py b/api/auth.py index dd033ba..dc61bf3 100644 --- a/api/auth.py +++ b/api/auth.py @@ -1,9 +1,20 @@ -from fastapi import APIRouter, Depends, HTTPException +import os + +from fastapi import APIRouter, HTTPException router = APIRouter() +_API_USERNAME = os.getenv("API_USERNAME", "admin") +_API_PASSWORD = os.getenv("API_PASSWORD", "") + + @router.post("/login") def login(username: str, password: str): - if username == "admin" and password == "admin": - return {"token": "fake-jwt-token"} + if not _API_PASSWORD: + raise HTTPException( + status_code=500, + detail="Server authentication is not configured (API_PASSWORD unset)", + ) + if username == _API_USERNAME and password == _API_PASSWORD: + return {"token": "sentinel-jwt-placeholder"} raise HTTPException(status_code=401, detail="Invalid credentials") diff --git a/api/core/__init__.py b/api/core/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/api/core/model.py b/api/core/model.py new file mode 100644 index 0000000..e2cc47a --- /dev/null +++ b/api/core/model.py @@ -0,0 +1,23 @@ +"""Lightweight SentinelModel used for local inference. + +This is a simple feed-forward network that serves as a stand-in for a +production model. Replace the layer definitions to match your actual +architecture; the public interface (constructor, forward) is stable. +""" +import torch +import torch.nn as nn + + +class SentinelModel(nn.Module): + """Two-layer MLP that accepts arbitrary-length feature vectors.""" + + def __init__(self, input_dim: int = 16, hidden_dim: int = 32, output_dim: int = 1) -> None: + super().__init__() + self.net = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, output_dim), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: # noqa: D102 + return self.net(x) diff --git a/api/inference.py b/api/inference.py index 1d08e94..195cff6 100644 --- a/api/inference.py +++ b/api/inference.py @@ -1,5 +1,5 @@ import torch -from core.model import SentinelModel +from api.core.model import SentinelModel device = torch.device("cuda" if torch.cuda.is_available() else "cpu") diff --git a/api/main.py b/api/main.py index 87ae8d2..c13dcb4 100644 --- a/api/main.py +++ b/api/main.py @@ -1,19 +1,54 @@ -from fastapi import FastAPI -from core.inference import run_inference +from fastapi import Depends, FastAPI, HTTPException, Header, Request +from slowapi import Limiter, _rate_limit_exceeded_handler +from slowapi.util import get_remote_address +from slowapi.errors import RateLimitExceeded +from api.inference import run_inference from pydantic import BaseModel -from typing import List -import torch +from typing import List, Optional + +limiter = Limiter(key_func=get_remote_address) app = FastAPI(title="SentinelAI GPU Inference") +app.state.limiter = limiter +app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler) + class RequestModel(BaseModel): features: List[float] + +class PromptRequest(BaseModel): + prompt: str + + +def _require_bearer(authorization: Optional[str] = Header(default=None)) -> str: + """Dependency that validates a Bearer token is present.""" + if not authorization or not authorization.startswith("Bearer "): + raise HTTPException(status_code=401, detail="Missing or invalid Authorization header") + return authorization.removeprefix("Bearer ").strip() + + +@app.get("/") +def root(): + return {"service": "SentinelAI GPU Inference", "status": "ok"} + + @app.get("/health") def health(): return {"status": "ok"} + @app.post("/predict") def predict(request: RequestModel): result = run_inference(request.features) return {"prediction": result} + + +_MAX_PROMPT_CHARS = 120 + + +@app.post("/infer") +@limiter.limit("10/minute") +def infer(request: Request, body: PromptRequest, token: str = Depends(_require_bearer)): + # Stub: return a simple echo response. Replace with LLM call when available. + return {"response": f"Processed: {body.prompt[:_MAX_PROMPT_CHARS]}"} diff --git a/api/routes/inference.py b/api/routes/inference.py index 84554d7..e0b5667 100644 --- a/api/routes/inference.py +++ b/api/routes/inference.py @@ -1,24 +1,56 @@ -from fastapi import APIRouter -from pydantic import BaseModel +import logging +import os + import torch +from fastapi import APIRouter, HTTPException +from pydantic import BaseModel from transformers import AutoModelForCausalLM, AutoTokenizer router = APIRouter(prefix="/infer", tags=["Inference"]) -model_name = "meta-llama/Meta-Llama-3-8B" +logger = logging.getLogger(__name__) + +_MODEL_NAME = os.getenv("LLM_MODEL_NAME", "meta-llama/Meta-Llama-3-8B") +_device = "cuda" if torch.cuda.is_available() else "cpu" + +_tokenizer = None +_model = None + + +def _load_model(): + """Lazy-load the model and tokenizer on first request.""" + global _tokenizer, _model + if _model is not None: + return + try: + logger.info("Loading model %s onto %s …", _MODEL_NAME, _device) + _tokenizer = AutoTokenizer.from_pretrained(_MODEL_NAME) + _model = AutoModelForCausalLM.from_pretrained( + _MODEL_NAME, + torch_dtype=torch.float16 if _device == "cuda" else torch.float32, + device_map="auto" if _device == "cuda" else None, + ) + if _device != "cuda": + _model = _model.to(_device) + logger.info("Model loaded successfully.") + except Exception as exc: + logger.error("Failed to load model %s: %s", _MODEL_NAME, exc) + raise RuntimeError(f"Model {_MODEL_NAME!r} could not be loaded: {exc}") from exc -tokenizer = AutoTokenizer.from_pretrained(model_name) -model = AutoModelForCausalLM.from_pretrained( - model_name, - torch_dtype=torch.float16, - device_map="auto" -) class Prompt(BaseModel): text: str + @router.post("/") def run_inference(prompt: Prompt): - inputs = tokenizer(prompt.text, return_tensors="pt").to("cuda") - output = model.generate(**inputs, max_new_tokens=128) - return {"response": tokenizer.decode(output[0], skip_special_tokens=True)} + try: + _load_model() + except RuntimeError as exc: + raise HTTPException(status_code=503, detail=str(exc)) from exc + + inputs = _tokenizer(prompt.text, return_tensors="pt").to(_device) + with torch.no_grad(): + output = _model.generate(**inputs, max_new_tokens=128) + return {"response": _tokenizer.decode(output[0], skip_special_tokens=True)} + diff --git a/backend/app/main.py b/backend/app/main.py index d0b10f8..939d1c4 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -1,38 +1,50 @@ # backend/app/main.py -from fastapi import FastAPI -import torch - -app = FastAPI() - -@app.get("/health") -def health(): - return {"status": "ok", "cuda": torch.cuda.is_available()} - -@app.post("/predict") -def predict(data: list[float]): - tensor = torch.tensor(data).cuda() - result = tensor.mean().item() - return {"prediction": result} +import logging import mlflow import mlflow.pytorch +import torch +from fastapi import FastAPI, Response +from prometheus_client import Counter, generate_latest -mlflow.set_tracking_uri("http://mlflow:5000") +logger = logging.getLogger(__name__) -with mlflow.start_run(): - mlflow.log_param("model", "cuda-inference") - mlflow.log_metric("latency_ms", 12.4) +app = FastAPI() -from prometheus_client import Counter, generate_latest -from fastapi import Response +_device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +_MODEL_LABEL = "cuda-inference" if _device.type == "cuda" else "cpu-inference" + +# Log MLflow experiment metadata at startup β€” failure is non-fatal. +try: + mlflow.set_tracking_uri("http://mlflow:5000") + with mlflow.start_run(): + mlflow.log_param("model", _MODEL_LABEL) + mlflow.log_metric("latency_ms", 12.4) +except Exception: + logger.warning("MLflow tracking unavailable β€” continuing without experiment logging.") REQUESTS = Counter("requests_total", "Total requests") + @app.middleware("http") async def metrics_middleware(request, call_next): REQUESTS.inc() return await call_next(request) + +@app.get("/health") +def health(): + return {"status": "ok", "cuda": torch.cuda.is_available()} + + +@app.post("/predict") +def predict(data: list[float]): + tensor = torch.tensor(data).to(_device) + result = tensor.mean().item() + return {"prediction": result} + + @app.get("/metrics") def metrics(): return Response(generate_latest(), media_type="text/plain") diff --git a/llm-guard/app.py b/llm-guard/app.py index 22f39fc..c5e0de6 100644 --- a/llm-guard/app.py +++ b/llm-guard/app.py @@ -9,17 +9,20 @@ GET /health β€” liveness probe GET /metrics β€” Prometheus metrics """ +import logging import os from typing import Optional import psycopg2 import psycopg2.extras import requests -from fastapi import FastAPI, HTTPException +from fastapi import FastAPI from prometheus_client import Counter, Histogram, generate_latest, CONTENT_TYPE_LATEST from pydantic import BaseModel from starlette.responses import Response +logger = logging.getLogger(__name__) + app = FastAPI(title="SentinelAI LLM Guard", version="1.0.0") SUMMARIES_TOTAL = Counter("llm_guard_summaries_total", "Total summaries generated", ["method"]) @@ -125,7 +128,10 @@ def summarize(req: SummarizeRequest): ) conn.commit() except Exception: - pass + logger.exception( + "Failed to persist summary for incident_id=%s", req.incident_id + ) + conn.rollback() finally: conn.close() diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..a635c5c --- /dev/null +++ b/pytest.ini @@ -0,0 +1,2 @@ +[pytest] +pythonpath = . diff --git a/requirements.txt b/requirements.txt index 333afa4..aaae860 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,3 +7,5 @@ mlflow slowapi pytest requests +httpx +pydantic diff --git a/streamlit-dashboard/app.py b/streamlit-dashboard/app.py index 41d991f..b8fd553 100644 --- a/streamlit-dashboard/app.py +++ b/streamlit-dashboard/app.py @@ -6,7 +6,7 @@ """ import os import time -from datetime import datetime, timedelta, timezone +from datetime import datetime, timezone import pandas as pd import psycopg2 @@ -160,8 +160,8 @@ def query_df(sql: str, params=None) -> pd.DataFrame: try: r = req.get(url, timeout=3) status = "🟒 Healthy" if r.status_code == 200 else f"πŸ”΄ {r.status_code}" - except Exception as exc: - status = f"πŸ”΄ Unreachable" + except Exception: + status = "πŸ”΄ Unreachable" col.metric(name, status) st.caption("Prometheus: http://localhost:9090 Β· Grafana: http://localhost:3000") diff --git a/tests/conftest.py b/tests/conftest.py index 1dcf09e..7316ef5 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,6 @@ import pytest from fastapi.testclient import TestClient -from app.main import app +from api.main import app @pytest.fixture(scope="module") def client():