From 552d154af7ac015c5c6140ac59d533e6f10d339f Mon Sep 17 00:00:00 2001 From: Himanshu Kumar Date: Sun, 19 Jul 2026 15:37:52 +0530 Subject: [PATCH] feat(ollama): detect and report GPU vs CPU execution (issue #11) Add a RuntimeStatus domain model with a device field (gpu/cpu/unknown) and a status() method on RuntimePort. OllamaRuntime detects the device from ollama ps VRAM usage (size_vram > 0 => GPU). The load CLI now prints the device after a successful load. Co-Authored-By: Claude --- src/modeldock/adapters/runtimes/base.py | 23 +++++++- src/modeldock/adapters/runtimes/ollama.py | 44 +++++++++++++- src/modeldock/cli/commands/load.py | 2 + src/modeldock/core/manager.py | 4 ++ src/modeldock/domain/model.py | 29 +++++++++ src/modeldock/ports/runtime.py | 6 +- tests/conftest.py | 10 +++- tests/integration/test_ollama.py | 46 ++++++++++++++- tests/unit/test_ollama_runtime.py | 72 ++++++++++++++++++++++- tests/unit/test_port_contract.py | 9 ++- 10 files changed, 238 insertions(+), 7 deletions(-) diff --git a/src/modeldock/adapters/runtimes/base.py b/src/modeldock/adapters/runtimes/base.py index 5a22c6f..646a55d 100644 --- a/src/modeldock/adapters/runtimes/base.py +++ b/src/modeldock/adapters/runtimes/base.py @@ -13,7 +13,7 @@ RuntimeUnavailableError, ) from modeldock.common.logging import get_logger -from modeldock.domain.model import ModelRef, ModelSpec, RuntimeBackend +from modeldock.domain.model import Device, ModelRef, ModelSpec, RuntimeBackend, RuntimeStatus from modeldock.ports.runtime import PullResult, RunResult _AVAILABILITY_TTL = 5.0 @@ -55,6 +55,27 @@ def default_tag_for(self, spec: ModelSpec) -> str: """Resolve the default tag for a spec (runtime may override).""" return spec.default_tag + def status(self) -> RuntimeStatus: + """Report availability and execution device. + + The base default reports availability and an ``UNKNOWN`` device; concrete + runtimes override ``_detect_device`` to populate the real device. + """ + available = self.is_available() + device = self._detect_device() if available else Device.UNKNOWN + return RuntimeStatus( + backend=self.backend, + available=available, + device=device, + ) + + def _detect_device(self) -> Device: + """Best-effort device detection; ``UNKNOWN`` by default. + + Concrete runtimes override this to inspect loaded-model metadata. + """ + return Device.UNKNOWN + def pull(self, ref: ModelRef, progress: Any = None) -> PullResult: """Normalized pull: checks availability, delegates to ``_do_pull``.""" if not self.is_available(): diff --git a/src/modeldock/adapters/runtimes/ollama.py b/src/modeldock/adapters/runtimes/ollama.py index f47dd95..8ded235 100644 --- a/src/modeldock/adapters/runtimes/ollama.py +++ b/src/modeldock/adapters/runtimes/ollama.py @@ -17,7 +17,7 @@ ModelNotInstalledError, RuntimeUnavailableError, ) -from modeldock.domain.model import ModelRef, RuntimeBackend +from modeldock.domain.model import Device, ModelRef, RuntimeBackend from modeldock.ports.runtime import PullResult, RunResult _OLLAMA_DEFAULT_HOST = "http://localhost:11434" @@ -271,5 +271,47 @@ def _write(text: str) -> None: sys.stdout.write(text) sys.stdout.flush() + # --- device / status -------------------------------------------------- + + def _detect_device(self) -> Device: + """Detect GPU vs CPU execution from ``ollama ps`` VRAM usage. + + Ollama reports ``size_vram`` for each loaded model. A positive value + means the model is resident on the GPU; zero/None means CPU execution. + When no model is loaded we cannot tell, so report ``UNKNOWN``. + """ + try: + client = self._ensure_client() + response = client.ps() + except Exception: + # Device detection is best-effort; never fail the status probe. + return Device.UNKNOWN + models = self._extract_ps_models(response) + if not models: + return Device.UNKNOWN + for entry in models: + vram = self._extract_size_vram(entry) + if vram and int(vram) > 0: + return Device.GPU + return Device.CPU + + @staticmethod + def _extract_ps_models(response: Any) -> List[Any]: + """Pull the ``models`` list out of a ps() response (object or dict).""" + if hasattr(response, "models"): + return list(getattr(response, "models", []) or []) + if isinstance(response, dict): + return list(response.get("models", []) or []) + return [] + + @staticmethod + def _extract_size_vram(entry: Any) -> int: + """Extract ``size_vram`` (VRAM bytes) from a ps() model entry.""" + if hasattr(entry, "size_vram"): + return int(entry.size_vram or 0) + if isinstance(entry, dict): + return int(entry.get("size_vram") or 0) + return 0 + __all__ = ["OllamaRuntime"] diff --git a/src/modeldock/cli/commands/load.py b/src/modeldock/cli/commands/load.py index 7ebcf04..3e2072e 100644 --- a/src/modeldock/cli/commands/load.py +++ b/src/modeldock/cli/commands/load.py @@ -23,7 +23,9 @@ def load_cmd( mgr = ModelManager() name = f"{model}:{tag}" if tag != "latest" else model client = mgr.load(name, auto_install=auto_install) + status = mgr.runtime_status() typer.echo(f"Loaded {name} -> {type(client).__name__}") + typer.echo(f"Device: {status.device.value}") except Exception as exc: # noqa: BLE001 - top-level CLI boundary print_error(exc, debug) raise typer.Exit(code=1) # noqa: B904 diff --git a/src/modeldock/core/manager.py b/src/modeldock/core/manager.py index f456c2b..cecbfe1 100644 --- a/src/modeldock/core/manager.py +++ b/src/modeldock/core/manager.py @@ -118,6 +118,10 @@ def categories(self) -> List[Category]: """Return all catalog categories.""" return self._registry.categories() + def runtime_status(self) -> Any: + """Report the active runtime's availability and execution device.""" + return self._runtime.status() + def recommend(self, task: str) -> List[Any]: """Recommend models for a task.""" return self._registry.recommend(task) diff --git a/src/modeldock/domain/model.py b/src/modeldock/domain/model.py index 05f8ae5..ba2296c 100644 --- a/src/modeldock/domain/model.py +++ b/src/modeldock/domain/model.py @@ -144,6 +144,35 @@ def from_spec(cls, spec: ModelSpec, installed_tags: List[str]) -> ModelInfo: ) +class Device(str, Enum): + """Execution device a runtime reports for a loaded model. + + ``UNKNOWN`` is used when the runtime cannot determine the device (e.g. the + model is not loaded, or the runtime exposes no device metadata). + """ + + GPU = "gpu" + CPU = "cpu" + UNKNOWN = "unknown" + + +class RuntimeStatus(BaseModel): + """Runtime execution status, including the device a model runs on. + + Pure data — no I/O. Adapters populate ``device`` from runtime metadata + (e.g. Ollama ``ps`` VRAM usage). See issue #11. + """ + + backend: RuntimeBackend + available: bool = False + device: Device = Device.UNKNOWN + loaded_models: List[str] = Field(default_factory=list) + details: str = "" + + def __repr__(self) -> str: + return f"RuntimeStatus({self.backend.value}, {self.device.value})" + + class ModelRef(BaseModel): """A concrete reference to a model: name plus optional tag and backend.""" diff --git a/src/modeldock/ports/runtime.py b/src/modeldock/ports/runtime.py index c44c38b..56f0479 100644 --- a/src/modeldock/ports/runtime.py +++ b/src/modeldock/ports/runtime.py @@ -9,7 +9,7 @@ from pathlib import Path from typing import Any, List, Optional, Protocol, runtime_checkable -from modeldock.domain.model import ModelRef, ModelSpec, RuntimeBackend +from modeldock.domain.model import ModelRef, ModelSpec, RuntimeBackend, RuntimeStatus @runtime_checkable @@ -49,6 +49,10 @@ def default_tag_for(self, spec: ModelSpec) -> str: """Resolve the default variant tag for a model spec.""" ... + def status(self) -> RuntimeStatus: + """Report runtime availability and the execution device (GPU/CPU).""" + ... + def run(self, ref: ModelRef, prompt: Optional[str] = None, **opts: Any) -> RunResult: """Run an interactive session for ``ref``, streaming output. diff --git a/tests/conftest.py b/tests/conftest.py index 027aabd..43fce64 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -13,6 +13,7 @@ from modeldock.domain.model import ( Capability, Category, + Device, ModelRef, ModelSpec, RuntimeBackend, @@ -21,7 +22,7 @@ from modeldock.ports.events import EventPort from modeldock.ports.progress import ProgressPort from modeldock.ports.registry import RegistryPort -from modeldock.ports.runtime import PullResult, RuntimePort +from modeldock.ports.runtime import PullResult, RuntimePort, RuntimeStatus class FakeProgress(ProgressPort): @@ -142,6 +143,13 @@ def get_model_client(self, ref: ModelRef) -> Any: def default_tag_for(self, spec: ModelSpec) -> str: return spec.default_tag + def status(self) -> RuntimeStatus: + return RuntimeStatus( + backend=self.backend, + available=self._available, + device=Device.UNKNOWN, + ) + def run(self, ref: ModelRef, prompt: Optional[str] = None, **opts: Any) -> Any: from modeldock.ports.runtime import RunResult diff --git a/tests/integration/test_ollama.py b/tests/integration/test_ollama.py index d056cf8..11a22ab 100644 --- a/tests/integration/test_ollama.py +++ b/tests/integration/test_ollama.py @@ -11,7 +11,7 @@ import pytest from modeldock.adapters.runtimes.registry import RuntimeRegistry -from modeldock.domain.model import ModelRef, RuntimeBackend +from modeldock.domain.model import Device, ModelRef, RuntimeBackend, RuntimeStatus ollama_present = shutil.which("ollama") is not None try: @@ -57,3 +57,47 @@ def test_ollama_pull_and_remove() -> None: assert runtime.is_installed(ref) runtime.remove(ref) assert not runtime.is_installed(ref) + + +@pytest.mark.skipif( + not (ollama_present and sdk_present), + reason="Ollama CLI and Python SDK not both available", +) +def test_ollama_status_device_detection() -> None: + """status() reports a valid device and matches `ollama ps` (issue #11).""" + import subprocess + + runtime = _runtime() + if not runtime.is_available(): + pytest.skip("Ollama daemon not running") + + status = runtime.status() + assert isinstance(status, RuntimeStatus) + assert status.available is True + assert status.device in (Device.GPU, Device.CPU, Device.UNKNOWN) + + # Load a small installed model so it becomes resident, then re-check. + # Skip cloud/subscription models (tag contains "cloud") which can't run locally. + installed = [r for r in runtime.list_installed() if "cloud" not in r.tag] + if not installed: + pytest.skip("No locally-runnable models installed to load") + ref = installed[0] + client = runtime._ensure_client() + try: + for _ in client.generate(ref.qualified_name(), prompt="ping", stream=True): + pass + except Exception: + pytest.skip("Could not load a model for device detection") + + status = runtime.status() + # Cross-check against `ollama ps` PROCESSOR column when available. + try: + out = subprocess.run( + ["ollama", "ps"], capture_output=True, text=True, timeout=10, check=True + ).stdout + except Exception: + out = "" + if ref.name in out and "CPU" in out: + assert status.device is Device.CPU + elif ref.name in out and "GPU" in out: + assert status.device is Device.GPU diff --git a/tests/unit/test_ollama_runtime.py b/tests/unit/test_ollama_runtime.py index 0c64a0f..c9e0c1e 100644 --- a/tests/unit/test_ollama_runtime.py +++ b/tests/unit/test_ollama_runtime.py @@ -9,7 +9,7 @@ from modeldock.adapters.runtimes.ollama import OllamaRuntime from modeldock.common.errors import DownloadError, ModelNotInstalledError, RuntimeUnavailableError -from modeldock.domain.model import ModelRef +from modeldock.domain.model import Device, ModelRef, RuntimeStatus class _PullClient: @@ -290,3 +290,73 @@ def generate(self, model: str, prompt: str = "", stream: bool = False, **opts: A result = runtime.run(ModelRef.parse("llama3"), prompt="hi") assert not result.success assert "boom" in result.error + + +# --- device / GPU vs CPU detection (issue #11) ------------------------------ + + +class _PsClient: + """Fake ollama.Client exposing list() + ps() for device detection.""" + + def __init__(self, ps_models: List[dict[str, Any]]) -> None: + self._ps_models = ps_models + + def list(self) -> dict[str, Any]: + return {"models": [{"name": "llama3:latest"}]} + + def ps(self) -> dict[str, Any]: + return {"models": self._ps_models} + + +def _runtime_with_ps(ps_models: List[dict[str, Any]]) -> OllamaRuntime: + runtime = OllamaRuntime() + runtime._client = _PsClient(ps_models=ps_models) + return runtime + + +def test_status_reports_gpu_when_model_has_vram() -> None: + runtime = _runtime_with_ps([{"name": "llama3:latest", "size_vram": 4_500_000_000}]) + status = runtime.status() + assert isinstance(status, RuntimeStatus) + assert status.available is True + assert status.device is Device.GPU + + +def test_status_reports_cpu_when_no_vram() -> None: + runtime = _runtime_with_ps([{"name": "llama3:latest", "size_vram": 0}]) + status = runtime.status() + assert status.available is True + assert status.device is Device.CPU + + +def test_status_unknown_when_no_models_loaded() -> None: + runtime = _runtime_with_ps([]) + status = runtime.status() + assert status.device is Device.UNKNOWN + + +def test_status_unknown_when_ps_raises() -> None: + class _PsBoomClient: + def list(self) -> dict[str, Any]: + return {"models": [{"name": "llama3:latest"}]} + + def ps(self) -> dict[str, Any]: + raise RuntimeError("ps unavailable") + + runtime = OllamaRuntime() + runtime._client = _PsBoomClient() + status = runtime.status() + # Best-effort: device unknown, but availability still reported. + assert status.available is True + assert status.device is Device.UNKNOWN + + +def test_status_unavailable_when_daemon_down() -> None: + class _DownClient: + def list(self) -> dict[str, Any]: + raise RuntimeError("connection refused") + + runtime = _runtime_with(_DownClient()) + status = runtime.status() + assert status.available is False + assert status.device is Device.UNKNOWN diff --git a/tests/unit/test_port_contract.py b/tests/unit/test_port_contract.py index c3fd212..bf6cda0 100644 --- a/tests/unit/test_port_contract.py +++ b/tests/unit/test_port_contract.py @@ -14,7 +14,7 @@ from modeldock.domain.model import Category, ModelRef, ModelSpec from modeldock.ports.cache import CachePort -from modeldock.ports.runtime import PullResult, RuntimePort +from modeldock.ports.runtime import PullResult, RuntimePort, RuntimeStatus # --- RuntimePort contract --------------------------------------------------- @@ -67,6 +67,13 @@ def test_runtime_default_tag(runtime_impl: RuntimePort) -> None: assert runtime_impl.default_tag_for(spec) == spec.default_tag +def test_runtime_status_contract(runtime_impl: RuntimePort) -> None: + status = runtime_impl.status() + assert isinstance(status, RuntimeStatus) + assert status.backend is not None + assert status.device.value in {"gpu", "cpu", "unknown"} + + # --- CachePort contract -----------------------------------------------------