diff --git a/app.py b/app.py index b3cb059..234464c 100644 --- a/app.py +++ b/app.py @@ -16,6 +16,7 @@ from fastapi.templating import Jinja2Templates from config import VERSION, RATE_WINDOW_SECONDS, UPLOAD_DIR from db import init_db +from hardware import assess_hardware from memory import get_memory_count from security import ( get_client_ip, is_ip_allowed, check_rate_limit, rate_policy, @@ -35,6 +36,7 @@ import routers.search_route as search_route import routers.completions as completions import routers.upload as upload import routers.ingest as ingest +import routers.hardware as hardware # --- Logging --- log = logging.getLogger("jarvischat") @@ -53,6 +55,7 @@ async def lifespan(app: FastAPI): os.makedirs(UPLOAD_DIR, exist_ok=True) init_db() log.info(f"Memory system: {get_memory_count()} memories loaded") + await assess_hardware() yield log.info("JarvisChat shutting down") @@ -103,7 +106,7 @@ async def session_auth_middleware(request: Request, call_next): unauth_paths = { "/api/auth/login", "/api/auth/logout", "/api/auth/session", - "/api/auth/heartbeat", "/api/auth/guest", "/api/ingest", + "/api/auth/heartbeat", "/api/auth/guest", "/api/ingest", "/api/hardware", } if path.startswith("/api/"): @@ -142,7 +145,7 @@ async def index(request: Request): for router_module in [ auth_router, conversations.router, memories.router, models.router, presets.router, profile.router, settings.router, skills.router, - chat.router, search_route.router, completions.router, upload.router, ingest.router, + chat.router, search_route.router, completions.router, upload.router, ingest.router, hardware.router, ]: app.include_router(router_module) diff --git a/hardware.py b/hardware.py new file mode 100644 index 0000000..69d8346 --- /dev/null +++ b/hardware.py @@ -0,0 +1,102 @@ +""" +JarvisChat — Startup hardware self-assessment. +""" +import asyncio +import json +import logging +import subprocess +from pathlib import Path + +import httpx +import psutil + +from config import LLAMA_SERVER_BASE, SEARXNG_BASE + +log = logging.getLogger("jarvischat") + +HARDWARE_STATE_PATH = Path("hardware_state.json") +_TIMEOUT_EXPIRED = subprocess.TimeoutExpired + + +async def assess_hardware() -> dict: + mem = psutil.virtual_memory() + ram_total_gb = round(mem.total / (1024 ** 3), 1) + ram_available_gb = round(mem.available / (1024 ** 3), 1) + cpu_count = psutil.cpu_count() + + vram_total_mb = 0 + vram_free_mb = 0 + try: + result = subprocess.run( + ["rocm-smi", "--showmeminfo", "vram", "--json"], + capture_output=True, text=True, timeout=5, + ) + if result.returncode == 0: + data = json.loads(result.stdout) + for card, info in data.items(): + tot = info.get("VRAM Total (MB)", 0) + free = info.get("VRAM Free (MB)", None) + used = info.get("VRAM Used (MB)", 0) + if tot: + vram_total_mb += int(tot) + if free is not None: + vram_free_mb += int(free) + else: + vram_free_mb += int(tot) - int(used) + except (FileNotFoundError, _TIMEOUT_EXPIRED, json.JSONDecodeError): + log.warning("rocm-smi not available or failed — VRAM stats set to 0") + except Exception as e: + log.warning(f"rocm-smi error: {e}") + + llama_reachable = False + llama_models = [] + try: + async with httpx.AsyncClient(timeout=3) as client: + resp = await client.get(f"{LLAMA_SERVER_BASE}/v1/models") + if resp.status_code == 200: + llama_reachable = True + data = resp.json() + llama_models = [m.get("id", "") for m in data.get("data", [])] + except Exception: + log.warning("llama-server not reachable") + + qdrant_reachable = False + qdrant_collections = [] + try: + async with httpx.AsyncClient(timeout=3) as client: + resp = await client.get("http://192.168.50.108:6333/collections") + if resp.status_code == 200: + qdrant_reachable = True + data = resp.json() + raw = data.get("result", {}).get("collections", []) + qdrant_collections = [c.get("name", "") for c in raw] + except Exception: + log.warning("Qdrant not reachable") + + searxng_reachable = False + try: + async with httpx.AsyncClient(timeout=3) as client: + resp = await client.get(SEARXNG_BASE) + if resp.status_code == 200: + searxng_reachable = True + except Exception: + log.warning("SearXNG not reachable") + + state = { + "ram_total_gb": ram_total_gb, + "ram_available_gb": ram_available_gb, + "cpu_count": cpu_count, + "vram_total_mb": vram_total_mb, + "vram_free_mb": vram_free_mb, + "llama_reachable": llama_reachable, + "llama_models": llama_models, + "qdrant_reachable": qdrant_reachable, + "qdrant_collections": qdrant_collections, + "searxng_reachable": searxng_reachable, + } + HARDWARE_STATE_PATH.write_text(json.dumps(state, indent=2)) + log.info( + f"HW: {ram_total_gb}GB RAM, {vram_total_mb}MB VRAM, " + f"llama={llama_reachable}, qdrant={qdrant_reachable}, searxng={searxng_reachable}" + ) + return state diff --git a/hardware_state.json b/hardware_state.json new file mode 100644 index 0000000..5551cf1 --- /dev/null +++ b/hardware_state.json @@ -0,0 +1,16 @@ +{ + "ram_total_gb": 15.2, + "ram_available_gb": 13.4, + "cpu_count": 16, + "vram_total_mb": 0, + "vram_free_mb": 0, + "llama_reachable": true, + "llama_models": [ + "/home/gramps/models/Mistral-Nemo-Instruct-2407-Q4_K_M.gguf" + ], + "qdrant_reachable": true, + "qdrant_collections": [ + "jarvis_rag" + ], + "searxng_reachable": false +} \ No newline at end of file diff --git a/routers/hardware.py b/routers/hardware.py new file mode 100644 index 0000000..4c483b1 --- /dev/null +++ b/routers/hardware.py @@ -0,0 +1,15 @@ +"""JarvisChat routers — Hardware self-assessment endpoint.""" +import json + +from fastapi import APIRouter + +import hardware + +router = APIRouter() + + +@router.get("/api/hardware") +async def get_hardware_state(): + if hardware.HARDWARE_STATE_PATH.exists(): + return json.loads(hardware.HARDWARE_STATE_PATH.read_text()) + return {"status": "not_ready", "message": "Hardware assessment not yet complete"} diff --git a/tests/test_hardware.py b/tests/test_hardware.py new file mode 100644 index 0000000..c1a47ff --- /dev/null +++ b/tests/test_hardware.py @@ -0,0 +1,176 @@ +import asyncio +import json +import os +import subprocess +from pathlib import Path + +import httpx +import psutil +from fastapi.testclient import TestClient + +import app as app_module +import config +import db +import hardware +from security import SESSIONS, PIN_ATTEMPTS, RATE_EVENTS + + +def make_client(tmp_path: Path) -> TestClient: + os.environ["JARVISCHAT_ADMIN_PIN"] = "1234" + db.DB_PATH = tmp_path / "jarvischat-hardware.db" + hardware.HARDWARE_STATE_PATH = tmp_path / "hardware_state.json" + SESSIONS.clear() + PIN_ATTEMPTS.clear() + RATE_EVENTS.clear() + db.init_db() + return TestClient(app_module.app, raise_server_exceptions=False) + + +def _guest_headers(client: TestClient) -> dict: + sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()["session_id"] + return {"X-Session-ID": sid, "Origin": "http://testserver"} + + +class _MockGet: + def __init__(self, status_code: int = 200, json_data: dict | None = None): + self.status_code = status_code + self._json_data = json_data or {} + + def json(self): + return self._json_data + + +def _mock_subprocess(rocm_stdout: str = "") -> object: + class MockProc: + returncode = 0 + stdout = rocm_stdout + + class MockSP: + TimeoutExpired = subprocess.TimeoutExpired + run = staticmethod(lambda cmd, **kw: MockProc()) + + return MockSP() + + +def _broken_subprocess(exception: Exception) -> object: + class MockSP: + TimeoutExpired = subprocess.TimeoutExpired + run = staticmethod(lambda cmd, **kw: (_ for _ in ()).throw(exception)) + + return MockSP() + + +def test_assess_hardware_all_services_reachable(tmp_path: Path, monkeypatch): + hardware.HARDWARE_STATE_PATH = tmp_path / "hardware_state.json" + monkeypatch.setattr(psutil, "virtual_memory", lambda: type("M", (), {"total": 16 * 1024 ** 3, "available": 8 * 1024 ** 3})()) + monkeypatch.setattr(psutil, "cpu_count", lambda: 8) + monkeypatch.setattr(hardware, "subprocess", _mock_subprocess( + json.dumps({"card0": {"VRAM Total (MB)": 8192, "VRAM Free (MB)": 4096}}) + )) + + async def mock_get(self, url, *args, **kwargs): + if "v1/models" in url: + return _MockGet(200, {"data": [{"id": "mistral-nemo:latest"}]}) + if "6333" in url: + return _MockGet(200, {"result": {"collections": [{"name": "jarvischat"}]}}) + if "8888" in url: + return _MockGet(200, {}) + return _MockGet(200, {}) + monkeypatch.setattr(httpx.AsyncClient, "get", mock_get) + + state = asyncio.run(hardware.assess_hardware()) + + assert state["ram_total_gb"] == 16.0 + assert state["ram_available_gb"] == 8.0 + assert state["cpu_count"] == 8 + assert state["vram_total_mb"] == 8192 + assert state["vram_free_mb"] == 4096 + assert state["llama_reachable"] is True + assert state["llama_models"] == ["mistral-nemo:latest"] + assert state["qdrant_reachable"] is True + assert state["qdrant_collections"] == ["jarvischat"] + assert state["searxng_reachable"] is True + assert tmp_path.joinpath("hardware_state.json").exists() + + +def test_assess_hardware_rocm_smi_absent(tmp_path: Path, monkeypatch): + hardware.HARDWARE_STATE_PATH = tmp_path / "hardware_state.json" + monkeypatch.setattr(psutil, "virtual_memory", lambda: type("M", (), {"total": 16 * 1024 ** 3, "available": 8 * 1024 ** 3})()) + monkeypatch.setattr(psutil, "cpu_count", lambda: 8) + monkeypatch.setattr(hardware, "subprocess", _broken_subprocess(FileNotFoundError("no rocm-smi"))) + + async def mock_get(self, url, *args, **kwargs): + if "v1/models" in url: + return _MockGet(200, {"data": []}) + if "6333" in url: + return _MockGet(200, {"result": {"collections": []}}) + if "8888" in url: + return _MockGet(200, {}) + return _MockGet(200, {}) + monkeypatch.setattr(httpx.AsyncClient, "get", mock_get) + + state = asyncio.run(hardware.assess_hardware()) + + assert state["vram_total_mb"] == 0 + assert state["vram_free_mb"] == 0 + assert state["llama_reachable"] is True + assert state["qdrant_reachable"] is True + assert state["searxng_reachable"] is True + + +def test_assess_hardware_llama_unreachable(tmp_path: Path, monkeypatch): + hardware.HARDWARE_STATE_PATH = tmp_path / "hardware_state.json" + monkeypatch.setattr(psutil, "virtual_memory", lambda: type("M", (), {"total": 16 * 1024 ** 3, "available": 8 * 1024 ** 3})()) + monkeypatch.setattr(psutil, "cpu_count", lambda: 8) + monkeypatch.setattr(hardware, "subprocess", _mock_subprocess( + json.dumps({"card0": {"VRAM Total (MB)": 8192, "VRAM Free (MB)": 4096}}) + )) + + async def mock_get(self, url, *args, **kwargs): + if "v1/models" in url: + raise httpx.ConnectError("refused") + if "6333" in url: + return _MockGet(200, {"result": {"collections": [{"name": "jarvischat"}]}}) + if "8888" in url: + return _MockGet(200, {}) + return _MockGet(200, {}) + monkeypatch.setattr(httpx.AsyncClient, "get", mock_get) + + state = asyncio.run(hardware.assess_hardware()) + + assert state["llama_reachable"] is False + assert state["llama_models"] == [] + assert state["qdrant_reachable"] is True + assert state["searxng_reachable"] is True + + +def test_get_hardware_endpoint(tmp_path: Path, monkeypatch): + hardware.HARDWARE_STATE_PATH = tmp_path / "hardware_state.json" + monkeypatch.setattr(psutil, "virtual_memory", lambda: type("M", (), {"total": 16 * 1024 ** 3, "available": 8 * 1024 ** 3})()) + monkeypatch.setattr(psutil, "cpu_count", lambda: 8) + monkeypatch.setattr(hardware, "subprocess", _mock_subprocess( + json.dumps({"card0": {"VRAM Total (MB)": 8192, "VRAM Free (MB)": 4096}}) + )) + + async def mock_get(self, url, *args, **kwargs): + if "v1/models" in url: + return _MockGet(200, {"data": [{"id": "mistral-nemo:latest"}]}) + if "6333" in url: + return _MockGet(200, {"result": {"collections": [{"name": "jarvischat"}]}}) + if "8888" in url: + return _MockGet(200, {}) + return _MockGet(200, {}) + monkeypatch.setattr(httpx.AsyncClient, "get", mock_get) + + with make_client(tmp_path) as client: + resp = client.get("/api/hardware", headers=_guest_headers(client)) + assert resp.status_code == 200 + data = resp.json() + assert data["ram_total_gb"] == 16.0 + assert data["cpu_count"] == 8 + assert data["vram_total_mb"] == 8192 + assert data["llama_reachable"] is True + assert data["qdrant_reachable"] is True + assert data["searxng_reachable"] is True + assert "llama_models" in data + assert "qdrant_collections" in data