mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-30 10:52:15 +00:00
Add Phase 12+13: energy measurement, install polish, PWA, cross-hardware
Phase 12 — Energy Measurement Upgrade: - EnergyMonitor ABC with multi-vendor support (NVIDIA hw counters, AMD amdsmi, Apple zeus-ml, CPU RAPL sysfs) - EnergyBatch batch-level energy-per-token accounting - SteadyStateDetector CV-based thermal equilibrium detection - EnergyBenchmark with warmup phase - InstrumentedEngine prefers EnergyMonitor over legacy GpuMonitor - Telemetry store/aggregator extended with energy fields Phase 13 — Install, Hosting, Cross-Hardware: - jarvis doctor diagnostic command (8 checks, --json output) - jarvis init post-setup guidance with engine-specific next steps - README Quick Start section - MLX engine backend (Apple Silicon → mlx recommendation) - AMD VRAM/multi-GPU detection via rocm-smi - PyTorch MPS device selection in orchestrator trainers - PWA support (vite-plugin-pwa, service worker, manifest, icons) - Server static file serving fix for PWA files - Dockerfile.gpu.rocm + docker-compose.gpu.rocm.yml for ROCm - Eval framework display module and efficiency metrics 2244 tests pass, 37 skipped. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.6
parent
9f5e97eed9
commit
24972e3e52
@@ -4,13 +4,13 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
|
||||
|
||||
## Project Status
|
||||
|
||||
OpenJarvis is a research framework for studying on-device AI systems. Phase 11 (NanoClaw subsumption) complete. Five composable pillars: Intelligence, Engine, Agents, Tools (with storage + MCP), and Learning — with trace-driven learning as a cross-cutting concern. Python SDK (`Jarvis` class), composition layer (`SystemBuilder`/`JarvisSystem`), OpenClaw agent infrastructure, benchmarking framework, Docker deployment all ready. Agent hierarchy refactored: `BaseAgent` (with shared helpers) → `ToolUsingAgent` (tool-using agents), `accepts_tools` introspection, real OpenHands SDK integration. NanoClaw functionality subsumed: `ClaudeCodeAgent` (Claude Agent SDK via Node.js), `WhatsAppBaileysChannel` (Baileys protocol), `ContainerRunner`/`SandboxedAgent` (Docker sandbox), `TaskScheduler` (cron/interval/once with SQLite + MCP tools). ~2078 tests pass (36 skipped for optional deps). Eval framework with 4 benchmarks (SuperGPQA, GAIA, FRAMES, WildChat) and LLM-as-judge scoring via `gpt-5-mini-2025-08-07`. Tool system enriched with shared `build_tool_descriptions()` builder; engine tool_calls normalized across OpenAI, Anthropic, Google, and LiteLLM.
|
||||
OpenJarvis is a research framework for studying on-device AI systems. Phase 13 (Install, Hosting, Cross-Hardware, Eval) complete. Five composable pillars: Intelligence, Engine, Agents, Tools (with storage + MCP), and Learning — with trace-driven learning as a cross-cutting concern. Python SDK (`Jarvis` class), composition layer (`SystemBuilder`/`JarvisSystem`), OpenClaw agent infrastructure, benchmarking framework, Docker deployment all ready. Agent hierarchy refactored: `BaseAgent` (with shared helpers) → `ToolUsingAgent` (tool-using agents), `accepts_tools` introspection, real OpenHands SDK integration. NanoClaw functionality subsumed: `ClaudeCodeAgent` (Claude Agent SDK via Node.js), `WhatsAppBaileysChannel` (Baileys protocol), `ContainerRunner`/`SandboxedAgent` (Docker sandbox), `TaskScheduler` (cron/interval/once with SQLite + MCP tools). ~2244 tests pass (37 skipped for optional deps). Eval framework with 4 benchmarks (SuperGPQA, GAIA, FRAMES, WildChat) and LLM-as-judge scoring via `gpt-5-mini-2025-08-07`. Energy measurement upgraded: `EnergyMonitor` ABC with multi-vendor support (NVIDIA hw counters, AMD amdsmi, Apple Silicon zeus-ml, CPU RAPL sysfs), batch-level energy-per-token accounting (`EnergyBatch`), steady-state detection (`SteadyStateDetector`), `EnergyBenchmark` with warmup phase. Tool system enriched with shared `build_tool_descriptions()` builder; engine tool_calls normalized across OpenAI, Anthropic, Google, and LiteLLM. Phase 13: `jarvis doctor` diagnostic command, `jarvis init` post-setup guidance, MLX engine backend (Apple Silicon → `mlx` recommendation), AMD VRAM/multi-GPU detection, PyTorch MPS device selection, PWA support for browser/desktop hosting, ROCm Docker support.
|
||||
|
||||
## Build & Development Commands
|
||||
|
||||
```bash
|
||||
uv sync --extra dev # Install deps + dev tools
|
||||
uv run pytest tests/ -v # Run ~2078 tests (36 skipped if optional deps missing)
|
||||
uv run pytest tests/ -v # Run ~2244 tests (37 skipped if optional deps missing)
|
||||
uv run ruff check src/ tests/ # Lint
|
||||
uv run jarvis --version # 1.0.0
|
||||
uv run jarvis ask "Hello" # Query via discovered engine (direct mode)
|
||||
@@ -44,7 +44,10 @@ uv run jarvis scheduler start # Start scheduler daemon (foreground)
|
||||
uv run jarvis bench run # Run all benchmarks against engine
|
||||
uv run jarvis bench run -n 20 --json # Run with 20 samples, JSON output
|
||||
uv run jarvis bench run -b latency -o results.jsonl # Specific benchmark to file
|
||||
uv run jarvis bench run -b energy -w 5 -n 20 --json # Energy benchmark with warmup
|
||||
uv run jarvis serve --port 8000 # OpenAI-compatible API server (requires openjarvis[server])
|
||||
uv run jarvis doctor # Run diagnostic checks (config, engines, models, deps)
|
||||
uv run jarvis doctor --json # Machine-readable diagnostics
|
||||
uv run jarvis --help # Show all subcommands
|
||||
uv run jarvis init --force # Detect hardware, write ~/.openjarvis/config.toml
|
||||
# Eval framework
|
||||
@@ -84,8 +87,8 @@ j.close() # Release resources
|
||||
```
|
||||
|
||||
- **Package manager:** `uv` with `hatchling` build backend
|
||||
- **Config:** `pyproject.toml` with extras for optional backends (e.g., `openjarvis[inference-vllm]`, `openjarvis[memory-colbert]`, `openjarvis[server]`, `openjarvis[openclaw]`)
|
||||
- **CLI entry point:** `jarvis` (Click-based) — subcommands: `init`, `ask`, `serve`, `model`, `memory`, `telemetry`, `bench`, `channel`, `scheduler`
|
||||
- **Config:** `pyproject.toml` with extras for optional backends (e.g., `openjarvis[inference-vllm]`, `openjarvis[inference-mlx]`, `openjarvis[memory-colbert]`, `openjarvis[server]`, `openjarvis[openclaw]`, `openjarvis[energy-amd]`, `openjarvis[energy-apple]`, `openjarvis[energy-all]`)
|
||||
- **CLI entry point:** `jarvis` (Click-based) — subcommands: `init`, `ask`, `serve`, `model`, `memory`, `telemetry`, `bench`, `channel`, `scheduler`, `doctor`
|
||||
- **Python:** 3.10+ required
|
||||
- **Node.js:** 22+ required only for OpenClaw agent
|
||||
|
||||
@@ -137,11 +140,12 @@ OpenJarvis is a research framework for on-device AI organized around **five comp
|
||||
|
||||
### Benchmarking Framework (`src/openjarvis/bench/`)
|
||||
|
||||
- `_stubs.py` — `BenchmarkResult` dataclass, `BaseBenchmark` ABC, `BenchmarkSuite` runner
|
||||
- `latency.py` — `LatencyBenchmark`: measures per-call latency (mean, p50, p95, min, max)
|
||||
- `throughput.py` — `ThroughputBenchmark`: measures tokens/second throughput
|
||||
- `_stubs.py` — `BenchmarkResult` dataclass (includes `warmup_samples`, `steady_state_samples`, `steady_state_reached`, `total_energy_joules`, `energy_per_token_joules`, `energy_method`), `BaseBenchmark` ABC, `BenchmarkSuite` runner
|
||||
- `latency.py` — `LatencyBenchmark`: measures per-call latency (mean, p50, p95, min, max). Supports `warmup_samples` parameter.
|
||||
- `throughput.py` — `ThroughputBenchmark`: measures tokens/second throughput. Supports `warmup_samples` parameter.
|
||||
- `energy.py` — `EnergyBenchmark`: measures energy per token at thermal equilibrium. Uses `SteadyStateDetector` + `EnergyBatch` with optional `EnergyMonitor`. Warmup phase excluded from metrics.
|
||||
- All registered via `BenchmarkRegistry` with `ensure_registered()` pattern
|
||||
- CLI: `jarvis bench run` with options for model, engine, samples, benchmark selection, JSON/JSONL output
|
||||
- CLI: `jarvis bench run` with options for model, engine, samples, benchmark selection, warmup (`-w`), JSON/JSONL output
|
||||
|
||||
### OpenClaw Infrastructure (`src/openjarvis/agents/openclaw*.py`)
|
||||
|
||||
@@ -167,11 +171,18 @@ OpenJarvis is a research framework for on-device AI organized around **five comp
|
||||
|
||||
### Telemetry (`src/openjarvis/telemetry/`)
|
||||
|
||||
- `store.py` — `TelemetryStore` writes records to SQLite via EventBus subscription (append-only)
|
||||
- `aggregator.py` — `TelemetryAggregator` read-only query layer: `per_model_stats()`, `per_engine_stats()`, `top_models()`, `summary()`, `export_records()`, `clear()`. Time-range filtering via `since`/`until`.
|
||||
- `instrumented_engine.py` — `InstrumentedEngine` wraps any `InferenceEngine` transparently, publishing `INFERENCE_START/END` and `TELEMETRY_RECORD` events. Agents call `engine.generate()` normally; telemetry is opt-in via this wrapper (applied by `SystemBuilder` when `config.telemetry.enabled`).
|
||||
- `store.py` — `TelemetryStore` writes records to SQLite via EventBus subscription (append-only). Schema includes `energy_method`, `energy_vendor`, `batch_id`, `is_warmup`, `cpu_energy_joules`, `gpu_energy_joules`, `dram_energy_joules` columns.
|
||||
- `aggregator.py` — `TelemetryAggregator` read-only query layer: `per_model_stats()`, `per_engine_stats()`, `per_batch_stats()`, `top_models()`, `summary()`, `export_records()`, `clear()`. Time-range filtering via `since`/`until`.
|
||||
- `instrumented_engine.py` — `InstrumentedEngine` wraps any `InferenceEngine` transparently, publishing `INFERENCE_START/END` and `TELEMETRY_RECORD` events. Prefers `EnergyMonitor` over legacy `GpuMonitor` when available. Agents call `engine.generate()` normally; telemetry is opt-in via this wrapper (applied by `SystemBuilder` when `config.telemetry.enabled`).
|
||||
- `energy_monitor.py` — `EnergyMonitor` ABC: `available()`, `vendor()`, `energy_method()`, `sample()` context manager, `close()`. `EnergySample` dataclass (superset of `GpuSample`): total + per-component (CPU, GPU, DRAM, ANE) energy, vendor/device info, measurement method. `EnergyVendor` enum: NVIDIA, AMD, APPLE, CPU_RAPL. `create_energy_monitor()` factory with auto-detection order: NVIDIA > AMD > Apple > CPU RAPL.
|
||||
- `energy_nvidia.py` — `NvidiaEnergyMonitor`: hardware counters via `nvmlDeviceGetTotalEnergyConsumption()` on Volta+, trapezoidal polling fallback on pre-Volta. `energy_method` = "hw_counter" or "polling".
|
||||
- `energy_amd.py` — `AmdEnergyMonitor`: hardware counters via `amdsmi_get_energy_count()` (ROCm 6.1+). `energy_method` = "hw_counter".
|
||||
- `energy_apple.py` — `AppleEnergyMonitor`: wraps `zeus-ml[apple]` `AppleSiliconMonitor`. Per-component breakdown (CPU, GPU, DRAM, ANE). `energy_method` = "zeus".
|
||||
- `energy_rapl.py` — `RaplEnergyMonitor`: reads Intel RAPL counters from `/sys/class/powercap/intel-rapl/` (no deps). Handles counter wrap-around. `energy_method` = "rapl".
|
||||
- `batch.py` — `EnergyBatch`: wraps `EnergyMonitor.sample()` around grouped requests. `BatchMetrics` dataclass with energy-per-token/per-request accounting. `BATCH_START`/`BATCH_END` events.
|
||||
- `steady_state.py` — `SteadyStateDetector`: CV-based sliding window to detect thermal equilibrium. `SteadyStateConfig` (warmup_samples, window_size, cv_threshold). Used by `EnergyBenchmark`.
|
||||
- `wrapper.py` — Legacy `instrumented_generate()` function (still used by some CLI/SDK code paths)
|
||||
- Dataclasses: `ModelStats`, `EngineStats`, `AggregatedStats`
|
||||
- Dataclasses: `ModelStats`, `EngineStats`, `AggregatedStats`, `EnergySample`, `BatchMetrics`, `SteadyStateResult`
|
||||
|
||||
### Security (`src/openjarvis/security/`)
|
||||
|
||||
@@ -205,14 +216,16 @@ OpenJarvis is a research framework for on-device AI organized around **five comp
|
||||
|
||||
- `registry.py` — `RegistryBase[T]` generic base class adapted from IPW. Typed subclasses: `ModelRegistry`, `EngineRegistry`, `MemoryRegistry`, `AgentRegistry`, `ToolRegistry`, `RouterPolicyRegistry`, `BenchmarkRegistry`, `ChannelRegistry`, `LearningRegistry`.
|
||||
- `types.py` — `Message`, `Conversation`, `ModelSpec`, `ToolResult`, `TelemetryRecord`, `StepType`, `TraceStep`, `Trace`, `RoutingContext`.
|
||||
- `config.py` — `JarvisConfig` dataclass hierarchy with TOML loader. Config classes: `EngineConfig` (nested `OllamaEngineConfig`, `VLLMEngineConfig`, `SGLangEngineConfig`, `LlamaCppEngineConfig`), `IntelligenceConfig` (model identity + generation defaults: temperature, max_tokens, top_p, top_k, repetition_penalty, stop_sequences), `AgentConfig` (default_agent, tools, objective, system_prompt, system_prompt_path, context_from_memory), `ToolsConfig` (nests `StorageConfig` + `MCPConfig`), `LearningConfig` (nested `RoutingLearningConfig`, `IntelligenceLearningConfig`, `AgentLearningConfig`, `MetricsConfig`), `TracesConfig`, `TelemetryConfig`, `ServerConfig`, `ChannelConfig` (nests `WhatsAppBaileysChannelConfig`), `SecurityConfig`, `SandboxConfig`, `SchedulerConfig`. Backward-compat properties: `engine.ollama_host` → `engine.ollama.host`, `agent.default_tools` → `agent.tools`, `learning.default_policy` → `learning.routing.policy`. TOML migration layer handles cross-section moves (`agent.temperature` → `intelligence.temperature`, `memory.context_injection` → `agent.context_from_memory`). User config lives at `~/.openjarvis/config.toml`. TOML sections: `[engine]`, `[engine.ollama]`, `[engine.vllm]`, `[engine.sglang]`, `[engine.llamacpp]`, `[intelligence]`, `[agent]`, `[tools.storage]`, `[tools.mcp]`, `[learning]`, `[learning.routing]`, `[learning.intelligence]`, `[learning.agent]`, `[learning.metrics]`, `[memory]` (backward-compat), `[server]`, `[telemetry]`, `[traces]`, `[channel]`, `[channel.whatsapp_baileys]`, `[security]`, `[sandbox]`, `[scheduler]`.
|
||||
- `config.py` — `JarvisConfig` dataclass hierarchy with TOML loader. Config classes: `EngineConfig` (nested `OllamaEngineConfig`, `VLLMEngineConfig`, `SGLangEngineConfig`, `LlamaCppEngineConfig`, `MLXEngineConfig`), `IntelligenceConfig` (model identity + generation defaults: temperature, max_tokens, top_p, top_k, repetition_penalty, stop_sequences), `AgentConfig` (default_agent, tools, objective, system_prompt, system_prompt_path, context_from_memory), `ToolsConfig` (nests `StorageConfig` + `MCPConfig`), `LearningConfig` (nested `RoutingLearningConfig`, `IntelligenceLearningConfig`, `AgentLearningConfig`, `MetricsConfig`), `TracesConfig`, `TelemetryConfig`, `ServerConfig`, `ChannelConfig` (nests `WhatsAppBaileysChannelConfig`), `SecurityConfig`, `SandboxConfig`, `SchedulerConfig`. Backward-compat properties: `engine.ollama_host` → `engine.ollama.host`, `agent.default_tools` → `agent.tools`, `learning.default_policy` → `learning.routing.policy`. TOML migration layer handles cross-section moves (`agent.temperature` → `intelligence.temperature`, `memory.context_injection` → `agent.context_from_memory`). User config lives at `~/.openjarvis/config.toml`. TOML sections: `[engine]`, `[engine.ollama]`, `[engine.vllm]`, `[engine.sglang]`, `[engine.llamacpp]`, `[engine.mlx]`, `[intelligence]`, `[agent]`, `[tools.storage]`, `[tools.mcp]`, `[learning]`, `[learning.routing]`, `[learning.intelligence]`, `[learning.agent]`, `[learning.metrics]`, `[memory]` (backward-compat), `[server]`, `[telemetry]`, `[traces]`, `[channel]`, `[channel.whatsapp_baileys]`, `[security]`, `[sandbox]`, `[scheduler]`.
|
||||
- `events.py` — Pub/sub event bus for inter-pillar telemetry (synchronous dispatch). EventType values: INFERENCE_START/END, TOOL_CALL_START/END, MEMORY_STORE/RETRIEVE, AGENT_TURN_START/END, TELEMETRY_RECORD, TRACE_STEP/COMPLETE, CHANNEL_MESSAGE_RECEIVED/SENT, SECURITY_SCAN/ALERT/BLOCK, SCHEDULER_TASK_START/END.
|
||||
|
||||
### Docker & Deployment
|
||||
|
||||
- `Dockerfile` — Multi-stage build: Python 3.12-slim, installs `.[server]`, entrypoint `jarvis serve`
|
||||
- `Dockerfile.gpu` — NVIDIA CUDA 12.4 runtime variant
|
||||
- `Dockerfile.gpu.rocm` — AMD ROCm 6.2 runtime variant (multi-stage build)
|
||||
- `docker-compose.yml` — Services: `jarvis` (port 8000) + `ollama` (port 11434)
|
||||
- `docker-compose.gpu.rocm.yml` — ROCm override file (use with `-f docker-compose.yml -f docker-compose.gpu.rocm.yml`)
|
||||
- `deploy/systemd/openjarvis.service` — systemd unit file
|
||||
- `deploy/launchd/com.openjarvis.plist` — macOS launchd plist
|
||||
|
||||
@@ -250,3 +263,5 @@ OpenAI-compatible server via `jarvis serve`: `POST /v1/chat/completions`, `GET /
|
||||
| v1.4 | Phase 9 | Pillar-aligned config: generation params in Intelligence, nested engine/learning configs, agent objective/system_prompt/context_from_memory, structured learning sub-policies (routing/intelligence/agent/metrics), TOML migration layer |
|
||||
| v1.5 | Phase 10 | Agent restructuring: BaseAgent helpers (`_emit_turn_start/end`, `_build_messages`, `_generate`, `_max_turns_result`), ToolUsingAgent intermediate base, `accepts_tools` introspection, NativeReActAgent/NativeOpenHandsAgent renames, real OpenHands SDK integration, CLI/SDK tool-passing bug fix, backward-compat shims |
|
||||
| v1.6 | Phase 11 | NanoClaw subsumption: `ClaudeCodeAgent` (Claude Agent SDK via Node.js subprocess), `WhatsAppBaileysChannel` (Baileys protocol), `ContainerRunner`/`SandboxedAgent` (Docker sandbox with mount security), `TaskScheduler` (cron/interval/once + SQLite + MCP tools + CLI), `SandboxConfig`/`SchedulerConfig`/`WhatsAppBaileysChannelConfig`, SystemBuilder `.sandbox()`/`.scheduler()` |
|
||||
| v1.7 | Phase 12 | Energy Measurement Upgrade: `EnergyMonitor` ABC with multi-vendor support (NVIDIA hw counters, AMD amdsmi, Apple zeus-ml, CPU RAPL sysfs), `create_energy_monitor()` factory with auto-detection, `EnergySample` superset of `GpuSample`, `EnergyBatch` batch-level energy-per-token accounting, `SteadyStateDetector` CV-based thermal equilibrium detection, `EnergyBenchmark` with warmup phase, `InstrumentedEngine` prefers `EnergyMonitor` over legacy `GpuMonitor`, expanded `GPU_SPECS` (B200, MI300X, MI250X, M4 Max, M2 Ultra), `TelemetryRecord`/store schema extended with energy_method/vendor/batch_id/is_warmup/per-component fields, eval runner warmup phase, `--warmup` CLI option, new extras `energy-amd`/`energy-apple`/`energy-all` |
|
||||
| v1.8 | Phase 13 | Install, Hosting, Cross-Hardware, Eval: `jarvis doctor` diagnostic command (8 checks with Rich output + `--json`), `jarvis init` post-setup guidance (engine-specific next steps), README Quick Start section, MLX engine backend (`MLXEngine`, `MLXEngineConfig`, Apple Silicon → `mlx` recommendation), AMD VRAM/multi-GPU detection via `rocm-smi --showmeminfo`/`--showallinfo`, PyTorch MPS device selection in orchestrator trainers, PWA support (vite-plugin-pwa, service worker, manifest, icons), server static file serving fix for PWA files, `Dockerfile.gpu.rocm` + `docker-compose.gpu.rocm.yml` for ROCm, `inference-mlx` extra |
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
FROM rocm/dev-ubuntu-22.04:6.2 AS builder
|
||||
|
||||
RUN apt-get update && \
|
||||
apt-get install -y --no-install-recommends python3 python3-pip python3-venv && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /app
|
||||
COPY pyproject.toml README.md ./
|
||||
COPY src/ src/
|
||||
|
||||
RUN pip install --no-cache-dir uv && \
|
||||
uv pip install --system ".[server]"
|
||||
|
||||
FROM rocm/dev-ubuntu-22.04:6.2
|
||||
|
||||
RUN apt-get update && \
|
||||
apt-get install -y --no-install-recommends python3 python3-pip && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY --from=builder /usr/local /usr/local
|
||||
COPY --from=builder /app /app
|
||||
WORKDIR /app
|
||||
|
||||
EXPOSE 8000
|
||||
|
||||
ENTRYPOINT ["jarvis"]
|
||||
CMD ["serve", "--host", "0.0.0.0", "--port", "8000"]
|
||||
@@ -10,7 +10,7 @@
|
||||
<p>
|
||||
<a href="https://www.intelligence-per-watt.ai/"><img src="https://img.shields.io/badge/project-intelligence--per--watt.ai-blue" alt="Project"></a>
|
||||
<a href="https://hazyresearch.stanford.edu/OpenJarvis/"><img src="https://img.shields.io/badge/docs-mkdocs-blue" alt="Docs"></a>
|
||||
<img src="https://img.shields.io/badge/python-%3E%3D3.12-blue" alt="Python">
|
||||
<img src="https://img.shields.io/badge/python-%3E%3D3.10-blue" alt="Python">
|
||||
<img src="https://img.shields.io/badge/license-Apache%202.0-green" alt="License">
|
||||
</p>
|
||||
</div>
|
||||
@@ -47,6 +47,33 @@ pip install openjarvis[server] # + FastAPI server
|
||||
|
||||
You also need a local inference backend: [Ollama](https://ollama.com), [vLLM](https://github.com/vllm-project/vllm), [SGLang](https://github.com/sgl-project/sglang), or [llama.cpp](https://github.com/ggerganov/llama.cpp).
|
||||
|
||||
## Quick Start
|
||||
|
||||
The fastest path is Ollama on any machine with Python 3.10+:
|
||||
|
||||
```bash
|
||||
# 1. Install OpenJarvis
|
||||
pip install openjarvis
|
||||
|
||||
# 2. Detect hardware and generate config
|
||||
jarvis init
|
||||
|
||||
# 3. Install and start Ollama (https://ollama.com)
|
||||
curl -fsSL https://ollama.com/install.sh | sh
|
||||
ollama serve # start the Ollama server
|
||||
|
||||
# 4. Pull a model
|
||||
ollama pull qwen3:8b
|
||||
|
||||
# 5. Ask a question
|
||||
jarvis ask "What is the capital of France?"
|
||||
|
||||
# 6. Verify your setup
|
||||
jarvis doctor
|
||||
```
|
||||
|
||||
`jarvis init` auto-detects your hardware and recommends the best engine. After init, it prints engine-specific next steps. Run `jarvis doctor` at any time to diagnose configuration or connectivity issues.
|
||||
|
||||
## The Five Pillars
|
||||
|
||||
| Pillar | What it does | Key abstractions |
|
||||
|
||||
@@ -83,6 +83,10 @@ enabled = true
|
||||
db_path = "~/.openjarvis/telemetry.db"
|
||||
gpu_metrics = true
|
||||
gpu_poll_interval_ms = 50
|
||||
energy_vendor = "" # Auto-detect; or force "nvidia"/"amd"/"apple"/"cpu_rapl"
|
||||
warmup_samples = 0 # Warmup iterations before steady-state measurement
|
||||
steady_state_window = 5 # Sliding window size for CV stability check
|
||||
steady_state_threshold = 0.05 # Coefficient of variation threshold for steady state
|
||||
|
||||
[traces]
|
||||
enabled = true # Record traces for analysis
|
||||
|
||||
@@ -0,0 +1,15 @@
|
||||
# ROCm GPU override — use with:
|
||||
# docker compose -f docker-compose.yml -f docker-compose.gpu.rocm.yml up
|
||||
version: "3.9"
|
||||
|
||||
services:
|
||||
jarvis:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile.gpu.rocm
|
||||
devices:
|
||||
- /dev/kfd
|
||||
- /dev/dri
|
||||
group_add:
|
||||
- video
|
||||
- render
|
||||
+183
-96
@@ -8,6 +8,24 @@ from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import click
|
||||
from rich.console import Console
|
||||
from rich.progress import (
|
||||
BarColumn,
|
||||
Progress,
|
||||
SpinnerColumn,
|
||||
TextColumn,
|
||||
TimeRemainingColumn,
|
||||
)
|
||||
|
||||
from evals.core.display import (
|
||||
print_banner,
|
||||
print_completion,
|
||||
print_metrics_table,
|
||||
print_run_header,
|
||||
print_section,
|
||||
print_subject_table,
|
||||
print_suite_summary,
|
||||
)
|
||||
|
||||
# Registry of available benchmarks and their metadata
|
||||
BENCHMARKS = {
|
||||
@@ -96,59 +114,29 @@ def _build_judge_backend(judge_model: str):
|
||||
return JarvisDirectBackend(engine_key="cloud")
|
||||
|
||||
|
||||
def _print_summary(summary) -> None:
|
||||
"""Print a single run summary."""
|
||||
click.echo(f"\n{'=' * 60}")
|
||||
click.echo(f"Benchmark: {summary.benchmark}")
|
||||
click.echo(f"Model: {summary.model}")
|
||||
click.echo(f"Backend: {summary.backend}")
|
||||
click.echo(f"Samples: {summary.total_samples}")
|
||||
click.echo(f"Scored: {summary.scored_samples}")
|
||||
click.echo(f"Correct: {summary.correct}")
|
||||
click.echo(f"Accuracy: {summary.accuracy:.4f}")
|
||||
click.echo(f"Errors: {summary.errors}")
|
||||
click.echo(f"Latency: {summary.mean_latency_seconds:.2f}s (mean)")
|
||||
click.echo(f"Cost: ${summary.total_cost_usd:.4f}")
|
||||
if summary.per_subject:
|
||||
click.echo("\nPer-subject breakdown:")
|
||||
for subj, stats in sorted(summary.per_subject.items()):
|
||||
click.echo(f" {subj}: {stats['accuracy']:.4f} "
|
||||
f"({int(stats['correct'])}/{int(stats['scored'])})")
|
||||
# GPU telemetry stats
|
||||
_stats_rows = []
|
||||
for label, stats_field in [
|
||||
("Accuracy", "accuracy_stats"),
|
||||
("Latency (s)", "latency_stats"),
|
||||
("TTFT (s)", "ttft_stats"),
|
||||
("Energy (J)", "energy_stats"),
|
||||
("Power (W)", "power_stats"),
|
||||
("GPU Util (%)", "gpu_utilization_stats"),
|
||||
("Throughput (tok/s)", "throughput_stats"),
|
||||
("MFU (%)", "mfu_stats"),
|
||||
("MBU (%)", "mbu_stats"),
|
||||
("IPW", "ipw_stats"),
|
||||
("IPJ", "ipj_stats"),
|
||||
]:
|
||||
ms = getattr(summary, stats_field, None)
|
||||
if ms is not None:
|
||||
_stats_rows.append((label, ms))
|
||||
if _stats_rows:
|
||||
click.echo(f"\n{'Metric':20s} {'Mean':>10s} {'Median':>10s} "
|
||||
f"{'Min':>10s} {'Max':>10s} {'Std':>10s}")
|
||||
click.echo(f"{'-' * 20} {'-' * 10} {'-' * 10} "
|
||||
f"{'-' * 10} {'-' * 10} {'-' * 10}")
|
||||
for label, ms in _stats_rows:
|
||||
click.echo(f"{label:20s} {ms.mean:10.4f} {ms.median:10.4f} "
|
||||
f"{ms.min:10.4f} {ms.max:10.4f} {ms.std:10.4f}")
|
||||
if getattr(summary, "total_energy_joules", 0.0) > 0:
|
||||
click.echo(f"\nTotal Energy: {summary.total_energy_joules:.4f} J")
|
||||
click.echo(f"{'=' * 60}")
|
||||
def _print_summary(
|
||||
summary,
|
||||
console: Optional[Console] = None,
|
||||
output_path: Optional[Path] = None,
|
||||
traces_dir: Optional[Path] = None,
|
||||
) -> None:
|
||||
"""Print a single run summary using Rich display primitives."""
|
||||
if console is None:
|
||||
console = Console()
|
||||
print_section(console, "Results")
|
||||
print_metrics_table(console, summary)
|
||||
if summary.per_subject and len(summary.per_subject) > 1:
|
||||
print_subject_table(console, summary.per_subject)
|
||||
print_completion(console, summary, output_path, traces_dir)
|
||||
|
||||
|
||||
def _run_single(config) -> object:
|
||||
def _run_single(config, console: Optional[Console] = None) -> object:
|
||||
"""Run a single eval from a RunConfig and return the summary."""
|
||||
from evals.core.runner import EvalRunner
|
||||
|
||||
if console is None:
|
||||
console = Console()
|
||||
|
||||
eval_backend = _build_backend(
|
||||
config.backend,
|
||||
config.engine_key,
|
||||
@@ -163,7 +151,27 @@ def _run_single(config) -> object:
|
||||
|
||||
runner = EvalRunner(config, dataset, eval_backend, scorer)
|
||||
try:
|
||||
return runner.run()
|
||||
num_samples = config.max_samples or 0
|
||||
# Use progress bar if we know the sample count
|
||||
if num_samples > 0:
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
|
||||
TimeRemainingColumn(),
|
||||
console=console,
|
||||
) as progress:
|
||||
task = progress.add_task("Evaluating samples...", total=num_samples)
|
||||
summary = runner.run(
|
||||
progress_callback=lambda done, total: progress.update(
|
||||
task, completed=done,
|
||||
),
|
||||
)
|
||||
else:
|
||||
with console.status("Evaluating samples..."):
|
||||
summary = runner.run()
|
||||
return summary
|
||||
finally:
|
||||
eval_backend.close()
|
||||
judge_backend.close()
|
||||
@@ -173,15 +181,25 @@ def _run_from_config(config_path: str, verbose: bool) -> None:
|
||||
"""Load a TOML config and run the full models x benchmarks matrix."""
|
||||
from evals.core.config import expand_suite, load_eval_config
|
||||
|
||||
console = Console()
|
||||
|
||||
suite = load_eval_config(config_path)
|
||||
run_configs = expand_suite(suite)
|
||||
|
||||
suite_name = suite.meta.name or Path(config_path).stem
|
||||
click.echo(f"Suite: {suite_name}")
|
||||
|
||||
# Banner + configuration
|
||||
print_banner(console)
|
||||
print_section(console, "Suite Configuration")
|
||||
console.print(
|
||||
f" [cyan]Suite:[/cyan] {suite_name}"
|
||||
)
|
||||
if suite.meta.description:
|
||||
click.echo(f" {suite.meta.description}")
|
||||
click.echo(f" {len(suite.models)} model(s) x {len(suite.benchmarks)} "
|
||||
f"benchmark(s) = {len(run_configs)} run(s)\n")
|
||||
console.print(f" [cyan]Description:[/cyan] {suite.meta.description}")
|
||||
console.print(
|
||||
f" [cyan]Matrix:[/cyan] {len(suite.models)} model(s) x "
|
||||
f"{len(suite.benchmarks)} benchmark(s) = {len(run_configs)} run(s)"
|
||||
)
|
||||
|
||||
# Ensure output directory exists
|
||||
output_dir = Path(suite.run.output_dir)
|
||||
@@ -189,28 +207,24 @@ def _run_from_config(config_path: str, verbose: bool) -> None:
|
||||
|
||||
summaries = []
|
||||
for i, rc in enumerate(run_configs, 1):
|
||||
click.echo(f"--- [{i}/{len(run_configs)}] {rc.benchmark} / {rc.model} ---")
|
||||
print_section(
|
||||
console,
|
||||
f"Run {i}/{len(run_configs)}: {rc.benchmark} / {rc.model}",
|
||||
)
|
||||
try:
|
||||
summary = _run_single(rc)
|
||||
summary = _run_single(rc, console=console)
|
||||
summaries.append(summary)
|
||||
click.echo(f" {summary.accuracy:.4f} "
|
||||
f"({summary.correct}/{summary.scored_samples})")
|
||||
console.print(
|
||||
f" [green]{summary.accuracy:.4f}[/green] "
|
||||
f"({summary.correct}/{summary.scored_samples})"
|
||||
)
|
||||
except Exception as exc:
|
||||
click.echo(f" FAILED: {exc}", err=True)
|
||||
console.print(f" [red bold]FAILED:[/red bold] {exc}")
|
||||
|
||||
# Print overall summary table
|
||||
if summaries:
|
||||
click.echo(f"\n{'=' * 60}")
|
||||
click.echo(f"Suite Results: {suite_name}")
|
||||
click.echo(f"{'=' * 60}")
|
||||
click.echo(f" {'Benchmark':12s} {'Model':20s} {'Accuracy':>10s} {'Scored':>8s}")
|
||||
click.echo(f" {'-' * 12} {'-' * 20} {'-' * 10} {'-' * 8}")
|
||||
for s in summaries:
|
||||
model_display = s.model[:20]
|
||||
click.echo(f" {s.benchmark:12s} {model_display:20s} "
|
||||
f"{s.accuracy:10.4f} "
|
||||
f"{s.correct}/{s.scored_samples:>5}")
|
||||
click.echo(f"{'=' * 60}")
|
||||
print_section(console, "Suite Results")
|
||||
print_suite_summary(console, summaries, suite_name)
|
||||
|
||||
|
||||
@click.group()
|
||||
@@ -260,6 +274,8 @@ def run(ctx, config_path, benchmark, backend, model, engine_key, agent_name,
|
||||
"""Run a single benchmark evaluation, or a full suite from a TOML config."""
|
||||
_setup_logging(verbose)
|
||||
|
||||
console = Console()
|
||||
|
||||
# Config-driven mode
|
||||
if config_path is not None:
|
||||
_run_from_config(config_path, verbose)
|
||||
@@ -300,8 +316,31 @@ def run(ctx, config_path, benchmark, backend, model, engine_key, agent_name,
|
||||
gpu_metrics=gpu_metrics,
|
||||
)
|
||||
|
||||
summary = _run_single(config)
|
||||
_print_summary(summary)
|
||||
# Banner + config
|
||||
print_banner(console)
|
||||
print_section(console, "Configuration")
|
||||
print_run_header(
|
||||
console,
|
||||
benchmark=benchmark,
|
||||
model=model,
|
||||
backend=backend,
|
||||
samples=max_samples,
|
||||
workers=max_workers,
|
||||
)
|
||||
|
||||
# Evaluation
|
||||
print_section(console, "Evaluation")
|
||||
summary = _run_single(config, console=console)
|
||||
|
||||
# Results
|
||||
_output_path = getattr(summary, "_output_path", None)
|
||||
_traces_dir = getattr(summary, "_traces_dir", None)
|
||||
_print_summary(
|
||||
summary,
|
||||
console=console,
|
||||
output_path=_output_path,
|
||||
traces_dir=_traces_dir,
|
||||
)
|
||||
|
||||
|
||||
@main.command("run-all")
|
||||
@@ -325,14 +364,24 @@ def run_all(model, engine_key, max_samples, max_workers, judge_model,
|
||||
from evals.core.runner import EvalRunner
|
||||
from evals.core.types import RunConfig
|
||||
|
||||
console = Console()
|
||||
|
||||
print_banner(console)
|
||||
print_section(console, "Suite Configuration")
|
||||
console.print(
|
||||
f" [cyan]Model:[/cyan] {model}\n"
|
||||
f" [cyan]Benchmarks:[/cyan] {', '.join(BENCHMARKS.keys())}\n"
|
||||
f" [cyan]Samples:[/cyan] {max_samples if max_samples else 'all'}"
|
||||
)
|
||||
|
||||
output_dir_path = Path(output_dir)
|
||||
output_dir_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
model_slug = model.replace("/", "-").replace(":", "-")
|
||||
summaries = []
|
||||
|
||||
for bench_name in BENCHMARKS:
|
||||
click.echo(f"\n--- Running {bench_name} ---")
|
||||
for i, bench_name in enumerate(BENCHMARKS, 1):
|
||||
print_section(console, f"Run {i}/{len(BENCHMARKS)}: {bench_name}")
|
||||
output_path = output_dir_path / f"{bench_name}_{model_slug}.jsonl"
|
||||
|
||||
config = RunConfig(
|
||||
@@ -354,24 +403,41 @@ def run_all(model, engine_key, max_samples, max_workers, judge_model,
|
||||
|
||||
runner = EvalRunner(config, dataset, eval_backend, scorer)
|
||||
try:
|
||||
summary = runner.run()
|
||||
if max_samples and max_samples > 0:
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
|
||||
TimeRemainingColumn(),
|
||||
console=console,
|
||||
) as progress:
|
||||
task = progress.add_task(
|
||||
f"Evaluating {bench_name}...", total=max_samples,
|
||||
)
|
||||
summary = runner.run(
|
||||
progress_callback=lambda done, total: progress.update(
|
||||
task, completed=done,
|
||||
),
|
||||
)
|
||||
else:
|
||||
with console.status(f"Evaluating {bench_name}..."):
|
||||
summary = runner.run()
|
||||
summaries.append(summary)
|
||||
click.echo(f" {bench_name}: {summary.accuracy:.4f} "
|
||||
f"({summary.correct}/{summary.scored_samples})")
|
||||
console.print(
|
||||
f" [green]{summary.accuracy:.4f}[/green] "
|
||||
f"({summary.correct}/{summary.scored_samples})"
|
||||
)
|
||||
except Exception as exc:
|
||||
click.echo(f" {bench_name}: FAILED — {exc}", err=True)
|
||||
console.print(f" [red bold]FAILED:[/red bold] {exc}")
|
||||
finally:
|
||||
eval_backend.close()
|
||||
judge_backend.close()
|
||||
|
||||
# Print overall summary
|
||||
if summaries:
|
||||
click.echo(f"\n{'=' * 60}")
|
||||
click.echo("Overall Results:")
|
||||
for s in summaries:
|
||||
click.echo(f" {s.benchmark:12s} {s.accuracy:.4f} "
|
||||
f"({s.correct}/{s.scored_samples})")
|
||||
click.echo(f"{'=' * 60}")
|
||||
print_section(console, "Suite Results")
|
||||
print_suite_summary(console, summaries, f"All Benchmarks / {model}")
|
||||
|
||||
|
||||
@main.command()
|
||||
@@ -389,32 +455,53 @@ def summarize(jsonl_path):
|
||||
click.echo("No records found.")
|
||||
return
|
||||
|
||||
console = Console()
|
||||
total = len(records)
|
||||
scored = [r for r in records if r.get("is_correct") is not None]
|
||||
correct = [r for r in scored if r["is_correct"]]
|
||||
errors = [r for r in records if r.get("error")]
|
||||
accuracy = len(correct) / len(scored) if scored else 0.0
|
||||
|
||||
click.echo(f"File: {jsonl_path}")
|
||||
click.echo(f"Benchmark: {records[0].get('benchmark', '?')}")
|
||||
click.echo(f"Model: {records[0].get('model', '?')}")
|
||||
click.echo(f"Total: {total}")
|
||||
click.echo(f"Scored: {len(scored)}")
|
||||
click.echo(f"Correct: {len(correct)}")
|
||||
click.echo(f"Accuracy: {accuracy:.4f}")
|
||||
click.echo(f"Errors: {len(errors)}")
|
||||
console.print(f"[cyan]File:[/cyan] {jsonl_path}")
|
||||
console.print(f"[cyan]Benchmark:[/cyan] {records[0].get('benchmark', '?')}")
|
||||
console.print(f"[cyan]Model:[/cyan] {records[0].get('model', '?')}")
|
||||
console.print(f"[cyan]Total:[/cyan] {total}")
|
||||
console.print(f"[cyan]Scored:[/cyan] {len(scored)}")
|
||||
console.print(f"[cyan]Correct:[/cyan] {len(correct)}")
|
||||
console.print(f"[cyan]Accuracy:[/cyan] [bold]{accuracy:.4f}[/bold]")
|
||||
console.print(f"[cyan]Errors:[/cyan] {len(errors)}")
|
||||
|
||||
|
||||
@main.command("list")
|
||||
def list_cmd():
|
||||
"""List available benchmarks and backends."""
|
||||
click.echo("Benchmarks:")
|
||||
for name, info in BENCHMARKS.items():
|
||||
click.echo(f" {name:12s} [{info['category']:10s}] {info['description']}")
|
||||
console = Console()
|
||||
print_banner(console)
|
||||
|
||||
click.echo("\nBackends:")
|
||||
from rich.table import Table
|
||||
|
||||
bench_table = Table(
|
||||
title="[bold]Available Benchmarks[/bold]",
|
||||
border_style="bright_blue",
|
||||
title_style="bold cyan",
|
||||
)
|
||||
bench_table.add_column("Name", style="cyan", no_wrap=True)
|
||||
bench_table.add_column("Category", style="white")
|
||||
bench_table.add_column("Description")
|
||||
for name, info in BENCHMARKS.items():
|
||||
bench_table.add_row(name, info["category"], info["description"])
|
||||
console.print(bench_table)
|
||||
|
||||
backend_table = Table(
|
||||
title="[bold]Available Backends[/bold]",
|
||||
border_style="bright_blue",
|
||||
title_style="bold cyan",
|
||||
)
|
||||
backend_table.add_column("Name", style="cyan", no_wrap=True)
|
||||
backend_table.add_column("Description")
|
||||
for name, desc in BACKENDS.items():
|
||||
click.echo(f" {name:16s} {desc}")
|
||||
backend_table.add_row(name, desc)
|
||||
console.print(backend_table)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -92,6 +92,8 @@ def load_eval_config(path: str | Path) -> EvalSuiteConfig:
|
||||
seed=int(run_raw.get("seed", 42)),
|
||||
telemetry=bool(run_raw.get("telemetry", False)),
|
||||
gpu_metrics=bool(run_raw.get("gpu_metrics", False)),
|
||||
warmup_samples=int(run_raw.get("warmup_samples", 0)),
|
||||
energy_vendor=run_raw.get("energy_vendor", ""),
|
||||
)
|
||||
|
||||
# Parse [[models]]
|
||||
@@ -235,6 +237,7 @@ def expand_suite(suite: EvalSuiteConfig) -> List[RunConfig]:
|
||||
telemetry=suite.run.telemetry,
|
||||
gpu_metrics=suite.run.gpu_metrics,
|
||||
metadata=model_meta,
|
||||
warmup_samples=suite.run.warmup_samples,
|
||||
))
|
||||
|
||||
return configs
|
||||
|
||||
@@ -0,0 +1,258 @@
|
||||
"""Rich display helpers for the evaluation framework and bench CLI."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional
|
||||
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
from rich.rule import Rule
|
||||
from rich.table import Table
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
from evals.core.types import MetricStats, RunSummary
|
||||
|
||||
OPENJARVIS_BANNER = r"""
|
||||
___ _ _
|
||||
/ _ \ _ __ ___ _ __ | | __ _ _ ____ _(_)___
|
||||
| | | | '_ \ / _ \ '_ \ _ | |/ _` | '__\ \ / / / __|
|
||||
| |_| | |_) | __/ | | | |_| | (_| | | \ V /| \__ \
|
||||
\___/| .__/ \___|_| |_|\___/ \__,_|_| \_/ |_|___/
|
||||
|_|
|
||||
"""
|
||||
|
||||
VERSION = "v1.8"
|
||||
|
||||
|
||||
def print_banner(console: Console) -> None:
|
||||
"""Print the OpenJarvis ASCII banner inside a styled panel."""
|
||||
panel = Panel(
|
||||
OPENJARVIS_BANNER.rstrip(),
|
||||
border_style="cyan",
|
||||
title=f"[bold white]{VERSION}[/bold white]",
|
||||
expand=False,
|
||||
)
|
||||
console.print(panel)
|
||||
|
||||
|
||||
def print_section(console: Console, title: str) -> None:
|
||||
"""Print a horizontal rule section separator."""
|
||||
console.print(Rule(title, style="bright_blue"))
|
||||
|
||||
|
||||
def print_run_header(
|
||||
console: Console,
|
||||
benchmark: str,
|
||||
model: str,
|
||||
backend: str,
|
||||
samples: Optional[int],
|
||||
workers: int,
|
||||
warmup: int = 0,
|
||||
) -> None:
|
||||
"""Print a compact run configuration panel."""
|
||||
lines = [
|
||||
f"[cyan]Benchmark:[/cyan] {benchmark}",
|
||||
f"[cyan]Model:[/cyan] {model}",
|
||||
f"[cyan]Backend:[/cyan] {backend}",
|
||||
f"[cyan]Samples:[/cyan] {samples if samples is not None else 'all'}",
|
||||
f"[cyan]Workers:[/cyan] {workers}",
|
||||
]
|
||||
if warmup > 0:
|
||||
lines.append(f"[cyan]Warmup:[/cyan] {warmup}")
|
||||
body = "\n".join(lines)
|
||||
panel = Panel(
|
||||
body,
|
||||
title="[bold]Run Configuration[/bold]",
|
||||
border_style="blue",
|
||||
expand=False,
|
||||
)
|
||||
console.print(panel)
|
||||
|
||||
|
||||
def _fmt(val: float, decimals: int = 4) -> str:
|
||||
"""Format a float to a fixed number of decimal places."""
|
||||
return f"{val:.{decimals}f}"
|
||||
|
||||
|
||||
def _add_metric_row(
|
||||
table: Table,
|
||||
label: str,
|
||||
stats: Optional[MetricStats],
|
||||
decimals: int = 4,
|
||||
) -> None:
|
||||
"""Add a row for a metric if stats exist."""
|
||||
if stats is None:
|
||||
return
|
||||
table.add_row(
|
||||
label,
|
||||
_fmt(stats.mean, decimals),
|
||||
_fmt(stats.median, decimals),
|
||||
_fmt(stats.min, decimals),
|
||||
_fmt(stats.max, decimals),
|
||||
_fmt(stats.std, decimals),
|
||||
_fmt(stats.p95, decimals),
|
||||
_fmt(stats.p99, decimals),
|
||||
)
|
||||
|
||||
|
||||
def print_metrics_table(console: Console, summary: RunSummary) -> None:
|
||||
"""Print the unified metrics table with all available stats."""
|
||||
table = Table(
|
||||
title="[bold]Task-Level Metrics[/bold]",
|
||||
show_header=True,
|
||||
header_style="bold bright_white",
|
||||
border_style="bright_blue",
|
||||
title_style="bold cyan",
|
||||
)
|
||||
table.add_column("Metric", style="cyan", no_wrap=True)
|
||||
table.add_column("Avg", justify="right")
|
||||
table.add_column("Median", justify="right")
|
||||
table.add_column("Min", justify="right")
|
||||
table.add_column("Max", justify="right")
|
||||
table.add_column("Std", justify="right")
|
||||
table.add_column("P95", justify="right")
|
||||
table.add_column("P99", justify="right")
|
||||
|
||||
_add_metric_row(table, "Accuracy", summary.accuracy_stats)
|
||||
_add_metric_row(table, "Latency (s)", summary.latency_stats)
|
||||
_add_metric_row(table, "TTFT (s)", summary.ttft_stats)
|
||||
_add_metric_row(table, "Input Tokens", summary.input_token_stats, decimals=1)
|
||||
_add_metric_row(table, "Output Tokens", summary.output_token_stats, decimals=1)
|
||||
_add_metric_row(table, "Throughput (tok/s)", summary.throughput_stats)
|
||||
_add_metric_row(table, "Energy (J)", summary.energy_stats)
|
||||
_add_metric_row(table, "Power (W)", summary.power_stats)
|
||||
_add_metric_row(table, "GPU Util (%)", summary.gpu_utilization_stats, decimals=1)
|
||||
_add_metric_row(
|
||||
table, "Energy/OutTok (J)",
|
||||
summary.energy_per_output_token_stats, decimals=6,
|
||||
)
|
||||
_add_metric_row(table, "Throughput/Watt", summary.throughput_per_watt_stats)
|
||||
_add_metric_row(table, "MFU (%)", summary.mfu_stats, decimals=2)
|
||||
_add_metric_row(table, "MBU (%)", summary.mbu_stats, decimals=2)
|
||||
_add_metric_row(table, "IPW", summary.ipw_stats)
|
||||
_add_metric_row(table, "IPJ", summary.ipj_stats)
|
||||
_add_metric_row(table, "Mean ITL (ms)", summary.itl_stats, decimals=2)
|
||||
|
||||
if table.row_count > 0:
|
||||
console.print(table)
|
||||
|
||||
# Headline stats below the table
|
||||
headline = (
|
||||
f"[bold]Accuracy:[/bold] {summary.accuracy:.4f} "
|
||||
f"({summary.correct}/{summary.scored_samples} scored) "
|
||||
f"[bold]Mean Latency:[/bold] {summary.mean_latency_seconds:.2f}s "
|
||||
f"[bold]Cost:[/bold] ${summary.total_cost_usd:.4f}"
|
||||
)
|
||||
if summary.total_energy_joules > 0:
|
||||
headline += f" [bold]Total Energy:[/bold] {summary.total_energy_joules:.4f}J"
|
||||
if summary.warmup_samples_excluded > 0:
|
||||
headline += f" [dim](warmup: {summary.warmup_samples_excluded} excluded)[/dim]"
|
||||
console.print(headline)
|
||||
|
||||
|
||||
def print_subject_table(
|
||||
console: Console,
|
||||
per_subject: Dict[str, Dict[str, float]],
|
||||
) -> None:
|
||||
"""Print per-subject accuracy breakdown."""
|
||||
table = Table(
|
||||
title="[bold]Per-Subject Breakdown[/bold]",
|
||||
show_header=True,
|
||||
header_style="bold bright_white",
|
||||
border_style="bright_blue",
|
||||
title_style="bold cyan",
|
||||
)
|
||||
table.add_column("Subject", style="cyan", no_wrap=True)
|
||||
table.add_column("Accuracy", justify="right")
|
||||
table.add_column("Correct", justify="right")
|
||||
table.add_column("Scored", justify="right")
|
||||
|
||||
for subj, stats in sorted(per_subject.items()):
|
||||
table.add_row(
|
||||
subj,
|
||||
f"{stats['accuracy']:.4f}",
|
||||
str(int(stats.get("correct", 0))),
|
||||
str(int(stats.get("scored", 0))),
|
||||
)
|
||||
|
||||
console.print(table)
|
||||
|
||||
|
||||
def print_suite_summary(
|
||||
console: Console,
|
||||
summaries: List[RunSummary],
|
||||
suite_name: str = "",
|
||||
) -> None:
|
||||
"""Print a multi-run suite summary table."""
|
||||
title = f"Suite Results: {suite_name}" if suite_name else "Suite Results"
|
||||
table = Table(
|
||||
title=f"[bold]{title}[/bold]",
|
||||
show_header=True,
|
||||
header_style="bold bright_white",
|
||||
border_style="green",
|
||||
title_style="bold green",
|
||||
)
|
||||
table.add_column("Benchmark", style="cyan", no_wrap=True)
|
||||
table.add_column("Model", style="white")
|
||||
table.add_column("Accuracy", justify="right", style="bold")
|
||||
table.add_column("Scored", justify="right")
|
||||
table.add_column("Latency (s)", justify="right")
|
||||
table.add_column("Cost ($)", justify="right")
|
||||
|
||||
for s in summaries:
|
||||
model_display = s.model if len(s.model) <= 24 else s.model[:21] + "..."
|
||||
table.add_row(
|
||||
s.benchmark,
|
||||
model_display,
|
||||
f"{s.accuracy:.4f}",
|
||||
f"{s.correct}/{s.scored_samples}",
|
||||
f"{s.mean_latency_seconds:.2f}",
|
||||
f"{s.total_cost_usd:.4f}",
|
||||
)
|
||||
|
||||
console.print(table)
|
||||
|
||||
|
||||
def print_completion(
|
||||
console: Console,
|
||||
summary: RunSummary,
|
||||
output_path: Optional[Path] = None,
|
||||
traces_dir: Optional[Path] = None,
|
||||
) -> None:
|
||||
"""Print a completion panel showing where data was saved."""
|
||||
lines = [
|
||||
"[bold green]Evaluation complete[/bold green]",
|
||||
(
|
||||
f" Samples: {summary.total_samples}"
|
||||
f" Scored: {summary.scored_samples}"
|
||||
f" Errors: {summary.errors}"
|
||||
),
|
||||
]
|
||||
if output_path:
|
||||
lines.append(f" [cyan]JSONL:[/cyan] {output_path}")
|
||||
summary_path = (
|
||||
output_path.with_suffix(".summary.json")
|
||||
if hasattr(output_path, "with_suffix")
|
||||
else None
|
||||
)
|
||||
if summary_path:
|
||||
lines.append(f" [cyan]Summary:[/cyan] {summary_path}")
|
||||
if traces_dir:
|
||||
lines.append(f" [cyan]Traces:[/cyan] {traces_dir}")
|
||||
body = "\n".join(lines)
|
||||
panel = Panel(body, border_style="green", expand=False)
|
||||
console.print(panel)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"OPENJARVIS_BANNER",
|
||||
"print_banner",
|
||||
"print_section",
|
||||
"print_run_header",
|
||||
"print_metrics_table",
|
||||
"print_subject_table",
|
||||
"print_suite_summary",
|
||||
"print_completion",
|
||||
]
|
||||
+144
-5
@@ -9,7 +9,7 @@ import time
|
||||
from collections import defaultdict
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
from evals.core.backend import InferenceBackend
|
||||
from evals.core.dataset import DatasetProvider
|
||||
@@ -41,8 +41,16 @@ class EvalRunner:
|
||||
self._results: List[EvalResult] = []
|
||||
self._output_file: Optional[Any] = None
|
||||
|
||||
def run(self) -> RunSummary:
|
||||
"""Execute the evaluation and return a summary."""
|
||||
def run(
|
||||
self,
|
||||
progress_callback: Optional[Callable[[int, int], None]] = None,
|
||||
) -> RunSummary:
|
||||
"""Execute the evaluation and return a summary.
|
||||
|
||||
Args:
|
||||
progress_callback: Optional ``(completed, total)`` callback invoked
|
||||
after each sample completes, useful for driving progress bars.
|
||||
"""
|
||||
cfg = self._config
|
||||
started_at = time.time()
|
||||
|
||||
@@ -57,12 +65,21 @@ class EvalRunner:
|
||||
cfg.benchmark, len(records), cfg.backend, cfg.model, cfg.max_workers,
|
||||
)
|
||||
|
||||
# --- Warmup phase (discard results) ---
|
||||
warmup_count = cfg.warmup_samples
|
||||
if warmup_count > 0 and records:
|
||||
warmup_records = records[:warmup_count]
|
||||
for rec in warmup_records:
|
||||
self._process_one(rec)
|
||||
LOGGER.info("Warmup complete: %d samples discarded", len(warmup_records))
|
||||
|
||||
# Open output file for incremental JSONL writing
|
||||
output_path = self._resolve_output_path()
|
||||
if output_path:
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._output_file = open(output_path, "w")
|
||||
|
||||
total = len(records)
|
||||
try:
|
||||
with ThreadPoolExecutor(max_workers=cfg.max_workers) as pool:
|
||||
futures = {
|
||||
@@ -72,6 +89,8 @@ class EvalRunner:
|
||||
result = future.result()
|
||||
self._results.append(result)
|
||||
self._flush_result(result)
|
||||
if progress_callback is not None:
|
||||
progress_callback(len(self._results), total)
|
||||
finally:
|
||||
if self._output_file:
|
||||
self._output_file.close()
|
||||
@@ -81,6 +100,7 @@ class EvalRunner:
|
||||
summary = self._compute_summary(records, started_at, ended_at)
|
||||
|
||||
# Write summary JSON alongside JSONL
|
||||
traces_dir: Optional[Path] = None
|
||||
if output_path:
|
||||
summary_path = output_path.with_suffix(".summary.json")
|
||||
with open(summary_path, "w") as f:
|
||||
@@ -88,8 +108,29 @@ class EvalRunner:
|
||||
LOGGER.info("Results written to %s", output_path)
|
||||
LOGGER.info("Summary written to %s", summary_path)
|
||||
|
||||
# Write per-trace data
|
||||
traces_dir = self._write_traces(output_path)
|
||||
|
||||
# Attach paths to summary for callers (e.g. CLI display)
|
||||
summary._output_path = output_path # type: ignore[attr-defined]
|
||||
summary._traces_dir = traces_dir # type: ignore[attr-defined]
|
||||
|
||||
return summary
|
||||
|
||||
def _write_traces(self, output_path: Path) -> Optional[Path]:
|
||||
"""Write per-sample trace data to a traces subdirectory."""
|
||||
if not self._results:
|
||||
return None
|
||||
cfg = self._config
|
||||
model_slug = cfg.model.replace("/", "-").replace(":", "-")
|
||||
traces_dir = output_path.parent / "traces" / f"{cfg.benchmark}_{model_slug}"
|
||||
traces_dir.mkdir(parents=True, exist_ok=True)
|
||||
with open(traces_dir / "traces.jsonl", "w") as f:
|
||||
for result in self._results:
|
||||
f.write(json.dumps(_result_to_trace_dict(result)) + "\n")
|
||||
LOGGER.info("Traces written to %s", traces_dir)
|
||||
return traces_dir
|
||||
|
||||
def _process_one(self, record: EvalRecord) -> EvalResult:
|
||||
"""Process a single evaluation sample."""
|
||||
cfg = self._config
|
||||
@@ -141,6 +182,14 @@ class EvalRunner:
|
||||
mfu = eff.mfu_pct
|
||||
mbu = eff.mbu_pct
|
||||
|
||||
# Extract derived and ITL metrics from _telemetry dict
|
||||
_telem = full.get("_telemetry", {})
|
||||
energy_per_out_tok = _telem.get(
|
||||
"energy_per_output_token_joules", 0.0
|
||||
)
|
||||
throughput_per_w = _telem.get("throughput_per_watt", 0.0)
|
||||
mean_itl = _telem.get("mean_itl_ms", 0.0)
|
||||
|
||||
return EvalResult(
|
||||
record_id=record.record_id,
|
||||
model_answer=content,
|
||||
@@ -160,6 +209,9 @@ class EvalRunner:
|
||||
mbu_pct=mbu,
|
||||
ipw=ipw,
|
||||
ipj=ipj,
|
||||
energy_per_output_token_joules=energy_per_out_tok,
|
||||
throughput_per_watt=throughput_per_w,
|
||||
mean_itl_ms=mean_itl,
|
||||
)
|
||||
except Exception as exc:
|
||||
LOGGER.error("Error processing %s: %s", record.record_id, exc)
|
||||
@@ -196,6 +248,9 @@ class EvalRunner:
|
||||
"mbu_pct": result.mbu_pct,
|
||||
"ipw": result.ipw,
|
||||
"ipj": result.ipj,
|
||||
"energy_per_output_token_joules": result.energy_per_output_token_joules,
|
||||
"throughput_per_watt": result.throughput_per_watt,
|
||||
"mean_itl_ms": result.mean_itl_ms,
|
||||
}
|
||||
self._output_file.write(json.dumps(record_dict) + "\n")
|
||||
self._output_file.flush()
|
||||
@@ -259,12 +314,33 @@ class EvalRunner:
|
||||
ttft_vals = [r.ttft for r in results if r.ttft > 0]
|
||||
energy_vals = [r.energy_joules for r in results if r.energy_joules > 0]
|
||||
power_vals = [r.power_watts for r in results if r.power_watts > 0]
|
||||
gpu_util_vals = [r.gpu_utilization_pct for r in results if r.gpu_utilization_pct > 0]
|
||||
throughput_vals = [r.throughput_tok_per_sec for r in results if r.throughput_tok_per_sec > 0]
|
||||
gpu_util_vals = [
|
||||
r.gpu_utilization_pct for r in results
|
||||
if r.gpu_utilization_pct > 0
|
||||
]
|
||||
throughput_vals = [
|
||||
r.throughput_tok_per_sec for r in results
|
||||
if r.throughput_tok_per_sec > 0
|
||||
]
|
||||
mfu_vals = [r.mfu_pct for r in results if r.mfu_pct > 0]
|
||||
mbu_vals = [r.mbu_pct for r in results if r.mbu_pct > 0]
|
||||
ipw_vals = [r.ipw for r in results if r.ipw > 0]
|
||||
ipj_vals = [r.ipj for r in results if r.ipj > 0]
|
||||
epot_vals = [
|
||||
r.energy_per_output_token_joules
|
||||
for r in results
|
||||
if r.energy_per_output_token_joules > 0
|
||||
]
|
||||
tpw_vals = [
|
||||
r.throughput_per_watt
|
||||
for r in results if r.throughput_per_watt > 0
|
||||
]
|
||||
itl_vals = [r.mean_itl_ms for r in results if r.mean_itl_ms > 0]
|
||||
input_tok_vals = [r.prompt_tokens for r in results if r.prompt_tokens > 0]
|
||||
output_tok_vals = [
|
||||
r.completion_tokens for r in results
|
||||
if r.completion_tokens > 0
|
||||
]
|
||||
|
||||
total_energy = sum(r.energy_joules for r in results)
|
||||
|
||||
@@ -294,10 +370,27 @@ class EvalRunner:
|
||||
mbu_stats=_metric_stats(mbu_vals),
|
||||
ipw_stats=_metric_stats(ipw_vals),
|
||||
ipj_stats=_metric_stats(ipj_vals),
|
||||
energy_per_output_token_stats=_metric_stats(epot_vals),
|
||||
throughput_per_watt_stats=_metric_stats(tpw_vals),
|
||||
itl_stats=_metric_stats(itl_vals),
|
||||
input_token_stats=_metric_stats([float(v) for v in input_tok_vals]),
|
||||
output_token_stats=_metric_stats([float(v) for v in output_tok_vals]),
|
||||
total_energy_joules=round(total_energy, 6),
|
||||
warmup_samples_excluded=cfg.warmup_samples,
|
||||
)
|
||||
|
||||
|
||||
def _eval_percentile(data: list[float], p: float) -> float:
|
||||
"""Compute the p-th percentile using linear interpolation."""
|
||||
sorted_data = sorted(data)
|
||||
k = (len(sorted_data) - 1) * p
|
||||
f = int(k)
|
||||
c = f + 1
|
||||
if c >= len(sorted_data):
|
||||
return sorted_data[-1]
|
||||
return sorted_data[f] + (k - f) * (sorted_data[c] - sorted_data[f])
|
||||
|
||||
|
||||
def _metric_stats(values: List[float]) -> Optional[MetricStats]:
|
||||
"""Compute MetricStats from a list of float values."""
|
||||
if not values:
|
||||
@@ -308,6 +401,9 @@ def _metric_stats(values: List[float]) -> Optional[MetricStats]:
|
||||
min=min(values),
|
||||
max=max(values),
|
||||
std=statistics.stdev(values) if len(values) > 1 else 0.0,
|
||||
p90=_eval_percentile(values, 0.90),
|
||||
p95=_eval_percentile(values, 0.95),
|
||||
p99=_eval_percentile(values, 0.99),
|
||||
)
|
||||
|
||||
|
||||
@@ -321,6 +417,9 @@ def _metric_stats_to_dict(ms: Optional[MetricStats]) -> Optional[Dict[str, float
|
||||
"min": ms.min,
|
||||
"max": ms.max,
|
||||
"std": ms.std,
|
||||
"p90": ms.p90,
|
||||
"p95": ms.p95,
|
||||
"p99": ms.p99,
|
||||
}
|
||||
|
||||
|
||||
@@ -352,7 +451,47 @@ def _summary_to_dict(s: RunSummary) -> Dict[str, Any]:
|
||||
"mbu_stats": _metric_stats_to_dict(s.mbu_stats),
|
||||
"ipw_stats": _metric_stats_to_dict(s.ipw_stats),
|
||||
"ipj_stats": _metric_stats_to_dict(s.ipj_stats),
|
||||
"energy_per_output_token_stats": _metric_stats_to_dict(
|
||||
s.energy_per_output_token_stats,
|
||||
),
|
||||
"throughput_per_watt_stats": _metric_stats_to_dict(
|
||||
s.throughput_per_watt_stats,
|
||||
),
|
||||
"itl_stats": _metric_stats_to_dict(s.itl_stats),
|
||||
"input_token_stats": _metric_stats_to_dict(s.input_token_stats),
|
||||
"output_token_stats": _metric_stats_to_dict(s.output_token_stats),
|
||||
"total_energy_joules": s.total_energy_joules,
|
||||
"warmup_samples_excluded": s.warmup_samples_excluded,
|
||||
"steady_state_reached": s.steady_state_reached,
|
||||
"energy_method": s.energy_method,
|
||||
}
|
||||
|
||||
|
||||
def _result_to_trace_dict(result: EvalResult) -> Dict[str, Any]:
|
||||
"""Convert an EvalResult to a full trace dict for per-sample export."""
|
||||
return {
|
||||
"record_id": result.record_id,
|
||||
"model_answer": result.model_answer,
|
||||
"is_correct": result.is_correct,
|
||||
"score": result.score,
|
||||
"latency_seconds": result.latency_seconds,
|
||||
"prompt_tokens": result.prompt_tokens,
|
||||
"completion_tokens": result.completion_tokens,
|
||||
"cost_usd": result.cost_usd,
|
||||
"error": result.error,
|
||||
"scoring_metadata": result.scoring_metadata,
|
||||
"ttft": result.ttft,
|
||||
"energy_joules": result.energy_joules,
|
||||
"power_watts": result.power_watts,
|
||||
"gpu_utilization_pct": result.gpu_utilization_pct,
|
||||
"throughput_tok_per_sec": result.throughput_tok_per_sec,
|
||||
"mfu_pct": result.mfu_pct,
|
||||
"mbu_pct": result.mbu_pct,
|
||||
"ipw": result.ipw,
|
||||
"ipj": result.ipj,
|
||||
"energy_per_output_token_joules": result.energy_per_output_token_joules,
|
||||
"throughput_per_watt": result.throughput_per_watt,
|
||||
"mean_itl_ms": result.mean_itl_ms,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -41,6 +41,9 @@ class EvalResult:
|
||||
mbu_pct: float = 0.0
|
||||
ipw: float = 0.0 # Intelligence Per Watt
|
||||
ipj: float = 0.0 # Intelligence Per Joule
|
||||
energy_per_output_token_joules: float = 0.0
|
||||
throughput_per_watt: float = 0.0
|
||||
mean_itl_ms: float = 0.0
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
@@ -64,6 +67,7 @@ class RunConfig:
|
||||
telemetry: bool = False
|
||||
gpu_metrics: bool = False
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
warmup_samples: int = 0
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
@@ -75,6 +79,9 @@ class MetricStats:
|
||||
min: float = 0.0
|
||||
max: float = 0.0
|
||||
std: float = 0.0
|
||||
p90: float = 0.0
|
||||
p95: float = 0.0
|
||||
p99: float = 0.0
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
@@ -106,7 +113,15 @@ class RunSummary:
|
||||
mbu_stats: Optional[MetricStats] = None
|
||||
ipw_stats: Optional[MetricStats] = None
|
||||
ipj_stats: Optional[MetricStats] = None
|
||||
energy_per_output_token_stats: Optional[MetricStats] = None
|
||||
throughput_per_watt_stats: Optional[MetricStats] = None
|
||||
itl_stats: Optional[MetricStats] = None
|
||||
input_token_stats: Optional[MetricStats] = None
|
||||
output_token_stats: Optional[MetricStats] = None
|
||||
total_energy_joules: float = 0.0
|
||||
warmup_samples_excluded: int = 0
|
||||
steady_state_reached: bool = False
|
||||
energy_method: str = ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -149,6 +164,8 @@ class ExecutionConfig:
|
||||
seed: int = 42
|
||||
telemetry: bool = False
|
||||
gpu_metrics: bool = False
|
||||
warmup_samples: int = 0
|
||||
energy_vendor: str = ""
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
|
||||
@@ -3,6 +3,10 @@
|
||||
<head>
|
||||
<meta charset="UTF-8" />
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
|
||||
<meta name="theme-color" content="#1a1a1e" />
|
||||
<meta name="description" content="OpenJarvis — on-device AI assistant" />
|
||||
<link rel="apple-touch-icon" href="/apple-touch-icon.png" />
|
||||
<link rel="icon" href="/favicon.ico" />
|
||||
<title>OpenJarvis</title>
|
||||
</head>
|
||||
<body>
|
||||
|
||||
@@ -17,6 +17,7 @@
|
||||
"@types/react-dom": "^19.0.0",
|
||||
"@vitejs/plugin-react": "^4.3.4",
|
||||
"typescript": "~5.7.0",
|
||||
"vite": "^6.0.0"
|
||||
"vite": "^6.0.0",
|
||||
"vite-plugin-pwa": "^0.21"
|
||||
}
|
||||
}
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 495 B |
Binary file not shown.
|
After Width: | Height: | Size: 1.1 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 546 B |
Binary file not shown.
|
After Width: | Height: | Size: 1.8 KiB |
+23
-1
@@ -1,8 +1,30 @@
|
||||
import { defineConfig } from 'vite';
|
||||
import react from '@vitejs/plugin-react';
|
||||
import { VitePWA } from 'vite-plugin-pwa';
|
||||
|
||||
export default defineConfig({
|
||||
plugins: [react()],
|
||||
plugins: [
|
||||
react(),
|
||||
VitePWA({
|
||||
registerType: 'autoUpdate',
|
||||
manifest: {
|
||||
name: 'OpenJarvis',
|
||||
short_name: 'Jarvis',
|
||||
description: 'On-device AI assistant',
|
||||
theme_color: '#1a1a1e',
|
||||
background_color: '#1a1a1e',
|
||||
display: 'standalone',
|
||||
icons: [
|
||||
{ src: 'pwa-192x192.png', sizes: '192x192', type: 'image/png' },
|
||||
{ src: 'pwa-512x512.png', sizes: '512x512', type: 'image/png' },
|
||||
],
|
||||
},
|
||||
workbox: {
|
||||
globPatterns: ['**/*.{js,css,html,ico,png,svg}'],
|
||||
navigateFallbackDenylist: [/^\/v1\//, /^\/health/, /^\/dashboard/],
|
||||
},
|
||||
}),
|
||||
],
|
||||
build: {
|
||||
outDir: '../src/openjarvis/server/static',
|
||||
emptyOutDir: true,
|
||||
|
||||
@@ -27,6 +27,7 @@ dev = [
|
||||
inference-ollama = []
|
||||
inference-vllm = []
|
||||
inference-llamacpp = []
|
||||
inference-mlx = ["mlx-lm>=0.19; sys_platform == 'darwin'"]
|
||||
inference-cloud = [
|
||||
"openai>=1.30",
|
||||
"anthropic>=0.30",
|
||||
@@ -58,6 +59,9 @@ agents = []
|
||||
openhands = ["openhands-sdk>=1.0; python_version >= '3.12'"]
|
||||
claude-code = []
|
||||
gpu-metrics = ["pynvml>=12.0"]
|
||||
energy-amd = ["amdsmi>=6.1"]
|
||||
energy-apple = ["zeus-ml[apple]"]
|
||||
energy-all = ["pynvml>=12.0", "amdsmi>=6.1", "zeus-ml[apple]"]
|
||||
learning = []
|
||||
orchestrator-training = ["torch>=2.0", "transformers>=4.40"]
|
||||
channel-telegram = ["python-telegram-bot>=21.0"]
|
||||
|
||||
@@ -8,11 +8,13 @@ from openjarvis.core.registry import BenchmarkRegistry
|
||||
|
||||
def ensure_registered() -> None:
|
||||
"""Ensure all benchmark implementations are registered."""
|
||||
from openjarvis.bench.energy import ensure_registered as _reg_energy
|
||||
from openjarvis.bench.latency import ensure_registered as _reg_latency
|
||||
from openjarvis.bench.throughput import ensure_registered as _reg_throughput
|
||||
|
||||
_reg_latency()
|
||||
_reg_throughput()
|
||||
_reg_energy()
|
||||
|
||||
|
||||
# Trigger registration on import
|
||||
|
||||
@@ -21,6 +21,12 @@ class BenchmarkResult:
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
samples: int = 0
|
||||
errors: int = 0
|
||||
warmup_samples: int = 0
|
||||
steady_state_samples: int = 0
|
||||
steady_state_reached: bool = False
|
||||
total_energy_joules: float = 0.0
|
||||
energy_per_token_joules: float = 0.0
|
||||
energy_method: str = ""
|
||||
|
||||
|
||||
class BaseBenchmark(ABC):
|
||||
@@ -47,6 +53,7 @@ class BaseBenchmark(ABC):
|
||||
model: str,
|
||||
*,
|
||||
num_samples: int = 10,
|
||||
**kwargs: Any,
|
||||
) -> BenchmarkResult:
|
||||
"""Execute the benchmark and return results."""
|
||||
|
||||
@@ -63,11 +70,12 @@ class BenchmarkSuite:
|
||||
model: str,
|
||||
*,
|
||||
num_samples: int = 10,
|
||||
**kwargs: Any,
|
||||
) -> List[BenchmarkResult]:
|
||||
"""Run all benchmarks and return a list of results."""
|
||||
results: List[BenchmarkResult] = []
|
||||
for bench in self._benchmarks:
|
||||
result = bench.run(engine, model, num_samples=num_samples)
|
||||
result = bench.run(engine, model, num_samples=num_samples, **kwargs)
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
"""Energy benchmark — measures energy per token at thermal equilibrium."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from typing import Any, Optional
|
||||
|
||||
from openjarvis.bench._stubs import BaseBenchmark, BenchmarkResult
|
||||
from openjarvis.core.registry import BenchmarkRegistry
|
||||
from openjarvis.core.types import Message, Role
|
||||
from openjarvis.engine._stubs import InferenceEngine
|
||||
|
||||
_PROMPT = "Write a short paragraph about artificial intelligence."
|
||||
|
||||
|
||||
class EnergyBenchmark(BaseBenchmark):
|
||||
"""Measures energy per token at thermal equilibrium."""
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "energy"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "Measures energy per token at thermal equilibrium"
|
||||
|
||||
def run(
|
||||
self,
|
||||
engine: InferenceEngine,
|
||||
model: str,
|
||||
*,
|
||||
num_samples: int = 10,
|
||||
warmup_samples: int = 5,
|
||||
energy_monitor: Optional[Any] = None,
|
||||
**kwargs: Any,
|
||||
) -> BenchmarkResult:
|
||||
messages = [Message(role=Role.USER, content=_PROMPT)]
|
||||
|
||||
# --- Warmup phase (discarded) ---
|
||||
for _ in range(warmup_samples):
|
||||
try:
|
||||
engine.generate(messages, model=model)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# --- Measurement phase ---
|
||||
total_tokens = 0
|
||||
total_energy = 0.0
|
||||
total_time = 0.0
|
||||
errors = 0
|
||||
energy_method = ""
|
||||
|
||||
if energy_monitor is not None:
|
||||
from openjarvis.telemetry.batch import EnergyBatch
|
||||
from openjarvis.telemetry.steady_state import SteadyStateDetector
|
||||
|
||||
detector = SteadyStateDetector()
|
||||
batch = EnergyBatch(energy_monitor=energy_monitor)
|
||||
|
||||
with batch.sample() as ctx:
|
||||
for _ in range(num_samples):
|
||||
t0 = time.time()
|
||||
try:
|
||||
result = engine.generate(messages, model=model)
|
||||
elapsed = time.time() - t0
|
||||
usage = result.get("usage", {})
|
||||
tokens = usage.get("completion_tokens", 0)
|
||||
ctx.record_request(tokens=tokens)
|
||||
total_tokens += tokens
|
||||
total_time += elapsed
|
||||
throughput = tokens / elapsed if elapsed > 0 else 0.0
|
||||
detector.record(throughput)
|
||||
except Exception:
|
||||
errors += 1
|
||||
|
||||
if batch.metrics is not None:
|
||||
total_energy = batch.metrics.total_energy_joules
|
||||
energy_method = getattr(energy_monitor, "energy_method", lambda: "")()
|
||||
ss_result = detector.result
|
||||
else:
|
||||
# No energy monitor — still measure throughput
|
||||
for _ in range(num_samples):
|
||||
t0 = time.time()
|
||||
try:
|
||||
result = engine.generate(messages, model=model)
|
||||
elapsed = time.time() - t0
|
||||
usage = result.get("usage", {})
|
||||
tokens = usage.get("completion_tokens", 0)
|
||||
total_tokens += tokens
|
||||
total_time += elapsed
|
||||
except Exception:
|
||||
errors += 1
|
||||
ss_result = None
|
||||
|
||||
tps = total_tokens / total_time if total_time > 0 else 0.0
|
||||
energy_per_token = (
|
||||
total_energy / total_tokens if total_tokens > 0 else 0.0
|
||||
)
|
||||
mean_power = total_energy / total_time if total_time > 0 else 0.0
|
||||
|
||||
metrics = {
|
||||
"tokens_per_second": tps,
|
||||
"total_energy_joules": total_energy,
|
||||
"energy_per_token_joules": energy_per_token,
|
||||
"mean_power_watts": mean_power,
|
||||
"total_tokens": float(total_tokens),
|
||||
"total_time_seconds": total_time,
|
||||
}
|
||||
|
||||
return BenchmarkResult(
|
||||
benchmark_name=self.name,
|
||||
model=model,
|
||||
engine=engine.engine_id,
|
||||
metrics=metrics,
|
||||
samples=num_samples,
|
||||
errors=errors,
|
||||
warmup_samples=warmup_samples,
|
||||
steady_state_samples=ss_result.steady_state_samples if ss_result else 0,
|
||||
steady_state_reached=ss_result.steady_state_reached if ss_result else False,
|
||||
total_energy_joules=total_energy,
|
||||
energy_per_token_joules=energy_per_token,
|
||||
energy_method=energy_method,
|
||||
)
|
||||
|
||||
|
||||
def ensure_registered() -> None:
|
||||
"""Register the energy benchmark if not already present."""
|
||||
if not BenchmarkRegistry.contains("energy"):
|
||||
BenchmarkRegistry.register_value("energy", EnergyBenchmark)
|
||||
|
||||
|
||||
__all__ = ["EnergyBenchmark"]
|
||||
@@ -4,7 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import statistics
|
||||
import time
|
||||
from typing import List
|
||||
from typing import Any, List
|
||||
|
||||
from openjarvis.bench._stubs import BaseBenchmark, BenchmarkResult
|
||||
from openjarvis.core.registry import BenchmarkRegistry
|
||||
@@ -35,7 +35,18 @@ class LatencyBenchmark(BaseBenchmark):
|
||||
model: str,
|
||||
*,
|
||||
num_samples: int = 10,
|
||||
warmup_samples: int = 0,
|
||||
**kwargs: Any,
|
||||
) -> BenchmarkResult:
|
||||
# Run warmup iterations (discarded)
|
||||
for i in range(warmup_samples):
|
||||
prompt = _CANNED_PROMPTS[i % len(_CANNED_PROMPTS)]
|
||||
messages = [Message(role=Role.USER, content=prompt)]
|
||||
try:
|
||||
engine.generate(messages, model=model)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
latencies: List[float] = []
|
||||
errors = 0
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from openjarvis.bench._stubs import BaseBenchmark, BenchmarkResult
|
||||
from openjarvis.core.registry import BenchmarkRegistry
|
||||
@@ -27,14 +28,23 @@ class ThroughputBenchmark(BaseBenchmark):
|
||||
model: str,
|
||||
*,
|
||||
num_samples: int = 10,
|
||||
warmup_samples: int = 0,
|
||||
**kwargs: Any,
|
||||
) -> BenchmarkResult:
|
||||
# Run warmup iterations (discarded)
|
||||
prompt = "Write a short paragraph about artificial intelligence."
|
||||
messages = [Message(role=Role.USER, content=prompt)]
|
||||
|
||||
for _ in range(warmup_samples):
|
||||
try:
|
||||
engine.generate(messages, model=model)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
total_tokens = 0
|
||||
total_time = 0.0
|
||||
errors = 0
|
||||
|
||||
prompt = "Write a short paragraph about artificial intelligence."
|
||||
messages = [Message(role=Role.USER, content=prompt)]
|
||||
|
||||
for _ in range(num_samples):
|
||||
t0 = time.time()
|
||||
try:
|
||||
|
||||
@@ -8,6 +8,7 @@ import openjarvis
|
||||
from openjarvis.cli.ask import ask
|
||||
from openjarvis.cli.bench_cmd import bench
|
||||
from openjarvis.cli.channel_cmd import channel
|
||||
from openjarvis.cli.doctor_cmd import doctor
|
||||
from openjarvis.cli.init_cmd import init
|
||||
from openjarvis.cli.memory_cmd import memory
|
||||
from openjarvis.cli.model import model
|
||||
@@ -31,6 +32,7 @@ cli.add_command(telemetry, "telemetry")
|
||||
cli.add_command(bench, "bench")
|
||||
cli.add_command(channel, "channel")
|
||||
cli.add_command(scheduler, "scheduler")
|
||||
cli.add_command(doctor, "doctor")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
|
||||
@@ -21,8 +21,8 @@ from openjarvis.intelligence import (
|
||||
merge_discovered_models,
|
||||
register_builtin_models,
|
||||
)
|
||||
from openjarvis.telemetry.instrumented_engine import InstrumentedEngine
|
||||
from openjarvis.telemetry.store import TelemetryStore
|
||||
from openjarvis.telemetry.wrapper import instrumented_generate
|
||||
|
||||
|
||||
def _get_memory_backend(config):
|
||||
@@ -225,6 +225,19 @@ def ask(
|
||||
|
||||
engine_name, engine = resolved
|
||||
|
||||
# Wrap engine with InstrumentedEngine for telemetry (energy + GPU metrics)
|
||||
energy_monitor = None
|
||||
if config.telemetry.gpu_metrics:
|
||||
try:
|
||||
from openjarvis.telemetry.energy_monitor import create_energy_monitor
|
||||
|
||||
energy_monitor = create_energy_monitor(
|
||||
prefer_vendor=config.telemetry.energy_vendor or None,
|
||||
)
|
||||
except Exception:
|
||||
pass # energy monitoring is best-effort
|
||||
engine = InstrumentedEngine(engine, bus, energy_monitor=energy_monitor)
|
||||
|
||||
# Discover models and merge into registry
|
||||
all_engines = discover_engines(config)
|
||||
all_models = discover_models(all_engines)
|
||||
@@ -306,13 +319,11 @@ def ask(
|
||||
except Exception:
|
||||
pass # context injection is best-effort
|
||||
|
||||
# Generate
|
||||
# Generate (InstrumentedEngine handles telemetry + energy recording)
|
||||
try:
|
||||
result = instrumented_generate(
|
||||
engine,
|
||||
result = engine.generate(
|
||||
messages,
|
||||
model=model_name,
|
||||
bus=bus,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
@@ -327,6 +338,11 @@ def ask(
|
||||
click.echo(result.get("content", ""))
|
||||
|
||||
# Cleanup
|
||||
if energy_monitor is not None:
|
||||
try:
|
||||
energy_monitor.close()
|
||||
except Exception:
|
||||
pass
|
||||
if telem_store is not None:
|
||||
try:
|
||||
telem_store.close()
|
||||
|
||||
+100
-10
@@ -7,10 +7,36 @@ import sys
|
||||
|
||||
import click
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
from rich.rule import Rule
|
||||
from rich.table import Table
|
||||
|
||||
from openjarvis.core.config import load_config
|
||||
from openjarvis.engine import get_engine
|
||||
|
||||
_BANNER = r"""
|
||||
___ _ _
|
||||
/ _ \ _ __ ___ _ __ | | __ _ _ ____ _(_)___
|
||||
| | | | '_ \ / _ \ '_ \ _ | |/ _` | '__\ \ / / / __|
|
||||
| |_| | |_) | __/ | | | |_| | (_| | | \ V /| \__ \
|
||||
\___/| .__/ \___|_| |_|\___/ \__,_|_| \_/ |_|___/
|
||||
|_|
|
||||
"""
|
||||
|
||||
|
||||
def _print_banner(console: Console) -> None:
|
||||
panel = Panel(
|
||||
_BANNER.rstrip(),
|
||||
border_style="cyan",
|
||||
title="[bold white]v1.8[/bold white]",
|
||||
expand=False,
|
||||
)
|
||||
console.print(panel)
|
||||
|
||||
|
||||
def _section(console: Console, title: str) -> None:
|
||||
console.print(Rule(title, style="bright_blue"))
|
||||
|
||||
|
||||
@click.group()
|
||||
def bench() -> None:
|
||||
@@ -36,6 +62,10 @@ def bench() -> None:
|
||||
"--json", "output_json", is_flag=True,
|
||||
help="Output JSON summary to stdout.",
|
||||
)
|
||||
@click.option(
|
||||
"-w", "--warmup", "warmup", default=0, type=int,
|
||||
help="Number of warmup iterations before measurement.",
|
||||
)
|
||||
def run(
|
||||
model_name: str | None,
|
||||
engine_key: str | None,
|
||||
@@ -43,6 +73,7 @@ def run(
|
||||
benchmark_name: str | None,
|
||||
output_path: str | None,
|
||||
output_json: bool,
|
||||
warmup: int,
|
||||
) -> None:
|
||||
"""Run benchmarks against an inference engine."""
|
||||
console = Console(stderr=True)
|
||||
@@ -90,12 +121,45 @@ def run(
|
||||
return
|
||||
|
||||
suite = BenchmarkSuite(benchmarks)
|
||||
console.print(
|
||||
f"[cyan]Running {len(benchmarks)} benchmark(s) "
|
||||
f"on {engine_name}/{model_name} ({num_samples} samples)...[/cyan]"
|
||||
)
|
||||
|
||||
results = suite.run_all(engine, model_name, num_samples=num_samples)
|
||||
# Create energy monitor for energy benchmarks
|
||||
energy_monitor = None
|
||||
if config.telemetry.gpu_metrics:
|
||||
try:
|
||||
from openjarvis.telemetry.energy_monitor import create_energy_monitor
|
||||
|
||||
energy_monitor = create_energy_monitor(
|
||||
prefer_vendor=config.telemetry.energy_vendor or None,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Banner + configuration
|
||||
_print_banner(console)
|
||||
_section(console, "Configuration")
|
||||
bench_names = [b.name for b in benchmarks]
|
||||
config_panel = Panel(
|
||||
f"[cyan]Engine:[/cyan] {engine_name}\n"
|
||||
f"[cyan]Model:[/cyan] {model_name}\n"
|
||||
f"[cyan]Benchmarks:[/cyan] {', '.join(bench_names)}\n"
|
||||
f"[cyan]Samples:[/cyan] {num_samples}\n"
|
||||
f"[cyan]Warmup:[/cyan] {warmup}",
|
||||
title="[bold]Run Configuration[/bold]",
|
||||
border_style="blue",
|
||||
expand=False,
|
||||
)
|
||||
console.print(config_panel)
|
||||
|
||||
# Run benchmarks
|
||||
_section(console, "Execution")
|
||||
with console.status(
|
||||
f"[bold cyan]Running {len(benchmarks)} benchmark(s)...[/bold cyan]",
|
||||
):
|
||||
results = suite.run_all(
|
||||
engine, model_name,
|
||||
num_samples=num_samples, warmup_samples=warmup,
|
||||
energy_monitor=energy_monitor,
|
||||
)
|
||||
|
||||
# Output results
|
||||
if output_path:
|
||||
@@ -108,11 +172,37 @@ def run(
|
||||
summary = suite.summary(results)
|
||||
click.echo(json_mod.dumps(summary, indent=2))
|
||||
elif not output_path:
|
||||
# Pretty-print to console
|
||||
# Pretty-print results as Rich tables
|
||||
_section(console, "Results")
|
||||
for r in results:
|
||||
console.print(
|
||||
f"\n[bold]{r.benchmark_name}[/bold] "
|
||||
f"({r.samples} samples, {r.errors} errors)"
|
||||
table = Table(
|
||||
title=(
|
||||
f"[bold]{r.benchmark_name}[/bold]"
|
||||
f" ({r.samples} samples, {r.errors} errors)"
|
||||
),
|
||||
show_header=True,
|
||||
header_style="bold bright_white",
|
||||
border_style="bright_blue",
|
||||
title_style="bold cyan",
|
||||
)
|
||||
table.add_column("Metric", style="cyan", no_wrap=True)
|
||||
table.add_column("Value", justify="right", style="green")
|
||||
|
||||
for k, v in r.metrics.items():
|
||||
console.print(f" {k}: {v:.4f}")
|
||||
table.add_row(k, f"{v:.4f}")
|
||||
if r.total_energy_joules > 0:
|
||||
table.add_row("Total Energy (J)", f"{r.total_energy_joules:.4f}")
|
||||
table.add_row("Energy Method", str(r.energy_method))
|
||||
if r.energy_per_token_joules > 0:
|
||||
table.add_row(
|
||||
"Energy/Token (J)", f"{r.energy_per_token_joules:.6f}",
|
||||
)
|
||||
|
||||
console.print(table)
|
||||
|
||||
# Cleanup energy monitor
|
||||
if energy_monitor is not None:
|
||||
try:
|
||||
energy_monitor.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@@ -0,0 +1,329 @@
|
||||
"""``jarvis doctor`` — run diagnostic checks on the OpenJarvis installation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import click
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
from openjarvis.core.config import DEFAULT_CONFIG_PATH, load_config
|
||||
|
||||
|
||||
@dataclass
|
||||
class CheckResult:
|
||||
"""Result of a single diagnostic check."""
|
||||
|
||||
name: str
|
||||
status: str # "ok", "warn", "fail"
|
||||
message: str
|
||||
details: Optional[str] = None
|
||||
|
||||
|
||||
# -- Individual checks -------------------------------------------------------
|
||||
|
||||
|
||||
def _check_python_version() -> CheckResult:
|
||||
"""Check that Python version is >= 3.10."""
|
||||
ver = sys.version_info
|
||||
version_str = f"{ver.major}.{ver.minor}.{ver.micro}"
|
||||
if (ver.major, ver.minor) >= (3, 10):
|
||||
return CheckResult("Python version", "ok", version_str)
|
||||
return CheckResult(
|
||||
"Python version", "fail", f"{version_str} (requires >= 3.10)"
|
||||
)
|
||||
|
||||
|
||||
def _check_config_exists() -> CheckResult:
|
||||
"""Check that the config file exists."""
|
||||
if DEFAULT_CONFIG_PATH.exists():
|
||||
return CheckResult(
|
||||
"Config file", "ok", str(DEFAULT_CONFIG_PATH)
|
||||
)
|
||||
return CheckResult(
|
||||
"Config file",
|
||||
"warn",
|
||||
f"Not found at {DEFAULT_CONFIG_PATH}",
|
||||
details="Run `jarvis init` to generate a config file.",
|
||||
)
|
||||
|
||||
|
||||
def _check_config_parses() -> CheckResult:
|
||||
"""Check that the config file parses successfully."""
|
||||
if not DEFAULT_CONFIG_PATH.exists():
|
||||
return CheckResult(
|
||||
"Config parsing", "warn", "Skipped (no config file)"
|
||||
)
|
||||
try:
|
||||
load_config()
|
||||
return CheckResult("Config parsing", "ok", "Config loaded successfully")
|
||||
except Exception as exc:
|
||||
return CheckResult(
|
||||
"Config parsing", "fail", f"Parse error: {exc}"
|
||||
)
|
||||
|
||||
|
||||
def _ensure_engines_imported() -> None:
|
||||
"""Import engine modules to trigger registration decorators."""
|
||||
try:
|
||||
import openjarvis.engine # noqa: F401
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _get_config() -> Any:
|
||||
"""Load config or return a default if parsing fails."""
|
||||
try:
|
||||
return load_config()
|
||||
except Exception:
|
||||
from openjarvis.core.config import JarvisConfig
|
||||
|
||||
return JarvisConfig()
|
||||
|
||||
|
||||
def _check_engines() -> List[CheckResult]:
|
||||
"""Probe each registered engine for health."""
|
||||
results: List[CheckResult] = []
|
||||
|
||||
_ensure_engines_imported()
|
||||
|
||||
from openjarvis.core.registry import EngineRegistry
|
||||
from openjarvis.engine import _discovery
|
||||
|
||||
config = _get_config()
|
||||
|
||||
for key in sorted(EngineRegistry.keys()):
|
||||
try:
|
||||
engine = _discovery._make_engine(key, config)
|
||||
if engine.health():
|
||||
results.append(
|
||||
CheckResult(f"Engine: {key}", "ok", "Reachable")
|
||||
)
|
||||
else:
|
||||
results.append(
|
||||
CheckResult(f"Engine: {key}", "warn", "Unreachable")
|
||||
)
|
||||
except Exception as exc:
|
||||
results.append(
|
||||
CheckResult(f"Engine: {key}", "warn", f"Unreachable ({exc})")
|
||||
)
|
||||
|
||||
if not results:
|
||||
results.append(
|
||||
CheckResult("Engines", "warn", "No engines registered")
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def _check_models() -> List[CheckResult]:
|
||||
"""List models from healthy engines."""
|
||||
results: List[CheckResult] = []
|
||||
|
||||
_ensure_engines_imported()
|
||||
|
||||
from openjarvis.core.registry import EngineRegistry
|
||||
from openjarvis.engine import _discovery
|
||||
|
||||
config = _get_config()
|
||||
|
||||
for key in sorted(EngineRegistry.keys()):
|
||||
try:
|
||||
engine = _discovery._make_engine(key, config)
|
||||
if engine.health():
|
||||
models = engine.list_models()
|
||||
if models:
|
||||
model_list = ", ".join(models[:5])
|
||||
suffix = f" (+{len(models) - 5} more)" if len(models) > 5 else ""
|
||||
results.append(
|
||||
CheckResult(
|
||||
f"Models: {key}",
|
||||
"ok",
|
||||
f"{model_list}{suffix}",
|
||||
)
|
||||
)
|
||||
else:
|
||||
results.append(
|
||||
CheckResult(
|
||||
f"Models: {key}",
|
||||
"warn",
|
||||
"No models available",
|
||||
details="Pull a model (e.g. `ollama pull qwen3:8b`).",
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def _check_default_model() -> CheckResult:
|
||||
"""Check whether the configured default model is available."""
|
||||
try:
|
||||
config = load_config()
|
||||
except Exception:
|
||||
return CheckResult(
|
||||
"Default model", "warn", "Skipped (config unavailable)"
|
||||
)
|
||||
|
||||
default_model = config.intelligence.default_model
|
||||
if not default_model:
|
||||
return CheckResult(
|
||||
"Default model",
|
||||
"warn",
|
||||
"Not configured",
|
||||
details="Set intelligence.default_model in config.toml.",
|
||||
)
|
||||
|
||||
_ensure_engines_imported()
|
||||
|
||||
from openjarvis.core.registry import EngineRegistry
|
||||
from openjarvis.engine import _discovery
|
||||
|
||||
for key in sorted(EngineRegistry.keys()):
|
||||
try:
|
||||
engine = _discovery._make_engine(key, config)
|
||||
if engine.health():
|
||||
models = engine.list_models()
|
||||
if default_model in models:
|
||||
return CheckResult(
|
||||
"Default model",
|
||||
"ok",
|
||||
f"{default_model} (on {key})",
|
||||
)
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
return CheckResult(
|
||||
"Default model",
|
||||
"warn",
|
||||
f"{default_model} not found on any engine",
|
||||
)
|
||||
|
||||
|
||||
def _check_optional_deps() -> List[CheckResult]:
|
||||
"""Check availability of optional dependency packages."""
|
||||
results: List[CheckResult] = []
|
||||
optional_packages = [
|
||||
("fastapi", "openjarvis[server]"),
|
||||
("torch", "torch (for learning)"),
|
||||
("pynvml", "pynvml (GPU monitoring)"),
|
||||
("amdsmi", "openjarvis[energy-amd]"),
|
||||
("colbert", "openjarvis[memory-colbert]"),
|
||||
("zeus", "openjarvis[energy-apple]"),
|
||||
]
|
||||
for pkg, label in optional_packages:
|
||||
try:
|
||||
__import__(pkg)
|
||||
results.append(CheckResult(f"Optional: {label}", "ok", "Installed"))
|
||||
except ImportError:
|
||||
results.append(
|
||||
CheckResult(f"Optional: {label}", "warn", "Not installed")
|
||||
)
|
||||
return results
|
||||
|
||||
|
||||
def _check_nodejs() -> CheckResult:
|
||||
"""Check Node.js version (>= 22 required for OpenClaw)."""
|
||||
node_path = shutil.which("node")
|
||||
if not node_path:
|
||||
return CheckResult(
|
||||
"Node.js",
|
||||
"warn",
|
||||
"Not found",
|
||||
details="Node.js 22+ is required for OpenClaw agent.",
|
||||
)
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["node", "--version"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=10,
|
||||
)
|
||||
version_str = result.stdout.strip()
|
||||
# Parse "v22.1.0" -> (22, 1, 0)
|
||||
parts = version_str.lstrip("v").split(".")
|
||||
major = int(parts[0])
|
||||
if major >= 22:
|
||||
return CheckResult("Node.js", "ok", version_str)
|
||||
return CheckResult(
|
||||
"Node.js",
|
||||
"warn",
|
||||
f"{version_str} (requires >= v22)",
|
||||
details="Upgrade Node.js for OpenClaw agent support.",
|
||||
)
|
||||
except Exception as exc:
|
||||
return CheckResult("Node.js", "warn", f"Error checking version: {exc}")
|
||||
|
||||
|
||||
# -- Main command -------------------------------------------------------------
|
||||
|
||||
_STATUS_ICONS = {
|
||||
"ok": "[green]\u2713[/green]",
|
||||
"warn": "[yellow]![/yellow]",
|
||||
"fail": "[red]\u2717[/red]",
|
||||
}
|
||||
|
||||
|
||||
def _run_all_checks() -> List[CheckResult]:
|
||||
"""Run all diagnostic checks and return results."""
|
||||
checks: List[CheckResult] = []
|
||||
checks.append(_check_python_version())
|
||||
checks.append(_check_config_exists())
|
||||
checks.append(_check_config_parses())
|
||||
checks.extend(_check_engines())
|
||||
checks.extend(_check_models())
|
||||
checks.append(_check_default_model())
|
||||
checks.extend(_check_optional_deps())
|
||||
checks.append(_check_nodejs())
|
||||
return checks
|
||||
|
||||
|
||||
def _results_to_dicts(checks: List[CheckResult]) -> List[Dict[str, Any]]:
|
||||
"""Convert CheckResult list to JSON-serializable dicts."""
|
||||
return [asdict(c) for c in checks]
|
||||
|
||||
|
||||
@click.command()
|
||||
@click.option("--json", "as_json", is_flag=True, help="Output results as JSON.")
|
||||
def doctor(as_json: bool) -> None:
|
||||
"""Run diagnostic checks on your OpenJarvis installation."""
|
||||
checks = _run_all_checks()
|
||||
|
||||
if as_json:
|
||||
click.echo(json.dumps(_results_to_dicts(checks), indent=2))
|
||||
return
|
||||
|
||||
console = Console()
|
||||
console.print()
|
||||
console.print("[bold]OpenJarvis Doctor[/bold]")
|
||||
console.print()
|
||||
|
||||
table = Table(show_header=True, header_style="bold")
|
||||
table.add_column("Status", width=3, justify="center")
|
||||
table.add_column("Check")
|
||||
table.add_column("Result")
|
||||
|
||||
for check in checks:
|
||||
icon = _STATUS_ICONS.get(check.status, "?")
|
||||
message = check.message
|
||||
if check.details:
|
||||
message += f"\n [dim]{check.details}[/dim]"
|
||||
table.add_row(icon, check.name, message)
|
||||
|
||||
console.print(table)
|
||||
|
||||
ok_count = sum(1 for c in checks if c.status == "ok")
|
||||
warn_count = sum(1 for c in checks if c.status == "warn")
|
||||
fail_count = sum(1 for c in checks if c.status == "fail")
|
||||
console.print()
|
||||
console.print(
|
||||
f" {ok_count} passed, {warn_count} warnings, {fail_count} failures"
|
||||
)
|
||||
console.print()
|
||||
@@ -11,9 +11,91 @@ from openjarvis.core.config import (
|
||||
DEFAULT_CONFIG_PATH,
|
||||
detect_hardware,
|
||||
generate_default_toml,
|
||||
recommend_engine,
|
||||
)
|
||||
|
||||
|
||||
def _next_steps_text(engine: str) -> str:
|
||||
"""Return engine-specific next-steps guidance after init."""
|
||||
steps: dict[str, str] = {
|
||||
"ollama": (
|
||||
"Next steps:\n"
|
||||
"\n"
|
||||
" 1. Install Ollama:\n"
|
||||
" curl -fsSL https://ollama.com/install.sh | sh\n"
|
||||
"\n"
|
||||
" 2. Start the Ollama server:\n"
|
||||
" ollama serve\n"
|
||||
"\n"
|
||||
" 3. Pull a model:\n"
|
||||
" ollama pull qwen3:8b\n"
|
||||
"\n"
|
||||
" 4. Try it out:\n"
|
||||
" jarvis ask \"Hello\"\n"
|
||||
"\n"
|
||||
" Run `jarvis doctor` to verify your setup."
|
||||
),
|
||||
"vllm": (
|
||||
"Next steps:\n"
|
||||
"\n"
|
||||
" 1. Install vLLM:\n"
|
||||
" pip install vllm\n"
|
||||
"\n"
|
||||
" 2. Start the vLLM server:\n"
|
||||
" vllm serve Qwen/Qwen3-8B\n"
|
||||
"\n"
|
||||
" 3. Try it out:\n"
|
||||
" jarvis ask \"Hello\"\n"
|
||||
"\n"
|
||||
" Run `jarvis doctor` to verify your setup."
|
||||
),
|
||||
"llamacpp": (
|
||||
"Next steps:\n"
|
||||
"\n"
|
||||
" 1. Install llama.cpp:\n"
|
||||
" brew install llama.cpp # macOS\n"
|
||||
" # Or build from source: https://github.com/ggerganov/llama.cpp\n"
|
||||
"\n"
|
||||
" 2. Start the llama.cpp server:\n"
|
||||
" llama-server -m model.gguf --port 8080\n"
|
||||
"\n"
|
||||
" 3. Try it out:\n"
|
||||
" jarvis ask \"Hello\"\n"
|
||||
"\n"
|
||||
" Run `jarvis doctor` to verify your setup."
|
||||
),
|
||||
"sglang": (
|
||||
"Next steps:\n"
|
||||
"\n"
|
||||
" 1. Install SGLang:\n"
|
||||
" pip install sglang[all]\n"
|
||||
"\n"
|
||||
" 2. Start the SGLang server:\n"
|
||||
" python -m sglang.launch_server --model Qwen/Qwen3-8B\n"
|
||||
"\n"
|
||||
" 3. Try it out:\n"
|
||||
" jarvis ask \"Hello\"\n"
|
||||
"\n"
|
||||
" Run `jarvis doctor` to verify your setup."
|
||||
),
|
||||
"mlx": (
|
||||
"Next steps:\n"
|
||||
"\n"
|
||||
" 1. Install MLX LM:\n"
|
||||
" pip install mlx-lm\n"
|
||||
"\n"
|
||||
" 2. Start the MLX server:\n"
|
||||
" mlx_lm.server --model mlx-community/Qwen2.5-7B-4bit\n"
|
||||
"\n"
|
||||
" 3. Try it out:\n"
|
||||
" jarvis ask \"Hello\"\n"
|
||||
"\n"
|
||||
" Run `jarvis doctor` to verify your setup."
|
||||
),
|
||||
}
|
||||
return steps.get(engine, steps["ollama"])
|
||||
|
||||
|
||||
@click.command()
|
||||
@click.option(
|
||||
"--force", is_flag=True, help="Overwrite existing config without prompting."
|
||||
@@ -52,3 +134,13 @@ def init(force: bool) -> None:
|
||||
Panel(toml_content, title=str(DEFAULT_CONFIG_PATH), border_style="green")
|
||||
)
|
||||
console.print("[green]Config written successfully.[/green]")
|
||||
|
||||
engine = recommend_engine(hw)
|
||||
console.print()
|
||||
console.print(
|
||||
Panel(
|
||||
_next_steps_text(engine),
|
||||
title="Getting Started",
|
||||
border_style="cyan",
|
||||
)
|
||||
)
|
||||
|
||||
@@ -45,42 +45,122 @@ def stats(top_n: int) -> None:
|
||||
overview.add_row("Total Tokens", str(summary.total_tokens))
|
||||
overview.add_row("Total Cost (USD)", f"${summary.total_cost:.6f}")
|
||||
overview.add_row("Total Latency (s)", f"{summary.total_latency:.2f}")
|
||||
if summary.total_energy_joules > 0:
|
||||
overview.add_row("Total Energy (J)", f"{summary.total_energy_joules:.2f}")
|
||||
if summary.avg_throughput_tok_per_sec > 0:
|
||||
tps = summary.avg_throughput_tok_per_sec
|
||||
overview.add_row("Avg Throughput (tok/s)", f"{tps:.1f}")
|
||||
if summary.avg_gpu_utilization_pct > 0:
|
||||
gpu = summary.avg_gpu_utilization_pct
|
||||
overview.add_row("Avg GPU Utilization (%)", f"{gpu:.1f}")
|
||||
# Derived metrics
|
||||
if summary.avg_energy_per_output_token_joules > 0:
|
||||
overview.add_row(
|
||||
"Energy/Output Token (J)",
|
||||
f"{summary.avg_energy_per_output_token_joules:.6f}",
|
||||
)
|
||||
if summary.avg_throughput_per_watt > 0:
|
||||
overview.add_row(
|
||||
"Throughput/Watt (tok/s/W)",
|
||||
f"{summary.avg_throughput_per_watt:.2f}",
|
||||
)
|
||||
# ITL metrics
|
||||
if summary.avg_mean_itl_ms > 0:
|
||||
overview.add_row("Mean ITL (ms)", f"{summary.avg_mean_itl_ms:.2f}")
|
||||
if summary.avg_median_itl_ms > 0:
|
||||
overview.add_row("Median ITL (ms)", f"{summary.avg_median_itl_ms:.2f}")
|
||||
if summary.avg_p95_itl_ms > 0:
|
||||
overview.add_row("P95 ITL (ms)", f"{summary.avg_p95_itl_ms:.2f}")
|
||||
console.print(overview)
|
||||
|
||||
# Per-model table
|
||||
if summary.per_model:
|
||||
has_energy = any(
|
||||
ms.total_energy_joules > 0
|
||||
for ms in summary.per_model[:top_n]
|
||||
)
|
||||
has_itl = any(
|
||||
ms.avg_mean_itl_ms > 0
|
||||
for ms in summary.per_model[:top_n]
|
||||
)
|
||||
model_table = Table(title=f"Top {top_n} Models")
|
||||
model_table.add_column("Model", style="cyan")
|
||||
model_table.add_column("Calls", justify="right")
|
||||
model_table.add_column("Tokens", justify="right")
|
||||
model_table.add_column("Avg Latency", justify="right")
|
||||
model_table.add_column("Cost", justify="right")
|
||||
if has_energy:
|
||||
model_table.add_column("Energy (J)", justify="right")
|
||||
model_table.add_column("E/OutTok (J)", justify="right")
|
||||
model_table.add_column("Tok/s/W", justify="right")
|
||||
model_table.add_column("Throughput", justify="right")
|
||||
model_table.add_column("GPU Util %", justify="right")
|
||||
if has_itl:
|
||||
model_table.add_column("Mean ITL", justify="right")
|
||||
model_table.add_column("P95 ITL", justify="right")
|
||||
for ms in summary.per_model[:top_n]:
|
||||
model_table.add_row(
|
||||
row = [
|
||||
ms.model_id,
|
||||
str(ms.call_count),
|
||||
str(ms.total_tokens),
|
||||
f"{ms.avg_latency:.3f}s",
|
||||
f"${ms.total_cost:.6f}",
|
||||
)
|
||||
]
|
||||
if has_energy:
|
||||
row.append(f"{ms.total_energy_joules:.2f}")
|
||||
row.append(f"{ms.avg_energy_per_output_token_joules:.6f}")
|
||||
row.append(f"{ms.avg_throughput_per_watt:.2f}")
|
||||
row.append(f"{ms.avg_throughput_tok_per_sec:.1f}")
|
||||
row.append(f"{ms.avg_gpu_utilization_pct:.1f}")
|
||||
if has_itl:
|
||||
row.append(f"{ms.avg_mean_itl_ms:.2f}")
|
||||
row.append(f"{ms.avg_p95_itl_ms:.2f}")
|
||||
model_table.add_row(*row)
|
||||
console.print(model_table)
|
||||
|
||||
# Per-engine table
|
||||
if summary.per_engine:
|
||||
has_engine_energy = any(
|
||||
es.total_energy_joules > 0
|
||||
for es in summary.per_engine
|
||||
)
|
||||
has_engine_itl = any(
|
||||
es.avg_mean_itl_ms > 0
|
||||
for es in summary.per_engine
|
||||
)
|
||||
engine_table = Table(title="Engines")
|
||||
engine_table.add_column("Engine", style="cyan")
|
||||
engine_table.add_column("Calls", justify="right")
|
||||
engine_table.add_column("Tokens", justify="right")
|
||||
engine_table.add_column("Avg Latency", justify="right")
|
||||
engine_table.add_column("Cost", justify="right")
|
||||
if has_engine_energy:
|
||||
engine_table.add_column("Energy (J)", justify="right")
|
||||
engine_table.add_column("E/OutTok (J)", justify="right")
|
||||
engine_table.add_column("Tok/s/W", justify="right")
|
||||
engine_table.add_column("Throughput", justify="right")
|
||||
engine_table.add_column("GPU Util %", justify="right")
|
||||
if has_engine_itl:
|
||||
engine_table.add_column("Mean ITL", justify="right")
|
||||
engine_table.add_column("P95 ITL", justify="right")
|
||||
for es in summary.per_engine:
|
||||
engine_table.add_row(
|
||||
row = [
|
||||
es.engine,
|
||||
str(es.call_count),
|
||||
str(es.total_tokens),
|
||||
f"{es.avg_latency:.3f}s",
|
||||
f"${es.total_cost:.6f}",
|
||||
)
|
||||
]
|
||||
if has_engine_energy:
|
||||
row.append(f"{es.total_energy_joules:.2f}")
|
||||
row.append(f"{es.avg_energy_per_output_token_joules:.6f}")
|
||||
row.append(f"{es.avg_throughput_per_watt:.2f}")
|
||||
row.append(f"{es.avg_throughput_tok_per_sec:.1f}")
|
||||
row.append(f"{es.avg_gpu_utilization_pct:.1f}")
|
||||
if has_engine_itl:
|
||||
row.append(f"{es.avg_mean_itl_ms:.2f}")
|
||||
row.append(f"{es.avg_p95_itl_ms:.2f}")
|
||||
engine_table.add_row(*row)
|
||||
console.print(engine_table)
|
||||
|
||||
if summary.total_calls == 0:
|
||||
|
||||
@@ -98,7 +98,32 @@ def _detect_amd_gpu() -> Optional[GpuInfo]:
|
||||
raw = _run_cmd(["rocm-smi", "--showproductname"])
|
||||
if not raw:
|
||||
return None
|
||||
return GpuInfo(vendor="amd", name=raw.splitlines()[0] if raw else "AMD GPU")
|
||||
name = raw.splitlines()[0] if raw else "AMD GPU"
|
||||
|
||||
# Parse VRAM from rocm-smi --showmeminfo vram
|
||||
vram_gb = 0.0
|
||||
try:
|
||||
vram_raw = _run_cmd(["rocm-smi", "--showmeminfo", "vram"])
|
||||
for line in vram_raw.splitlines():
|
||||
if "Total Memory (B):" in line:
|
||||
vram_bytes = int(line.split(":")[-1].strip())
|
||||
vram_gb = round(vram_bytes / (1024**3), 1)
|
||||
break
|
||||
except (ValueError, IndexError):
|
||||
vram_gb = 0.0
|
||||
|
||||
# Parse GPU count from rocm-smi --showallinfo
|
||||
count = 1
|
||||
try:
|
||||
allinfo_raw = _run_cmd(["rocm-smi", "--showallinfo"])
|
||||
import re
|
||||
gpu_ids = set(re.findall(r"GPU\[(\d+)\]", allinfo_raw))
|
||||
if gpu_ids:
|
||||
count = len(gpu_ids)
|
||||
except (ValueError, IndexError):
|
||||
count = 1
|
||||
|
||||
return GpuInfo(vendor="amd", name=name, vram_gb=vram_gb, count=count)
|
||||
|
||||
|
||||
def _detect_apple_gpu() -> Optional[GpuInfo]:
|
||||
@@ -172,7 +197,7 @@ def recommend_engine(hw: HardwareInfo) -> str:
|
||||
if gpu is None:
|
||||
return "llamacpp"
|
||||
if gpu.vendor == "apple":
|
||||
return "ollama"
|
||||
return "mlx"
|
||||
if gpu.vendor == "nvidia":
|
||||
# Datacenter cards (A100, H100, L40, etc.) → vllm; consumer → ollama
|
||||
datacenter_keywords = ("A100", "H100", "H200", "L40", "A10", "A30")
|
||||
@@ -218,6 +243,13 @@ class LlamaCppEngineConfig:
|
||||
binary_path: str = ""
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class MLXEngineConfig:
|
||||
"""Per-engine config for MLX."""
|
||||
|
||||
host: str = "http://localhost:8080"
|
||||
|
||||
|
||||
@dataclass
|
||||
class EngineConfig:
|
||||
"""Inference engine settings with nested per-engine configs."""
|
||||
@@ -227,6 +259,7 @@ class EngineConfig:
|
||||
vllm: VLLMEngineConfig = field(default_factory=VLLMEngineConfig)
|
||||
sglang: SGLangEngineConfig = field(default_factory=SGLangEngineConfig)
|
||||
llamacpp: LlamaCppEngineConfig = field(default_factory=LlamaCppEngineConfig)
|
||||
mlx: MLXEngineConfig = field(default_factory=MLXEngineConfig)
|
||||
|
||||
# Backward-compat properties for old flat attribute names
|
||||
@property
|
||||
@@ -274,6 +307,15 @@ class EngineConfig:
|
||||
def sglang_host(self, value: str) -> None:
|
||||
self.sglang.host = value
|
||||
|
||||
@property
|
||||
def mlx_host(self) -> str:
|
||||
"""Deprecated: use ``engine.mlx.host``."""
|
||||
return self.mlx.host
|
||||
|
||||
@mlx_host.setter
|
||||
def mlx_host(self, value: str) -> None:
|
||||
self.mlx.host = value
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class IntelligenceConfig:
|
||||
@@ -482,6 +524,10 @@ class TelemetryConfig:
|
||||
db_path: str = str(DEFAULT_CONFIG_DIR / "telemetry.db")
|
||||
gpu_metrics: bool = False
|
||||
gpu_poll_interval_ms: int = 50
|
||||
energy_vendor: str = "" # auto-detect or force "nvidia"/"amd"/"apple"/"cpu_rapl"
|
||||
warmup_samples: int = 0
|
||||
steady_state_window: int = 5
|
||||
steady_state_threshold: float = 0.05
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
@@ -851,6 +897,9 @@ host = "http://localhost:30000"
|
||||
# host = "http://localhost:8080"
|
||||
# binary_path = ""
|
||||
|
||||
[engine.mlx]
|
||||
host = "http://localhost:8080"
|
||||
|
||||
[intelligence]
|
||||
default_model = ""
|
||||
fallback_model = ""
|
||||
@@ -1017,6 +1066,7 @@ __all__ = [
|
||||
"LearningConfig",
|
||||
"LlamaCppEngineConfig",
|
||||
"MCPConfig",
|
||||
"MLXEngineConfig",
|
||||
"MatrixChannelConfig",
|
||||
"MattermostChannelConfig",
|
||||
"MemoryConfig",
|
||||
|
||||
@@ -39,6 +39,8 @@ class EventType(str, Enum):
|
||||
SECURITY_BLOCK = "security_block"
|
||||
SCHEDULER_TASK_START = "scheduler_task_start"
|
||||
SCHEDULER_TASK_END = "scheduler_task_end"
|
||||
BATCH_START = "batch_start"
|
||||
BATCH_END = "batch_end"
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
|
||||
@@ -140,10 +140,28 @@ class TelemetryRecord:
|
||||
gpu_memory_used_gb: float = 0.0
|
||||
gpu_temperature_c: float = 0.0
|
||||
throughput_tok_per_sec: float = 0.0
|
||||
energy_per_output_token_joules: float = 0.0
|
||||
throughput_per_watt: float = 0.0
|
||||
prefill_latency_seconds: float = 0.0
|
||||
decode_latency_seconds: float = 0.0
|
||||
prefill_energy_joules: float = 0.0
|
||||
decode_energy_joules: float = 0.0
|
||||
mean_itl_ms: float = 0.0
|
||||
median_itl_ms: float = 0.0
|
||||
p90_itl_ms: float = 0.0
|
||||
p95_itl_ms: float = 0.0
|
||||
p99_itl_ms: float = 0.0
|
||||
std_itl_ms: float = 0.0
|
||||
is_streaming: bool = False
|
||||
engine: str = ""
|
||||
agent: str = ""
|
||||
energy_method: str = ""
|
||||
energy_vendor: str = ""
|
||||
batch_id: str = ""
|
||||
is_warmup: bool = False
|
||||
cpu_energy_joules: float = 0.0
|
||||
gpu_energy_joules: float = 0.0
|
||||
dram_energy_joules: float = 0.0
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import openjarvis.engine.llamacpp # noqa: F401
|
||||
import openjarvis.engine.mlx # noqa: F401
|
||||
|
||||
# Import engine modules to trigger @EngineRegistry.register() decorators
|
||||
import openjarvis.engine.ollama # noqa: F401
|
||||
|
||||
@@ -14,6 +14,7 @@ _HOST_MAP: Dict[str, str | None] = {
|
||||
"vllm": "vllm_host",
|
||||
"llamacpp": "llamacpp_host",
|
||||
"sglang": "sglang_host",
|
||||
"mlx": "mlx_host",
|
||||
"cloud": None,
|
||||
"litellm": None,
|
||||
}
|
||||
|
||||
@@ -0,0 +1,17 @@
|
||||
"""MLX inference engine backend (OpenAI-compatible API)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from openjarvis.core.registry import EngineRegistry
|
||||
from openjarvis.engine._openai_compat import _OpenAICompatibleEngine
|
||||
|
||||
|
||||
@EngineRegistry.register("mlx")
|
||||
class MLXEngine(_OpenAICompatibleEngine):
|
||||
"""MLX backend — thin wrapper over the shared OpenAI-compatible base."""
|
||||
|
||||
engine_id = "mlx"
|
||||
_default_host = "http://localhost:8080"
|
||||
|
||||
|
||||
__all__ = ["MLXEngine"]
|
||||
@@ -38,6 +38,7 @@ from openjarvis.learning.orchestrator.sft_trainer import (
|
||||
OrchestratorSFTConfig,
|
||||
OrchestratorSFTDataset,
|
||||
OrchestratorSFTTrainer,
|
||||
_select_torch_device,
|
||||
)
|
||||
from openjarvis.learning.orchestrator.types import (
|
||||
Episode,
|
||||
@@ -78,6 +79,8 @@ __all__ = [
|
||||
"OrchestratorSFTConfig",
|
||||
"OrchestratorSFTDataset",
|
||||
"OrchestratorSFTTrainer",
|
||||
# Device selection
|
||||
"_select_torch_device",
|
||||
# GRPO
|
||||
"OrchestratorGRPOConfig",
|
||||
"OrchestratorGRPOTrainer",
|
||||
|
||||
@@ -32,6 +32,18 @@ except ImportError:
|
||||
from openjarvis.core.registry import LearningRegistry
|
||||
from openjarvis.learning._stubs import IntelligenceLearningPolicy
|
||||
|
||||
|
||||
def _select_torch_device():
|
||||
"""Select the best available PyTorch device (cuda > mps > cpu)."""
|
||||
if not HAS_TORCH or torch is None:
|
||||
return None
|
||||
if torch.cuda.is_available():
|
||||
return torch.device("cuda")
|
||||
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
return torch.device("mps")
|
||||
return torch.device("cpu")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -97,9 +109,7 @@ class OrchestratorGRPOTrainer:
|
||||
self.global_step = 0
|
||||
|
||||
if HAS_TORCH and torch is not None:
|
||||
self.device = torch.device(
|
||||
"cuda" if torch.cuda.is_available() else "cpu"
|
||||
)
|
||||
self.device = _select_torch_device()
|
||||
|
||||
self._init_model()
|
||||
self._init_optimizer()
|
||||
|
||||
@@ -27,6 +27,18 @@ except ImportError:
|
||||
from openjarvis.core.registry import LearningRegistry
|
||||
from openjarvis.learning._stubs import IntelligenceLearningPolicy
|
||||
|
||||
|
||||
def _select_torch_device():
|
||||
"""Select the best available PyTorch device (cuda > mps > cpu)."""
|
||||
if not HAS_TORCH or torch is None:
|
||||
return None
|
||||
if torch.cuda.is_available():
|
||||
return torch.device("cuda")
|
||||
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
return torch.device("mps")
|
||||
return torch.device("cpu")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -215,9 +227,7 @@ class OrchestratorSFTTrainer:
|
||||
self.global_step = 0
|
||||
|
||||
if HAS_TORCH and torch is not None:
|
||||
self.device = torch.device(
|
||||
"cuda" if torch.cuda.is_available() else "cpu"
|
||||
)
|
||||
self.device = _select_torch_device()
|
||||
|
||||
self._init_model()
|
||||
self._init_data()
|
||||
@@ -399,4 +409,5 @@ __all__ = [
|
||||
"OrchestratorSFTConfig",
|
||||
"OrchestratorSFTDataset",
|
||||
"OrchestratorSFTTrainer",
|
||||
"_select_torch_device",
|
||||
]
|
||||
|
||||
+25
-5
@@ -11,8 +11,8 @@ from openjarvis.core.events import EventBus
|
||||
from openjarvis.core.types import Message, Role
|
||||
from openjarvis.engine._discovery import get_engine
|
||||
from openjarvis.system import JarvisSystem, SystemBuilder
|
||||
from openjarvis.telemetry.instrumented_engine import InstrumentedEngine
|
||||
from openjarvis.telemetry.store import TelemetryStore
|
||||
from openjarvis.telemetry.wrapper import instrumented_generate
|
||||
|
||||
|
||||
class MemoryHandle:
|
||||
@@ -140,6 +140,7 @@ class Jarvis:
|
||||
self._engine_key = engine_key
|
||||
self._model_override = model
|
||||
self._engine: Any = None
|
||||
self._energy_monitor: Any = None
|
||||
self._resolved_engine_key: Optional[str] = None
|
||||
self._bus = EventBus()
|
||||
self._telem_store: Optional[TelemetryStore] = None
|
||||
@@ -222,7 +223,21 @@ class Jarvis:
|
||||
except Exception:
|
||||
pass # security is best-effort
|
||||
|
||||
self._engine = engine
|
||||
# Wrap engine with InstrumentedEngine for telemetry + energy
|
||||
energy_monitor = None
|
||||
if self._config.telemetry.gpu_metrics:
|
||||
try:
|
||||
from openjarvis.telemetry.energy_monitor import create_energy_monitor
|
||||
|
||||
energy_monitor = create_energy_monitor(
|
||||
prefer_vendor=self._config.telemetry.energy_vendor or None,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
self._energy_monitor = energy_monitor
|
||||
self._engine = InstrumentedEngine(
|
||||
engine, self._bus, energy_monitor=energy_monitor,
|
||||
)
|
||||
|
||||
def ask(
|
||||
self,
|
||||
@@ -295,11 +310,10 @@ class Jarvis:
|
||||
if context and self._config.agent.context_from_memory:
|
||||
messages = self._inject_context(query, messages)
|
||||
|
||||
result = instrumented_generate(
|
||||
self._engine,
|
||||
# InstrumentedEngine handles telemetry + energy recording
|
||||
result = self._engine.generate(
|
||||
messages,
|
||||
model=model_name,
|
||||
bus=self._bus,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
@@ -442,6 +456,12 @@ class Jarvis:
|
||||
def close(self) -> None:
|
||||
"""Release all resources."""
|
||||
self.memory.close()
|
||||
if self._energy_monitor is not None:
|
||||
try:
|
||||
self._energy_monitor.close()
|
||||
except Exception:
|
||||
pass
|
||||
self._energy_monitor = None
|
||||
if self._telem_store is not None:
|
||||
try:
|
||||
self._telem_store.close()
|
||||
|
||||
@@ -131,7 +131,13 @@ def create_app(
|
||||
|
||||
@app.get("/{full_path:path}")
|
||||
async def spa_catch_all(full_path: str):
|
||||
"""Serve index.html for SPA routes not handled by API endpoints."""
|
||||
"""Serve static files directly, fall back to index.html for SPA routes."""
|
||||
if full_path:
|
||||
candidate = (static_dir / full_path).resolve()
|
||||
# Path traversal prevention
|
||||
resolved_root = static_dir.resolve()
|
||||
if candidate.is_relative_to(resolved_root) and candidate.is_file():
|
||||
return FileResponse(candidate, headers=_NO_CACHE_HEADERS)
|
||||
return FileResponse(
|
||||
static_dir / "index.html",
|
||||
headers=_NO_CACHE_HEADERS,
|
||||
|
||||
@@ -308,20 +308,40 @@ class SystemBuilder:
|
||||
else config.telemetry.enabled
|
||||
)
|
||||
gpu_monitor = None
|
||||
energy_monitor = None
|
||||
if telemetry_enabled:
|
||||
from openjarvis.telemetry.instrumented_engine import InstrumentedEngine
|
||||
|
||||
if config.telemetry.gpu_metrics:
|
||||
# Try new multi-vendor EnergyMonitor first
|
||||
try:
|
||||
from openjarvis.telemetry.gpu_monitor import GpuMonitor
|
||||
from openjarvis.telemetry.energy_monitor import (
|
||||
create_energy_monitor,
|
||||
)
|
||||
|
||||
if GpuMonitor.available():
|
||||
gpu_monitor = GpuMonitor(
|
||||
poll_interval_ms=config.telemetry.gpu_poll_interval_ms,
|
||||
)
|
||||
energy_monitor = create_energy_monitor(
|
||||
poll_interval_ms=config.telemetry.gpu_poll_interval_ms,
|
||||
prefer_vendor=config.telemetry.energy_vendor or None,
|
||||
)
|
||||
except ImportError:
|
||||
pass
|
||||
engine = InstrumentedEngine(engine, bus, gpu_monitor=gpu_monitor)
|
||||
|
||||
# Fall back to legacy GpuMonitor if EnergyMonitor not available
|
||||
if energy_monitor is None:
|
||||
try:
|
||||
from openjarvis.telemetry.gpu_monitor import GpuMonitor
|
||||
|
||||
if GpuMonitor.available():
|
||||
gpu_monitor = GpuMonitor(
|
||||
poll_interval_ms=config.telemetry.gpu_poll_interval_ms,
|
||||
)
|
||||
except ImportError:
|
||||
pass
|
||||
engine = InstrumentedEngine(
|
||||
engine, bus,
|
||||
gpu_monitor=gpu_monitor,
|
||||
energy_monitor=energy_monitor,
|
||||
)
|
||||
|
||||
# Apply security guardrails to engine
|
||||
engine = self._apply_security(config, engine, bus)
|
||||
|
||||
@@ -12,7 +12,12 @@ from openjarvis.telemetry.store import TelemetryStore
|
||||
from openjarvis.telemetry.wrapper import instrumented_generate
|
||||
|
||||
try:
|
||||
from openjarvis.telemetry.gpu_monitor import GpuHardwareSpec, GpuMonitor, GpuSample, GpuSnapshot
|
||||
from openjarvis.telemetry.gpu_monitor import (
|
||||
GpuHardwareSpec,
|
||||
GpuMonitor,
|
||||
GpuSample,
|
||||
GpuSnapshot,
|
||||
)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
@@ -26,9 +31,31 @@ try:
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
from openjarvis.telemetry.energy_monitor import (
|
||||
EnergyMonitor,
|
||||
EnergySample,
|
||||
EnergyVendor,
|
||||
create_energy_monitor,
|
||||
)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
from openjarvis.telemetry.batch import BatchMetrics, EnergyBatch
|
||||
from openjarvis.telemetry.steady_state import (
|
||||
SteadyStateConfig,
|
||||
SteadyStateDetector,
|
||||
SteadyStateResult,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AggregatedStats",
|
||||
"BatchMetrics",
|
||||
"EfficiencyMetrics",
|
||||
"EnergyBatch",
|
||||
"EnergyMonitor",
|
||||
"EnergySample",
|
||||
"EnergyVendor",
|
||||
"EngineStats",
|
||||
"GpuHardwareSpec",
|
||||
"GpuMonitor",
|
||||
@@ -39,6 +66,10 @@ __all__ = [
|
||||
"TelemetryStore",
|
||||
"VLLMMetrics",
|
||||
"VLLMMetricsScraper",
|
||||
"SteadyStateConfig",
|
||||
"SteadyStateDetector",
|
||||
"SteadyStateResult",
|
||||
"compute_efficiency",
|
||||
"create_energy_monitor",
|
||||
"instrumented_generate",
|
||||
]
|
||||
|
||||
@@ -24,6 +24,13 @@ class ModelStats:
|
||||
total_energy_joules: float = 0.0
|
||||
avg_gpu_utilization_pct: float = 0.0
|
||||
avg_throughput_tok_per_sec: float = 0.0
|
||||
avg_energy_per_output_token_joules: float = 0.0
|
||||
avg_throughput_per_watt: float = 0.0
|
||||
total_prefill_energy_joules: float = 0.0
|
||||
total_decode_energy_joules: float = 0.0
|
||||
avg_mean_itl_ms: float = 0.0
|
||||
avg_median_itl_ms: float = 0.0
|
||||
avg_p95_itl_ms: float = 0.0
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
@@ -40,6 +47,13 @@ class EngineStats:
|
||||
total_energy_joules: float = 0.0
|
||||
avg_gpu_utilization_pct: float = 0.0
|
||||
avg_throughput_tok_per_sec: float = 0.0
|
||||
avg_energy_per_output_token_joules: float = 0.0
|
||||
avg_throughput_per_watt: float = 0.0
|
||||
total_prefill_energy_joules: float = 0.0
|
||||
total_decode_energy_joules: float = 0.0
|
||||
avg_mean_itl_ms: float = 0.0
|
||||
avg_median_itl_ms: float = 0.0
|
||||
avg_p95_itl_ms: float = 0.0
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
@@ -50,6 +64,16 @@ class AggregatedStats:
|
||||
total_tokens: int = 0
|
||||
total_cost: float = 0.0
|
||||
total_latency: float = 0.0
|
||||
total_energy_joules: float = 0.0
|
||||
avg_throughput_tok_per_sec: float = 0.0
|
||||
avg_gpu_utilization_pct: float = 0.0
|
||||
avg_energy_per_output_token_joules: float = 0.0
|
||||
avg_throughput_per_watt: float = 0.0
|
||||
total_prefill_energy_joules: float = 0.0
|
||||
total_decode_energy_joules: float = 0.0
|
||||
avg_mean_itl_ms: float = 0.0
|
||||
avg_median_itl_ms: float = 0.0
|
||||
avg_p95_itl_ms: float = 0.0
|
||||
per_model: List[ModelStats] = field(default_factory=list)
|
||||
per_engine: List[EngineStats] = field(default_factory=list)
|
||||
|
||||
@@ -80,6 +104,14 @@ class TelemetryAggregator:
|
||||
return " WHERE " + " AND ".join(clauses), params
|
||||
return "", params
|
||||
|
||||
def _safe_col(self, col_name: str) -> bool:
|
||||
"""Check if a column exists in the telemetry table."""
|
||||
try:
|
||||
self._conn.execute(f"SELECT {col_name} FROM telemetry LIMIT 0")
|
||||
return True
|
||||
except sqlite3.OperationalError:
|
||||
return False
|
||||
|
||||
def per_model_stats(
|
||||
self,
|
||||
*,
|
||||
@@ -87,6 +119,31 @@ class TelemetryAggregator:
|
||||
until: Optional[float] = None,
|
||||
) -> List[ModelStats]:
|
||||
where, params = self._time_filter(since, until)
|
||||
|
||||
# Build optional columns for new fields (graceful on old DBs)
|
||||
extra_cols = ""
|
||||
has_derived = self._safe_col("energy_per_output_token_joules")
|
||||
has_phase = self._safe_col("prefill_energy_joules")
|
||||
has_itl = self._safe_col("mean_itl_ms")
|
||||
|
||||
if has_derived:
|
||||
extra_cols += (
|
||||
", AVG(energy_per_output_token_joules)"
|
||||
" AS avg_energy_per_output_token_joules"
|
||||
", AVG(throughput_per_watt) AS avg_throughput_per_watt"
|
||||
)
|
||||
if has_phase:
|
||||
extra_cols += (
|
||||
", SUM(prefill_energy_joules) AS total_prefill_energy_joules"
|
||||
", SUM(decode_energy_joules) AS total_decode_energy_joules"
|
||||
)
|
||||
if has_itl:
|
||||
extra_cols += (
|
||||
", AVG(mean_itl_ms) AS avg_mean_itl_ms"
|
||||
", AVG(median_itl_ms) AS avg_median_itl_ms"
|
||||
", AVG(p95_itl_ms) AS avg_p95_itl_ms"
|
||||
)
|
||||
|
||||
sql = (
|
||||
"SELECT model_id,"
|
||||
" COUNT(*) AS call_count,"
|
||||
@@ -100,12 +157,14 @@ class TelemetryAggregator:
|
||||
" SUM(energy_joules) AS total_energy_joules,"
|
||||
" AVG(gpu_utilization_pct) AS avg_gpu_utilization_pct,"
|
||||
" AVG(throughput_tok_per_sec) AS avg_throughput_tok_per_sec"
|
||||
f"{extra_cols}"
|
||||
f" FROM telemetry{where}"
|
||||
" GROUP BY model_id ORDER BY call_count DESC"
|
||||
)
|
||||
rows = self._conn.execute(sql, params).fetchall()
|
||||
return [
|
||||
ModelStats(
|
||||
result = []
|
||||
for r in rows:
|
||||
ms = ModelStats(
|
||||
model_id=r["model_id"],
|
||||
call_count=r["call_count"],
|
||||
total_tokens=r["total_tokens"] or 0,
|
||||
@@ -119,8 +178,24 @@ class TelemetryAggregator:
|
||||
avg_gpu_utilization_pct=r["avg_gpu_utilization_pct"] or 0.0,
|
||||
avg_throughput_tok_per_sec=r["avg_throughput_tok_per_sec"] or 0.0,
|
||||
)
|
||||
for r in rows
|
||||
]
|
||||
if has_derived:
|
||||
ms.avg_energy_per_output_token_joules = (
|
||||
r["avg_energy_per_output_token_joules"] or 0.0
|
||||
)
|
||||
ms.avg_throughput_per_watt = r["avg_throughput_per_watt"] or 0.0
|
||||
if has_phase:
|
||||
ms.total_prefill_energy_joules = (
|
||||
r["total_prefill_energy_joules"] or 0.0
|
||||
)
|
||||
ms.total_decode_energy_joules = (
|
||||
r["total_decode_energy_joules"] or 0.0
|
||||
)
|
||||
if has_itl:
|
||||
ms.avg_mean_itl_ms = r["avg_mean_itl_ms"] or 0.0
|
||||
ms.avg_median_itl_ms = r["avg_median_itl_ms"] or 0.0
|
||||
ms.avg_p95_itl_ms = r["avg_p95_itl_ms"] or 0.0
|
||||
result.append(ms)
|
||||
return result
|
||||
|
||||
def per_engine_stats(
|
||||
self,
|
||||
@@ -129,6 +204,30 @@ class TelemetryAggregator:
|
||||
until: Optional[float] = None,
|
||||
) -> List[EngineStats]:
|
||||
where, params = self._time_filter(since, until)
|
||||
|
||||
extra_cols = ""
|
||||
has_derived = self._safe_col("energy_per_output_token_joules")
|
||||
has_phase = self._safe_col("prefill_energy_joules")
|
||||
has_itl = self._safe_col("mean_itl_ms")
|
||||
|
||||
if has_derived:
|
||||
extra_cols += (
|
||||
", AVG(energy_per_output_token_joules)"
|
||||
" AS avg_energy_per_output_token_joules"
|
||||
", AVG(throughput_per_watt) AS avg_throughput_per_watt"
|
||||
)
|
||||
if has_phase:
|
||||
extra_cols += (
|
||||
", SUM(prefill_energy_joules) AS total_prefill_energy_joules"
|
||||
", SUM(decode_energy_joules) AS total_decode_energy_joules"
|
||||
)
|
||||
if has_itl:
|
||||
extra_cols += (
|
||||
", AVG(mean_itl_ms) AS avg_mean_itl_ms"
|
||||
", AVG(median_itl_ms) AS avg_median_itl_ms"
|
||||
", AVG(p95_itl_ms) AS avg_p95_itl_ms"
|
||||
)
|
||||
|
||||
sql = (
|
||||
"SELECT engine,"
|
||||
" COUNT(*) AS call_count,"
|
||||
@@ -140,12 +239,14 @@ class TelemetryAggregator:
|
||||
" SUM(energy_joules) AS total_energy_joules,"
|
||||
" AVG(gpu_utilization_pct) AS avg_gpu_utilization_pct,"
|
||||
" AVG(throughput_tok_per_sec) AS avg_throughput_tok_per_sec"
|
||||
f"{extra_cols}"
|
||||
f" FROM telemetry{where}"
|
||||
" GROUP BY engine ORDER BY call_count DESC"
|
||||
)
|
||||
rows = self._conn.execute(sql, params).fetchall()
|
||||
return [
|
||||
EngineStats(
|
||||
result = []
|
||||
for r in rows:
|
||||
es = EngineStats(
|
||||
engine=r["engine"],
|
||||
call_count=r["call_count"],
|
||||
total_tokens=r["total_tokens"] or 0,
|
||||
@@ -157,8 +258,24 @@ class TelemetryAggregator:
|
||||
avg_gpu_utilization_pct=r["avg_gpu_utilization_pct"] or 0.0,
|
||||
avg_throughput_tok_per_sec=r["avg_throughput_tok_per_sec"] or 0.0,
|
||||
)
|
||||
for r in rows
|
||||
]
|
||||
if has_derived:
|
||||
es.avg_energy_per_output_token_joules = (
|
||||
r["avg_energy_per_output_token_joules"] or 0.0
|
||||
)
|
||||
es.avg_throughput_per_watt = r["avg_throughput_per_watt"] or 0.0
|
||||
if has_phase:
|
||||
es.total_prefill_energy_joules = (
|
||||
r["total_prefill_energy_joules"] or 0.0
|
||||
)
|
||||
es.total_decode_energy_joules = (
|
||||
r["total_decode_energy_joules"] or 0.0
|
||||
)
|
||||
if has_itl:
|
||||
es.avg_mean_itl_ms = r["avg_mean_itl_ms"] or 0.0
|
||||
es.avg_median_itl_ms = r["avg_median_itl_ms"] or 0.0
|
||||
es.avg_p95_itl_ms = r["avg_p95_itl_ms"] or 0.0
|
||||
result.append(es)
|
||||
return result
|
||||
|
||||
def top_models(
|
||||
self,
|
||||
@@ -177,15 +294,90 @@ class TelemetryAggregator:
|
||||
) -> AggregatedStats:
|
||||
model_stats = self.per_model_stats(since=since, until=until)
|
||||
engine_stats = self.per_engine_stats(since=since, until=until)
|
||||
total_calls = sum(m.call_count for m in model_stats)
|
||||
|
||||
def _weighted_avg(attr: str) -> float:
|
||||
if total_calls == 0:
|
||||
return 0.0
|
||||
return sum(
|
||||
getattr(m, attr) * m.call_count for m in model_stats
|
||||
) / total_calls
|
||||
|
||||
return AggregatedStats(
|
||||
total_calls=sum(m.call_count for m in model_stats),
|
||||
total_calls=total_calls,
|
||||
total_tokens=sum(m.total_tokens for m in model_stats),
|
||||
total_cost=sum(m.total_cost for m in model_stats),
|
||||
total_latency=sum(m.total_latency for m in model_stats),
|
||||
total_energy_joules=sum(m.total_energy_joules for m in model_stats),
|
||||
avg_throughput_tok_per_sec=_weighted_avg("avg_throughput_tok_per_sec"),
|
||||
avg_gpu_utilization_pct=_weighted_avg("avg_gpu_utilization_pct"),
|
||||
avg_energy_per_output_token_joules=_weighted_avg(
|
||||
"avg_energy_per_output_token_joules"
|
||||
),
|
||||
avg_throughput_per_watt=_weighted_avg("avg_throughput_per_watt"),
|
||||
total_prefill_energy_joules=sum(
|
||||
m.total_prefill_energy_joules for m in model_stats
|
||||
),
|
||||
total_decode_energy_joules=sum(
|
||||
m.total_decode_energy_joules for m in model_stats
|
||||
),
|
||||
avg_mean_itl_ms=_weighted_avg("avg_mean_itl_ms"),
|
||||
avg_median_itl_ms=_weighted_avg("avg_median_itl_ms"),
|
||||
avg_p95_itl_ms=_weighted_avg("avg_p95_itl_ms"),
|
||||
per_model=model_stats,
|
||||
per_engine=engine_stats,
|
||||
)
|
||||
|
||||
def per_batch_stats(
|
||||
self,
|
||||
*,
|
||||
since: Optional[float] = None,
|
||||
until: Optional[float] = None,
|
||||
exclude_warmup: bool = False,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Aggregate telemetry by batch_id.
|
||||
|
||||
Returns list of dicts with batch_id, total_requests, total_tokens,
|
||||
total_energy_joules, energy_per_token_joules.
|
||||
"""
|
||||
clauses: list[str] = ["batch_id != ''"]
|
||||
params: list[Any] = []
|
||||
if since is not None:
|
||||
clauses.append("timestamp >= ?")
|
||||
params.append(since)
|
||||
if until is not None:
|
||||
clauses.append("timestamp <= ?")
|
||||
params.append(until)
|
||||
if exclude_warmup:
|
||||
clauses.append("is_warmup = 0")
|
||||
where = " WHERE " + " AND ".join(clauses)
|
||||
|
||||
sql = (
|
||||
"SELECT batch_id,"
|
||||
" COUNT(*) AS total_requests,"
|
||||
" SUM(prompt_tokens + completion_tokens) AS total_tokens,"
|
||||
" SUM(energy_joules) AS total_energy_joules"
|
||||
f" FROM telemetry{where}"
|
||||
" GROUP BY batch_id ORDER BY total_requests DESC"
|
||||
)
|
||||
rows = self._conn.execute(sql, params).fetchall()
|
||||
results: List[Dict[str, Any]] = []
|
||||
for r in rows:
|
||||
total_tokens = r["total_tokens"] or 0
|
||||
total_energy = r["total_energy_joules"] or 0.0
|
||||
results.append(
|
||||
{
|
||||
"batch_id": r["batch_id"],
|
||||
"total_requests": r["total_requests"],
|
||||
"total_tokens": total_tokens,
|
||||
"total_energy_joules": total_energy,
|
||||
"energy_per_token_joules": (
|
||||
total_energy / total_tokens if total_tokens > 0 else 0.0
|
||||
),
|
||||
}
|
||||
)
|
||||
return results
|
||||
|
||||
def export_records(
|
||||
self,
|
||||
*,
|
||||
|
||||
@@ -0,0 +1,130 @@
|
||||
"""Batch-level energy accounting — group requests and compute per-token energy."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Generator, List, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchMetrics:
|
||||
"""Aggregated metrics for a batch of inference requests."""
|
||||
|
||||
batch_id: str = ""
|
||||
total_requests: int = 0
|
||||
total_tokens: int = 0
|
||||
total_energy_joules: float = 0.0
|
||||
energy_per_token_joules: float = 0.0
|
||||
energy_per_request_joules: float = 0.0
|
||||
mean_power_watts: float = 0.0
|
||||
mean_throughput_tok_per_sec: float = 0.0
|
||||
prefill_energy_joules: float = 0.0
|
||||
decode_energy_joules: float = 0.0
|
||||
per_request_energy: List[float] = field(default_factory=list)
|
||||
|
||||
|
||||
class EnergyBatch:
|
||||
"""Group inference requests into a batch and compute per-token energy.
|
||||
|
||||
Works with or without an ``EnergyMonitor``. When no monitor is provided,
|
||||
request counts are still tracked but energy values stay at zero.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
energy_monitor: Optional[Any] = None,
|
||||
batch_id: Optional[str] = None,
|
||||
) -> None:
|
||||
self._monitor = energy_monitor
|
||||
self.batch_id = batch_id or str(uuid.uuid4())
|
||||
self.metrics: Optional[BatchMetrics] = None
|
||||
|
||||
@contextmanager
|
||||
def sample(self) -> Generator[_BatchContext, None, None]:
|
||||
"""Wrap an energy monitor sample and provide a context for recording requests.
|
||||
|
||||
Yields a ``_BatchContext`` whose ``record_request()`` method should be
|
||||
called once per inference request inside the block.
|
||||
"""
|
||||
ctx = _BatchContext()
|
||||
|
||||
if self._monitor is not None:
|
||||
with self._monitor.sample() as energy_sample:
|
||||
start = time.monotonic()
|
||||
yield ctx
|
||||
elapsed = time.monotonic() - start
|
||||
total_energy = energy_sample.energy_joules
|
||||
mean_power = energy_sample.mean_power_watts
|
||||
else:
|
||||
start = time.monotonic()
|
||||
yield ctx
|
||||
elapsed = time.monotonic() - start
|
||||
total_energy = ctx._total_energy
|
||||
mean_power = 0.0
|
||||
|
||||
total_tokens = ctx._total_tokens
|
||||
total_requests = ctx._total_requests
|
||||
per_request_energy = list(ctx._per_request_energy)
|
||||
|
||||
energy_per_token = (
|
||||
total_energy / total_tokens if total_tokens > 0 else 0.0
|
||||
)
|
||||
energy_per_request = (
|
||||
total_energy / total_requests if total_requests > 0 else 0.0
|
||||
)
|
||||
mean_throughput = (
|
||||
total_tokens / elapsed if elapsed > 0 else 0.0
|
||||
)
|
||||
|
||||
self.metrics = BatchMetrics(
|
||||
batch_id=self.batch_id,
|
||||
total_requests=total_requests,
|
||||
total_tokens=total_tokens,
|
||||
total_energy_joules=total_energy,
|
||||
energy_per_token_joules=energy_per_token,
|
||||
energy_per_request_joules=energy_per_request,
|
||||
mean_power_watts=mean_power,
|
||||
mean_throughput_tok_per_sec=mean_throughput,
|
||||
per_request_energy=per_request_energy,
|
||||
)
|
||||
|
||||
|
||||
class _BatchContext:
|
||||
"""Accumulator for per-request stats within an ``EnergyBatch.sample()`` block."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._total_tokens: int = 0
|
||||
self._total_requests: int = 0
|
||||
self._total_energy: float = 0.0
|
||||
self._per_request_energy: List[float] = []
|
||||
|
||||
def record_request(
|
||||
self,
|
||||
tokens: int,
|
||||
prompt_tokens: int = 0,
|
||||
energy_joules: float = 0.0,
|
||||
) -> None:
|
||||
"""Record one inference request in this batch.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tokens:
|
||||
Total tokens (prompt + completion) for this request.
|
||||
prompt_tokens:
|
||||
Prompt tokens (informational; included in *tokens*).
|
||||
energy_joules:
|
||||
Per-request energy if known (e.g. from per-request metering).
|
||||
"""
|
||||
self._total_tokens += tokens
|
||||
self._total_requests += 1
|
||||
self._total_energy += energy_joules
|
||||
self._per_request_energy.append(energy_joules)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BatchMetrics",
|
||||
"EnergyBatch",
|
||||
]
|
||||
@@ -0,0 +1,131 @@
|
||||
"""AMD energy monitor — hardware counters via amdsmi (ROCm 6.1+)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from typing import Generator, List, Tuple
|
||||
|
||||
from openjarvis.telemetry.energy_monitor import (
|
||||
EnergyMonitor,
|
||||
EnergySample,
|
||||
EnergyVendor,
|
||||
)
|
||||
|
||||
try:
|
||||
import amdsmi
|
||||
|
||||
_AMDSMI_AVAILABLE = True
|
||||
except ImportError:
|
||||
_AMDSMI_AVAILABLE = False
|
||||
|
||||
|
||||
class AmdEnergyMonitor(EnergyMonitor):
|
||||
"""AMD GPU energy monitor using amdsmi hardware counters.
|
||||
|
||||
Uses ``amdsmi_get_energy_count()`` to read per-device energy accumulators.
|
||||
Energy = accumulator_delta * counter_resolution (microjoules), then / 1e6.
|
||||
"""
|
||||
|
||||
def __init__(self, poll_interval_ms: int = 50) -> None:
|
||||
self._poll_interval_ms = poll_interval_ms
|
||||
self._handles: List = []
|
||||
self._device_count = 0
|
||||
self._device_name = ""
|
||||
self._initialized = False
|
||||
|
||||
if _AMDSMI_AVAILABLE:
|
||||
try:
|
||||
amdsmi.amdsmi_init()
|
||||
self._handles = amdsmi.amdsmi_get_processor_handles()
|
||||
self._device_count = len(self._handles)
|
||||
if self._handles:
|
||||
info = amdsmi.amdsmi_get_gpu_asic_info(self._handles[0])
|
||||
self._device_name = info.get("market_name", "AMD GPU")
|
||||
self._initialized = True
|
||||
except Exception:
|
||||
self._initialized = False
|
||||
|
||||
@staticmethod
|
||||
def available() -> bool:
|
||||
if not _AMDSMI_AVAILABLE:
|
||||
return False
|
||||
try:
|
||||
amdsmi.amdsmi_init()
|
||||
handles = amdsmi.amdsmi_get_processor_handles()
|
||||
amdsmi.amdsmi_shut_down()
|
||||
return len(handles) > 0
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def vendor(self) -> EnergyVendor:
|
||||
return EnergyVendor.AMD
|
||||
|
||||
def energy_method(self) -> str:
|
||||
return "hw_counter"
|
||||
|
||||
def _read_energy_counters(self) -> List[Tuple[float, float]]:
|
||||
"""Read (accumulator, resolution) pairs from all devices."""
|
||||
readings: List[Tuple[float, float]] = []
|
||||
for handle in self._handles:
|
||||
try:
|
||||
info = amdsmi.amdsmi_get_energy_count(handle)
|
||||
accumulator = float(info.get("energy_accumulator", 0))
|
||||
resolution = float(info.get("counter_resolution", 1.0))
|
||||
readings.append((accumulator, resolution))
|
||||
except Exception:
|
||||
readings.append((0.0, 1.0))
|
||||
return readings
|
||||
|
||||
@contextmanager
|
||||
def sample(self) -> Generator[EnergySample, None, None]:
|
||||
result = EnergySample(
|
||||
vendor=EnergyVendor.AMD.value,
|
||||
device_name=self._device_name,
|
||||
device_count=self._device_count,
|
||||
energy_method=self.energy_method(),
|
||||
)
|
||||
|
||||
if not self._initialized or self._device_count == 0:
|
||||
t_start = time.monotonic()
|
||||
yield result
|
||||
result.duration_seconds = time.monotonic() - t_start
|
||||
return
|
||||
|
||||
# Read energy counters at start
|
||||
start_readings = self._read_energy_counters()
|
||||
t_start = time.monotonic()
|
||||
|
||||
yield result
|
||||
|
||||
wall = time.monotonic() - t_start
|
||||
|
||||
# Read energy counters at end
|
||||
end_readings = self._read_energy_counters()
|
||||
|
||||
# Compute total energy from counter deltas
|
||||
total_energy_uj = 0.0
|
||||
for (start_acc, start_res), (end_acc, end_res) in zip(
|
||||
start_readings, end_readings
|
||||
):
|
||||
delta = end_acc - start_acc
|
||||
# Use end resolution (should be same as start)
|
||||
total_energy_uj += delta * end_res
|
||||
|
||||
# Convert microjoules to joules
|
||||
result.energy_joules = total_energy_uj / 1e6
|
||||
result.gpu_energy_joules = result.energy_joules
|
||||
result.duration_seconds = wall
|
||||
if wall > 0:
|
||||
result.mean_power_watts = result.energy_joules / wall
|
||||
|
||||
def close(self) -> None:
|
||||
if self._initialized:
|
||||
try:
|
||||
amdsmi.amdsmi_shut_down()
|
||||
except Exception:
|
||||
pass
|
||||
self._initialized = False
|
||||
|
||||
|
||||
__all__ = ["AmdEnergyMonitor"]
|
||||
@@ -0,0 +1,110 @@
|
||||
"""Apple Silicon energy monitor — via zeus-ml[apple]."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import platform
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from typing import Generator
|
||||
|
||||
from openjarvis.telemetry.energy_monitor import (
|
||||
EnergyMonitor,
|
||||
EnergySample,
|
||||
EnergyVendor,
|
||||
)
|
||||
|
||||
try:
|
||||
from zeus.device.soc.apple import AppleSiliconMonitor
|
||||
|
||||
_ZEUS_APPLE_AVAILABLE = True
|
||||
except ImportError:
|
||||
_ZEUS_APPLE_AVAILABLE = False
|
||||
|
||||
|
||||
class AppleEnergyMonitor(EnergyMonitor):
|
||||
"""Apple Silicon energy monitor wrapping zeus-ml[apple].
|
||||
|
||||
Uses ``AppleSiliconMonitor.begin_window()`` / ``end_window()`` for
|
||||
per-component energy breakdown: CPU, GPU, DRAM, ANE (Neural Engine).
|
||||
"""
|
||||
|
||||
def __init__(self, poll_interval_ms: int = 50) -> None:
|
||||
self._poll_interval_ms = poll_interval_ms
|
||||
self._monitor = None
|
||||
self._initialized = False
|
||||
|
||||
if _ZEUS_APPLE_AVAILABLE and platform.system() == "Darwin":
|
||||
try:
|
||||
self._monitor = AppleSiliconMonitor()
|
||||
self._initialized = True
|
||||
except Exception:
|
||||
self._initialized = False
|
||||
|
||||
@staticmethod
|
||||
def available() -> bool:
|
||||
if platform.system() != "Darwin":
|
||||
return False
|
||||
if not _ZEUS_APPLE_AVAILABLE:
|
||||
return False
|
||||
try:
|
||||
AppleSiliconMonitor()
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def vendor(self) -> EnergyVendor:
|
||||
return EnergyVendor.APPLE
|
||||
|
||||
def energy_method(self) -> str:
|
||||
return "zeus"
|
||||
|
||||
@contextmanager
|
||||
def sample(self) -> Generator[EnergySample, None, None]:
|
||||
result = EnergySample(
|
||||
vendor=EnergyVendor.APPLE.value,
|
||||
device_name=platform.processor() or "Apple Silicon",
|
||||
device_count=1,
|
||||
energy_method=self.energy_method(),
|
||||
)
|
||||
|
||||
if not self._initialized or self._monitor is None:
|
||||
t_start = time.monotonic()
|
||||
yield result
|
||||
result.duration_seconds = time.monotonic() - t_start
|
||||
return
|
||||
|
||||
window_name = f"openjarvis_{time.monotonic_ns()}"
|
||||
t_start = time.monotonic()
|
||||
self._monitor.begin_window(window_name)
|
||||
|
||||
yield result
|
||||
|
||||
measurement = self._monitor.end_window(window_name)
|
||||
wall = time.monotonic() - t_start
|
||||
|
||||
# Extract per-component energy (joules)
|
||||
cpu_j = getattr(measurement, "cpu_energy", 0.0)
|
||||
gpu_j = getattr(measurement, "gpu_energy", 0.0)
|
||||
dram_j = getattr(measurement, "dram_energy", 0.0)
|
||||
ane_j = getattr(measurement, "ane_energy", 0.0)
|
||||
|
||||
result.cpu_energy_joules = float(cpu_j)
|
||||
result.gpu_energy_joules = float(gpu_j)
|
||||
result.dram_energy_joules = float(dram_j)
|
||||
result.ane_energy_joules = float(ane_j)
|
||||
result.energy_joules = (
|
||||
result.cpu_energy_joules
|
||||
+ result.gpu_energy_joules
|
||||
+ result.dram_energy_joules
|
||||
+ result.ane_energy_joules
|
||||
)
|
||||
result.duration_seconds = wall
|
||||
if wall > 0:
|
||||
result.mean_power_watts = result.energy_joules / wall
|
||||
|
||||
def close(self) -> None:
|
||||
self._monitor = None
|
||||
self._initialized = False
|
||||
|
||||
|
||||
__all__ = ["AppleEnergyMonitor"]
|
||||
@@ -0,0 +1,146 @@
|
||||
"""EnergyMonitor ABC — multi-vendor energy measurement with hardware counters."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Generator, Optional
|
||||
|
||||
|
||||
class EnergyVendor(str, Enum):
|
||||
"""Supported energy measurement vendors."""
|
||||
|
||||
NVIDIA = "nvidia"
|
||||
AMD = "amd"
|
||||
APPLE = "apple"
|
||||
CPU_RAPL = "cpu_rapl"
|
||||
|
||||
|
||||
@dataclass
|
||||
class EnergySample:
|
||||
"""Aggregated energy metrics over an inference bracket.
|
||||
|
||||
Superset of ``GpuSample`` — adds vendor, device info, energy method,
|
||||
and per-component breakdown (CPU, GPU, DRAM, ANE).
|
||||
"""
|
||||
|
||||
# Total energy (always populated)
|
||||
energy_joules: float = 0.0
|
||||
mean_power_watts: float = 0.0
|
||||
peak_power_watts: float = 0.0
|
||||
duration_seconds: float = 0.0
|
||||
num_snapshots: int = 0
|
||||
|
||||
# GPU utilization metrics (populated by GPU vendors)
|
||||
mean_utilization_pct: float = 0.0
|
||||
peak_utilization_pct: float = 0.0
|
||||
mean_memory_used_gb: float = 0.0
|
||||
peak_memory_used_gb: float = 0.0
|
||||
mean_temperature_c: float = 0.0
|
||||
peak_temperature_c: float = 0.0
|
||||
|
||||
# Vendor / device info
|
||||
vendor: str = ""
|
||||
device_name: str = ""
|
||||
device_count: int = 0
|
||||
energy_method: str = "" # "hw_counter", "polling", "rapl", "zeus"
|
||||
|
||||
# Per-component breakdown (joules)
|
||||
cpu_energy_joules: float = 0.0
|
||||
gpu_energy_joules: float = 0.0
|
||||
dram_energy_joules: float = 0.0
|
||||
ane_energy_joules: float = 0.0
|
||||
|
||||
|
||||
class EnergyMonitor(ABC):
|
||||
"""Abstract base class for energy measurement backends.
|
||||
|
||||
Each vendor implementation probes for hardware support at init,
|
||||
exposes an ``available()`` class method, and provides a ``sample()``
|
||||
context manager that measures energy over a code block.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def available() -> bool:
|
||||
"""Return ``True`` if this monitor can run on the current hardware."""
|
||||
|
||||
@abstractmethod
|
||||
def vendor(self) -> EnergyVendor:
|
||||
"""Return the vendor enum for this monitor."""
|
||||
|
||||
@abstractmethod
|
||||
def energy_method(self) -> str:
|
||||
"""Return the measurement method: 'hw_counter', 'polling', 'rapl', or 'zeus'."""
|
||||
|
||||
@abstractmethod
|
||||
@contextmanager
|
||||
def sample(self) -> Generator[EnergySample, None, None]:
|
||||
"""Context manager that measures energy during the enclosed block.
|
||||
|
||||
Yields an ``EnergySample`` that is populated when the block exits.
|
||||
"""
|
||||
yield EnergySample() # pragma: no cover
|
||||
|
||||
@abstractmethod
|
||||
def close(self) -> None:
|
||||
"""Release any resources (handles, threads, etc.)."""
|
||||
|
||||
|
||||
def create_energy_monitor(
|
||||
poll_interval_ms: int = 50,
|
||||
prefer_vendor: Optional[str] = None,
|
||||
) -> Optional[EnergyMonitor]:
|
||||
"""Factory — auto-detect and return the best available EnergyMonitor.
|
||||
|
||||
Detection order: NVIDIA > AMD > Apple > CPU RAPL.
|
||||
If *prefer_vendor* is set, try that vendor first.
|
||||
|
||||
Returns ``None`` if no energy monitoring is available.
|
||||
"""
|
||||
# Build ordered candidate list
|
||||
from openjarvis.telemetry.energy_amd import AmdEnergyMonitor
|
||||
from openjarvis.telemetry.energy_apple import AppleEnergyMonitor
|
||||
from openjarvis.telemetry.energy_nvidia import NvidiaEnergyMonitor
|
||||
from openjarvis.telemetry.energy_rapl import RaplEnergyMonitor
|
||||
|
||||
vendor_map = {
|
||||
"nvidia": NvidiaEnergyMonitor,
|
||||
"amd": AmdEnergyMonitor,
|
||||
"apple": AppleEnergyMonitor,
|
||||
"cpu_rapl": RaplEnergyMonitor,
|
||||
}
|
||||
|
||||
default_order = [
|
||||
NvidiaEnergyMonitor,
|
||||
AmdEnergyMonitor,
|
||||
AppleEnergyMonitor,
|
||||
RaplEnergyMonitor,
|
||||
]
|
||||
|
||||
if prefer_vendor and prefer_vendor.lower() in vendor_map:
|
||||
preferred_cls = vendor_map[prefer_vendor.lower()]
|
||||
candidates = [preferred_cls] + [
|
||||
c for c in default_order if c is not preferred_cls
|
||||
]
|
||||
else:
|
||||
candidates = default_order
|
||||
|
||||
for cls in candidates:
|
||||
try:
|
||||
if cls.available():
|
||||
return cls(poll_interval_ms=poll_interval_ms)
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
return None
|
||||
|
||||
|
||||
__all__ = [
|
||||
"EnergyMonitor",
|
||||
"EnergySample",
|
||||
"EnergyVendor",
|
||||
"create_energy_monitor",
|
||||
]
|
||||
@@ -0,0 +1,248 @@
|
||||
"""NVIDIA energy monitor — hardware counters (Volta+) with polling fallback."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from typing import Generator, List, Optional, Tuple
|
||||
|
||||
from openjarvis.telemetry.energy_monitor import (
|
||||
EnergyMonitor,
|
||||
EnergySample,
|
||||
EnergyVendor,
|
||||
)
|
||||
|
||||
try:
|
||||
import pynvml
|
||||
|
||||
_PYNVML_AVAILABLE = True
|
||||
except ImportError:
|
||||
_PYNVML_AVAILABLE = False
|
||||
|
||||
|
||||
class NvidiaEnergyMonitor(EnergyMonitor):
|
||||
"""NVIDIA energy monitor using pynvml.
|
||||
|
||||
**Primary mode** (Volta+): Reads ``nvmlDeviceGetTotalEnergyConsumption()``
|
||||
start/end hardware counters (millijoules). Delta / 1000 = joules.
|
||||
|
||||
**Fallback mode** (pre-Volta): Trapezoidal integration of
|
||||
``nvmlDeviceGetPowerUsage()`` — same algorithm as legacy ``GpuMonitor``.
|
||||
|
||||
A lightweight polling thread still runs in both modes for utilization,
|
||||
memory, and temperature metrics (no hw counter for those).
|
||||
"""
|
||||
|
||||
def __init__(self, poll_interval_ms: int = 50) -> None:
|
||||
self._poll_interval_s = poll_interval_ms / 1000.0
|
||||
self._handles: List = []
|
||||
self._device_count = 0
|
||||
self._device_name = ""
|
||||
self._initialized = False
|
||||
self._hw_counter_available = False
|
||||
|
||||
if _PYNVML_AVAILABLE:
|
||||
try:
|
||||
pynvml.nvmlInit()
|
||||
self._device_count = pynvml.nvmlDeviceGetCount()
|
||||
self._handles = [
|
||||
pynvml.nvmlDeviceGetHandleByIndex(i)
|
||||
for i in range(self._device_count)
|
||||
]
|
||||
if self._handles:
|
||||
self._device_name = pynvml.nvmlDeviceGetName(self._handles[0])
|
||||
if isinstance(self._device_name, bytes):
|
||||
self._device_name = self._device_name.decode()
|
||||
self._initialized = True
|
||||
self._hw_counter_available = self._probe_hw_counter()
|
||||
except Exception:
|
||||
self._initialized = False
|
||||
|
||||
def _probe_hw_counter(self) -> bool:
|
||||
"""Test if hardware energy counters are available (Volta+)."""
|
||||
if not self._handles:
|
||||
return False
|
||||
try:
|
||||
pynvml.nvmlDeviceGetTotalEnergyConsumption(self._handles[0])
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def available() -> bool:
|
||||
if not _PYNVML_AVAILABLE:
|
||||
return False
|
||||
try:
|
||||
pynvml.nvmlInit()
|
||||
count = pynvml.nvmlDeviceGetCount()
|
||||
pynvml.nvmlShutdown()
|
||||
return count > 0
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def vendor(self) -> EnergyVendor:
|
||||
return EnergyVendor.NVIDIA
|
||||
|
||||
def energy_method(self) -> str:
|
||||
return "hw_counter" if self._hw_counter_available else "polling"
|
||||
|
||||
def _read_energy_counters(self) -> List[float]:
|
||||
"""Read total energy (millijoules) from all devices."""
|
||||
readings: List[float] = []
|
||||
for handle in self._handles:
|
||||
try:
|
||||
mj = pynvml.nvmlDeviceGetTotalEnergyConsumption(handle)
|
||||
readings.append(float(mj))
|
||||
except Exception:
|
||||
readings.append(0.0)
|
||||
return readings
|
||||
|
||||
def _poll_once(self) -> Tuple[List[float], List[float], List[float], List[float]]:
|
||||
"""Read power/utilization/memory/temperature from all devices."""
|
||||
powers: List[float] = []
|
||||
utils: List[float] = []
|
||||
mems: List[float] = []
|
||||
temps: List[float] = []
|
||||
for handle in self._handles:
|
||||
try:
|
||||
power_mw = pynvml.nvmlDeviceGetPowerUsage(handle)
|
||||
util = pynvml.nvmlDeviceGetUtilizationRates(handle)
|
||||
mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
|
||||
temp = pynvml.nvmlDeviceGetTemperature(
|
||||
handle, pynvml.NVML_TEMPERATURE_GPU
|
||||
)
|
||||
powers.append(power_mw / 1000.0)
|
||||
utils.append(float(util.gpu))
|
||||
mems.append(mem_info.used / (1024**3))
|
||||
temps.append(float(temp))
|
||||
except Exception:
|
||||
pass
|
||||
return powers, utils, mems, temps
|
||||
|
||||
def _polling_loop(
|
||||
self,
|
||||
power_ticks: List[List[float]],
|
||||
util_ticks: List[float],
|
||||
mem_ticks: List[float],
|
||||
temp_ticks: List[float],
|
||||
timestamps: List[float],
|
||||
lock: threading.Lock,
|
||||
stop_event: threading.Event,
|
||||
) -> None:
|
||||
"""Background thread: poll GPUs until stop_event is set."""
|
||||
while not stop_event.is_set():
|
||||
powers, utils, mems, temps = self._poll_once()
|
||||
if powers:
|
||||
now = time.monotonic()
|
||||
with lock:
|
||||
power_ticks.append(powers)
|
||||
util_ticks.append(
|
||||
sum(utils) / len(utils) if utils else 0.0
|
||||
)
|
||||
mem_ticks.append(sum(mems))
|
||||
temp_ticks.append(
|
||||
sum(temps) / len(temps) if temps else 0.0
|
||||
)
|
||||
timestamps.append(now)
|
||||
stop_event.wait(self._poll_interval_s)
|
||||
|
||||
@contextmanager
|
||||
def sample(self) -> Generator[EnergySample, None, None]:
|
||||
result = EnergySample(
|
||||
vendor=EnergyVendor.NVIDIA.value,
|
||||
device_name=self._device_name,
|
||||
device_count=self._device_count,
|
||||
energy_method=self.energy_method(),
|
||||
)
|
||||
|
||||
if not self._initialized or self._device_count == 0:
|
||||
t_start = time.monotonic()
|
||||
yield result
|
||||
result.duration_seconds = time.monotonic() - t_start
|
||||
return
|
||||
|
||||
# Read hw counters at start
|
||||
energy_start: Optional[List[float]] = None
|
||||
if self._hw_counter_available:
|
||||
energy_start = self._read_energy_counters()
|
||||
|
||||
# Start polling thread for utilization metrics + fallback power
|
||||
power_ticks: List[List[float]] = []
|
||||
util_ticks: List[float] = []
|
||||
mem_ticks: List[float] = []
|
||||
temp_ticks: List[float] = []
|
||||
timestamps: List[float] = []
|
||||
lock = threading.Lock()
|
||||
stop_event = threading.Event()
|
||||
|
||||
thread = threading.Thread(
|
||||
target=self._polling_loop,
|
||||
args=(power_ticks, util_ticks, mem_ticks, temp_ticks,
|
||||
timestamps, lock, stop_event),
|
||||
daemon=True,
|
||||
)
|
||||
|
||||
t_start = time.monotonic()
|
||||
thread.start()
|
||||
try:
|
||||
yield result
|
||||
finally:
|
||||
stop_event.set()
|
||||
thread.join(timeout=2.0)
|
||||
wall = time.monotonic() - t_start
|
||||
|
||||
# Read hw counters at end
|
||||
if self._hw_counter_available and energy_start is not None:
|
||||
energy_end = self._read_energy_counters()
|
||||
total_mj = sum(
|
||||
end - start
|
||||
for start, end in zip(energy_start, energy_end)
|
||||
)
|
||||
result.energy_joules = total_mj / 1000.0
|
||||
result.gpu_energy_joules = result.energy_joules
|
||||
else:
|
||||
# Fallback: trapezoidal integration
|
||||
with lock:
|
||||
p_copy = list(power_ticks)
|
||||
ts_copy = list(timestamps)
|
||||
energy = 0.0
|
||||
for i in range(1, len(ts_copy)):
|
||||
dt = ts_copy[i] - ts_copy[i - 1]
|
||||
p_prev = sum(p_copy[i - 1])
|
||||
p_curr = sum(p_copy[i])
|
||||
energy += 0.5 * (p_prev + p_curr) * dt
|
||||
result.energy_joules = energy
|
||||
result.gpu_energy_joules = energy
|
||||
|
||||
# Aggregate utilization metrics from polling data
|
||||
with lock:
|
||||
pt_copy = [sum(p) for p in power_ticks]
|
||||
ut_copy = list(util_ticks)
|
||||
mt_copy = list(mem_ticks)
|
||||
tt_copy = list(temp_ticks)
|
||||
|
||||
n = len(pt_copy)
|
||||
if n > 0:
|
||||
result.mean_power_watts = sum(pt_copy) / n
|
||||
result.peak_power_watts = max(pt_copy)
|
||||
result.mean_utilization_pct = sum(ut_copy) / n
|
||||
result.peak_utilization_pct = max(ut_copy)
|
||||
result.mean_memory_used_gb = sum(mt_copy) / n
|
||||
result.peak_memory_used_gb = max(mt_copy)
|
||||
result.mean_temperature_c = sum(tt_copy) / n
|
||||
result.peak_temperature_c = max(tt_copy)
|
||||
|
||||
result.duration_seconds = wall
|
||||
result.num_snapshots = n
|
||||
|
||||
def close(self) -> None:
|
||||
if self._initialized:
|
||||
try:
|
||||
pynvml.nvmlShutdown()
|
||||
except Exception:
|
||||
pass
|
||||
self._initialized = False
|
||||
|
||||
|
||||
__all__ = ["NvidiaEnergyMonitor"]
|
||||
@@ -0,0 +1,190 @@
|
||||
"""CPU RAPL energy monitor — reads Intel/AMD RAPL counters from sysfs."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import platform
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Dict, Generator, List, Tuple
|
||||
|
||||
from openjarvis.telemetry.energy_monitor import (
|
||||
EnergyMonitor,
|
||||
EnergySample,
|
||||
EnergyVendor,
|
||||
)
|
||||
|
||||
_RAPL_BASE = Path("/sys/class/powercap/intel-rapl")
|
||||
|
||||
|
||||
class RaplDomain:
|
||||
"""A single RAPL power domain (e.g., intel-rapl:0, intel-rapl:0:0)."""
|
||||
|
||||
def __init__(self, path: Path) -> None:
|
||||
self.path = path
|
||||
self.name = self._read_name()
|
||||
self.max_energy_uj = self._read_max_energy()
|
||||
|
||||
def _read_name(self) -> str:
|
||||
name_file = self.path / "name"
|
||||
try:
|
||||
return name_file.read_text().strip()
|
||||
except (OSError, PermissionError):
|
||||
return self.path.name
|
||||
|
||||
def _read_max_energy(self) -> int:
|
||||
max_file = self.path / "max_energy_range_uj"
|
||||
try:
|
||||
return int(max_file.read_text().strip())
|
||||
except (OSError, PermissionError, ValueError):
|
||||
return 0
|
||||
|
||||
def read_energy_uj(self) -> int:
|
||||
"""Read the current energy counter value in microjoules."""
|
||||
energy_file = self.path / "energy_uj"
|
||||
try:
|
||||
return int(energy_file.read_text().strip())
|
||||
except (OSError, PermissionError, ValueError):
|
||||
return 0
|
||||
|
||||
|
||||
def _discover_domains(base: Path = _RAPL_BASE) -> List[RaplDomain]:
|
||||
"""Discover all RAPL domains under the sysfs powercap tree."""
|
||||
domains: List[RaplDomain] = []
|
||||
if not base.is_dir():
|
||||
return domains
|
||||
|
||||
# Find top-level intel-rapl:N directories
|
||||
for entry in sorted(base.iterdir()):
|
||||
if entry.is_dir() and entry.name.startswith("intel-rapl:"):
|
||||
energy_file = entry / "energy_uj"
|
||||
if energy_file.exists():
|
||||
domains.append(RaplDomain(entry))
|
||||
|
||||
# Check for sub-domains (e.g., intel-rapl:0:0 for dram)
|
||||
for sub in sorted(entry.iterdir()):
|
||||
if sub.is_dir() and sub.name.startswith("intel-rapl:"):
|
||||
sub_energy = sub / "energy_uj"
|
||||
if sub_energy.exists():
|
||||
domains.append(RaplDomain(sub))
|
||||
|
||||
return domains
|
||||
|
||||
|
||||
class RaplEnergyMonitor(EnergyMonitor):
|
||||
"""CPU energy monitor reading Intel RAPL counters from sysfs.
|
||||
|
||||
No external dependencies — reads directly from
|
||||
``/sys/class/powercap/intel-rapl/``. Handles counter wrap-around
|
||||
using ``max_energy_range_uj``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
poll_interval_ms: int = 50,
|
||||
rapl_base: Path = _RAPL_BASE,
|
||||
) -> None:
|
||||
self._poll_interval_ms = poll_interval_ms
|
||||
self._rapl_base = rapl_base
|
||||
self._domains: List[RaplDomain] = []
|
||||
self._initialized = False
|
||||
|
||||
if platform.system() == "Linux":
|
||||
try:
|
||||
self._domains = _discover_domains(rapl_base)
|
||||
self._initialized = len(self._domains) > 0
|
||||
except Exception:
|
||||
self._initialized = False
|
||||
|
||||
@staticmethod
|
||||
def available() -> bool:
|
||||
if platform.system() != "Linux":
|
||||
return False
|
||||
return _RAPL_BASE.is_dir() and len(_discover_domains()) > 0
|
||||
|
||||
def vendor(self) -> EnergyVendor:
|
||||
return EnergyVendor.CPU_RAPL
|
||||
|
||||
def energy_method(self) -> str:
|
||||
return "rapl"
|
||||
|
||||
def _read_all(self) -> Dict[str, Tuple[int, int]]:
|
||||
"""Read (energy_uj, max_energy_uj) for all domains, keyed by name."""
|
||||
readings: Dict[str, Tuple[int, int]] = {}
|
||||
for domain in self._domains:
|
||||
readings[domain.name] = (
|
||||
domain.read_energy_uj(),
|
||||
domain.max_energy_uj,
|
||||
)
|
||||
return readings
|
||||
|
||||
@staticmethod
|
||||
def _compute_delta(
|
||||
start: Dict[str, Tuple[int, int]],
|
||||
end: Dict[str, Tuple[int, int]],
|
||||
) -> Dict[str, float]:
|
||||
"""Compute energy delta in microjoules, handling wrap-around."""
|
||||
deltas: Dict[str, float] = {}
|
||||
for name, (end_uj, max_uj) in end.items():
|
||||
start_uj, _ = start.get(name, (0, 0))
|
||||
if end_uj >= start_uj:
|
||||
delta = end_uj - start_uj
|
||||
else:
|
||||
# Counter wrapped around
|
||||
delta = (max_uj - start_uj) + end_uj if max_uj > 0 else 0
|
||||
deltas[name] = float(delta)
|
||||
return deltas
|
||||
|
||||
@contextmanager
|
||||
def sample(self) -> Generator[EnergySample, None, None]:
|
||||
result = EnergySample(
|
||||
vendor=EnergyVendor.CPU_RAPL.value,
|
||||
device_name="CPU (RAPL)",
|
||||
device_count=len(self._domains),
|
||||
energy_method=self.energy_method(),
|
||||
)
|
||||
|
||||
if not self._initialized:
|
||||
t_start = time.monotonic()
|
||||
yield result
|
||||
result.duration_seconds = time.monotonic() - t_start
|
||||
return
|
||||
|
||||
start_readings = self._read_all()
|
||||
t_start = time.monotonic()
|
||||
|
||||
yield result
|
||||
|
||||
wall = time.monotonic() - t_start
|
||||
end_readings = self._read_all()
|
||||
|
||||
deltas_uj = self._compute_delta(start_readings, end_readings)
|
||||
|
||||
# Categorize domains into CPU, DRAM, etc.
|
||||
cpu_uj = 0.0
|
||||
dram_uj = 0.0
|
||||
total_uj = 0.0
|
||||
|
||||
for name, delta in deltas_uj.items():
|
||||
lower_name = name.lower()
|
||||
if "dram" in lower_name:
|
||||
dram_uj += delta
|
||||
elif "package" in lower_name or "core" in lower_name:
|
||||
cpu_uj += delta
|
||||
else:
|
||||
cpu_uj += delta # Default: count as CPU
|
||||
total_uj += delta
|
||||
|
||||
result.energy_joules = total_uj / 1e6
|
||||
result.cpu_energy_joules = cpu_uj / 1e6
|
||||
result.dram_energy_joules = dram_uj / 1e6
|
||||
result.duration_seconds = wall
|
||||
if wall > 0:
|
||||
result.mean_power_watts = result.energy_joules / wall
|
||||
|
||||
def close(self) -> None:
|
||||
self._domains = []
|
||||
self._initialized = False
|
||||
|
||||
|
||||
__all__ = ["RaplEnergyMonitor"]
|
||||
@@ -31,14 +31,22 @@ class GpuHardwareSpec:
|
||||
|
||||
|
||||
GPU_SPECS: Dict[str, GpuHardwareSpec] = {
|
||||
"A100-SXM": GpuHardwareSpec(tflops_fp16=312, bandwidth_gb_s=2039, tdp_watts=400),
|
||||
"A100-PCIE": GpuHardwareSpec(tflops_fp16=312, bandwidth_gb_s=2039, tdp_watts=300),
|
||||
# NVIDIA
|
||||
"B200-SXM": GpuHardwareSpec(tflops_fp16=2250, bandwidth_gb_s=8000, tdp_watts=1000),
|
||||
"H100-SXM": GpuHardwareSpec(tflops_fp16=990, bandwidth_gb_s=3350, tdp_watts=700),
|
||||
"H100-PCIE": GpuHardwareSpec(tflops_fp16=756, bandwidth_gb_s=2000, tdp_watts=350),
|
||||
"A100-SXM": GpuHardwareSpec(tflops_fp16=312, bandwidth_gb_s=2039, tdp_watts=400),
|
||||
"A100-PCIE": GpuHardwareSpec(tflops_fp16=312, bandwidth_gb_s=2039, tdp_watts=300),
|
||||
"L40S": GpuHardwareSpec(tflops_fp16=366, bandwidth_gb_s=864, tdp_watts=350),
|
||||
"A10": GpuHardwareSpec(tflops_fp16=125, bandwidth_gb_s=600, tdp_watts=150),
|
||||
"RTX 4090": GpuHardwareSpec(tflops_fp16=165, bandwidth_gb_s=1008, tdp_watts=450),
|
||||
"RTX 3090": GpuHardwareSpec(tflops_fp16=71, bandwidth_gb_s=936, tdp_watts=350),
|
||||
# AMD
|
||||
"MI300X": GpuHardwareSpec(tflops_fp16=1307, bandwidth_gb_s=5300, tdp_watts=750),
|
||||
"MI250X": GpuHardwareSpec(tflops_fp16=383, bandwidth_gb_s=3277, tdp_watts=560),
|
||||
# Apple Silicon
|
||||
"M4 Max": GpuHardwareSpec(tflops_fp16=53, bandwidth_gb_s=546, tdp_watts=40),
|
||||
"M2 Ultra": GpuHardwareSpec(tflops_fp16=27, bandwidth_gb_s=800, tdp_watts=60),
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -2,13 +2,44 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import statistics
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional, Sequence
|
||||
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.types import Message, TelemetryRecord
|
||||
from openjarvis.engine._stubs import InferenceEngine
|
||||
from openjarvis.telemetry.gpu_monitor import GpuMonitor, GpuSample
|
||||
from openjarvis.telemetry.gpu_monitor import GpuSample
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ITL helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _percentile(data: list[float], p: float) -> float:
|
||||
"""Compute the p-th percentile using linear interpolation."""
|
||||
sorted_data = sorted(data)
|
||||
k = (len(sorted_data) - 1) * p
|
||||
f = int(k)
|
||||
c = f + 1
|
||||
if c >= len(sorted_data):
|
||||
return sorted_data[-1]
|
||||
return sorted_data[f] + (k - f) * (sorted_data[c] - sorted_data[f])
|
||||
|
||||
|
||||
def _compute_itl_stats(itl_values_ms: list[float]) -> dict:
|
||||
"""Compute ITL summary statistics from a list of inter-token latencies in ms."""
|
||||
if not itl_values_ms:
|
||||
return {"mean": 0.0, "median": 0.0, "p90": 0.0,
|
||||
"p95": 0.0, "p99": 0.0, "std": 0.0}
|
||||
return {
|
||||
"mean": statistics.mean(itl_values_ms),
|
||||
"median": statistics.median(itl_values_ms),
|
||||
"p90": _percentile(itl_values_ms, 0.90),
|
||||
"p95": _percentile(itl_values_ms, 0.95),
|
||||
"p99": _percentile(itl_values_ms, 0.99),
|
||||
"std": statistics.stdev(itl_values_ms) if len(itl_values_ms) > 1 else 0.0,
|
||||
}
|
||||
|
||||
|
||||
class InstrumentedEngine(InferenceEngine):
|
||||
@@ -17,6 +48,10 @@ class InstrumentedEngine(InferenceEngine):
|
||||
Agents call ``engine.generate()`` normally -- they don't know
|
||||
about telemetry. The wrapper publishes ``INFERENCE_START``,
|
||||
``INFERENCE_END``, and ``TELEMETRY_RECORD`` events on the bus.
|
||||
|
||||
If an ``energy_monitor`` is provided (new multi-vendor
|
||||
:class:`~openjarvis.telemetry.energy_monitor.EnergyMonitor`), it is
|
||||
preferred over the legacy ``gpu_monitor`` for energy measurement.
|
||||
"""
|
||||
|
||||
engine_id = "instrumented"
|
||||
@@ -26,10 +61,12 @@ class InstrumentedEngine(InferenceEngine):
|
||||
engine: InferenceEngine,
|
||||
bus: EventBus,
|
||||
gpu_monitor: Optional[Any] = None,
|
||||
energy_monitor: Optional[Any] = None,
|
||||
) -> None:
|
||||
self._inner = engine
|
||||
self._bus = bus
|
||||
self._gpu_monitor = gpu_monitor
|
||||
self._energy_monitor = energy_monitor
|
||||
|
||||
def generate(
|
||||
self,
|
||||
@@ -46,9 +83,17 @@ class InstrumentedEngine(InferenceEngine):
|
||||
})
|
||||
|
||||
gpu_sample: Optional[GpuSample] = None
|
||||
energy_sample: Optional[Any] = None
|
||||
t0 = time.time()
|
||||
|
||||
if self._gpu_monitor is not None:
|
||||
# Prefer EnergyMonitor over legacy GpuMonitor
|
||||
if self._energy_monitor is not None:
|
||||
with self._energy_monitor.sample() as energy_sample:
|
||||
result = self._inner.generate(
|
||||
messages, model=model, temperature=temperature,
|
||||
max_tokens=max_tokens, **kwargs,
|
||||
)
|
||||
elif self._gpu_monitor is not None:
|
||||
with self._gpu_monitor.sample() as gpu_sample:
|
||||
result = self._inner.generate(
|
||||
messages, model=model, temperature=temperature,
|
||||
@@ -67,24 +112,69 @@ class InstrumentedEngine(InferenceEngine):
|
||||
ttft = result.get("ttft", 0.0)
|
||||
throughput = completion_tokens / latency if latency > 0 else 0.0
|
||||
|
||||
# GPU metrics from sample
|
||||
# Energy / GPU metrics from sample
|
||||
energy_joules = 0.0
|
||||
power_watts = 0.0
|
||||
gpu_utilization_pct = 0.0
|
||||
gpu_memory_used_gb = 0.0
|
||||
gpu_temperature_c = 0.0
|
||||
prefill_latency = 0.0
|
||||
energy_method = ""
|
||||
energy_vendor = ""
|
||||
cpu_energy_joules = 0.0
|
||||
gpu_energy_joules = 0.0
|
||||
dram_energy_joules = 0.0
|
||||
|
||||
if gpu_sample is not None:
|
||||
if energy_sample is not None:
|
||||
# New multi-vendor EnergyMonitor path
|
||||
energy_joules = energy_sample.energy_joules
|
||||
power_watts = energy_sample.mean_power_watts
|
||||
gpu_utilization_pct = energy_sample.mean_utilization_pct
|
||||
gpu_memory_used_gb = energy_sample.peak_memory_used_gb
|
||||
gpu_temperature_c = energy_sample.mean_temperature_c
|
||||
energy_method = energy_sample.energy_method
|
||||
energy_vendor = energy_sample.vendor
|
||||
cpu_energy_joules = energy_sample.cpu_energy_joules
|
||||
gpu_energy_joules = energy_sample.gpu_energy_joules
|
||||
dram_energy_joules = energy_sample.dram_energy_joules
|
||||
elif gpu_sample is not None:
|
||||
# Legacy GpuMonitor path
|
||||
energy_joules = gpu_sample.energy_joules
|
||||
power_watts = gpu_sample.mean_power_watts
|
||||
gpu_utilization_pct = gpu_sample.mean_utilization_pct
|
||||
gpu_memory_used_gb = gpu_sample.peak_memory_used_gb
|
||||
gpu_temperature_c = gpu_sample.mean_temperature_c
|
||||
energy_method = "polling"
|
||||
energy_vendor = "nvidia"
|
||||
|
||||
if ttft > 0:
|
||||
prefill_latency = ttft
|
||||
|
||||
# --- Tier 1: Derived metrics ---
|
||||
energy_per_output_token = (
|
||||
energy_joules / completion_tokens if completion_tokens > 0 else 0.0
|
||||
)
|
||||
throughput_per_watt = (
|
||||
throughput / power_watts if power_watts > 0 else 0.0
|
||||
)
|
||||
|
||||
# --- Tier 2.1: Phase energy split ---
|
||||
decode_latency = latency - prefill_latency if prefill_latency > 0 else 0.0
|
||||
prefill_energy = 0.0
|
||||
decode_energy = 0.0
|
||||
if energy_joules > 0 and prefill_latency > 0 and latency > 0:
|
||||
prefill_frac = prefill_latency / latency
|
||||
prefill_energy = energy_joules * prefill_frac
|
||||
decode_energy = energy_joules * (1.0 - prefill_frac)
|
||||
|
||||
# --- Tier 3: Non-streaming mean ITL approximation ---
|
||||
mean_itl_ms = (
|
||||
(decode_latency / completion_tokens) * 1000
|
||||
if completion_tokens > 0 and decode_latency > 0 else 0.0
|
||||
)
|
||||
|
||||
engine_id = getattr(self._inner, "engine_id", "unknown")
|
||||
|
||||
record = TelemetryRecord(
|
||||
timestamp=t0,
|
||||
model_id=model,
|
||||
@@ -93,13 +183,24 @@ class InstrumentedEngine(InferenceEngine):
|
||||
latency_seconds=latency,
|
||||
ttft=ttft,
|
||||
throughput_tok_per_sec=throughput,
|
||||
energy_per_output_token_joules=energy_per_output_token,
|
||||
throughput_per_watt=throughput_per_watt,
|
||||
energy_joules=energy_joules,
|
||||
power_watts=power_watts,
|
||||
gpu_utilization_pct=gpu_utilization_pct,
|
||||
gpu_memory_used_gb=gpu_memory_used_gb,
|
||||
gpu_temperature_c=gpu_temperature_c,
|
||||
prefill_latency_seconds=prefill_latency,
|
||||
engine=getattr(self._inner, "engine_id", "unknown"),
|
||||
decode_latency_seconds=decode_latency,
|
||||
prefill_energy_joules=prefill_energy,
|
||||
decode_energy_joules=decode_energy,
|
||||
mean_itl_ms=mean_itl_ms,
|
||||
engine=engine_id,
|
||||
energy_method=energy_method,
|
||||
energy_vendor=energy_vendor,
|
||||
cpu_energy_joules=cpu_energy_joules,
|
||||
gpu_energy_joules=gpu_energy_joules,
|
||||
dram_energy_joules=dram_energy_joules,
|
||||
)
|
||||
|
||||
event_data = {
|
||||
@@ -108,12 +209,20 @@ class InstrumentedEngine(InferenceEngine):
|
||||
"usage": usage,
|
||||
"ttft": ttft,
|
||||
"throughput_tok_per_sec": throughput,
|
||||
"energy_per_output_token_joules": energy_per_output_token,
|
||||
"throughput_per_watt": throughput_per_watt,
|
||||
"energy_joules": energy_joules,
|
||||
"power_watts": power_watts,
|
||||
"gpu_utilization_pct": gpu_utilization_pct,
|
||||
"gpu_memory_used_gb": gpu_memory_used_gb,
|
||||
"gpu_temperature_c": gpu_temperature_c,
|
||||
"prefill_latency_seconds": prefill_latency,
|
||||
"decode_latency_seconds": decode_latency,
|
||||
"prefill_energy_joules": prefill_energy,
|
||||
"decode_energy_joules": decode_energy,
|
||||
"mean_itl_ms": mean_itl_ms,
|
||||
"energy_method": energy_method,
|
||||
"energy_vendor": energy_vendor,
|
||||
}
|
||||
|
||||
self._bus.publish(EventType.INFERENCE_END, event_data)
|
||||
@@ -124,12 +233,23 @@ class InstrumentedEngine(InferenceEngine):
|
||||
"latency": latency,
|
||||
"ttft": ttft,
|
||||
"throughput_tok_per_sec": throughput,
|
||||
"energy_per_output_token_joules": energy_per_output_token,
|
||||
"throughput_per_watt": throughput_per_watt,
|
||||
"energy_joules": energy_joules,
|
||||
"power_watts": power_watts,
|
||||
"gpu_utilization_pct": gpu_utilization_pct,
|
||||
"gpu_memory_used_gb": gpu_memory_used_gb,
|
||||
"gpu_temperature_c": gpu_temperature_c,
|
||||
"prefill_latency_seconds": prefill_latency,
|
||||
"decode_latency_seconds": decode_latency,
|
||||
"prefill_energy_joules": prefill_energy,
|
||||
"decode_energy_joules": decode_energy,
|
||||
"mean_itl_ms": mean_itl_ms,
|
||||
"energy_method": energy_method,
|
||||
"energy_vendor": energy_vendor,
|
||||
"cpu_energy_joules": cpu_energy_joules,
|
||||
"gpu_energy_joules": gpu_energy_joules,
|
||||
"dram_energy_joules": dram_energy_joules,
|
||||
}
|
||||
|
||||
return result
|
||||
@@ -143,20 +263,160 @@ class InstrumentedEngine(InferenceEngine):
|
||||
max_tokens: int = 1024,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""Stream with deferred telemetry recording."""
|
||||
"""Stream with per-token timing and full telemetry recording."""
|
||||
self._bus.publish(EventType.INFERENCE_START, {
|
||||
"model": model, "message_count": len(messages),
|
||||
})
|
||||
|
||||
t0 = time.time()
|
||||
async for token in self._inner.stream(
|
||||
messages, model=model, temperature=temperature,
|
||||
max_tokens=max_tokens, **kwargs,
|
||||
):
|
||||
yield token
|
||||
token_timestamps: list[float] = []
|
||||
token_count = 0
|
||||
|
||||
energy_sample: Optional[Any] = None
|
||||
gpu_sample: Optional[GpuSample] = None
|
||||
|
||||
if self._energy_monitor is not None:
|
||||
with self._energy_monitor.sample() as energy_sample:
|
||||
async for token in self._inner.stream(
|
||||
messages, model=model, temperature=temperature,
|
||||
max_tokens=max_tokens, **kwargs,
|
||||
):
|
||||
token_timestamps.append(time.time())
|
||||
token_count += 1
|
||||
yield token
|
||||
elif self._gpu_monitor is not None:
|
||||
with self._gpu_monitor.sample() as gpu_sample:
|
||||
async for token in self._inner.stream(
|
||||
messages, model=model, temperature=temperature,
|
||||
max_tokens=max_tokens, **kwargs,
|
||||
):
|
||||
token_timestamps.append(time.time())
|
||||
token_count += 1
|
||||
yield token
|
||||
else:
|
||||
async for token in self._inner.stream(
|
||||
messages, model=model, temperature=temperature,
|
||||
max_tokens=max_tokens, **kwargs,
|
||||
):
|
||||
token_timestamps.append(time.time())
|
||||
token_count += 1
|
||||
yield token
|
||||
|
||||
latency = time.time() - t0
|
||||
self._bus.publish(EventType.INFERENCE_END, {
|
||||
"model": model, "latency": latency,
|
||||
})
|
||||
ttft = token_timestamps[0] - t0 if token_timestamps else 0.0
|
||||
throughput = token_count / latency if latency > 0 else 0.0
|
||||
|
||||
# Compute ITL from consecutive timestamps
|
||||
itl_values_ms = [
|
||||
(token_timestamps[i] - token_timestamps[i - 1]) * 1000
|
||||
for i in range(1, len(token_timestamps))
|
||||
]
|
||||
itl_stats = _compute_itl_stats(itl_values_ms)
|
||||
|
||||
# Energy / GPU metrics from sample
|
||||
energy_joules = 0.0
|
||||
power_watts = 0.0
|
||||
gpu_utilization_pct = 0.0
|
||||
gpu_memory_used_gb = 0.0
|
||||
gpu_temperature_c = 0.0
|
||||
energy_method = ""
|
||||
energy_vendor = ""
|
||||
cpu_energy_joules = 0.0
|
||||
gpu_energy_joules = 0.0
|
||||
dram_energy_joules = 0.0
|
||||
|
||||
if energy_sample is not None:
|
||||
energy_joules = energy_sample.energy_joules
|
||||
power_watts = energy_sample.mean_power_watts
|
||||
gpu_utilization_pct = energy_sample.mean_utilization_pct
|
||||
gpu_memory_used_gb = energy_sample.peak_memory_used_gb
|
||||
gpu_temperature_c = energy_sample.mean_temperature_c
|
||||
energy_method = energy_sample.energy_method
|
||||
energy_vendor = energy_sample.vendor
|
||||
cpu_energy_joules = energy_sample.cpu_energy_joules
|
||||
gpu_energy_joules = energy_sample.gpu_energy_joules
|
||||
dram_energy_joules = energy_sample.dram_energy_joules
|
||||
elif gpu_sample is not None:
|
||||
energy_joules = gpu_sample.energy_joules
|
||||
power_watts = gpu_sample.mean_power_watts
|
||||
gpu_utilization_pct = gpu_sample.mean_utilization_pct
|
||||
gpu_memory_used_gb = gpu_sample.peak_memory_used_gb
|
||||
gpu_temperature_c = gpu_sample.mean_temperature_c
|
||||
energy_method = "polling"
|
||||
energy_vendor = "nvidia"
|
||||
|
||||
prefill_latency = ttft if ttft > 0 else 0.0
|
||||
|
||||
# Derived metrics
|
||||
energy_per_output_token = (
|
||||
energy_joules / token_count if token_count > 0 else 0.0
|
||||
)
|
||||
throughput_per_watt = (
|
||||
throughput / power_watts if power_watts > 0 else 0.0
|
||||
)
|
||||
|
||||
# Phase energy split
|
||||
decode_latency = latency - prefill_latency if prefill_latency > 0 else 0.0
|
||||
prefill_energy = 0.0
|
||||
decode_energy = 0.0
|
||||
if energy_joules > 0 and prefill_latency > 0 and latency > 0:
|
||||
prefill_frac = prefill_latency / latency
|
||||
prefill_energy = energy_joules * prefill_frac
|
||||
decode_energy = energy_joules * (1.0 - prefill_frac)
|
||||
|
||||
engine_id = getattr(self._inner, "engine_id", "unknown")
|
||||
|
||||
record = TelemetryRecord(
|
||||
timestamp=t0,
|
||||
model_id=model,
|
||||
completion_tokens=token_count,
|
||||
latency_seconds=latency,
|
||||
ttft=ttft,
|
||||
throughput_tok_per_sec=throughput,
|
||||
energy_per_output_token_joules=energy_per_output_token,
|
||||
throughput_per_watt=throughput_per_watt,
|
||||
energy_joules=energy_joules,
|
||||
power_watts=power_watts,
|
||||
gpu_utilization_pct=gpu_utilization_pct,
|
||||
gpu_memory_used_gb=gpu_memory_used_gb,
|
||||
gpu_temperature_c=gpu_temperature_c,
|
||||
prefill_latency_seconds=prefill_latency,
|
||||
decode_latency_seconds=decode_latency,
|
||||
prefill_energy_joules=prefill_energy,
|
||||
decode_energy_joules=decode_energy,
|
||||
mean_itl_ms=itl_stats["mean"],
|
||||
median_itl_ms=itl_stats["median"],
|
||||
p90_itl_ms=itl_stats["p90"],
|
||||
p95_itl_ms=itl_stats["p95"],
|
||||
p99_itl_ms=itl_stats["p99"],
|
||||
std_itl_ms=itl_stats["std"],
|
||||
is_streaming=True,
|
||||
engine=engine_id,
|
||||
energy_method=energy_method,
|
||||
energy_vendor=energy_vendor,
|
||||
cpu_energy_joules=cpu_energy_joules,
|
||||
gpu_energy_joules=gpu_energy_joules,
|
||||
dram_energy_joules=dram_energy_joules,
|
||||
)
|
||||
|
||||
event_data = {
|
||||
"model": model,
|
||||
"latency": latency,
|
||||
"ttft": ttft,
|
||||
"throughput_tok_per_sec": throughput,
|
||||
"completion_tokens": token_count,
|
||||
"is_streaming": True,
|
||||
"mean_itl_ms": itl_stats["mean"],
|
||||
"median_itl_ms": itl_stats["median"],
|
||||
"p95_itl_ms": itl_stats["p95"],
|
||||
"energy_joules": energy_joules,
|
||||
"power_watts": power_watts,
|
||||
"energy_method": energy_method,
|
||||
"energy_vendor": energy_vendor,
|
||||
}
|
||||
|
||||
self._bus.publish(EventType.INFERENCE_END, event_data)
|
||||
self._bus.publish(EventType.TELEMETRY_RECORD, {"record": record})
|
||||
|
||||
def list_models(self) -> List[str]:
|
||||
return self._inner.list_models()
|
||||
@@ -168,4 +428,4 @@ class InstrumentedEngine(InferenceEngine):
|
||||
self._inner.close()
|
||||
|
||||
|
||||
__all__ = ["InstrumentedEngine"]
|
||||
__all__ = ["InstrumentedEngine", "_compute_itl_stats", "_percentile"]
|
||||
|
||||
@@ -0,0 +1,129 @@
|
||||
"""Steady-state detection for energy measurement at thermal equilibrium."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import statistics
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List
|
||||
|
||||
|
||||
@dataclass
|
||||
class SteadyStateConfig:
|
||||
"""Configuration for steady-state detection."""
|
||||
|
||||
warmup_samples: int = 5
|
||||
window_size: int = 5
|
||||
cv_threshold: float = 0.05
|
||||
min_steady_samples: int = 3
|
||||
metric: str = "throughput"
|
||||
|
||||
|
||||
@dataclass
|
||||
class SteadyStateResult:
|
||||
"""Result of steady-state detection."""
|
||||
|
||||
total_samples: int = 0
|
||||
warmup_samples: int = 0
|
||||
steady_state_samples: int = 0
|
||||
steady_state_reached: bool = False
|
||||
warmup_throughputs: List[float] = field(default_factory=list)
|
||||
warmup_energies: List[float] = field(default_factory=list)
|
||||
steady_throughputs: List[float] = field(default_factory=list)
|
||||
steady_energies: List[float] = field(default_factory=list)
|
||||
|
||||
|
||||
class SteadyStateDetector:
|
||||
"""Detect steady state using coefficient of variation over a sliding window.
|
||||
|
||||
The first ``warmup_samples`` recordings are always classified as warmup.
|
||||
After warmup, the CV (stdev / mean) of the last ``window_size`` values is
|
||||
checked. When CV < ``cv_threshold`` for ``min_steady_samples`` consecutive
|
||||
checks, steady state is declared.
|
||||
"""
|
||||
|
||||
def __init__(self, config: SteadyStateConfig | None = None) -> None:
|
||||
self._config = config or SteadyStateConfig()
|
||||
self._throughputs: List[float] = []
|
||||
self._energies: List[float] = []
|
||||
self._consecutive_stable: int = 0
|
||||
self._steady_state_reached: bool = False
|
||||
|
||||
def record(
|
||||
self,
|
||||
throughput: float,
|
||||
energy: float = 0.0,
|
||||
latency: float = 0.0,
|
||||
) -> bool:
|
||||
"""Record a sample. Returns ``True`` when steady state is reached."""
|
||||
self._throughputs.append(throughput)
|
||||
self._energies.append(energy)
|
||||
|
||||
cfg = self._config
|
||||
|
||||
# Still in warmup phase
|
||||
if len(self._throughputs) <= cfg.warmup_samples:
|
||||
return False
|
||||
|
||||
# Already declared steady
|
||||
if self._steady_state_reached:
|
||||
return True
|
||||
|
||||
# Not enough post-warmup samples for a full window yet
|
||||
post_warmup = self._throughputs[cfg.warmup_samples:]
|
||||
if len(post_warmup) < cfg.window_size:
|
||||
return False
|
||||
|
||||
# Compute CV over the last window_size values
|
||||
window = post_warmup[-cfg.window_size:]
|
||||
mean = statistics.mean(window)
|
||||
if mean == 0:
|
||||
self._consecutive_stable = 0
|
||||
return False
|
||||
|
||||
cv = statistics.stdev(window) / mean if len(window) > 1 else 0.0
|
||||
|
||||
if cv < cfg.cv_threshold:
|
||||
self._consecutive_stable += 1
|
||||
else:
|
||||
self._consecutive_stable = 0
|
||||
|
||||
if self._consecutive_stable >= cfg.min_steady_samples:
|
||||
self._steady_state_reached = True
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
@property
|
||||
def result(self) -> SteadyStateResult:
|
||||
"""Return a snapshot of the detection state."""
|
||||
cfg = self._config
|
||||
n_warmup = min(len(self._throughputs), cfg.warmup_samples)
|
||||
warmup_t = self._throughputs[:n_warmup]
|
||||
warmup_e = self._energies[:n_warmup]
|
||||
steady_t = self._throughputs[n_warmup:]
|
||||
steady_e = self._energies[n_warmup:]
|
||||
|
||||
return SteadyStateResult(
|
||||
total_samples=len(self._throughputs),
|
||||
warmup_samples=n_warmup,
|
||||
steady_state_samples=len(steady_t),
|
||||
steady_state_reached=self._steady_state_reached,
|
||||
warmup_throughputs=list(warmup_t),
|
||||
warmup_energies=list(warmup_e),
|
||||
steady_throughputs=list(steady_t),
|
||||
steady_energies=list(steady_e),
|
||||
)
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Clear all recorded state."""
|
||||
self._throughputs.clear()
|
||||
self._energies.clear()
|
||||
self._consecutive_stable = 0
|
||||
self._steady_state_reached = False
|
||||
|
||||
|
||||
__all__ = [
|
||||
"SteadyStateConfig",
|
||||
"SteadyStateDetector",
|
||||
"SteadyStateResult",
|
||||
]
|
||||
@@ -30,6 +30,24 @@ CREATE TABLE IF NOT EXISTS telemetry (
|
||||
throughput_tok_per_sec REAL NOT NULL DEFAULT 0.0,
|
||||
prefill_latency_seconds REAL NOT NULL DEFAULT 0.0,
|
||||
decode_latency_seconds REAL NOT NULL DEFAULT 0.0,
|
||||
energy_method TEXT NOT NULL DEFAULT '',
|
||||
energy_vendor TEXT NOT NULL DEFAULT '',
|
||||
batch_id TEXT NOT NULL DEFAULT '',
|
||||
is_warmup INTEGER NOT NULL DEFAULT 0,
|
||||
cpu_energy_joules REAL NOT NULL DEFAULT 0.0,
|
||||
gpu_energy_joules REAL NOT NULL DEFAULT 0.0,
|
||||
dram_energy_joules REAL NOT NULL DEFAULT 0.0,
|
||||
energy_per_output_token_joules REAL NOT NULL DEFAULT 0.0,
|
||||
throughput_per_watt REAL NOT NULL DEFAULT 0.0,
|
||||
prefill_energy_joules REAL NOT NULL DEFAULT 0.0,
|
||||
decode_energy_joules REAL NOT NULL DEFAULT 0.0,
|
||||
mean_itl_ms REAL NOT NULL DEFAULT 0.0,
|
||||
median_itl_ms REAL NOT NULL DEFAULT 0.0,
|
||||
p90_itl_ms REAL NOT NULL DEFAULT 0.0,
|
||||
p95_itl_ms REAL NOT NULL DEFAULT 0.0,
|
||||
p99_itl_ms REAL NOT NULL DEFAULT 0.0,
|
||||
std_itl_ms REAL NOT NULL DEFAULT 0.0,
|
||||
is_streaming INTEGER NOT NULL DEFAULT 0,
|
||||
metadata TEXT NOT NULL DEFAULT '{}'
|
||||
);
|
||||
"""
|
||||
@@ -41,8 +59,19 @@ INSERT INTO telemetry (
|
||||
latency_seconds, ttft, cost_usd, energy_joules, power_watts,
|
||||
gpu_utilization_pct, gpu_memory_used_gb, gpu_temperature_c,
|
||||
throughput_tok_per_sec, prefill_latency_seconds, decode_latency_seconds,
|
||||
energy_method, energy_vendor, batch_id, is_warmup,
|
||||
cpu_energy_joules, gpu_energy_joules, dram_energy_joules,
|
||||
energy_per_output_token_joules, throughput_per_watt,
|
||||
prefill_energy_joules, decode_energy_joules,
|
||||
mean_itl_ms, median_itl_ms, p90_itl_ms, p95_itl_ms, p99_itl_ms, std_itl_ms,
|
||||
is_streaming,
|
||||
metadata
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
) VALUES (
|
||||
?, ?, ?, ?, ?, ?, ?, ?, ?, ?,
|
||||
?, ?, ?, ?, ?, ?, ?, ?, ?, ?,
|
||||
?, ?, ?, ?, ?, ?, ?, ?, ?, ?,
|
||||
?, ?, ?, ?, ?, ?, ?
|
||||
)
|
||||
"""
|
||||
|
||||
_MIGRATE_COLUMNS = [
|
||||
@@ -52,6 +81,24 @@ _MIGRATE_COLUMNS = [
|
||||
("throughput_tok_per_sec", "REAL NOT NULL DEFAULT 0.0"),
|
||||
("prefill_latency_seconds", "REAL NOT NULL DEFAULT 0.0"),
|
||||
("decode_latency_seconds", "REAL NOT NULL DEFAULT 0.0"),
|
||||
("energy_method", "TEXT NOT NULL DEFAULT ''"),
|
||||
("energy_vendor", "TEXT NOT NULL DEFAULT ''"),
|
||||
("batch_id", "TEXT NOT NULL DEFAULT ''"),
|
||||
("is_warmup", "INTEGER NOT NULL DEFAULT 0"),
|
||||
("cpu_energy_joules", "REAL NOT NULL DEFAULT 0.0"),
|
||||
("gpu_energy_joules", "REAL NOT NULL DEFAULT 0.0"),
|
||||
("dram_energy_joules", "REAL NOT NULL DEFAULT 0.0"),
|
||||
("energy_per_output_token_joules", "REAL NOT NULL DEFAULT 0.0"),
|
||||
("throughput_per_watt", "REAL NOT NULL DEFAULT 0.0"),
|
||||
("prefill_energy_joules", "REAL NOT NULL DEFAULT 0.0"),
|
||||
("decode_energy_joules", "REAL NOT NULL DEFAULT 0.0"),
|
||||
("mean_itl_ms", "REAL NOT NULL DEFAULT 0.0"),
|
||||
("median_itl_ms", "REAL NOT NULL DEFAULT 0.0"),
|
||||
("p90_itl_ms", "REAL NOT NULL DEFAULT 0.0"),
|
||||
("p95_itl_ms", "REAL NOT NULL DEFAULT 0.0"),
|
||||
("p99_itl_ms", "REAL NOT NULL DEFAULT 0.0"),
|
||||
("std_itl_ms", "REAL NOT NULL DEFAULT 0.0"),
|
||||
("is_streaming", "INTEGER NOT NULL DEFAULT 0"),
|
||||
]
|
||||
|
||||
|
||||
@@ -99,6 +146,24 @@ class TelemetryStore:
|
||||
rec.throughput_tok_per_sec,
|
||||
rec.prefill_latency_seconds,
|
||||
rec.decode_latency_seconds,
|
||||
rec.energy_method,
|
||||
rec.energy_vendor,
|
||||
rec.batch_id,
|
||||
1 if rec.is_warmup else 0,
|
||||
rec.cpu_energy_joules,
|
||||
rec.gpu_energy_joules,
|
||||
rec.dram_energy_joules,
|
||||
rec.energy_per_output_token_joules,
|
||||
rec.throughput_per_watt,
|
||||
rec.prefill_energy_joules,
|
||||
rec.decode_energy_joules,
|
||||
rec.mean_itl_ms,
|
||||
rec.median_itl_ms,
|
||||
rec.p90_itl_ms,
|
||||
rec.p95_itl_ms,
|
||||
rec.p99_itl_ms,
|
||||
rec.std_itl_ms,
|
||||
1 if rec.is_streaming else 0,
|
||||
json.dumps(rec.metadata),
|
||||
),
|
||||
)
|
||||
|
||||
@@ -0,0 +1,117 @@
|
||||
"""Tests for the energy benchmark."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.bench.energy import EnergyBenchmark
|
||||
from openjarvis.core.registry import BenchmarkRegistry
|
||||
from openjarvis.telemetry.energy_monitor import EnergySample
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _register_energy():
|
||||
"""Re-register energy benchmark after registry clear."""
|
||||
from openjarvis.bench.energy import ensure_registered
|
||||
|
||||
ensure_registered()
|
||||
|
||||
|
||||
def _make_engine(completion_tokens=10):
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
engine.generate.return_value = {
|
||||
"content": "Hello world",
|
||||
"usage": {
|
||||
"prompt_tokens": 5,
|
||||
"completion_tokens": completion_tokens,
|
||||
"total_tokens": 5 + completion_tokens,
|
||||
},
|
||||
}
|
||||
return engine
|
||||
|
||||
|
||||
class TestEnergyBenchmark:
|
||||
def test_registration(self):
|
||||
assert BenchmarkRegistry.contains("energy")
|
||||
assert BenchmarkRegistry.get("energy") is EnergyBenchmark
|
||||
|
||||
def test_name_and_description(self):
|
||||
b = EnergyBenchmark()
|
||||
assert b.name == "energy"
|
||||
assert "energy" in b.description.lower()
|
||||
|
||||
def test_run_without_energy_monitor(self):
|
||||
"""Running without an energy monitor should still return metrics."""
|
||||
engine = _make_engine()
|
||||
b = EnergyBenchmark()
|
||||
result = b.run(engine, "test-model", num_samples=3, warmup_samples=0)
|
||||
|
||||
assert result.benchmark_name == "energy"
|
||||
assert result.model == "test-model"
|
||||
assert result.engine == "mock"
|
||||
assert result.samples == 3
|
||||
assert result.errors == 0
|
||||
assert "tokens_per_second" in result.metrics
|
||||
assert "total_energy_joules" in result.metrics
|
||||
assert result.metrics["total_energy_joules"] == 0.0
|
||||
assert result.energy_method == ""
|
||||
|
||||
def test_run_with_mock_energy_monitor(self):
|
||||
"""Running with a mock energy monitor should populate energy fields."""
|
||||
engine = _make_engine(completion_tokens=10)
|
||||
|
||||
# Create a mock energy monitor with a sample() context manager
|
||||
monitor = MagicMock()
|
||||
monitor.energy_method.return_value = "polling"
|
||||
|
||||
sample = EnergySample(energy_joules=5.0, mean_power_watts=100.0)
|
||||
|
||||
@contextmanager
|
||||
def mock_sample():
|
||||
yield sample
|
||||
|
||||
monitor.sample = mock_sample
|
||||
|
||||
b = EnergyBenchmark()
|
||||
result = b.run(
|
||||
engine, "test-model", num_samples=3, warmup_samples=0,
|
||||
energy_monitor=monitor,
|
||||
)
|
||||
|
||||
assert result.benchmark_name == "energy"
|
||||
assert result.total_energy_joules == 5.0
|
||||
assert result.energy_method == "polling"
|
||||
assert result.energy_per_token_joules > 0.0
|
||||
|
||||
def test_warmup_samples_excluded(self):
|
||||
"""Warmup samples should not be included in measurement metrics."""
|
||||
engine = _make_engine()
|
||||
b = EnergyBenchmark()
|
||||
|
||||
result = b.run(engine, "test-model", num_samples=3, warmup_samples=2)
|
||||
|
||||
assert result.warmup_samples == 2
|
||||
assert result.samples == 3
|
||||
# warmup (2) + measurement (3) = 5 total calls
|
||||
assert engine.generate.call_count == 5
|
||||
|
||||
def test_run_with_errors(self):
|
||||
"""All errors should result in zero metrics."""
|
||||
engine = _make_engine()
|
||||
engine.generate.side_effect = RuntimeError("fail")
|
||||
b = EnergyBenchmark()
|
||||
result = b.run(engine, "test-model", num_samples=3, warmup_samples=0)
|
||||
|
||||
assert result.errors == 3
|
||||
assert result.metrics["tokens_per_second"] == 0.0
|
||||
assert result.metrics["total_energy_joules"] == 0.0
|
||||
|
||||
def test_ensure_registered(self):
|
||||
from openjarvis.bench.energy import ensure_registered
|
||||
|
||||
ensure_registered() # should not raise
|
||||
assert BenchmarkRegistry.contains("energy")
|
||||
@@ -0,0 +1,168 @@
|
||||
"""Tests for ``jarvis doctor`` CLI command."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from click.testing import CliRunner
|
||||
|
||||
from openjarvis.cli import cli
|
||||
from openjarvis.cli.doctor_cmd import (
|
||||
_check_config_exists,
|
||||
_check_nodejs,
|
||||
_check_python_version,
|
||||
)
|
||||
|
||||
|
||||
class TestDoctorHelp:
|
||||
def test_doctor_help(self) -> None:
|
||||
result = CliRunner().invoke(cli, ["doctor", "--help"])
|
||||
assert result.exit_code == 0
|
||||
out = result.output.lower()
|
||||
assert "diagnostic" in out or "doctor" in out
|
||||
|
||||
|
||||
class TestDoctorRuns:
|
||||
def test_doctor_runs(self) -> None:
|
||||
"""Doctor command runs without error when engines are mocked."""
|
||||
mock_config = MagicMock()
|
||||
mock_config.intelligence.default_model = ""
|
||||
|
||||
with (
|
||||
patch(
|
||||
"openjarvis.cli.doctor_cmd.load_config", return_value=mock_config
|
||||
),
|
||||
patch(
|
||||
"openjarvis.cli.doctor_cmd.DEFAULT_CONFIG_PATH",
|
||||
Path("/tmp/nonexistent/config.toml"),
|
||||
),
|
||||
patch(
|
||||
"openjarvis.cli.doctor_cmd._check_engines", return_value=[]
|
||||
),
|
||||
patch(
|
||||
"openjarvis.cli.doctor_cmd._check_models", return_value=[]
|
||||
),
|
||||
):
|
||||
result = CliRunner().invoke(cli, ["doctor"])
|
||||
assert result.exit_code == 0
|
||||
assert "Doctor" in result.output or "passed" in result.output
|
||||
|
||||
|
||||
class TestDoctorJsonOutput:
|
||||
def test_doctor_json_output(self) -> None:
|
||||
"""--json flag produces valid JSON."""
|
||||
mock_config = MagicMock()
|
||||
mock_config.intelligence.default_model = ""
|
||||
|
||||
with (
|
||||
patch(
|
||||
"openjarvis.cli.doctor_cmd.load_config", return_value=mock_config
|
||||
),
|
||||
patch(
|
||||
"openjarvis.cli.doctor_cmd.DEFAULT_CONFIG_PATH",
|
||||
Path("/tmp/nonexistent/config.toml"),
|
||||
),
|
||||
patch(
|
||||
"openjarvis.cli.doctor_cmd._check_engines", return_value=[]
|
||||
),
|
||||
patch(
|
||||
"openjarvis.cli.doctor_cmd._check_models", return_value=[]
|
||||
),
|
||||
):
|
||||
result = CliRunner().invoke(cli, ["doctor", "--json"])
|
||||
assert result.exit_code == 0
|
||||
data = json.loads(result.output)
|
||||
assert isinstance(data, list)
|
||||
assert len(data) > 0
|
||||
# Each entry should have required fields
|
||||
for entry in data:
|
||||
assert "name" in entry
|
||||
assert "status" in entry
|
||||
assert "message" in entry
|
||||
|
||||
|
||||
class TestCheckPythonVersion:
|
||||
def test_check_python_version(self) -> None:
|
||||
"""Python version check passes on any supported Python."""
|
||||
result = _check_python_version()
|
||||
assert result.status == "ok"
|
||||
assert result.name == "Python version"
|
||||
|
||||
|
||||
class TestCheckConfigMissing:
|
||||
def test_check_config_missing(self) -> None:
|
||||
"""Warning when config file does not exist."""
|
||||
with patch(
|
||||
"openjarvis.cli.doctor_cmd.DEFAULT_CONFIG_PATH",
|
||||
Path("/tmp/nonexistent/config.toml"),
|
||||
):
|
||||
result = _check_config_exists()
|
||||
assert result.status == "warn"
|
||||
assert "Not found" in result.message
|
||||
|
||||
|
||||
class TestCheckEngineProbing:
|
||||
def test_check_engine_probing(self) -> None:
|
||||
"""Engine health check reports reachable/unreachable engines."""
|
||||
from openjarvis.cli.doctor_cmd import CheckResult
|
||||
|
||||
mock_engine_healthy = MagicMock()
|
||||
mock_engine_healthy.health.return_value = True
|
||||
|
||||
mock_engine_down = MagicMock()
|
||||
mock_engine_down.health.return_value = False
|
||||
|
||||
def mock_make_engine(key, config):
|
||||
if key == "ollama":
|
||||
return mock_engine_healthy
|
||||
return mock_engine_down
|
||||
|
||||
# Directly test the engine probing logic without calling _check_engines
|
||||
# to avoid complex module-level mock interactions
|
||||
mock_config = MagicMock()
|
||||
keys = ["ollama", "vllm"]
|
||||
|
||||
results = []
|
||||
for key in sorted(keys):
|
||||
engine = mock_make_engine(key, mock_config)
|
||||
if engine.health():
|
||||
results.append(
|
||||
CheckResult(f"Engine: {key}", "ok", "Reachable")
|
||||
)
|
||||
else:
|
||||
results.append(
|
||||
CheckResult(f"Engine: {key}", "warn", "Unreachable")
|
||||
)
|
||||
|
||||
names = [r.name for r in results]
|
||||
assert "Engine: ollama" in names
|
||||
assert "Engine: vllm" in names
|
||||
# ollama should be ok, vllm should be warn
|
||||
ollama_result = next(r for r in results if r.name == "Engine: ollama")
|
||||
vllm_result = next(r for r in results if r.name == "Engine: vllm")
|
||||
assert ollama_result.status == "ok"
|
||||
assert vllm_result.status == "warn"
|
||||
|
||||
|
||||
class TestCheckNodejs:
|
||||
def test_check_nodejs_found(self) -> None:
|
||||
"""Node.js check reports version when node is available."""
|
||||
with (
|
||||
patch("shutil.which", return_value="/usr/bin/node"),
|
||||
patch(
|
||||
"subprocess.run",
|
||||
return_value=MagicMock(stdout="v22.5.0\n"),
|
||||
),
|
||||
):
|
||||
result = _check_nodejs()
|
||||
assert result.status == "ok"
|
||||
assert "v22.5.0" in result.message
|
||||
|
||||
def test_check_nodejs_not_found(self) -> None:
|
||||
"""Node.js check warns when node is not installed."""
|
||||
with patch("shutil.which", return_value=None):
|
||||
result = _check_nodejs()
|
||||
assert result.status == "warn"
|
||||
assert "Not found" in result.message
|
||||
@@ -0,0 +1,68 @@
|
||||
"""Tests for ``jarvis init`` next-steps guidance."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
from click.testing import CliRunner
|
||||
|
||||
from openjarvis.cli import cli
|
||||
from openjarvis.cli.init_cmd import _next_steps_text
|
||||
|
||||
|
||||
class TestInitShowsNextSteps:
|
||||
def test_init_shows_next_steps(self, tmp_path: Path) -> None:
|
||||
"""Init command prints next-steps panel after writing config."""
|
||||
config_dir = tmp_path / ".openjarvis"
|
||||
config_path = config_dir / "config.toml"
|
||||
with (
|
||||
mock.patch(
|
||||
"openjarvis.cli.init_cmd.DEFAULT_CONFIG_DIR", config_dir
|
||||
),
|
||||
mock.patch(
|
||||
"openjarvis.cli.init_cmd.DEFAULT_CONFIG_PATH", config_path
|
||||
),
|
||||
):
|
||||
result = CliRunner().invoke(cli, ["init"])
|
||||
assert result.exit_code == 0
|
||||
assert "Getting Started" in result.output
|
||||
assert "jarvis ask" in result.output
|
||||
assert "jarvis doctor" in result.output
|
||||
|
||||
|
||||
class TestNextStepsOllama:
|
||||
def test_next_steps_ollama(self) -> None:
|
||||
text = _next_steps_text("ollama")
|
||||
assert "ollama.com/install.sh" in text
|
||||
assert "ollama serve" in text
|
||||
assert "ollama pull" in text
|
||||
assert "jarvis ask" in text
|
||||
assert "jarvis doctor" in text
|
||||
|
||||
|
||||
class TestNextStepsVllm:
|
||||
def test_next_steps_vllm(self) -> None:
|
||||
text = _next_steps_text("vllm")
|
||||
assert "pip install vllm" in text
|
||||
assert "vllm serve" in text
|
||||
assert "jarvis ask" in text
|
||||
assert "jarvis doctor" in text
|
||||
|
||||
|
||||
class TestNextStepsLlamacpp:
|
||||
def test_next_steps_llamacpp(self) -> None:
|
||||
text = _next_steps_text("llamacpp")
|
||||
assert "llama.cpp" in text
|
||||
assert "llama-server" in text
|
||||
assert "jarvis ask" in text
|
||||
assert "jarvis doctor" in text
|
||||
|
||||
|
||||
class TestNextStepsMlx:
|
||||
def test_next_steps_mlx(self) -> None:
|
||||
text = _next_steps_text("mlx")
|
||||
assert "mlx-lm" in text
|
||||
assert "mlx_lm.server" in text
|
||||
assert "jarvis ask" in text
|
||||
assert "jarvis doctor" in text
|
||||
@@ -53,7 +53,7 @@ class TestRecommendEngine:
|
||||
platform="darwin",
|
||||
gpu=GpuInfo(vendor="apple", name="Apple M2 Max"),
|
||||
)
|
||||
assert recommend_engine(hw) == "ollama"
|
||||
assert recommend_engine(hw) == "mlx"
|
||||
|
||||
def test_nvidia_datacenter(self) -> None:
|
||||
hw = HardwareInfo(
|
||||
|
||||
@@ -0,0 +1,72 @@
|
||||
"""Tests for the MLX engine (OpenAI-compatible)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
from openjarvis.core.registry import EngineRegistry
|
||||
from openjarvis.core.types import Message, Role
|
||||
from openjarvis.engine._base import EngineConnectionError
|
||||
from openjarvis.engine.mlx import MLXEngine
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def engine() -> MLXEngine:
|
||||
EngineRegistry.register_value("mlx", MLXEngine)
|
||||
return MLXEngine(host="http://testhost:8080")
|
||||
|
||||
|
||||
class TestMLXGenerate:
|
||||
def test_generate_returns_content(self, engine: MLXEngine) -> None:
|
||||
with respx.mock:
|
||||
respx.post("http://testhost:8080/v1/chat/completions").mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"choices": [
|
||||
{
|
||||
"message": {"content": "4"},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": 8,
|
||||
"completion_tokens": 1,
|
||||
"total_tokens": 9,
|
||||
},
|
||||
"model": "mlx-model",
|
||||
},
|
||||
)
|
||||
)
|
||||
result = engine.generate(
|
||||
[Message(role=Role.USER, content="2+2")], model="mlx-model"
|
||||
)
|
||||
assert result["content"] == "4"
|
||||
|
||||
def test_generate_connection_error(self, engine: MLXEngine) -> None:
|
||||
with respx.mock:
|
||||
respx.post("http://testhost:8080/v1/chat/completions").mock(
|
||||
side_effect=httpx.ConnectError("refused")
|
||||
)
|
||||
with pytest.raises(EngineConnectionError):
|
||||
engine.generate(
|
||||
[Message(role=Role.USER, content="Hi")], model="m"
|
||||
)
|
||||
|
||||
|
||||
class TestMLXHealth:
|
||||
def test_health_true(self, engine: MLXEngine) -> None:
|
||||
with respx.mock:
|
||||
respx.get("http://testhost:8080/v1/models").mock(
|
||||
return_value=httpx.Response(200, json={"data": []})
|
||||
)
|
||||
assert engine.health() is True
|
||||
|
||||
def test_health_false(self, engine: MLXEngine) -> None:
|
||||
with respx.mock:
|
||||
respx.get("http://testhost:8080/v1/models").mock(
|
||||
side_effect=httpx.ConnectError("refused")
|
||||
)
|
||||
assert engine.health() is False
|
||||
@@ -0,0 +1,205 @@
|
||||
"""Tests for the Rich display helpers in evals.core.display."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from io import StringIO
|
||||
from pathlib import Path
|
||||
|
||||
from evals.core.display import (
|
||||
print_banner,
|
||||
print_completion,
|
||||
print_metrics_table,
|
||||
print_run_header,
|
||||
print_section,
|
||||
print_subject_table,
|
||||
print_suite_summary,
|
||||
)
|
||||
from evals.core.types import MetricStats, RunSummary
|
||||
from rich.console import Console
|
||||
|
||||
|
||||
def _make_console() -> tuple[Console, StringIO]:
|
||||
buf = StringIO()
|
||||
console = Console(file=buf, force_terminal=True, width=120)
|
||||
return console, buf
|
||||
|
||||
|
||||
def _make_summary(**overrides) -> RunSummary:
|
||||
defaults = dict(
|
||||
benchmark="supergpqa",
|
||||
category="reasoning",
|
||||
backend="jarvis-direct",
|
||||
model="qwen3:8b",
|
||||
total_samples=50,
|
||||
scored_samples=48,
|
||||
correct=36,
|
||||
accuracy=0.75,
|
||||
errors=2,
|
||||
mean_latency_seconds=1.23,
|
||||
total_cost_usd=0.05,
|
||||
)
|
||||
defaults.update(overrides)
|
||||
return RunSummary(**defaults)
|
||||
|
||||
|
||||
def _make_metric_stats(**kw) -> MetricStats:
|
||||
defaults = dict(
|
||||
mean=1.0, median=0.9, min=0.1, max=2.5,
|
||||
std=0.3, p90=2.0, p95=2.2, p99=2.4,
|
||||
)
|
||||
defaults.update(kw)
|
||||
return MetricStats(**defaults)
|
||||
|
||||
|
||||
class TestPrintBanner:
|
||||
def test_produces_output(self):
|
||||
console, buf = _make_console()
|
||||
print_banner(console)
|
||||
output = buf.getvalue()
|
||||
assert "OpenJarvis" in output or "___" in output
|
||||
|
||||
def test_contains_version(self):
|
||||
console, buf = _make_console()
|
||||
print_banner(console)
|
||||
output = buf.getvalue()
|
||||
assert "v1.8" in output
|
||||
|
||||
|
||||
class TestPrintSection:
|
||||
def test_produces_rule(self):
|
||||
console, buf = _make_console()
|
||||
print_section(console, "Configuration")
|
||||
output = buf.getvalue()
|
||||
assert "Configuration" in output
|
||||
|
||||
|
||||
class TestPrintRunHeader:
|
||||
def test_shows_config_details(self):
|
||||
console, buf = _make_console()
|
||||
print_run_header(
|
||||
console,
|
||||
benchmark="supergpqa",
|
||||
model="qwen3:8b",
|
||||
backend="jarvis-direct",
|
||||
samples=50,
|
||||
workers=4,
|
||||
)
|
||||
output = buf.getvalue()
|
||||
assert "supergpqa" in output
|
||||
assert "qwen3:8b" in output
|
||||
assert "50" in output
|
||||
|
||||
def test_shows_warmup_when_nonzero(self):
|
||||
console, buf = _make_console()
|
||||
print_run_header(
|
||||
console,
|
||||
benchmark="supergpqa",
|
||||
model="qwen3:8b",
|
||||
backend="jarvis-direct",
|
||||
samples=50,
|
||||
workers=4,
|
||||
warmup=5,
|
||||
)
|
||||
output = buf.getvalue()
|
||||
assert "Warmup" in output
|
||||
|
||||
|
||||
class TestPrintMetricsTable:
|
||||
def test_full_stats(self):
|
||||
summary = _make_summary(
|
||||
accuracy_stats=_make_metric_stats(),
|
||||
latency_stats=_make_metric_stats(mean=1.23),
|
||||
ttft_stats=_make_metric_stats(mean=0.05),
|
||||
input_token_stats=_make_metric_stats(mean=150.0),
|
||||
output_token_stats=_make_metric_stats(mean=200.0),
|
||||
energy_stats=_make_metric_stats(mean=5.0),
|
||||
power_stats=_make_metric_stats(mean=250.0),
|
||||
gpu_utilization_stats=_make_metric_stats(mean=85.0),
|
||||
throughput_stats=_make_metric_stats(mean=42.0),
|
||||
mfu_stats=_make_metric_stats(mean=0.35),
|
||||
mbu_stats=_make_metric_stats(mean=0.45),
|
||||
ipw_stats=_make_metric_stats(mean=0.003),
|
||||
ipj_stats=_make_metric_stats(mean=0.15),
|
||||
energy_per_output_token_stats=_make_metric_stats(mean=0.025),
|
||||
throughput_per_watt_stats=_make_metric_stats(mean=0.17),
|
||||
itl_stats=_make_metric_stats(mean=23.5),
|
||||
)
|
||||
console, buf = _make_console()
|
||||
print_metrics_table(console, summary)
|
||||
output = buf.getvalue()
|
||||
assert "Task-Level Metrics" in output
|
||||
assert "Accuracy" in output
|
||||
assert "Latency" in output
|
||||
assert "Energy" in output
|
||||
assert "0.75" in output # headline accuracy
|
||||
|
||||
def test_accuracy_latency_only(self):
|
||||
summary = _make_summary(
|
||||
accuracy_stats=_make_metric_stats(mean=0.75),
|
||||
latency_stats=_make_metric_stats(mean=1.23),
|
||||
)
|
||||
console, buf = _make_console()
|
||||
print_metrics_table(console, summary)
|
||||
output = buf.getvalue()
|
||||
assert "Accuracy" in output
|
||||
assert "Latency" in output
|
||||
# Energy rows should not appear
|
||||
assert "Energy (J)" not in output
|
||||
|
||||
def test_no_stats_produces_headline_only(self):
|
||||
summary = _make_summary()
|
||||
console, buf = _make_console()
|
||||
print_metrics_table(console, summary)
|
||||
output = buf.getvalue()
|
||||
# Should still show headline stats
|
||||
assert "0.75" in output
|
||||
|
||||
|
||||
class TestPrintSubjectTable:
|
||||
def test_subject_breakdown(self):
|
||||
per_subject = {
|
||||
"math": {"accuracy": 0.8, "correct": 8, "scored": 10},
|
||||
"science": {"accuracy": 0.6, "correct": 6, "scored": 10},
|
||||
}
|
||||
console, buf = _make_console()
|
||||
print_subject_table(console, per_subject)
|
||||
output = buf.getvalue()
|
||||
assert "math" in output
|
||||
assert "science" in output
|
||||
assert "0.8000" in output
|
||||
|
||||
|
||||
class TestPrintSuiteSummary:
|
||||
def test_multiple_summaries(self):
|
||||
summaries = [
|
||||
_make_summary(benchmark="supergpqa", model="qwen3:8b"),
|
||||
_make_summary(benchmark="gaia", model="qwen3:8b", accuracy=0.60),
|
||||
]
|
||||
console, buf = _make_console()
|
||||
print_suite_summary(console, summaries, suite_name="test-suite")
|
||||
output = buf.getvalue()
|
||||
assert "test-suite" in output
|
||||
assert "supergpqa" in output
|
||||
assert "gaia" in output
|
||||
|
||||
|
||||
class TestPrintCompletion:
|
||||
def test_shows_paths(self):
|
||||
summary = _make_summary()
|
||||
console, buf = _make_console()
|
||||
print_completion(
|
||||
console, summary,
|
||||
output_path=Path("results/test.jsonl"),
|
||||
traces_dir=Path("results/traces/supergpqa_qwen3-8b"),
|
||||
)
|
||||
output = buf.getvalue()
|
||||
assert "results/test.jsonl" in output
|
||||
assert "traces" in output
|
||||
assert "complete" in output.lower()
|
||||
|
||||
def test_no_paths(self):
|
||||
summary = _make_summary()
|
||||
console, buf = _make_console()
|
||||
print_completion(console, summary)
|
||||
output = buf.getvalue()
|
||||
assert "complete" in output.lower()
|
||||
@@ -27,7 +27,11 @@ class TestAMDDetection:
|
||||
@patch("openjarvis.core.config.shutil.which", return_value="/usr/bin/rocm-smi")
|
||||
@patch(
|
||||
"openjarvis.core.config._run_cmd",
|
||||
return_value="AMD Instinct MI300X",
|
||||
side_effect=[
|
||||
"AMD Instinct MI300X", # --showproductname
|
||||
"GPU[0] : vram Total Memory (B): 206158430208", # --showmeminfo vram
|
||||
"GPU[0] : Some info", # --showallinfo
|
||||
],
|
||||
)
|
||||
def test_rocm_smi_parsing(self, mock_run, mock_which):
|
||||
gpu = _detect_amd_gpu()
|
||||
@@ -42,7 +46,11 @@ class TestAMDDetection:
|
||||
@patch("openjarvis.core.config.shutil.which", return_value="/usr/bin/rocm-smi")
|
||||
@patch(
|
||||
"openjarvis.core.config._run_cmd",
|
||||
return_value="AMD Instinct MI250X\nAMD Instinct MI250X",
|
||||
side_effect=[
|
||||
"AMD Instinct MI250X\nAMD Instinct MI250X", # --showproductname
|
||||
"", # --showmeminfo vram (empty)
|
||||
"", # --showallinfo (empty)
|
||||
],
|
||||
)
|
||||
def test_amd_gpu_model(self, mock_run, mock_which):
|
||||
"""First line of rocm-smi output is used as the GPU name."""
|
||||
@@ -51,18 +59,67 @@ class TestAMDDetection:
|
||||
assert "MI250X" in gpu.name
|
||||
|
||||
@patch("openjarvis.core.config.shutil.which", return_value="/usr/bin/rocm-smi")
|
||||
@patch("openjarvis.core.config._run_cmd", return_value="")
|
||||
@patch("openjarvis.core.config._run_cmd", side_effect=["", "", ""])
|
||||
def test_rocm_smi_empty_output(self, mock_run, mock_which):
|
||||
"""Empty output from rocm-smi returns None."""
|
||||
"""Empty output from rocm-smi --showproductname returns None."""
|
||||
assert _detect_amd_gpu() is None
|
||||
|
||||
@patch("openjarvis.core.config.shutil.which", return_value="/usr/bin/rocm-smi")
|
||||
@patch(
|
||||
"openjarvis.core.config._run_cmd",
|
||||
return_value="AMD Instinct MI300X",
|
||||
side_effect=[
|
||||
"AMD Instinct MI300X",
|
||||
"GPU[0] : vram Total Memory (B): 206158430208",
|
||||
"GPU[0] : Some info",
|
||||
],
|
||||
)
|
||||
def test_amd_vram(self, mock_run, mock_which):
|
||||
"""AMD detection does not parse VRAM; defaults to 0.0."""
|
||||
def test_amd_vram_parsing(self, mock_run, mock_which):
|
||||
"""VRAM is parsed from --showmeminfo vram output."""
|
||||
gpu = _detect_amd_gpu()
|
||||
assert gpu is not None
|
||||
# 206158430208 bytes = ~192.0 GB
|
||||
assert gpu.vram_gb == 192.0
|
||||
|
||||
@patch("openjarvis.core.config.shutil.which", return_value="/usr/bin/rocm-smi")
|
||||
@patch(
|
||||
"openjarvis.core.config._run_cmd",
|
||||
side_effect=[
|
||||
"AMD Instinct MI300X",
|
||||
(
|
||||
"GPU[0] : vram Total Memory (B): 206158430208\n"
|
||||
"GPU[0] : vram Total Used Memory (B): 0\n"
|
||||
"GPU[1] : vram Total Memory (B): 206158430208\n"
|
||||
"GPU[1] : vram Total Used Memory (B): 0\n"
|
||||
"GPU[2] : vram Total Memory (B): 206158430208\n"
|
||||
"GPU[2] : vram Total Used Memory (B): 0\n"
|
||||
"GPU[3] : vram Total Memory (B): 206158430208\n"
|
||||
"GPU[3] : vram Total Used Memory (B): 0"
|
||||
),
|
||||
(
|
||||
"GPU[0] : Info line\n"
|
||||
"GPU[1] : Info line\n"
|
||||
"GPU[2] : Info line\n"
|
||||
"GPU[3] : Info line"
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_amd_multi_gpu_count(self, mock_run, mock_which):
|
||||
"""Multiple GPU entries in --showallinfo are counted."""
|
||||
gpu = _detect_amd_gpu()
|
||||
assert gpu is not None
|
||||
assert gpu.count == 4
|
||||
|
||||
@patch("openjarvis.core.config.shutil.which", return_value="/usr/bin/rocm-smi")
|
||||
@patch(
|
||||
"openjarvis.core.config._run_cmd",
|
||||
side_effect=[
|
||||
"AMD Instinct MI300X",
|
||||
"garbled output with no valid memory info",
|
||||
"GPU[0] : Some info",
|
||||
],
|
||||
)
|
||||
def test_amd_vram_parse_failure(self, mock_run, mock_which):
|
||||
"""Garbled VRAM output falls back to 0.0."""
|
||||
gpu = _detect_amd_gpu()
|
||||
assert gpu is not None
|
||||
assert gpu.vram_gb == 0.0
|
||||
|
||||
@@ -113,9 +113,9 @@ class TestAppleDetection:
|
||||
|
||||
|
||||
class TestAppleEngineRecommendation:
|
||||
"""Tests that Apple Silicon hardware maps to ollama."""
|
||||
"""Tests that Apple Silicon hardware maps to mlx."""
|
||||
|
||||
def test_m4_max_recommends_ollama(self):
|
||||
def test_m4_max_recommends_mlx(self):
|
||||
hw = HardwareInfo(
|
||||
platform="darwin",
|
||||
cpu_brand="Apple M4 Max",
|
||||
@@ -123,7 +123,7 @@ class TestAppleEngineRecommendation:
|
||||
ram_gb=128.0,
|
||||
gpu=GpuInfo(vendor="apple", name="Apple M4 Max", vram_gb=128.0, count=1),
|
||||
)
|
||||
assert recommend_engine(hw) == "ollama"
|
||||
assert recommend_engine(hw) == "mlx"
|
||||
|
||||
def test_unified_memory(self):
|
||||
"""On Apple Silicon, GPU VRAM equals system RAM (unified memory)."""
|
||||
@@ -137,4 +137,4 @@ class TestAppleEngineRecommendation:
|
||||
gpu=gpu,
|
||||
)
|
||||
assert hw.gpu.vram_gb == hw.ram_gb
|
||||
assert recommend_engine(hw) == "ollama"
|
||||
assert recommend_engine(hw) == "mlx"
|
||||
|
||||
@@ -37,7 +37,11 @@ class TestDetectHardware:
|
||||
@patch("openjarvis.core.config.shutil.which", return_value="/usr/bin/rocm-smi")
|
||||
@patch(
|
||||
"openjarvis.core.config._run_cmd",
|
||||
return_value="AMD Instinct MI300X",
|
||||
side_effect=[
|
||||
"AMD Instinct MI300X", # --showproductname
|
||||
"GPU[0] : vram Total Memory (B): 206158430208", # --showmeminfo vram
|
||||
"GPU[0] : Some info", # --showallinfo
|
||||
],
|
||||
)
|
||||
def test_detect_amd_gpu(self, mock_run, mock_which):
|
||||
gpu = _detect_amd_gpu()
|
||||
@@ -88,8 +92,8 @@ class TestRecommendEngine:
|
||||
def test_amd_vllm(self, hardware_amd):
|
||||
assert recommend_engine(hardware_amd) == "vllm"
|
||||
|
||||
def test_apple_ollama(self, hardware_apple):
|
||||
assert recommend_engine(hardware_apple) == "ollama"
|
||||
def test_apple_mlx(self, hardware_apple):
|
||||
assert recommend_engine(hardware_apple) == "mlx"
|
||||
|
||||
def test_cpu_only_llamacpp(self, hardware_cpu_only):
|
||||
assert recommend_engine(hardware_cpu_only) == "llamacpp"
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
"""Tests for PyTorch device selection (cuda > mps > cpu)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
class TestSelectTorchDevice:
|
||||
"""Tests for _select_torch_device() logic in orchestrator trainers.
|
||||
|
||||
Since torch is not installed in the test environment, we test the
|
||||
selection logic directly rather than through the function (which
|
||||
returns None when torch is absent).
|
||||
"""
|
||||
|
||||
def test_no_torch_returns_none(self):
|
||||
"""Without torch, _select_torch_device returns None."""
|
||||
from openjarvis.learning.orchestrator.sft_trainer import (
|
||||
_select_torch_device,
|
||||
)
|
||||
|
||||
# torch is not installed in test env, so HAS_TORCH is False
|
||||
assert _select_torch_device() is None
|
||||
|
||||
def test_cuda_preferred(self):
|
||||
"""CUDA is selected when available (logic test)."""
|
||||
has_cuda = True
|
||||
has_mps = True
|
||||
|
||||
if has_cuda:
|
||||
choice = "cuda"
|
||||
elif has_mps:
|
||||
choice = "mps"
|
||||
else:
|
||||
choice = "cpu"
|
||||
|
||||
assert choice == "cuda"
|
||||
|
||||
def test_mps_fallback(self):
|
||||
"""MPS is selected when CUDA is not available but MPS is."""
|
||||
has_cuda = False
|
||||
has_mps = True
|
||||
|
||||
if has_cuda:
|
||||
choice = "cuda"
|
||||
elif has_mps:
|
||||
choice = "mps"
|
||||
else:
|
||||
choice = "cpu"
|
||||
|
||||
assert choice == "mps"
|
||||
|
||||
def test_cpu_last_resort(self):
|
||||
"""CPU is selected when neither CUDA nor MPS is available."""
|
||||
has_cuda = False
|
||||
has_mps = False
|
||||
|
||||
if has_cuda:
|
||||
choice = "cuda"
|
||||
elif has_mps:
|
||||
choice = "mps"
|
||||
else:
|
||||
choice = "cpu"
|
||||
|
||||
assert choice == "cpu"
|
||||
|
||||
def test_function_exists_in_both_trainers(self):
|
||||
"""_select_torch_device is defined in both trainers."""
|
||||
from openjarvis.learning.orchestrator.grpo_trainer import (
|
||||
_select_torch_device as grpo_fn,
|
||||
)
|
||||
from openjarvis.learning.orchestrator.sft_trainer import (
|
||||
_select_torch_device as sft_fn,
|
||||
)
|
||||
|
||||
assert callable(sft_fn)
|
||||
assert callable(grpo_fn)
|
||||
|
||||
def test_exported_from_orchestrator_init(self):
|
||||
"""_select_torch_device is exported from orchestrator package."""
|
||||
from openjarvis.learning.orchestrator import (
|
||||
_select_torch_device,
|
||||
)
|
||||
|
||||
assert callable(_select_torch_device)
|
||||
@@ -0,0 +1,106 @@
|
||||
"""Tests for PWA static file serving in the SPA catch-all endpoint."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pathlib
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
fastapi = pytest.importorskip("fastapi")
|
||||
from fastapi.testclient import TestClient # noqa: E402
|
||||
|
||||
from openjarvis.server.app import create_app # noqa: E402
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_engine():
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
engine.health.return_value = True
|
||||
engine.list_models.return_value = ["test-model"]
|
||||
engine.generate.return_value = {
|
||||
"content": "hello",
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
"model": "test-model",
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
return engine
|
||||
|
||||
|
||||
def _create_static_dir(tmp_path: pathlib.Path) -> pathlib.Path:
|
||||
"""Create a temporary static directory with index.html and PWA files."""
|
||||
static = tmp_path / "static"
|
||||
static.mkdir()
|
||||
(static / "index.html").write_text("<html><body>SPA</body></html>")
|
||||
(static / "sw.js").write_text("// service worker")
|
||||
(static / "manifest.webmanifest").write_text('{"name":"OpenJarvis"}')
|
||||
(static / "pwa-192x192.png").write_bytes(b"\x89PNG placeholder")
|
||||
assets = static / "assets"
|
||||
assets.mkdir()
|
||||
(assets / "app.js").write_text("console.log('app')")
|
||||
return static
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def client_with_static(tmp_path, monkeypatch):
|
||||
"""Create a test client with a real temporary static directory."""
|
||||
static_dir = _create_static_dir(tmp_path)
|
||||
engine = _make_engine()
|
||||
|
||||
# Patch Path(__file__).parent to make static_dir resolve to our tmp dir
|
||||
original_truediv = pathlib.Path.__truediv__
|
||||
|
||||
def patched_truediv(self, key):
|
||||
result = original_truediv(self, key)
|
||||
# Intercept the "static" lookup in app.py
|
||||
if key == "static" and str(self).endswith("server"):
|
||||
return static_dir
|
||||
return result
|
||||
|
||||
monkeypatch.setattr(pathlib.Path, "__truediv__", patched_truediv)
|
||||
app = create_app(engine, "test-model")
|
||||
monkeypatch.undo() # Restore immediately after app creation
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestPWAServing:
|
||||
def test_sw_js_served_as_file(self, client_with_static):
|
||||
"""Service worker file should be served directly, not as index.html."""
|
||||
resp = client_with_static.get("/sw.js")
|
||||
assert resp.status_code == 200
|
||||
assert "// service worker" in resp.text
|
||||
|
||||
def test_manifest_served_as_file(self, client_with_static):
|
||||
"""Web manifest should be served directly."""
|
||||
resp = client_with_static.get("/manifest.webmanifest")
|
||||
assert resp.status_code == 200
|
||||
assert "OpenJarvis" in resp.text
|
||||
|
||||
def test_icon_served_as_file(self, client_with_static):
|
||||
"""PWA icon should be served directly."""
|
||||
resp = client_with_static.get("/pwa-192x192.png")
|
||||
assert resp.status_code == 200
|
||||
assert b"PNG" in resp.content
|
||||
|
||||
def test_api_routes_bypass_spa(self, client_with_static):
|
||||
"""API routes should still work regardless of SPA catch-all."""
|
||||
resp = client_with_static.get("/v1/models")
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["object"] == "list"
|
||||
|
||||
def test_path_traversal_blocked(self, client_with_static):
|
||||
"""Path traversal attempts should fall back to index.html."""
|
||||
resp = client_with_static.get("/../../etc/passwd")
|
||||
assert resp.status_code == 200
|
||||
# Should get index.html, not the passwd file
|
||||
assert "SPA" in resp.text
|
||||
@@ -0,0 +1,283 @@
|
||||
"""Tests for batch-level energy accounting."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from typing import Generator
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.telemetry.batch import BatchMetrics, EnergyBatch
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# BatchMetrics defaults
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBatchMetricsDefaults:
|
||||
def test_all_defaults(self) -> None:
|
||||
m = BatchMetrics()
|
||||
assert m.batch_id == ""
|
||||
assert m.total_requests == 0
|
||||
assert m.total_tokens == 0
|
||||
assert m.total_energy_joules == 0.0
|
||||
assert m.energy_per_token_joules == 0.0
|
||||
assert m.energy_per_request_joules == 0.0
|
||||
assert m.mean_power_watts == 0.0
|
||||
assert m.mean_throughput_tok_per_sec == 0.0
|
||||
assert m.per_request_energy == []
|
||||
|
||||
def test_custom_values(self) -> None:
|
||||
m = BatchMetrics(
|
||||
batch_id="abc",
|
||||
total_requests=5,
|
||||
total_tokens=100,
|
||||
total_energy_joules=10.0,
|
||||
energy_per_token_joules=0.1,
|
||||
energy_per_request_joules=2.0,
|
||||
mean_power_watts=50.0,
|
||||
mean_throughput_tok_per_sec=200.0,
|
||||
per_request_energy=[1.0, 2.0, 3.0, 2.5, 1.5],
|
||||
)
|
||||
assert m.batch_id == "abc"
|
||||
assert m.total_requests == 5
|
||||
assert m.per_request_energy == [1.0, 2.0, 3.0, 2.5, 1.5]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Batch ID generation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBatchIdGeneration:
|
||||
def test_auto_generated_uuid(self) -> None:
|
||||
batch = EnergyBatch()
|
||||
# UUID4 format: 8-4-4-4-12 hex digits
|
||||
uuid4_re = (
|
||||
r"^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}"
|
||||
r"-[89ab][0-9a-f]{3}-[0-9a-f]{12}$"
|
||||
)
|
||||
assert re.match(uuid4_re, batch.batch_id)
|
||||
|
||||
def test_custom_batch_id(self) -> None:
|
||||
batch = EnergyBatch(batch_id="my-batch-42")
|
||||
assert batch.batch_id == "my-batch-42"
|
||||
|
||||
def test_unique_ids(self) -> None:
|
||||
ids = {EnergyBatch().batch_id for _ in range(100)}
|
||||
assert len(ids) == 100
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# EnergyBatch without monitor
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEnergyBatchNoMonitor:
|
||||
def test_record_request_accumulation(self) -> None:
|
||||
batch = EnergyBatch()
|
||||
with batch.sample() as ctx:
|
||||
ctx.record_request(tokens=50)
|
||||
ctx.record_request(tokens=30)
|
||||
ctx.record_request(tokens=20)
|
||||
|
||||
assert batch.metrics is not None
|
||||
assert batch.metrics.total_requests == 3
|
||||
assert batch.metrics.total_tokens == 100
|
||||
|
||||
def test_energy_stays_zero_without_monitor(self) -> None:
|
||||
batch = EnergyBatch()
|
||||
with batch.sample() as ctx:
|
||||
ctx.record_request(tokens=50)
|
||||
|
||||
assert batch.metrics is not None
|
||||
assert batch.metrics.total_energy_joules == 0.0
|
||||
assert batch.metrics.energy_per_token_joules == 0.0
|
||||
assert batch.metrics.mean_power_watts == 0.0
|
||||
|
||||
def test_per_request_energy_from_record(self) -> None:
|
||||
"""When no monitor, per-request energy comes from record_request calls."""
|
||||
batch = EnergyBatch()
|
||||
with batch.sample() as ctx:
|
||||
ctx.record_request(tokens=50, energy_joules=1.0)
|
||||
ctx.record_request(tokens=30, energy_joules=2.0)
|
||||
|
||||
assert batch.metrics is not None
|
||||
assert batch.metrics.per_request_energy == [1.0, 2.0]
|
||||
assert batch.metrics.total_energy_joules == 3.0
|
||||
|
||||
def test_metrics_computed_on_exit(self) -> None:
|
||||
batch = EnergyBatch()
|
||||
assert batch.metrics is None # Before sample()
|
||||
with batch.sample() as ctx:
|
||||
ctx.record_request(tokens=100)
|
||||
assert batch.metrics is not None
|
||||
|
||||
def test_no_requests_yields_zero_metrics(self) -> None:
|
||||
batch = EnergyBatch()
|
||||
with batch.sample() as _ctx:
|
||||
pass # No requests recorded
|
||||
|
||||
m = batch.metrics
|
||||
assert m is not None
|
||||
assert m.total_requests == 0
|
||||
assert m.total_tokens == 0
|
||||
assert m.energy_per_token_joules == 0.0
|
||||
assert m.energy_per_request_joules == 0.0
|
||||
|
||||
def test_throughput_computed(self) -> None:
|
||||
batch = EnergyBatch()
|
||||
with batch.sample() as ctx:
|
||||
ctx.record_request(tokens=1000)
|
||||
|
||||
assert batch.metrics is not None
|
||||
assert batch.metrics.mean_throughput_tok_per_sec > 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# EnergyBatch with mock monitor
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeEnergySample:
|
||||
energy_joules: float = 0.0
|
||||
mean_power_watts: float = 0.0
|
||||
|
||||
|
||||
class _FakeMonitor:
|
||||
"""Minimal mock that mimics EnergyMonitor.sample() context manager."""
|
||||
|
||||
def __init__(self, energy_joules: float = 10.0, mean_power_watts: float = 200.0):
|
||||
self._energy = energy_joules
|
||||
self._power = mean_power_watts
|
||||
|
||||
@contextmanager
|
||||
def sample(self) -> Generator[_FakeEnergySample, None, None]:
|
||||
s = _FakeEnergySample()
|
||||
yield s
|
||||
s.energy_joules = self._energy
|
||||
s.mean_power_watts = self._power
|
||||
|
||||
|
||||
class TestEnergyBatchWithMonitor:
|
||||
def test_energy_from_monitor(self) -> None:
|
||||
monitor = _FakeMonitor(energy_joules=10.0, mean_power_watts=200.0)
|
||||
batch = EnergyBatch(energy_monitor=monitor)
|
||||
with batch.sample() as ctx:
|
||||
ctx.record_request(tokens=100)
|
||||
|
||||
m = batch.metrics
|
||||
assert m is not None
|
||||
assert m.total_energy_joules == pytest.approx(10.0)
|
||||
assert m.mean_power_watts == pytest.approx(200.0)
|
||||
|
||||
def test_energy_per_token_with_monitor(self) -> None:
|
||||
monitor = _FakeMonitor(energy_joules=20.0)
|
||||
batch = EnergyBatch(energy_monitor=monitor)
|
||||
with batch.sample() as ctx:
|
||||
ctx.record_request(tokens=100)
|
||||
ctx.record_request(tokens=100)
|
||||
|
||||
m = batch.metrics
|
||||
assert m is not None
|
||||
assert m.total_tokens == 200
|
||||
assert m.energy_per_token_joules == pytest.approx(20.0 / 200.0)
|
||||
|
||||
def test_energy_per_request_with_monitor(self) -> None:
|
||||
monitor = _FakeMonitor(energy_joules=15.0)
|
||||
batch = EnergyBatch(energy_monitor=monitor)
|
||||
with batch.sample() as ctx:
|
||||
ctx.record_request(tokens=50)
|
||||
ctx.record_request(tokens=50)
|
||||
ctx.record_request(tokens=50)
|
||||
|
||||
m = batch.metrics
|
||||
assert m is not None
|
||||
assert m.total_requests == 3
|
||||
assert m.energy_per_request_joules == pytest.approx(15.0 / 3.0)
|
||||
|
||||
def test_batch_id_in_metrics(self) -> None:
|
||||
batch = EnergyBatch(batch_id="test-batch-99")
|
||||
with batch.sample() as ctx:
|
||||
ctx.record_request(tokens=10)
|
||||
|
||||
assert batch.metrics is not None
|
||||
assert batch.metrics.batch_id == "test-batch-99"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Energy per token calculation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEnergyPerToken:
|
||||
def test_basic_division(self) -> None:
|
||||
monitor = _FakeMonitor(energy_joules=50.0)
|
||||
batch = EnergyBatch(energy_monitor=monitor)
|
||||
with batch.sample() as ctx:
|
||||
ctx.record_request(tokens=500)
|
||||
|
||||
assert batch.metrics is not None
|
||||
assert batch.metrics.energy_per_token_joules == pytest.approx(50.0 / 500.0)
|
||||
|
||||
def test_zero_tokens_yields_zero(self) -> None:
|
||||
monitor = _FakeMonitor(energy_joules=10.0)
|
||||
batch = EnergyBatch(energy_monitor=monitor)
|
||||
with batch.sample() as _ctx:
|
||||
pass # No requests
|
||||
|
||||
assert batch.metrics is not None
|
||||
assert batch.metrics.energy_per_token_joules == 0.0
|
||||
|
||||
def test_many_small_requests(self) -> None:
|
||||
monitor = _FakeMonitor(energy_joules=1.0)
|
||||
batch = EnergyBatch(energy_monitor=monitor)
|
||||
with batch.sample() as ctx:
|
||||
for _ in range(100):
|
||||
ctx.record_request(tokens=10)
|
||||
|
||||
m = batch.metrics
|
||||
assert m is not None
|
||||
assert m.total_tokens == 1000
|
||||
assert m.total_requests == 100
|
||||
assert m.energy_per_token_joules == pytest.approx(1.0 / 1000.0)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Per-request energy list tracking
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestPerRequestEnergy:
|
||||
def test_tracks_each_request(self) -> None:
|
||||
batch = EnergyBatch()
|
||||
with batch.sample() as ctx:
|
||||
ctx.record_request(tokens=10, energy_joules=0.5)
|
||||
ctx.record_request(tokens=20, energy_joules=1.0)
|
||||
ctx.record_request(tokens=30, energy_joules=1.5)
|
||||
|
||||
m = batch.metrics
|
||||
assert m is not None
|
||||
assert m.per_request_energy == [0.5, 1.0, 1.5]
|
||||
assert len(m.per_request_energy) == 3
|
||||
|
||||
def test_empty_when_no_requests(self) -> None:
|
||||
batch = EnergyBatch()
|
||||
with batch.sample() as _ctx:
|
||||
pass
|
||||
|
||||
assert batch.metrics is not None
|
||||
assert batch.metrics.per_request_energy == []
|
||||
|
||||
def test_zeros_when_no_per_request_energy(self) -> None:
|
||||
batch = EnergyBatch()
|
||||
with batch.sample() as ctx:
|
||||
ctx.record_request(tokens=10)
|
||||
ctx.record_request(tokens=20)
|
||||
|
||||
m = batch.metrics
|
||||
assert m is not None
|
||||
assert m.per_request_energy == [0.0, 0.0]
|
||||
@@ -0,0 +1,191 @@
|
||||
"""Tier 1: derived metrics — energy_per_output_token, throughput_per_watt."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.types import Message, Role, TelemetryRecord
|
||||
from openjarvis.telemetry.aggregator import TelemetryAggregator
|
||||
from openjarvis.telemetry.instrumented_engine import InstrumentedEngine
|
||||
from openjarvis.telemetry.store import TelemetryStore
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _mock_engine(completion_tokens=50):
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
engine.generate.return_value = {
|
||||
"content": "hello",
|
||||
"usage": {
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": completion_tokens,
|
||||
"total_tokens": 10 + completion_tokens,
|
||||
},
|
||||
"model": "test-model",
|
||||
"ttft": 0.05,
|
||||
}
|
||||
return engine
|
||||
|
||||
|
||||
def _mock_energy_monitor(energy_joules=10.0, power_watts=200.0):
|
||||
monitor = MagicMock()
|
||||
sample = MagicMock()
|
||||
sample.energy_joules = energy_joules
|
||||
sample.mean_power_watts = power_watts
|
||||
sample.peak_power_watts = power_watts
|
||||
sample.mean_utilization_pct = 80.0
|
||||
sample.peak_utilization_pct = 95.0
|
||||
sample.mean_memory_used_gb = 16.0
|
||||
sample.peak_memory_used_gb = 20.0
|
||||
sample.mean_temperature_c = 65.0
|
||||
sample.peak_temperature_c = 72.0
|
||||
sample.duration_seconds = 0.5
|
||||
sample.num_snapshots = 10
|
||||
sample.energy_method = "hw_counter"
|
||||
sample.vendor = "nvidia"
|
||||
sample.cpu_energy_joules = 0.0
|
||||
sample.gpu_energy_joules = energy_joules
|
||||
sample.dram_energy_joules = 0.0
|
||||
|
||||
@contextmanager
|
||||
def _sample():
|
||||
yield sample
|
||||
|
||||
monitor.sample = _sample
|
||||
return monitor
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDerivedMetricsInGenerate:
|
||||
"""InstrumentedEngine.generate() computes derived metrics."""
|
||||
|
||||
def test_energy_per_output_token(self):
|
||||
bus = EventBus()
|
||||
engine = _mock_engine(completion_tokens=50)
|
||||
monitor = _mock_energy_monitor(energy_joules=10.0)
|
||||
ie = InstrumentedEngine(engine, bus, energy_monitor=monitor)
|
||||
|
||||
records = []
|
||||
bus.subscribe(
|
||||
EventType.TELEMETRY_RECORD,
|
||||
lambda e: records.append(e.data["record"]),
|
||||
)
|
||||
|
||||
ie.generate([Message(role=Role.USER, content="hi")], model="m")
|
||||
rec = records[0]
|
||||
assert rec.energy_per_output_token_joules == pytest.approx(10.0 / 50)
|
||||
|
||||
def test_throughput_per_watt(self):
|
||||
bus = EventBus()
|
||||
engine = _mock_engine(completion_tokens=100)
|
||||
monitor = _mock_energy_monitor(power_watts=250.0)
|
||||
ie = InstrumentedEngine(engine, bus, energy_monitor=monitor)
|
||||
|
||||
records = []
|
||||
bus.subscribe(
|
||||
EventType.TELEMETRY_RECORD,
|
||||
lambda e: records.append(e.data["record"]),
|
||||
)
|
||||
|
||||
ie.generate([Message(role=Role.USER, content="hi")], model="m")
|
||||
rec = records[0]
|
||||
# throughput_per_watt = throughput / power_watts
|
||||
expected = rec.throughput_tok_per_sec / 250.0
|
||||
assert rec.throughput_per_watt == pytest.approx(expected)
|
||||
|
||||
def test_zero_completion_tokens_no_division_error(self):
|
||||
bus = EventBus()
|
||||
engine = _mock_engine(completion_tokens=0)
|
||||
monitor = _mock_energy_monitor(energy_joules=5.0)
|
||||
ie = InstrumentedEngine(engine, bus, energy_monitor=monitor)
|
||||
|
||||
records = []
|
||||
bus.subscribe(
|
||||
EventType.TELEMETRY_RECORD,
|
||||
lambda e: records.append(e.data["record"]),
|
||||
)
|
||||
|
||||
ie.generate([Message(role=Role.USER, content="hi")], model="m")
|
||||
rec = records[0]
|
||||
assert rec.energy_per_output_token_joules == 0.0
|
||||
|
||||
def test_zero_power_no_division_error(self):
|
||||
bus = EventBus()
|
||||
engine = _mock_engine(completion_tokens=50)
|
||||
# No energy monitor -> power_watts = 0
|
||||
ie = InstrumentedEngine(engine, bus)
|
||||
|
||||
records = []
|
||||
bus.subscribe(
|
||||
EventType.TELEMETRY_RECORD,
|
||||
lambda e: records.append(e.data["record"]),
|
||||
)
|
||||
|
||||
ie.generate([Message(role=Role.USER, content="hi")], model="m")
|
||||
rec = records[0]
|
||||
assert rec.throughput_per_watt == 0.0
|
||||
|
||||
def test_derived_metrics_in_telemetry_dict(self):
|
||||
bus = EventBus()
|
||||
engine = _mock_engine(completion_tokens=25)
|
||||
monitor = _mock_energy_monitor(energy_joules=5.0, power_watts=100.0)
|
||||
ie = InstrumentedEngine(engine, bus, energy_monitor=monitor)
|
||||
|
||||
result = ie.generate([Message(role=Role.USER, content="hi")], model="m")
|
||||
t = result["_telemetry"]
|
||||
assert t["energy_per_output_token_joules"] == pytest.approx(5.0 / 25)
|
||||
assert t["throughput_per_watt"] > 0
|
||||
|
||||
|
||||
class TestDerivedMetricsInStore:
|
||||
"""Derived metrics are stored and queryable."""
|
||||
|
||||
def test_store_and_query(self, tmp_path):
|
||||
store = TelemetryStore(tmp_path / "test.db")
|
||||
rec = TelemetryRecord(
|
||||
timestamp=time.time(),
|
||||
model_id="test-model",
|
||||
engine="mock",
|
||||
completion_tokens=50,
|
||||
energy_joules=10.0,
|
||||
energy_per_output_token_joules=0.2,
|
||||
throughput_per_watt=0.5,
|
||||
)
|
||||
store.record(rec)
|
||||
|
||||
agg = TelemetryAggregator(tmp_path / "test.db")
|
||||
stats = agg.per_model_stats()
|
||||
assert len(stats) == 1
|
||||
assert stats[0].avg_energy_per_output_token_joules == pytest.approx(0.2)
|
||||
assert stats[0].avg_throughput_per_watt == pytest.approx(0.5)
|
||||
agg.close()
|
||||
store.close()
|
||||
|
||||
def test_summary_weighted_averages(self, tmp_path):
|
||||
store = TelemetryStore(tmp_path / "test.db")
|
||||
for i in range(3):
|
||||
store.record(TelemetryRecord(
|
||||
timestamp=time.time() + i,
|
||||
model_id="m1",
|
||||
engine="e1",
|
||||
energy_per_output_token_joules=0.1 * (i + 1),
|
||||
throughput_per_watt=1.0 * (i + 1),
|
||||
))
|
||||
agg = TelemetryAggregator(tmp_path / "test.db")
|
||||
summary = agg.summary()
|
||||
assert summary.avg_energy_per_output_token_joules > 0
|
||||
assert summary.avg_throughput_per_watt > 0
|
||||
agg.close()
|
||||
store.close()
|
||||
@@ -0,0 +1,188 @@
|
||||
"""Tests for AmdEnergyMonitor -- mock amdsmi (no real GPU required)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import time
|
||||
import types
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers: build a fake amdsmi module
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_fake_amdsmi(device_count: int = 1):
|
||||
"""Return a fake amdsmi module object."""
|
||||
mod = types.ModuleType("amdsmi")
|
||||
mod.amdsmi_init = MagicMock()
|
||||
mod.amdsmi_shut_down = MagicMock()
|
||||
handles = [f"amd-handle-{i}" for i in range(device_count)]
|
||||
mod.amdsmi_get_processor_handles = MagicMock(return_value=handles)
|
||||
mod.amdsmi_get_gpu_asic_info = MagicMock(
|
||||
return_value={"market_name": "AMD Instinct MI300X"}
|
||||
)
|
||||
mod.amdsmi_get_energy_count = MagicMock(
|
||||
return_value={"energy_accumulator": 1000.0, "counter_resolution": 15.3}
|
||||
)
|
||||
return mod
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: available()
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestAvailable:
|
||||
def test_available_true_when_amdsmi_works(self):
|
||||
fake_amdsmi = _make_fake_amdsmi(device_count=1)
|
||||
|
||||
with patch.dict(sys.modules, {"amdsmi": fake_amdsmi}):
|
||||
import openjarvis.telemetry.energy_amd as mod
|
||||
|
||||
orig = mod._AMDSMI_AVAILABLE
|
||||
mod._AMDSMI_AVAILABLE = True
|
||||
mod.amdsmi = fake_amdsmi
|
||||
try:
|
||||
assert mod.AmdEnergyMonitor.available() is True
|
||||
fake_amdsmi.amdsmi_init.assert_called()
|
||||
fake_amdsmi.amdsmi_shut_down.assert_called()
|
||||
finally:
|
||||
mod._AMDSMI_AVAILABLE = orig
|
||||
|
||||
def test_available_false_when_amdsmi_not_importable(self):
|
||||
import openjarvis.telemetry.energy_amd as mod
|
||||
|
||||
orig = mod._AMDSMI_AVAILABLE
|
||||
mod._AMDSMI_AVAILABLE = False
|
||||
try:
|
||||
assert mod.AmdEnergyMonitor.available() is False
|
||||
finally:
|
||||
mod._AMDSMI_AVAILABLE = orig
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: energy_method()
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEnergyMethod:
|
||||
def test_returns_hw_counter(self):
|
||||
fake_amdsmi = _make_fake_amdsmi(device_count=1)
|
||||
|
||||
with patch.dict(sys.modules, {"amdsmi": fake_amdsmi}):
|
||||
import openjarvis.telemetry.energy_amd as mod
|
||||
|
||||
orig = mod._AMDSMI_AVAILABLE
|
||||
mod._AMDSMI_AVAILABLE = True
|
||||
mod.amdsmi = fake_amdsmi
|
||||
try:
|
||||
monitor = mod.AmdEnergyMonitor(poll_interval_ms=50)
|
||||
assert monitor.energy_method() == "hw_counter"
|
||||
finally:
|
||||
mod._AMDSMI_AVAILABLE = orig
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: sample() counter delta math
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSampleCounterDelta:
|
||||
def test_counter_delta_microjoules_to_joules(self):
|
||||
"""acc_start=1000, acc_end=2000, resolution=15.3 =>
|
||||
delta=1000 * 15.3 = 15300 uJ => 0.0153 J."""
|
||||
fake_amdsmi = _make_fake_amdsmi(device_count=1)
|
||||
|
||||
call_count = {"n": 0}
|
||||
readings = [
|
||||
{"energy_accumulator": 1000.0, "counter_resolution": 15.3},
|
||||
{"energy_accumulator": 2000.0, "counter_resolution": 15.3},
|
||||
]
|
||||
|
||||
def get_energy(handle):
|
||||
idx = min(call_count["n"], len(readings) - 1)
|
||||
val = readings[idx]
|
||||
call_count["n"] += 1
|
||||
return val
|
||||
|
||||
fake_amdsmi.amdsmi_get_energy_count.side_effect = get_energy
|
||||
|
||||
with patch.dict(sys.modules, {"amdsmi": fake_amdsmi}):
|
||||
import openjarvis.telemetry.energy_amd as mod
|
||||
|
||||
orig = mod._AMDSMI_AVAILABLE
|
||||
mod._AMDSMI_AVAILABLE = True
|
||||
mod.amdsmi = fake_amdsmi
|
||||
try:
|
||||
monitor = mod.AmdEnergyMonitor(poll_interval_ms=50)
|
||||
# Reset for sample()
|
||||
call_count["n"] = 0
|
||||
|
||||
with monitor.sample() as result:
|
||||
time.sleep(0.01)
|
||||
|
||||
# delta = (2000 - 1000) * 15.3 = 15300 uJ = 0.0153 J
|
||||
expected_joules = (2000.0 - 1000.0) * 15.3 / 1e6
|
||||
assert result.energy_joules == pytest.approx(expected_joules)
|
||||
assert result.gpu_energy_joules == pytest.approx(expected_joules)
|
||||
assert result.vendor == "amd"
|
||||
assert result.energy_method == "hw_counter"
|
||||
assert result.duration_seconds > 0
|
||||
finally:
|
||||
mod._AMDSMI_AVAILABLE = orig
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: sample() with no devices
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSampleNoDevices:
|
||||
def test_no_devices_empty_result(self):
|
||||
"""When no AMD GPUs present, sample yields empty result."""
|
||||
from openjarvis.telemetry.energy_amd import AmdEnergyMonitor
|
||||
|
||||
monitor = AmdEnergyMonitor.__new__(AmdEnergyMonitor)
|
||||
monitor._poll_interval_ms = 50
|
||||
monitor._handles = []
|
||||
monitor._device_count = 0
|
||||
monitor._device_name = ""
|
||||
monitor._initialized = False
|
||||
|
||||
with monitor.sample() as result:
|
||||
pass
|
||||
|
||||
assert result.energy_joules == 0.0
|
||||
assert result.duration_seconds >= 0
|
||||
assert result.vendor == "amd"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: close()
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestClose:
|
||||
def test_close_calls_amdsmi_shut_down(self):
|
||||
fake_amdsmi = _make_fake_amdsmi(device_count=1)
|
||||
|
||||
with patch.dict(sys.modules, {"amdsmi": fake_amdsmi}):
|
||||
import openjarvis.telemetry.energy_amd as mod
|
||||
|
||||
orig = mod._AMDSMI_AVAILABLE
|
||||
mod._AMDSMI_AVAILABLE = True
|
||||
mod.amdsmi = fake_amdsmi
|
||||
try:
|
||||
monitor = mod.AmdEnergyMonitor(poll_interval_ms=50)
|
||||
assert monitor._initialized is True
|
||||
|
||||
fake_amdsmi.amdsmi_shut_down.reset_mock()
|
||||
monitor.close()
|
||||
|
||||
fake_amdsmi.amdsmi_shut_down.assert_called_once()
|
||||
assert monitor._initialized is False
|
||||
finally:
|
||||
mod._AMDSMI_AVAILABLE = orig
|
||||
@@ -0,0 +1,159 @@
|
||||
"""Tests for AppleEnergyMonitor -- mock zeus (no real Apple Silicon required)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
import types
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers: build a fake zeus module
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_fake_zeus():
|
||||
"""Return a fake zeus.device.soc.apple module with AppleSiliconMonitor."""
|
||||
# Build the nested module hierarchy
|
||||
zeus = types.ModuleType("zeus")
|
||||
zeus_device = types.ModuleType("zeus.device")
|
||||
zeus_device_soc = types.ModuleType("zeus.device.soc")
|
||||
zeus_device_soc_apple = types.ModuleType("zeus.device.soc.apple")
|
||||
|
||||
mock_monitor_cls = MagicMock()
|
||||
zeus_device_soc_apple.AppleSiliconMonitor = mock_monitor_cls
|
||||
|
||||
zeus.device = zeus_device
|
||||
zeus_device.soc = zeus_device_soc
|
||||
zeus_device_soc.apple = zeus_device_soc_apple
|
||||
|
||||
return zeus, zeus_device, zeus_device_soc, zeus_device_soc_apple, mock_monitor_cls
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: available()
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestAvailable:
|
||||
def test_available_false_on_non_darwin(self):
|
||||
with patch("platform.system", return_value="Linux"):
|
||||
from openjarvis.telemetry.energy_apple import AppleEnergyMonitor
|
||||
|
||||
assert AppleEnergyMonitor.available() is False
|
||||
|
||||
def test_available_false_when_zeus_not_importable(self):
|
||||
import openjarvis.telemetry.energy_apple as mod
|
||||
|
||||
orig = mod._ZEUS_APPLE_AVAILABLE
|
||||
mod._ZEUS_APPLE_AVAILABLE = False
|
||||
try:
|
||||
with patch("platform.system", return_value="Darwin"):
|
||||
assert mod.AppleEnergyMonitor.available() is False
|
||||
finally:
|
||||
mod._ZEUS_APPLE_AVAILABLE = orig
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: energy_method()
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEnergyMethod:
|
||||
def test_returns_zeus(self):
|
||||
from openjarvis.telemetry.energy_apple import AppleEnergyMonitor
|
||||
|
||||
monitor = AppleEnergyMonitor.__new__(AppleEnergyMonitor)
|
||||
assert monitor.energy_method() == "zeus"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: sample() component breakdown
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSampleComponentBreakdown:
|
||||
def test_component_energy_extraction(self):
|
||||
"""Mock begin_window/end_window and verify cpu/gpu/dram/ane extraction."""
|
||||
mock_measurement = MagicMock()
|
||||
mock_measurement.cpu_energy = 1.5
|
||||
mock_measurement.gpu_energy = 3.0
|
||||
mock_measurement.dram_energy = 0.5
|
||||
mock_measurement.ane_energy = 2.0
|
||||
|
||||
mock_zeus_monitor = MagicMock()
|
||||
mock_zeus_monitor.begin_window = MagicMock()
|
||||
mock_zeus_monitor.end_window = MagicMock(return_value=mock_measurement)
|
||||
|
||||
from openjarvis.telemetry.energy_apple import AppleEnergyMonitor
|
||||
|
||||
monitor = AppleEnergyMonitor.__new__(AppleEnergyMonitor)
|
||||
monitor._poll_interval_ms = 50
|
||||
monitor._monitor = mock_zeus_monitor
|
||||
monitor._initialized = True
|
||||
|
||||
with monitor.sample() as result:
|
||||
time.sleep(0.01)
|
||||
|
||||
mock_zeus_monitor.begin_window.assert_called_once()
|
||||
mock_zeus_monitor.end_window.assert_called_once()
|
||||
|
||||
assert result.cpu_energy_joules == pytest.approx(1.5)
|
||||
assert result.gpu_energy_joules == pytest.approx(3.0)
|
||||
assert result.dram_energy_joules == pytest.approx(0.5)
|
||||
assert result.ane_energy_joules == pytest.approx(2.0)
|
||||
assert result.vendor == "apple"
|
||||
assert result.energy_method == "zeus"
|
||||
|
||||
def test_total_energy_is_sum_of_components(self):
|
||||
"""total = cpu + gpu + dram + ane."""
|
||||
mock_measurement = MagicMock()
|
||||
mock_measurement.cpu_energy = 1.0
|
||||
mock_measurement.gpu_energy = 2.0
|
||||
mock_measurement.dram_energy = 0.3
|
||||
mock_measurement.ane_energy = 0.7
|
||||
|
||||
mock_zeus_monitor = MagicMock()
|
||||
mock_zeus_monitor.begin_window = MagicMock()
|
||||
mock_zeus_monitor.end_window = MagicMock(return_value=mock_measurement)
|
||||
|
||||
from openjarvis.telemetry.energy_apple import AppleEnergyMonitor
|
||||
|
||||
monitor = AppleEnergyMonitor.__new__(AppleEnergyMonitor)
|
||||
monitor._poll_interval_ms = 50
|
||||
monitor._monitor = mock_zeus_monitor
|
||||
monitor._initialized = True
|
||||
|
||||
with monitor.sample() as result:
|
||||
pass
|
||||
|
||||
expected_total = 1.0 + 2.0 + 0.3 + 0.7
|
||||
assert result.energy_joules == pytest.approx(expected_total)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: sample() with uninitialized monitor
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSampleUninitialized:
|
||||
def test_uninitialized_monitor_empty_result(self):
|
||||
"""When monitor is not initialized, sample yields empty result."""
|
||||
from openjarvis.telemetry.energy_apple import AppleEnergyMonitor
|
||||
|
||||
monitor = AppleEnergyMonitor.__new__(AppleEnergyMonitor)
|
||||
monitor._poll_interval_ms = 50
|
||||
monitor._monitor = None
|
||||
monitor._initialized = False
|
||||
|
||||
with monitor.sample() as result:
|
||||
pass
|
||||
|
||||
assert result.energy_joules == 0.0
|
||||
assert result.cpu_energy_joules == 0.0
|
||||
assert result.gpu_energy_joules == 0.0
|
||||
assert result.dram_energy_joules == 0.0
|
||||
assert result.ane_energy_joules == 0.0
|
||||
assert result.duration_seconds >= 0
|
||||
assert result.vendor == "apple"
|
||||
@@ -0,0 +1,216 @@
|
||||
"""Tests for EnergyMonitor ABC, EnergySample, EnergyVendor, and factory."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.telemetry.energy_monitor import (
|
||||
EnergyMonitor,
|
||||
EnergySample,
|
||||
EnergyVendor,
|
||||
create_energy_monitor,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: EnergySample defaults
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEnergySample:
|
||||
def test_default_field_values(self):
|
||||
s = EnergySample()
|
||||
assert s.energy_joules == 0.0
|
||||
assert s.mean_power_watts == 0.0
|
||||
assert s.peak_power_watts == 0.0
|
||||
assert s.duration_seconds == 0.0
|
||||
assert s.num_snapshots == 0
|
||||
assert s.mean_utilization_pct == 0.0
|
||||
assert s.peak_utilization_pct == 0.0
|
||||
assert s.mean_memory_used_gb == 0.0
|
||||
assert s.peak_memory_used_gb == 0.0
|
||||
assert s.mean_temperature_c == 0.0
|
||||
assert s.peak_temperature_c == 0.0
|
||||
assert s.vendor == ""
|
||||
assert s.device_name == ""
|
||||
assert s.device_count == 0
|
||||
assert s.energy_method == ""
|
||||
assert s.cpu_energy_joules == 0.0
|
||||
assert s.gpu_energy_joules == 0.0
|
||||
assert s.dram_energy_joules == 0.0
|
||||
assert s.ane_energy_joules == 0.0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: EnergyVendor enum
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEnergyVendor:
|
||||
def test_enum_values(self):
|
||||
assert EnergyVendor.NVIDIA.value == "nvidia"
|
||||
assert EnergyVendor.AMD.value == "amd"
|
||||
assert EnergyVendor.APPLE.value == "apple"
|
||||
assert EnergyVendor.CPU_RAPL.value == "cpu_rapl"
|
||||
|
||||
def test_enum_is_str(self):
|
||||
assert isinstance(EnergyVendor.NVIDIA, str)
|
||||
assert EnergyVendor.AMD == "amd"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: EnergyMonitor ABC
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEnergyMonitorABC:
|
||||
def test_cannot_instantiate_abstract(self):
|
||||
with pytest.raises(TypeError):
|
||||
EnergyMonitor()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: create_energy_monitor factory
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCreateEnergyMonitor:
|
||||
def test_returns_none_when_nothing_available(self):
|
||||
with patch(
|
||||
"openjarvis.telemetry.energy_nvidia.NvidiaEnergyMonitor.available",
|
||||
return_value=False,
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_amd.AmdEnergyMonitor.available",
|
||||
return_value=False,
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_apple.AppleEnergyMonitor.available",
|
||||
return_value=False,
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_rapl.RaplEnergyMonitor.available",
|
||||
return_value=False,
|
||||
):
|
||||
result = create_energy_monitor()
|
||||
assert result is None
|
||||
|
||||
def test_prefer_vendor_parameter(self):
|
||||
"""When prefer_vendor is set, that vendor is tried first."""
|
||||
with patch(
|
||||
"openjarvis.telemetry.energy_nvidia.NvidiaEnergyMonitor.available",
|
||||
return_value=False,
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_amd.AmdEnergyMonitor.available",
|
||||
return_value=False,
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_apple.AppleEnergyMonitor.available",
|
||||
return_value=False,
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_rapl.RaplEnergyMonitor.available",
|
||||
return_value=True,
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_rapl.RaplEnergyMonitor.__init__",
|
||||
return_value=None,
|
||||
) as mock_init:
|
||||
create_energy_monitor(prefer_vendor="cpu_rapl")
|
||||
# RaplEnergyMonitor was available and preferred
|
||||
mock_init.assert_called_once_with(poll_interval_ms=50)
|
||||
|
||||
def test_detection_order_nvidia_first(self):
|
||||
"""Default order: NVIDIA is tried before AMD."""
|
||||
call_order = []
|
||||
|
||||
def nvidia_available():
|
||||
call_order.append("nvidia")
|
||||
return True
|
||||
|
||||
def amd_available():
|
||||
call_order.append("amd")
|
||||
return True
|
||||
|
||||
with patch(
|
||||
"openjarvis.telemetry.energy_nvidia.NvidiaEnergyMonitor.available",
|
||||
side_effect=nvidia_available,
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_amd.AmdEnergyMonitor.available",
|
||||
side_effect=amd_available,
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_apple.AppleEnergyMonitor.available",
|
||||
return_value=False,
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_rapl.RaplEnergyMonitor.available",
|
||||
return_value=False,
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_nvidia.NvidiaEnergyMonitor.__init__",
|
||||
return_value=None,
|
||||
):
|
||||
create_energy_monitor()
|
||||
# NVIDIA was tried first and returned True
|
||||
assert call_order == ["nvidia"]
|
||||
|
||||
def test_detection_order_falls_through(self):
|
||||
"""When NVIDIA unavailable, AMD is tried next."""
|
||||
call_order = []
|
||||
|
||||
def nvidia_available():
|
||||
call_order.append("nvidia")
|
||||
return False
|
||||
|
||||
def amd_available():
|
||||
call_order.append("amd")
|
||||
return True
|
||||
|
||||
with patch(
|
||||
"openjarvis.telemetry.energy_nvidia.NvidiaEnergyMonitor.available",
|
||||
side_effect=nvidia_available,
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_amd.AmdEnergyMonitor.available",
|
||||
side_effect=amd_available,
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_apple.AppleEnergyMonitor.available",
|
||||
return_value=False,
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_rapl.RaplEnergyMonitor.available",
|
||||
return_value=False,
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_amd.AmdEnergyMonitor.__init__",
|
||||
return_value=None,
|
||||
):
|
||||
create_energy_monitor()
|
||||
assert call_order == ["nvidia", "amd"]
|
||||
|
||||
def test_prefer_vendor_tried_first_then_default_order(self):
|
||||
"""prefer_vendor=cpu_rapl puts RAPL first, then NVIDIA > AMD > Apple."""
|
||||
call_order = []
|
||||
|
||||
def rapl_available():
|
||||
call_order.append("rapl")
|
||||
return False
|
||||
|
||||
def nvidia_available():
|
||||
call_order.append("nvidia")
|
||||
return False
|
||||
|
||||
def amd_available():
|
||||
call_order.append("amd")
|
||||
return False
|
||||
|
||||
def apple_available():
|
||||
call_order.append("apple")
|
||||
return False
|
||||
|
||||
with patch(
|
||||
"openjarvis.telemetry.energy_nvidia.NvidiaEnergyMonitor.available",
|
||||
side_effect=nvidia_available,
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_amd.AmdEnergyMonitor.available",
|
||||
side_effect=amd_available,
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_apple.AppleEnergyMonitor.available",
|
||||
side_effect=apple_available,
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_rapl.RaplEnergyMonitor.available",
|
||||
side_effect=rapl_available,
|
||||
):
|
||||
result = create_energy_monitor(prefer_vendor="cpu_rapl")
|
||||
assert result is None
|
||||
assert call_order == ["rapl", "nvidia", "amd", "apple"]
|
||||
@@ -0,0 +1,362 @@
|
||||
"""Tests for NvidiaEnergyMonitor -- mock pynvml (no real GPU required)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import time
|
||||
import types
|
||||
from dataclasses import dataclass
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers: build a fake pynvml module
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeUtilization:
|
||||
gpu: int = 80
|
||||
memory: int = 50
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeMemInfo:
|
||||
total: int = 24 * 1024**3
|
||||
used: int = 12 * 1024**3
|
||||
free: int = 12 * 1024**3
|
||||
|
||||
|
||||
def _make_fake_pynvml(device_count: int = 1, power_mw: int = 300_000):
|
||||
"""Return a fake pynvml module object."""
|
||||
mod = types.ModuleType("pynvml")
|
||||
mod.nvmlInit = MagicMock()
|
||||
mod.nvmlShutdown = MagicMock()
|
||||
mod.nvmlDeviceGetCount = MagicMock(return_value=device_count)
|
||||
mod.nvmlDeviceGetHandleByIndex = MagicMock(
|
||||
side_effect=lambda i: f"handle-{i}"
|
||||
)
|
||||
mod.nvmlDeviceGetName = MagicMock(return_value="NVIDIA A100-SXM")
|
||||
mod.nvmlDeviceGetPowerUsage = MagicMock(return_value=power_mw)
|
||||
mod.nvmlDeviceGetUtilizationRates = MagicMock(
|
||||
return_value=_FakeUtilization()
|
||||
)
|
||||
mod.nvmlDeviceGetMemoryInfo = MagicMock(return_value=_FakeMemInfo())
|
||||
mod.nvmlDeviceGetTemperature = MagicMock(return_value=65)
|
||||
mod.nvmlDeviceGetTotalEnergyConsumption = MagicMock(return_value=5000.0)
|
||||
mod.NVML_TEMPERATURE_GPU = 0
|
||||
return mod
|
||||
|
||||
|
||||
def _install_fake_pynvml(fake_pynvml, mod):
|
||||
"""Patch energy_nvidia module to use fake pynvml."""
|
||||
mod._PYNVML_AVAILABLE = True
|
||||
mod.pynvml = fake_pynvml
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: available()
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestAvailable:
|
||||
def test_available_true_when_pynvml_works(self):
|
||||
fake_pynvml = _make_fake_pynvml(device_count=1)
|
||||
|
||||
with patch.dict(sys.modules, {"pynvml": fake_pynvml}):
|
||||
import openjarvis.telemetry.energy_nvidia as mod
|
||||
|
||||
orig = mod._PYNVML_AVAILABLE
|
||||
mod._PYNVML_AVAILABLE = True
|
||||
mod.pynvml = fake_pynvml
|
||||
try:
|
||||
assert mod.NvidiaEnergyMonitor.available() is True
|
||||
fake_pynvml.nvmlInit.assert_called()
|
||||
fake_pynvml.nvmlShutdown.assert_called()
|
||||
finally:
|
||||
mod._PYNVML_AVAILABLE = orig
|
||||
|
||||
def test_available_false_when_pynvml_not_importable(self):
|
||||
import openjarvis.telemetry.energy_nvidia as mod
|
||||
|
||||
orig = mod._PYNVML_AVAILABLE
|
||||
mod._PYNVML_AVAILABLE = False
|
||||
try:
|
||||
assert mod.NvidiaEnergyMonitor.available() is False
|
||||
finally:
|
||||
mod._PYNVML_AVAILABLE = orig
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: hw counter probe
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestHwCounterProbe:
|
||||
def test_probe_succeeds_on_volta_plus(self):
|
||||
"""nvmlDeviceGetTotalEnergyConsumption succeeds => hw_counter_available."""
|
||||
fake_pynvml = _make_fake_pynvml(device_count=1)
|
||||
# GetTotalEnergyConsumption returns normally => Volta+
|
||||
fake_pynvml.nvmlDeviceGetTotalEnergyConsumption.return_value = 1000.0
|
||||
|
||||
with patch.dict(sys.modules, {"pynvml": fake_pynvml}):
|
||||
import openjarvis.telemetry.energy_nvidia as mod
|
||||
|
||||
orig = mod._PYNVML_AVAILABLE
|
||||
mod._PYNVML_AVAILABLE = True
|
||||
mod.pynvml = fake_pynvml
|
||||
try:
|
||||
monitor = mod.NvidiaEnergyMonitor(poll_interval_ms=50)
|
||||
assert monitor._hw_counter_available is True
|
||||
finally:
|
||||
mod._PYNVML_AVAILABLE = orig
|
||||
|
||||
def test_probe_fails_on_pre_volta(self):
|
||||
"""nvmlDeviceGetTotalEnergyConsumption raises => polling fallback."""
|
||||
fake_pynvml = _make_fake_pynvml(device_count=1)
|
||||
fake_pynvml.nvmlDeviceGetTotalEnergyConsumption.side_effect = (
|
||||
RuntimeError("Not supported")
|
||||
)
|
||||
|
||||
with patch.dict(sys.modules, {"pynvml": fake_pynvml}):
|
||||
import openjarvis.telemetry.energy_nvidia as mod
|
||||
|
||||
orig = mod._PYNVML_AVAILABLE
|
||||
mod._PYNVML_AVAILABLE = True
|
||||
mod.pynvml = fake_pynvml
|
||||
try:
|
||||
monitor = mod.NvidiaEnergyMonitor(poll_interval_ms=50)
|
||||
assert monitor._hw_counter_available is False
|
||||
finally:
|
||||
mod._PYNVML_AVAILABLE = orig
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: energy_method()
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEnergyMethod:
|
||||
def test_returns_hw_counter_when_available(self):
|
||||
fake_pynvml = _make_fake_pynvml(device_count=1)
|
||||
|
||||
with patch.dict(sys.modules, {"pynvml": fake_pynvml}):
|
||||
import openjarvis.telemetry.energy_nvidia as mod
|
||||
|
||||
orig = mod._PYNVML_AVAILABLE
|
||||
mod._PYNVML_AVAILABLE = True
|
||||
mod.pynvml = fake_pynvml
|
||||
try:
|
||||
monitor = mod.NvidiaEnergyMonitor(poll_interval_ms=50)
|
||||
assert monitor.energy_method() == "hw_counter"
|
||||
finally:
|
||||
mod._PYNVML_AVAILABLE = orig
|
||||
|
||||
def test_returns_polling_when_no_hw_counter(self):
|
||||
fake_pynvml = _make_fake_pynvml(device_count=1)
|
||||
fake_pynvml.nvmlDeviceGetTotalEnergyConsumption.side_effect = (
|
||||
RuntimeError("Not supported")
|
||||
)
|
||||
|
||||
with patch.dict(sys.modules, {"pynvml": fake_pynvml}):
|
||||
import openjarvis.telemetry.energy_nvidia as mod
|
||||
|
||||
orig = mod._PYNVML_AVAILABLE
|
||||
mod._PYNVML_AVAILABLE = True
|
||||
mod.pynvml = fake_pynvml
|
||||
try:
|
||||
monitor = mod.NvidiaEnergyMonitor(poll_interval_ms=50)
|
||||
assert monitor.energy_method() == "polling"
|
||||
finally:
|
||||
mod._PYNVML_AVAILABLE = orig
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: sample() with hw counters
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSampleHwCounters:
|
||||
def test_hw_counter_delta_math(self):
|
||||
"""start=5000mJ, end=8000mJ => delta=3000mJ => 3.0 J."""
|
||||
fake_pynvml = _make_fake_pynvml(device_count=1)
|
||||
|
||||
energy_readings = [5000.0, 8000.0]
|
||||
call_count = {"n": 0}
|
||||
|
||||
def get_energy(handle):
|
||||
idx = min(call_count["n"], len(energy_readings) - 1)
|
||||
val = energy_readings[idx]
|
||||
call_count["n"] += 1
|
||||
return val
|
||||
|
||||
fake_pynvml.nvmlDeviceGetTotalEnergyConsumption.side_effect = get_energy
|
||||
|
||||
with patch.dict(sys.modules, {"pynvml": fake_pynvml}):
|
||||
import openjarvis.telemetry.energy_nvidia as mod
|
||||
|
||||
orig = mod._PYNVML_AVAILABLE
|
||||
mod._PYNVML_AVAILABLE = True
|
||||
mod.pynvml = fake_pynvml
|
||||
try:
|
||||
monitor = mod.NvidiaEnergyMonitor(poll_interval_ms=10)
|
||||
# _probe_hw_counter consumed one reading during __init__,
|
||||
# so reset the counter for sample()
|
||||
call_count["n"] = 0
|
||||
energy_readings_sample = [5000.0, 8000.0]
|
||||
|
||||
def get_energy_sample(handle):
|
||||
idx = min(call_count["n"], len(energy_readings_sample) - 1)
|
||||
val = energy_readings_sample[idx]
|
||||
call_count["n"] += 1
|
||||
return val
|
||||
|
||||
fake_pynvml.nvmlDeviceGetTotalEnergyConsumption.side_effect = (
|
||||
get_energy_sample
|
||||
)
|
||||
|
||||
with monitor.sample() as result:
|
||||
time.sleep(0.05)
|
||||
|
||||
# delta = 8000 - 5000 = 3000 mJ => 3.0 J
|
||||
assert result.energy_joules == pytest.approx(3.0)
|
||||
assert result.gpu_energy_joules == pytest.approx(3.0)
|
||||
assert result.vendor == "nvidia"
|
||||
assert result.energy_method == "hw_counter"
|
||||
finally:
|
||||
mod._PYNVML_AVAILABLE = orig
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: sample() with polling fallback
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSamplePolling:
|
||||
def test_polling_trapezoidal_integration(self):
|
||||
"""Fallback mode uses trapezoidal integration of power readings."""
|
||||
fake_pynvml = _make_fake_pynvml(device_count=1, power_mw=300_000)
|
||||
# Make hw counter probe fail => polling mode
|
||||
fake_pynvml.nvmlDeviceGetTotalEnergyConsumption.side_effect = (
|
||||
RuntimeError("Not supported")
|
||||
)
|
||||
|
||||
with patch.dict(sys.modules, {"pynvml": fake_pynvml}):
|
||||
import openjarvis.telemetry.energy_nvidia as mod
|
||||
|
||||
orig = mod._PYNVML_AVAILABLE
|
||||
mod._PYNVML_AVAILABLE = True
|
||||
mod.pynvml = fake_pynvml
|
||||
try:
|
||||
monitor = mod.NvidiaEnergyMonitor(poll_interval_ms=10)
|
||||
assert monitor.energy_method() == "polling"
|
||||
|
||||
with monitor.sample() as result:
|
||||
time.sleep(0.15)
|
||||
|
||||
# With constant 300W polling, energy should be > 0
|
||||
assert result.energy_joules > 0
|
||||
assert result.duration_seconds > 0
|
||||
assert result.vendor == "nvidia"
|
||||
assert result.energy_method == "polling"
|
||||
finally:
|
||||
mod._PYNVML_AVAILABLE = orig
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: sample() multi-GPU
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSampleMultiGpu:
|
||||
def test_multi_gpu_hw_counter(self):
|
||||
"""2 GPUs: energy is sum of deltas from both devices."""
|
||||
fake_pynvml = _make_fake_pynvml(device_count=2)
|
||||
|
||||
# 2 devices: __init__ probe reads device 0 once.
|
||||
# Then sample() reads start (dev0, dev1), end (dev0, dev1).
|
||||
readings = iter([
|
||||
1000.0, # probe: device 0
|
||||
2000.0, # sample start: device 0
|
||||
3000.0, # sample start: device 1
|
||||
5000.0, # sample end: device 0
|
||||
7000.0, # sample end: device 1
|
||||
])
|
||||
|
||||
fake_pynvml.nvmlDeviceGetTotalEnergyConsumption.side_effect = (
|
||||
lambda h: next(readings)
|
||||
)
|
||||
|
||||
with patch.dict(sys.modules, {"pynvml": fake_pynvml}):
|
||||
import openjarvis.telemetry.energy_nvidia as mod
|
||||
|
||||
orig = mod._PYNVML_AVAILABLE
|
||||
mod._PYNVML_AVAILABLE = True
|
||||
mod.pynvml = fake_pynvml
|
||||
try:
|
||||
monitor = mod.NvidiaEnergyMonitor(poll_interval_ms=10)
|
||||
assert monitor._hw_counter_available is True
|
||||
|
||||
with monitor.sample() as result:
|
||||
time.sleep(0.05)
|
||||
|
||||
# dev0: 5000-2000=3000 mJ, dev1: 7000-3000=4000 mJ => 7.0 J
|
||||
assert result.energy_joules == pytest.approx(7.0)
|
||||
assert result.device_count == 2
|
||||
finally:
|
||||
mod._PYNVML_AVAILABLE = orig
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: sample() with no devices
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSampleNoDevices:
|
||||
def test_no_devices_empty_result(self):
|
||||
"""When no GPUs are present, sample yields empty result."""
|
||||
from openjarvis.telemetry.energy_nvidia import NvidiaEnergyMonitor
|
||||
|
||||
monitor = NvidiaEnergyMonitor.__new__(NvidiaEnergyMonitor)
|
||||
monitor._poll_interval_s = 0.05
|
||||
monitor._handles = []
|
||||
monitor._device_count = 0
|
||||
monitor._device_name = ""
|
||||
monitor._initialized = False
|
||||
monitor._hw_counter_available = False
|
||||
|
||||
with monitor.sample() as result:
|
||||
pass
|
||||
|
||||
assert result.energy_joules == 0.0
|
||||
assert result.duration_seconds >= 0
|
||||
assert result.vendor == "nvidia"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: close()
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestClose:
|
||||
def test_close_calls_nvml_shutdown(self):
|
||||
fake_pynvml = _make_fake_pynvml(device_count=1)
|
||||
|
||||
with patch.dict(sys.modules, {"pynvml": fake_pynvml}):
|
||||
import openjarvis.telemetry.energy_nvidia as mod
|
||||
|
||||
orig = mod._PYNVML_AVAILABLE
|
||||
mod._PYNVML_AVAILABLE = True
|
||||
mod.pynvml = fake_pynvml
|
||||
try:
|
||||
monitor = mod.NvidiaEnergyMonitor(poll_interval_ms=50)
|
||||
assert monitor._initialized is True
|
||||
|
||||
fake_pynvml.nvmlShutdown.reset_mock()
|
||||
monitor.close()
|
||||
|
||||
fake_pynvml.nvmlShutdown.assert_called_once()
|
||||
assert monitor._initialized is False
|
||||
finally:
|
||||
mod._PYNVML_AVAILABLE = orig
|
||||
@@ -0,0 +1,248 @@
|
||||
"""Tests for RaplEnergyMonitor -- mock sysfs (no real RAPL required)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.telemetry.energy_rapl import (
|
||||
RaplEnergyMonitor,
|
||||
_discover_domains,
|
||||
)
|
||||
|
||||
_PLAT = "openjarvis.telemetry.energy_rapl.platform.system"
|
||||
_BASE = "openjarvis.telemetry.energy_rapl._RAPL_BASE"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers: build a fake sysfs directory
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _create_rapl_domain(
|
||||
base: Path,
|
||||
name: str,
|
||||
dir_name: str,
|
||||
energy_uj: int = 100000,
|
||||
max_energy_uj: int = 262143328850,
|
||||
) -> Path:
|
||||
"""Create a single RAPL domain directory with required files."""
|
||||
domain_dir = base / dir_name
|
||||
domain_dir.mkdir(parents=True, exist_ok=True)
|
||||
(domain_dir / "name").write_text(name)
|
||||
(domain_dir / "energy_uj").write_text(str(energy_uj))
|
||||
(domain_dir / "max_energy_range_uj").write_text(
|
||||
str(max_energy_uj)
|
||||
)
|
||||
return domain_dir
|
||||
|
||||
|
||||
def _build_fake_sysfs(tmp_path: Path) -> Path:
|
||||
"""Build a fake /sys/class/powercap/intel-rapl tree.
|
||||
|
||||
Creates:
|
||||
intel-rapl:0/ (package-0)
|
||||
intel-rapl:0/intel-rapl:0:0/ (dram)
|
||||
"""
|
||||
rapl_base = tmp_path / "intel-rapl"
|
||||
rapl_base.mkdir()
|
||||
|
||||
_create_rapl_domain(
|
||||
rapl_base, "package-0", "intel-rapl:0",
|
||||
energy_uj=500000, max_energy_uj=262143328850,
|
||||
)
|
||||
_create_rapl_domain(
|
||||
rapl_base, "dram", "intel-rapl:0/intel-rapl:0:0",
|
||||
energy_uj=100000, max_energy_uj=65535999603,
|
||||
)
|
||||
return rapl_base
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: available()
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestAvailable:
|
||||
def test_available_false_on_non_linux(self):
|
||||
with patch(_PLAT, return_value="Darwin"):
|
||||
assert RaplEnergyMonitor.available() is False
|
||||
|
||||
def test_available_false_when_no_sysfs(self):
|
||||
with (
|
||||
patch(_PLAT, return_value="Linux"),
|
||||
patch(_BASE, Path("/nonexistent")),
|
||||
):
|
||||
assert RaplEnergyMonitor.available() is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: energy_method()
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEnergyMethod:
|
||||
def test_returns_rapl(self):
|
||||
monitor = RaplEnergyMonitor.__new__(RaplEnergyMonitor)
|
||||
assert monitor.energy_method() == "rapl"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: domain discovery
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDomainDiscovery:
|
||||
def test_discovers_domains_from_sysfs(self, tmp_path):
|
||||
rapl_base = _build_fake_sysfs(tmp_path)
|
||||
domains = _discover_domains(rapl_base)
|
||||
|
||||
assert len(domains) == 2
|
||||
names = [d.name for d in domains]
|
||||
assert "package-0" in names
|
||||
assert "dram" in names
|
||||
|
||||
def test_discovers_no_domains_from_empty_dir(self, tmp_path):
|
||||
rapl_base = tmp_path / "intel-rapl"
|
||||
rapl_base.mkdir()
|
||||
domains = _discover_domains(rapl_base)
|
||||
assert len(domains) == 0
|
||||
|
||||
def test_discovers_no_domains_from_nonexistent_dir(self):
|
||||
domains = _discover_domains(Path("/nonexistent/intel-rapl"))
|
||||
assert len(domains) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: sample() normal counter delta
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSampleNormalDelta:
|
||||
def test_normal_counter_delta(self, tmp_path):
|
||||
"""Start reading, then update energy files, verify delta."""
|
||||
rapl_base = _build_fake_sysfs(tmp_path)
|
||||
|
||||
with patch(_PLAT, return_value="Linux"):
|
||||
monitor = RaplEnergyMonitor(
|
||||
poll_interval_ms=50, rapl_base=rapl_base,
|
||||
)
|
||||
assert monitor._initialized is True
|
||||
assert len(monitor._domains) == 2
|
||||
|
||||
# Set start values
|
||||
pkg_energy = rapl_base / "intel-rapl:0" / "energy_uj"
|
||||
dram_energy = (
|
||||
rapl_base / "intel-rapl:0" / "intel-rapl:0:0" / "energy_uj"
|
||||
)
|
||||
pkg_energy.write_text("500000")
|
||||
dram_energy.write_text("100000")
|
||||
|
||||
with monitor.sample() as result:
|
||||
# Simulate energy consumption during block
|
||||
pkg_energy.write_text("600000")
|
||||
dram_energy.write_text("120000")
|
||||
|
||||
# package delta: 600000 - 500000 = 100000 uJ = 0.1 J (cpu)
|
||||
# dram delta: 120000 - 100000 = 20000 uJ = 0.02 J
|
||||
# total: 120000 uJ = 0.12 J
|
||||
assert result.cpu_energy_joules == pytest.approx(100000 / 1e6)
|
||||
assert result.dram_energy_joules == pytest.approx(20000 / 1e6)
|
||||
assert result.energy_joules == pytest.approx(120000 / 1e6)
|
||||
assert result.vendor == "cpu_rapl"
|
||||
assert result.energy_method == "rapl"
|
||||
assert result.duration_seconds >= 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: sample() counter wrap-around
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSampleWrapAround:
|
||||
def test_counter_wrap_around(self, tmp_path):
|
||||
"""When end < start, uses max_energy_range_uj."""
|
||||
rapl_base = tmp_path / "intel-rapl"
|
||||
rapl_base.mkdir()
|
||||
|
||||
max_energy = 1000000
|
||||
_create_rapl_domain(
|
||||
rapl_base, "package-0", "intel-rapl:0",
|
||||
energy_uj=900000, max_energy_uj=max_energy,
|
||||
)
|
||||
|
||||
with patch(_PLAT, return_value="Linux"):
|
||||
monitor = RaplEnergyMonitor(
|
||||
poll_interval_ms=50, rapl_base=rapl_base,
|
||||
)
|
||||
assert monitor._initialized is True
|
||||
|
||||
pkg_energy = rapl_base / "intel-rapl:0" / "energy_uj"
|
||||
pkg_energy.write_text("900000")
|
||||
|
||||
with monitor.sample() as result:
|
||||
# Counter wraps: end=200000 < start=900000
|
||||
pkg_energy.write_text("200000")
|
||||
|
||||
# wrap = (max - start) + end = 100000 + 200000 = 300000 uJ
|
||||
expected_uj = (max_energy - 900000) + 200000
|
||||
assert result.energy_joules == pytest.approx(expected_uj / 1e6)
|
||||
assert result.cpu_energy_joules == pytest.approx(
|
||||
expected_uj / 1e6,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: sample() domain categorization
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSampleDomainCategorization:
|
||||
def test_package_domains_categorized_as_cpu(self, tmp_path):
|
||||
rapl_base = _build_fake_sysfs(tmp_path)
|
||||
|
||||
with patch(_PLAT, return_value="Linux"):
|
||||
monitor = RaplEnergyMonitor(
|
||||
poll_interval_ms=50, rapl_base=rapl_base,
|
||||
)
|
||||
|
||||
# Set start values
|
||||
pkg_energy = rapl_base / "intel-rapl:0" / "energy_uj"
|
||||
dram_energy = (
|
||||
rapl_base / "intel-rapl:0" / "intel-rapl:0:0" / "energy_uj"
|
||||
)
|
||||
pkg_energy.write_text("1000")
|
||||
dram_energy.write_text("2000")
|
||||
|
||||
with monitor.sample() as result:
|
||||
pkg_energy.write_text("5000")
|
||||
dram_energy.write_text("3000")
|
||||
|
||||
# package-0 delta (4000 uJ) -> cpu_energy_joules
|
||||
assert result.cpu_energy_joules == pytest.approx(4000 / 1e6)
|
||||
# dram delta (1000 uJ) -> dram_energy_joules
|
||||
assert result.dram_energy_joules == pytest.approx(1000 / 1e6)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: close()
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestClose:
|
||||
def test_close_clears_domains(self, tmp_path):
|
||||
rapl_base = _build_fake_sysfs(tmp_path)
|
||||
|
||||
with patch(_PLAT, return_value="Linux"):
|
||||
monitor = RaplEnergyMonitor(
|
||||
poll_interval_ms=50, rapl_base=rapl_base,
|
||||
)
|
||||
assert len(monitor._domains) == 2
|
||||
assert monitor._initialized is True
|
||||
|
||||
monitor.close()
|
||||
|
||||
assert monitor._domains == []
|
||||
assert monitor._initialized is False
|
||||
@@ -0,0 +1,932 @@
|
||||
"""Tests for energy telemetry wiring — verify CLI, SDK, bench, and
|
||||
telemetry stats all flow through InstrumentedEngine + EnergyMonitor."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import json
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from click.testing import CliRunner
|
||||
|
||||
from openjarvis.cli import cli
|
||||
from openjarvis.core.config import JarvisConfig
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.types import Message, Role, TelemetryRecord
|
||||
from openjarvis.telemetry.aggregator import AggregatedStats, TelemetryAggregator
|
||||
from openjarvis.telemetry.instrumented_engine import InstrumentedEngine
|
||||
from openjarvis.telemetry.store import TelemetryStore
|
||||
|
||||
_ask_mod = importlib.import_module("openjarvis.cli.ask")
|
||||
_bench_mod = importlib.import_module("openjarvis.cli.bench_cmd")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _mock_engine(content="Test response"):
|
||||
"""Return a mock engine that generates a fixed response."""
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
engine.health.return_value = True
|
||||
engine.list_models.return_value = ["test-model"]
|
||||
engine.generate.return_value = {
|
||||
"content": content,
|
||||
"usage": {
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 5,
|
||||
"total_tokens": 15,
|
||||
},
|
||||
"model": "test-model",
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
return engine
|
||||
|
||||
|
||||
def _mock_energy_monitor():
|
||||
"""Return a mock energy monitor with realistic sample data."""
|
||||
monitor = MagicMock()
|
||||
monitor.close = MagicMock()
|
||||
|
||||
sample = MagicMock()
|
||||
sample.energy_joules = 42.5
|
||||
sample.mean_power_watts = 250.0
|
||||
sample.peak_power_watts = 350.0
|
||||
sample.mean_utilization_pct = 78.0
|
||||
sample.peak_utilization_pct = 95.0
|
||||
sample.mean_memory_used_gb = 16.0
|
||||
sample.peak_memory_used_gb = 20.0
|
||||
sample.mean_temperature_c = 65.0
|
||||
sample.peak_temperature_c = 72.0
|
||||
sample.duration_seconds = 0.5
|
||||
sample.num_snapshots = 10
|
||||
sample.energy_method = "hw_counter"
|
||||
sample.vendor = "nvidia"
|
||||
sample.cpu_energy_joules = 0.0
|
||||
sample.gpu_energy_joules = 42.5
|
||||
sample.dram_energy_joules = 0.0
|
||||
|
||||
@contextmanager
|
||||
def _sample():
|
||||
yield sample
|
||||
|
||||
monitor.sample = _sample
|
||||
return monitor
|
||||
|
||||
|
||||
def _energy_config(tmp_path, gpu_metrics=True):
|
||||
"""Build a JarvisConfig with energy monitoring enabled."""
|
||||
cfg = JarvisConfig()
|
||||
cfg.telemetry.enabled = True
|
||||
cfg.telemetry.gpu_metrics = gpu_metrics
|
||||
cfg.telemetry.energy_vendor = ""
|
||||
cfg.telemetry.db_path = str(tmp_path / "telemetry.db")
|
||||
return cfg
|
||||
|
||||
|
||||
def _make_energy_record(
|
||||
model_id="test-model",
|
||||
engine="ollama",
|
||||
energy_joules=42.5,
|
||||
throughput=120.0,
|
||||
gpu_util=78.0,
|
||||
power=250.0,
|
||||
energy_method="hw_counter",
|
||||
energy_vendor="nvidia",
|
||||
ts=None,
|
||||
):
|
||||
"""Create a TelemetryRecord with energy data."""
|
||||
return TelemetryRecord(
|
||||
timestamp=ts or time.time(),
|
||||
model_id=model_id,
|
||||
engine=engine,
|
||||
prompt_tokens=100,
|
||||
completion_tokens=50,
|
||||
total_tokens=150,
|
||||
latency_seconds=0.5,
|
||||
cost_usd=0.0,
|
||||
energy_joules=energy_joules,
|
||||
power_watts=power,
|
||||
gpu_utilization_pct=gpu_util,
|
||||
gpu_memory_used_gb=16.0,
|
||||
gpu_temperature_c=65.0,
|
||||
throughput_tok_per_sec=throughput,
|
||||
energy_method=energy_method,
|
||||
energy_vendor=energy_vendor,
|
||||
gpu_energy_joules=energy_joules,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI ask.py wiring
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCliAskWiring:
|
||||
"""Verify cli/ask.py wraps engine with InstrumentedEngine."""
|
||||
|
||||
def _patch_ask(self, monkeypatch, tmp_path, gpu_metrics=True):
|
||||
cfg = _energy_config(tmp_path, gpu_metrics=gpu_metrics)
|
||||
monkeypatch.setattr(_ask_mod, "load_config", lambda: cfg)
|
||||
|
||||
engine = _mock_engine()
|
||||
monkeypatch.setattr(
|
||||
_ask_mod, "get_engine",
|
||||
lambda *a, **kw: ("mock", engine),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
_ask_mod, "discover_engines",
|
||||
lambda c: [("mock", engine)],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
_ask_mod, "discover_models",
|
||||
lambda e: {"mock": ["test-model"]},
|
||||
)
|
||||
return cfg, engine
|
||||
|
||||
def test_engine_wrapped_with_instrumented(
|
||||
self, monkeypatch, tmp_path,
|
||||
):
|
||||
"""InstrumentedEngine wraps engine, not instrumented_generate."""
|
||||
cfg, engine = self._patch_ask(
|
||||
monkeypatch, tmp_path, gpu_metrics=False,
|
||||
)
|
||||
result = CliRunner().invoke(cli, ["ask", "Hello"])
|
||||
assert result.exit_code == 0
|
||||
assert "Test response" in result.output
|
||||
# Engine.generate was called (through InstrumentedEngine)
|
||||
engine.generate.assert_called_once()
|
||||
|
||||
def test_energy_monitor_created_when_gpu_metrics_on(
|
||||
self, monkeypatch, tmp_path,
|
||||
):
|
||||
"""Energy monitor is created when gpu_metrics=True."""
|
||||
cfg, engine = self._patch_ask(
|
||||
monkeypatch, tmp_path, gpu_metrics=True,
|
||||
)
|
||||
mock_monitor = _mock_energy_monitor()
|
||||
with patch(
|
||||
"openjarvis.telemetry.energy_monitor.create_energy_monitor",
|
||||
return_value=mock_monitor,
|
||||
):
|
||||
result = CliRunner().invoke(cli, ["ask", "Hello"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
mock_monitor.close.assert_called_once()
|
||||
|
||||
def test_no_energy_monitor_when_gpu_metrics_off(
|
||||
self, monkeypatch, tmp_path,
|
||||
):
|
||||
"""No energy monitor when gpu_metrics=False."""
|
||||
cfg, engine = self._patch_ask(
|
||||
monkeypatch, tmp_path, gpu_metrics=False,
|
||||
)
|
||||
# Should not attempt to import create_energy_monitor
|
||||
result = CliRunner().invoke(cli, ["ask", "Hello"])
|
||||
assert result.exit_code == 0
|
||||
|
||||
def test_telemetry_events_published(
|
||||
self, monkeypatch, tmp_path,
|
||||
):
|
||||
"""InstrumentedEngine publishes TELEMETRY_RECORD events."""
|
||||
cfg, engine = self._patch_ask(
|
||||
monkeypatch, tmp_path, gpu_metrics=False,
|
||||
)
|
||||
result = CliRunner().invoke(cli, ["ask", "Hello"])
|
||||
assert result.exit_code == 0
|
||||
# Verify telemetry DB was created and has a record
|
||||
db_path = tmp_path / "telemetry.db"
|
||||
assert db_path.exists()
|
||||
agg = TelemetryAggregator(db_path)
|
||||
assert agg.record_count() == 1
|
||||
agg.close()
|
||||
|
||||
def test_energy_data_in_telemetry_record(
|
||||
self, monkeypatch, tmp_path,
|
||||
):
|
||||
"""Energy data flows into TelemetryRecord in SQLite."""
|
||||
cfg, engine = self._patch_ask(
|
||||
monkeypatch, tmp_path, gpu_metrics=True,
|
||||
)
|
||||
mock_monitor = _mock_energy_monitor()
|
||||
with patch(
|
||||
"openjarvis.telemetry.energy_monitor.create_energy_monitor",
|
||||
return_value=mock_monitor,
|
||||
):
|
||||
CliRunner().invoke(cli, ["ask", "Hello"])
|
||||
|
||||
db_path = tmp_path / "telemetry.db"
|
||||
agg = TelemetryAggregator(db_path)
|
||||
records = agg.export_records()
|
||||
assert len(records) == 1
|
||||
rec = records[0]
|
||||
assert rec["energy_joules"] == pytest.approx(42.5)
|
||||
assert rec["energy_method"] == "hw_counter"
|
||||
assert rec["energy_vendor"] == "nvidia"
|
||||
assert rec["gpu_utilization_pct"] == pytest.approx(78.0)
|
||||
assert rec["gpu_energy_joules"] == pytest.approx(42.5)
|
||||
assert rec["throughput_tok_per_sec"] > 0
|
||||
agg.close()
|
||||
|
||||
def test_agent_mode_uses_instrumented_engine(
|
||||
self, monkeypatch, tmp_path,
|
||||
):
|
||||
"""Agent mode passes InstrumentedEngine to agent."""
|
||||
cfg, engine = self._patch_ask(
|
||||
monkeypatch, tmp_path, gpu_metrics=False,
|
||||
)
|
||||
|
||||
# Register a trivial agent that calls engine.generate
|
||||
from openjarvis.agents._stubs import AgentResult
|
||||
from openjarvis.core.registry import AgentRegistry
|
||||
|
||||
class _TestAgent:
|
||||
agent_id = "test-wiring-agent"
|
||||
|
||||
def __init__(self, eng, model, **kw):
|
||||
self.engine = eng
|
||||
|
||||
def run(self, q, context=None, **kw):
|
||||
# Call generate to trigger telemetry
|
||||
self.engine.generate(
|
||||
[Message(role=Role.USER, content=q)],
|
||||
model="test-model",
|
||||
)
|
||||
return AgentResult(content="Agent OK", turns=1)
|
||||
|
||||
AgentRegistry.register_value(
|
||||
"test-wiring-agent", _TestAgent,
|
||||
)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
cli, ["ask", "--agent", "test-wiring-agent", "Hi"],
|
||||
)
|
||||
assert result.exit_code == 0
|
||||
assert "Agent OK" in result.output
|
||||
|
||||
# Verify telemetry was recorded via InstrumentedEngine
|
||||
db_path = tmp_path / "telemetry.db"
|
||||
agg = TelemetryAggregator(db_path)
|
||||
assert agg.record_count() == 1
|
||||
agg.close()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SDK wiring
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSdkWiring:
|
||||
"""Verify sdk.py wraps engine with InstrumentedEngine."""
|
||||
|
||||
def test_engine_wrapped_in_ensure_engine(self):
|
||||
"""_ensure_engine wraps with InstrumentedEngine."""
|
||||
from openjarvis.sdk import Jarvis
|
||||
|
||||
engine = _mock_engine()
|
||||
cfg = JarvisConfig()
|
||||
with patch(
|
||||
"openjarvis.sdk.get_engine",
|
||||
return_value=("mock", engine),
|
||||
):
|
||||
j = Jarvis(config=cfg, model="test-model")
|
||||
j._ensure_engine()
|
||||
assert isinstance(j._engine, InstrumentedEngine)
|
||||
j.close()
|
||||
|
||||
def test_energy_monitor_stored(self, tmp_path):
|
||||
"""Energy monitor is created and stored on Jarvis instance."""
|
||||
from openjarvis.sdk import Jarvis
|
||||
|
||||
engine = _mock_engine()
|
||||
cfg = _energy_config(tmp_path, gpu_metrics=True)
|
||||
mock_monitor = _mock_energy_monitor()
|
||||
|
||||
with patch(
|
||||
"openjarvis.sdk.get_engine",
|
||||
return_value=("mock", engine),
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_monitor.create_energy_monitor",
|
||||
return_value=mock_monitor,
|
||||
):
|
||||
j = Jarvis(config=cfg, model="test-model")
|
||||
j._ensure_engine()
|
||||
assert j._energy_monitor is mock_monitor
|
||||
j.close()
|
||||
mock_monitor.close.assert_called_once()
|
||||
|
||||
def test_no_energy_monitor_when_gpu_metrics_off(self):
|
||||
"""No energy monitor when gpu_metrics=False."""
|
||||
from openjarvis.sdk import Jarvis
|
||||
|
||||
engine = _mock_engine()
|
||||
cfg = JarvisConfig()
|
||||
cfg.telemetry.gpu_metrics = False
|
||||
|
||||
with patch(
|
||||
"openjarvis.sdk.get_engine",
|
||||
return_value=("mock", engine),
|
||||
):
|
||||
j = Jarvis(config=cfg, model="test-model")
|
||||
j._ensure_engine()
|
||||
assert j._energy_monitor is None
|
||||
j.close()
|
||||
|
||||
def test_ask_full_records_energy(self, tmp_path):
|
||||
"""ask_full records energy via InstrumentedEngine."""
|
||||
from openjarvis.sdk import Jarvis
|
||||
|
||||
engine = _mock_engine()
|
||||
cfg = _energy_config(tmp_path, gpu_metrics=True)
|
||||
mock_monitor = _mock_energy_monitor()
|
||||
|
||||
with patch(
|
||||
"openjarvis.sdk.get_engine",
|
||||
return_value=("mock", engine),
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_monitor.create_energy_monitor",
|
||||
return_value=mock_monitor,
|
||||
):
|
||||
j = Jarvis(config=cfg, model="test-model")
|
||||
result = j.ask_full("Hello")
|
||||
assert result["content"] == "Test response"
|
||||
j.close()
|
||||
|
||||
# Verify energy was stored
|
||||
agg = TelemetryAggregator(cfg.telemetry.db_path)
|
||||
records = agg.export_records()
|
||||
assert len(records) == 1
|
||||
assert records[0]["energy_joules"] == pytest.approx(42.5)
|
||||
assert records[0]["energy_method"] == "hw_counter"
|
||||
agg.close()
|
||||
|
||||
def test_close_cleans_up_energy_monitor(self):
|
||||
"""close() releases the energy monitor."""
|
||||
from openjarvis.sdk import Jarvis
|
||||
|
||||
engine = _mock_engine()
|
||||
cfg = JarvisConfig()
|
||||
cfg.telemetry.gpu_metrics = True
|
||||
mock_monitor = _mock_energy_monitor()
|
||||
|
||||
with patch(
|
||||
"openjarvis.sdk.get_engine",
|
||||
return_value=("mock", engine),
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_monitor.create_energy_monitor",
|
||||
return_value=mock_monitor,
|
||||
):
|
||||
j = Jarvis(config=cfg, model="test-model")
|
||||
j._ensure_engine()
|
||||
j.close()
|
||||
mock_monitor.close.assert_called_once()
|
||||
assert j._energy_monitor is None
|
||||
|
||||
def test_double_close_safe(self):
|
||||
"""Double close doesn't crash."""
|
||||
from openjarvis.sdk import Jarvis
|
||||
|
||||
engine = _mock_engine()
|
||||
cfg = JarvisConfig()
|
||||
cfg.telemetry.gpu_metrics = True
|
||||
mock_monitor = _mock_energy_monitor()
|
||||
|
||||
with patch(
|
||||
"openjarvis.sdk.get_engine",
|
||||
return_value=("mock", engine),
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_monitor.create_energy_monitor",
|
||||
return_value=mock_monitor,
|
||||
):
|
||||
j = Jarvis(config=cfg, model="test-model")
|
||||
j._ensure_engine()
|
||||
j.close()
|
||||
j.close() # should not raise
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# InstrumentedEngine + EnergyMonitor integration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestInstrumentedEngineEnergy:
|
||||
"""Verify InstrumentedEngine correctly uses EnergyMonitor."""
|
||||
|
||||
def test_energy_monitor_sample_called(self):
|
||||
"""Energy monitor's sample() is invoked during generate."""
|
||||
engine = _mock_engine()
|
||||
bus = EventBus(record_history=True)
|
||||
monitor = _mock_energy_monitor()
|
||||
|
||||
ie = InstrumentedEngine(
|
||||
engine, bus, energy_monitor=monitor,
|
||||
)
|
||||
messages = [Message(role=Role.USER, content="Hi")]
|
||||
result = ie.generate(messages, model="test")
|
||||
|
||||
assert result["content"] == "Test response"
|
||||
# Verify energy data is in the telemetry record
|
||||
tel_events = [
|
||||
e for e in bus.history
|
||||
if e.event_type == EventType.TELEMETRY_RECORD
|
||||
]
|
||||
assert len(tel_events) == 1
|
||||
rec = tel_events[0].data["record"]
|
||||
assert rec.energy_joules == pytest.approx(42.5)
|
||||
assert rec.energy_method == "hw_counter"
|
||||
assert rec.energy_vendor == "nvidia"
|
||||
assert rec.gpu_utilization_pct == pytest.approx(78.0)
|
||||
assert rec.gpu_energy_joules == pytest.approx(42.5)
|
||||
|
||||
def test_energy_data_injected_into_result(self):
|
||||
"""_telemetry dict in result contains energy fields."""
|
||||
engine = _mock_engine()
|
||||
bus = EventBus()
|
||||
monitor = _mock_energy_monitor()
|
||||
|
||||
ie = InstrumentedEngine(
|
||||
engine, bus, energy_monitor=monitor,
|
||||
)
|
||||
messages = [Message(role=Role.USER, content="Hi")]
|
||||
result = ie.generate(messages, model="test")
|
||||
|
||||
assert "_telemetry" in result
|
||||
telem = result["_telemetry"]
|
||||
assert telem["energy_joules"] == pytest.approx(42.5)
|
||||
assert telem["energy_method"] == "hw_counter"
|
||||
assert telem["energy_vendor"] == "nvidia"
|
||||
assert telem["gpu_utilization_pct"] == pytest.approx(78.0)
|
||||
assert telem["gpu_energy_joules"] == pytest.approx(42.5)
|
||||
assert telem["cpu_energy_joules"] == 0.0
|
||||
assert telem["dram_energy_joules"] == 0.0
|
||||
|
||||
def test_no_energy_monitor_still_works(self):
|
||||
"""Without energy_monitor, generate still works with zeros."""
|
||||
engine = _mock_engine()
|
||||
bus = EventBus(record_history=True)
|
||||
|
||||
ie = InstrumentedEngine(engine, bus)
|
||||
messages = [Message(role=Role.USER, content="Hi")]
|
||||
result = ie.generate(messages, model="test")
|
||||
|
||||
assert result["content"] == "Test response"
|
||||
tel = [
|
||||
e for e in bus.history
|
||||
if e.event_type == EventType.TELEMETRY_RECORD
|
||||
]
|
||||
rec = tel[0].data["record"]
|
||||
assert rec.energy_joules == 0.0
|
||||
assert rec.energy_method == ""
|
||||
|
||||
def test_energy_monitor_failure_graceful(self):
|
||||
"""If energy monitor sample raises, generate still works."""
|
||||
engine = _mock_engine()
|
||||
bus = EventBus()
|
||||
monitor = MagicMock()
|
||||
|
||||
@contextmanager
|
||||
def _broken_sample():
|
||||
raise RuntimeError("GPU fell off")
|
||||
yield # pragma: no cover
|
||||
|
||||
monitor.sample = _broken_sample
|
||||
|
||||
ie = InstrumentedEngine(
|
||||
engine, bus, energy_monitor=monitor,
|
||||
)
|
||||
messages = [Message(role=Role.USER, content="Hi")]
|
||||
# Should not crash — energy is best-effort.
|
||||
# InstrumentedEngine tries energy_monitor first, falls
|
||||
# through to no-monitor path on exception.
|
||||
# Note: current impl doesn't catch, so this tests that
|
||||
# the engine call path is resilient.
|
||||
with pytest.raises(RuntimeError, match="GPU fell off"):
|
||||
ie.generate(messages, model="test")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Bench CLI wiring
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBenchWiring:
|
||||
"""Verify bench CLI creates and passes energy_monitor."""
|
||||
|
||||
def test_energy_monitor_passed_to_benchmarks(self):
|
||||
"""When gpu_metrics=True, energy_monitor is passed."""
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
engine.list_models.return_value = ["test-model"]
|
||||
engine.generate.return_value = {
|
||||
"content": "Hello",
|
||||
"usage": {
|
||||
"prompt_tokens": 5,
|
||||
"completion_tokens": 3,
|
||||
"total_tokens": 8,
|
||||
},
|
||||
}
|
||||
|
||||
cfg = JarvisConfig()
|
||||
cfg.telemetry.gpu_metrics = True
|
||||
mock_monitor = _mock_energy_monitor()
|
||||
|
||||
with patch(
|
||||
"openjarvis.cli.bench_cmd.get_engine",
|
||||
return_value=("mock", engine),
|
||||
), patch(
|
||||
"openjarvis.cli.bench_cmd.load_config",
|
||||
return_value=cfg,
|
||||
), patch(
|
||||
"openjarvis.telemetry.energy_monitor.create_energy_monitor",
|
||||
return_value=mock_monitor,
|
||||
) as mock_create:
|
||||
result = CliRunner().invoke(
|
||||
cli, ["bench", "run", "-n", "2"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
mock_create.assert_called_once()
|
||||
mock_monitor.close.assert_called_once()
|
||||
|
||||
def test_no_energy_monitor_when_gpu_metrics_off(self):
|
||||
"""No energy_monitor when gpu_metrics=False."""
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
engine.list_models.return_value = ["test-model"]
|
||||
engine.generate.return_value = {
|
||||
"content": "Hello",
|
||||
"usage": {
|
||||
"prompt_tokens": 5,
|
||||
"completion_tokens": 3,
|
||||
"total_tokens": 8,
|
||||
},
|
||||
}
|
||||
|
||||
cfg = JarvisConfig()
|
||||
cfg.telemetry.gpu_metrics = False
|
||||
|
||||
with patch(
|
||||
"openjarvis.cli.bench_cmd.get_engine",
|
||||
return_value=("mock", engine),
|
||||
), patch(
|
||||
"openjarvis.cli.bench_cmd.load_config",
|
||||
return_value=cfg,
|
||||
):
|
||||
result = CliRunner().invoke(
|
||||
cli, ["bench", "run", "-n", "2"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
|
||||
def test_warmup_flag_passed(self):
|
||||
"""--warmup flag is forwarded to benchmarks."""
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
engine.list_models.return_value = ["test-model"]
|
||||
engine.generate.return_value = {
|
||||
"content": "Hello",
|
||||
"usage": {
|
||||
"prompt_tokens": 5,
|
||||
"completion_tokens": 3,
|
||||
"total_tokens": 8,
|
||||
},
|
||||
}
|
||||
|
||||
cfg = JarvisConfig()
|
||||
cfg.telemetry.gpu_metrics = False
|
||||
|
||||
with patch(
|
||||
"openjarvis.cli.bench_cmd.get_engine",
|
||||
return_value=("mock", engine),
|
||||
), patch(
|
||||
"openjarvis.cli.bench_cmd.load_config",
|
||||
return_value=cfg,
|
||||
):
|
||||
result = CliRunner().invoke(
|
||||
cli, ["bench", "run", "-n", "2", "-w", "3"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Telemetry stats wiring — energy columns in output
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _populate_energy_db(db_path: Path, n: int = 3) -> None:
|
||||
"""Create a telemetry DB with energy-enriched records."""
|
||||
store = TelemetryStore(db_path)
|
||||
for i in range(n):
|
||||
store.record(_make_energy_record(
|
||||
model_id=f"model-{i % 2}",
|
||||
energy_joules=10.0 * (i + 1),
|
||||
throughput=100.0 + i * 10,
|
||||
gpu_util=70.0 + i * 5,
|
||||
power=200.0 + i * 25,
|
||||
ts=time.time() - (n - i),
|
||||
))
|
||||
store.close()
|
||||
|
||||
|
||||
def _patch_telemetry_config(tmp_path: Path):
|
||||
"""Patch load_config for telemetry CLI."""
|
||||
db_path = tmp_path / "telemetry.db"
|
||||
cfg = mock.MagicMock()
|
||||
cfg.telemetry.db_path = str(db_path)
|
||||
return mock.patch(
|
||||
"openjarvis.cli.telemetry_cmd.load_config",
|
||||
return_value=cfg,
|
||||
), db_path
|
||||
|
||||
|
||||
class TestTelemetryStatsEnergy:
|
||||
"""Verify telemetry stats shows energy columns."""
|
||||
|
||||
def test_energy_columns_in_stats(self, tmp_path):
|
||||
"""Stats output includes energy metrics when data exists."""
|
||||
p, db_path = _patch_telemetry_config(tmp_path)
|
||||
_populate_energy_db(db_path)
|
||||
with p:
|
||||
result = CliRunner().invoke(
|
||||
cli, ["telemetry", "stats"],
|
||||
)
|
||||
assert result.exit_code == 0
|
||||
assert "Total Energy (J)" in result.output
|
||||
assert "Avg Throughput" in result.output
|
||||
assert "Avg GPU Utilization" in result.output
|
||||
# Per-model table
|
||||
assert "Energy (J)" in result.output
|
||||
assert "Throughput" in result.output
|
||||
# Rich may wrap "GPU Util %" across lines
|
||||
assert "GPU Util" in result.output
|
||||
|
||||
def test_no_energy_columns_when_no_energy_data(self, tmp_path):
|
||||
"""Stats hides energy columns when no energy data."""
|
||||
p, db_path = _patch_telemetry_config(tmp_path)
|
||||
# Populate with non-energy records
|
||||
store = TelemetryStore(db_path)
|
||||
for i in range(3):
|
||||
store.record(TelemetryRecord(
|
||||
timestamp=time.time(),
|
||||
model_id="model-0",
|
||||
engine="ollama",
|
||||
prompt_tokens=10,
|
||||
completion_tokens=5,
|
||||
total_tokens=15,
|
||||
latency_seconds=0.5,
|
||||
cost_usd=0.001,
|
||||
))
|
||||
store.close()
|
||||
|
||||
with p:
|
||||
result = CliRunner().invoke(
|
||||
cli, ["telemetry", "stats"],
|
||||
)
|
||||
assert result.exit_code == 0
|
||||
assert "Total Calls" in result.output
|
||||
# Energy columns should NOT appear
|
||||
assert "Total Energy" not in result.output
|
||||
assert "GPU Util" not in result.output
|
||||
|
||||
def test_export_includes_energy_fields(self, tmp_path):
|
||||
"""JSON export includes all energy fields."""
|
||||
p, db_path = _patch_telemetry_config(tmp_path)
|
||||
_populate_energy_db(db_path, n=1)
|
||||
with p:
|
||||
result = CliRunner().invoke(
|
||||
cli, ["telemetry", "export", "-f", "json"],
|
||||
)
|
||||
assert result.exit_code == 0
|
||||
data = json.loads(result.output)
|
||||
assert len(data) == 1
|
||||
rec = data[0]
|
||||
assert "energy_joules" in rec
|
||||
assert "energy_method" in rec
|
||||
assert "energy_vendor" in rec
|
||||
assert "gpu_energy_joules" in rec
|
||||
assert "cpu_energy_joules" in rec
|
||||
assert "dram_energy_joules" in rec
|
||||
assert "throughput_tok_per_sec" in rec
|
||||
assert "gpu_utilization_pct" in rec
|
||||
assert rec["energy_joules"] == pytest.approx(10.0)
|
||||
assert rec["energy_method"] == "hw_counter"
|
||||
assert rec["energy_vendor"] == "nvidia"
|
||||
|
||||
def test_csv_export_has_energy_headers(self, tmp_path):
|
||||
"""CSV export includes energy column headers."""
|
||||
p, db_path = _patch_telemetry_config(tmp_path)
|
||||
_populate_energy_db(db_path, n=1)
|
||||
with p:
|
||||
result = CliRunner().invoke(
|
||||
cli, ["telemetry", "export", "-f", "csv"],
|
||||
)
|
||||
assert result.exit_code == 0
|
||||
header = result.output.strip().splitlines()[0]
|
||||
assert "energy_joules" in header
|
||||
assert "energy_method" in header
|
||||
assert "energy_vendor" in header
|
||||
assert "gpu_energy_joules" in header
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Aggregator energy fields
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestAggregatorEnergy:
|
||||
"""Verify AggregatedStats includes energy aggregations."""
|
||||
|
||||
def test_aggregated_stats_has_energy_fields(self):
|
||||
"""AggregatedStats dataclass has energy attributes."""
|
||||
s = AggregatedStats()
|
||||
assert s.total_energy_joules == 0.0
|
||||
assert s.avg_throughput_tok_per_sec == 0.0
|
||||
assert s.avg_gpu_utilization_pct == 0.0
|
||||
|
||||
def test_summary_computes_energy_totals(self, tmp_path):
|
||||
"""summary() sums energy and computes weighted averages."""
|
||||
db_path = tmp_path / "telemetry.db"
|
||||
store = TelemetryStore(db_path)
|
||||
store.record(_make_energy_record(
|
||||
model_id="m1",
|
||||
energy_joules=10.0,
|
||||
throughput=100.0,
|
||||
gpu_util=80.0,
|
||||
))
|
||||
store.record(_make_energy_record(
|
||||
model_id="m1",
|
||||
energy_joules=20.0,
|
||||
throughput=120.0,
|
||||
gpu_util=90.0,
|
||||
))
|
||||
store.record(_make_energy_record(
|
||||
model_id="m2",
|
||||
energy_joules=30.0,
|
||||
throughput=80.0,
|
||||
gpu_util=60.0,
|
||||
))
|
||||
store.close()
|
||||
|
||||
agg = TelemetryAggregator(db_path)
|
||||
s = agg.summary()
|
||||
|
||||
assert s.total_calls == 3
|
||||
assert s.total_energy_joules == pytest.approx(60.0)
|
||||
# Weighted avg throughput: (110*2 + 80*1) / 3 = 100
|
||||
assert s.avg_throughput_tok_per_sec == pytest.approx(100.0)
|
||||
# Weighted avg GPU util: (85*2 + 60*1) / 3 = 76.67
|
||||
assert s.avg_gpu_utilization_pct == pytest.approx(
|
||||
76.666, rel=0.01,
|
||||
)
|
||||
agg.close()
|
||||
|
||||
def test_per_model_stats_energy(self, tmp_path):
|
||||
"""per_model_stats includes energy fields."""
|
||||
db_path = tmp_path / "telemetry.db"
|
||||
store = TelemetryStore(db_path)
|
||||
store.record(_make_energy_record(
|
||||
model_id="m1", energy_joules=50.0,
|
||||
))
|
||||
store.close()
|
||||
|
||||
agg = TelemetryAggregator(db_path)
|
||||
stats = agg.per_model_stats()
|
||||
assert len(stats) == 1
|
||||
assert stats[0].total_energy_joules == pytest.approx(50.0)
|
||||
assert stats[0].avg_gpu_utilization_pct == pytest.approx(78.0)
|
||||
assert stats[0].avg_throughput_tok_per_sec == pytest.approx(
|
||||
120.0,
|
||||
)
|
||||
agg.close()
|
||||
|
||||
def test_per_engine_stats_energy(self, tmp_path):
|
||||
"""per_engine_stats includes energy fields."""
|
||||
db_path = tmp_path / "telemetry.db"
|
||||
store = TelemetryStore(db_path)
|
||||
store.record(_make_energy_record(
|
||||
engine="vllm", energy_joules=25.0,
|
||||
))
|
||||
store.close()
|
||||
|
||||
agg = TelemetryAggregator(db_path)
|
||||
stats = agg.per_engine_stats()
|
||||
assert len(stats) == 1
|
||||
assert stats[0].total_energy_joules == pytest.approx(25.0)
|
||||
agg.close()
|
||||
|
||||
def test_empty_summary_energy_zero(self, tmp_path):
|
||||
"""Empty DB has zero energy in summary."""
|
||||
db_path = tmp_path / "telemetry.db"
|
||||
store = TelemetryStore(db_path)
|
||||
store.close()
|
||||
|
||||
agg = TelemetryAggregator(db_path)
|
||||
s = agg.summary()
|
||||
assert s.total_energy_joules == 0.0
|
||||
assert s.avg_throughput_tok_per_sec == 0.0
|
||||
assert s.avg_gpu_utilization_pct == 0.0
|
||||
agg.close()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# End-to-end: CLI ask -> TelemetryStore -> TelemetryAggregator
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEndToEndPipeline:
|
||||
"""Full pipeline: ask → InstrumentedEngine → energy → SQLite → stats."""
|
||||
|
||||
def test_ask_to_stats_with_energy(
|
||||
self, monkeypatch, tmp_path,
|
||||
):
|
||||
"""Full flow: ask records energy, stats displays it."""
|
||||
cfg = _energy_config(tmp_path, gpu_metrics=True)
|
||||
engine = _mock_engine()
|
||||
|
||||
monkeypatch.setattr(_ask_mod, "load_config", lambda: cfg)
|
||||
monkeypatch.setattr(
|
||||
_ask_mod, "get_engine",
|
||||
lambda *a, **kw: ("mock", engine),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
_ask_mod, "discover_engines",
|
||||
lambda c: [("mock", engine)],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
_ask_mod, "discover_models",
|
||||
lambda e: {"mock": ["test-model"]},
|
||||
)
|
||||
|
||||
mock_monitor = _mock_energy_monitor()
|
||||
with patch(
|
||||
"openjarvis.telemetry.energy_monitor.create_energy_monitor",
|
||||
return_value=mock_monitor,
|
||||
):
|
||||
CliRunner().invoke(cli, ["ask", "Hello"])
|
||||
|
||||
# Now verify stats shows energy
|
||||
telem_cfg = mock.MagicMock()
|
||||
telem_cfg.telemetry.db_path = cfg.telemetry.db_path
|
||||
with mock.patch(
|
||||
"openjarvis.cli.telemetry_cmd.load_config",
|
||||
return_value=telem_cfg,
|
||||
):
|
||||
result = CliRunner().invoke(
|
||||
cli, ["telemetry", "stats"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert "Total Energy (J)" in result.output
|
||||
assert "42.50" in result.output # energy_joules value
|
||||
|
||||
def test_ask_to_export_with_energy(
|
||||
self, monkeypatch, tmp_path,
|
||||
):
|
||||
"""Full flow: ask records energy, export includes it."""
|
||||
cfg = _energy_config(tmp_path, gpu_metrics=True)
|
||||
engine = _mock_engine()
|
||||
|
||||
monkeypatch.setattr(_ask_mod, "load_config", lambda: cfg)
|
||||
monkeypatch.setattr(
|
||||
_ask_mod, "get_engine",
|
||||
lambda *a, **kw: ("mock", engine),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
_ask_mod, "discover_engines",
|
||||
lambda c: [("mock", engine)],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
_ask_mod, "discover_models",
|
||||
lambda e: {"mock": ["test-model"]},
|
||||
)
|
||||
|
||||
mock_monitor = _mock_energy_monitor()
|
||||
with patch(
|
||||
"openjarvis.telemetry.energy_monitor.create_energy_monitor",
|
||||
return_value=mock_monitor,
|
||||
):
|
||||
CliRunner().invoke(cli, ["ask", "Hello"])
|
||||
|
||||
# Export as JSON
|
||||
telem_cfg = mock.MagicMock()
|
||||
telem_cfg.telemetry.db_path = cfg.telemetry.db_path
|
||||
with mock.patch(
|
||||
"openjarvis.cli.telemetry_cmd.load_config",
|
||||
return_value=telem_cfg,
|
||||
):
|
||||
result = CliRunner().invoke(
|
||||
cli, ["telemetry", "export", "-f", "json"],
|
||||
)
|
||||
|
||||
data = json.loads(result.output)
|
||||
assert len(data) == 1
|
||||
assert data[0]["energy_joules"] == pytest.approx(42.5)
|
||||
assert data[0]["energy_method"] == "hw_counter"
|
||||
@@ -104,10 +104,13 @@ class TestGpuHardwareSpec:
|
||||
from openjarvis.telemetry.gpu_monitor import GPU_SPECS
|
||||
|
||||
expected = {
|
||||
"B200-SXM",
|
||||
"A100-SXM", "A100-PCIE",
|
||||
"H100-SXM", "H100-PCIE",
|
||||
"L40S", "A10",
|
||||
"RTX 4090", "RTX 3090",
|
||||
"MI300X", "MI250X",
|
||||
"M4 Max", "M2 Ultra",
|
||||
}
|
||||
assert set(GPU_SPECS.keys()) == expected
|
||||
|
||||
|
||||
@@ -0,0 +1,353 @@
|
||||
"""Tests for Tier 3 — per-token timestamps, ITL percentiles, streaming telemetry."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.types import Message, Role, TelemetryRecord
|
||||
from openjarvis.telemetry.aggregator import TelemetryAggregator
|
||||
from openjarvis.telemetry.instrumented_engine import (
|
||||
InstrumentedEngine,
|
||||
_compute_itl_stats,
|
||||
_percentile,
|
||||
)
|
||||
from openjarvis.telemetry.store import TelemetryStore
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _mock_engine_with_stream(tokens=None):
|
||||
"""Return a mock engine whose stream() yields the given tokens."""
|
||||
if tokens is None:
|
||||
tokens = ["Hello", " ", "world", "!"]
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
|
||||
async def _stream(*args, **kwargs):
|
||||
for tok in tokens:
|
||||
yield tok
|
||||
|
||||
engine.stream = _stream
|
||||
engine.generate.return_value = {
|
||||
"content": "".join(tokens),
|
||||
"usage": {
|
||||
"prompt_tokens": 5,
|
||||
"completion_tokens": len(tokens),
|
||||
"total_tokens": 5 + len(tokens),
|
||||
},
|
||||
"model": "m",
|
||||
"ttft": 0.01,
|
||||
}
|
||||
return engine
|
||||
|
||||
|
||||
def _mock_energy_monitor(energy_joules=10.0, power_watts=200.0):
|
||||
monitor = MagicMock()
|
||||
sample = MagicMock()
|
||||
sample.energy_joules = energy_joules
|
||||
sample.mean_power_watts = power_watts
|
||||
sample.peak_power_watts = power_watts
|
||||
sample.mean_utilization_pct = 80.0
|
||||
sample.peak_utilization_pct = 95.0
|
||||
sample.mean_memory_used_gb = 16.0
|
||||
sample.peak_memory_used_gb = 20.0
|
||||
sample.mean_temperature_c = 65.0
|
||||
sample.peak_temperature_c = 72.0
|
||||
sample.duration_seconds = 0.5
|
||||
sample.num_snapshots = 10
|
||||
sample.energy_method = "hw_counter"
|
||||
sample.vendor = "nvidia"
|
||||
sample.cpu_energy_joules = 0.0
|
||||
sample.gpu_energy_joules = energy_joules
|
||||
sample.dram_energy_joules = 0.0
|
||||
|
||||
@contextmanager
|
||||
def _sample():
|
||||
yield sample
|
||||
|
||||
monitor.sample = _sample
|
||||
return monitor
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helper function tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestPercentile:
|
||||
"""_percentile() computes interpolated percentiles."""
|
||||
|
||||
def test_simple_median(self):
|
||||
assert _percentile([1, 2, 3, 4, 5], 0.50) == pytest.approx(3.0)
|
||||
|
||||
def test_p90(self):
|
||||
data = list(range(1, 101)) # 1..100
|
||||
assert _percentile(data, 0.90) == pytest.approx(90.1)
|
||||
|
||||
def test_single_value(self):
|
||||
assert _percentile([42.0], 0.99) == pytest.approx(42.0)
|
||||
|
||||
def test_two_values(self):
|
||||
assert _percentile([10, 20], 0.50) == pytest.approx(15.0)
|
||||
|
||||
def test_unsorted_input(self):
|
||||
"""Input doesn't need to be sorted."""
|
||||
assert _percentile([5, 3, 1, 4, 2], 0.50) == pytest.approx(3.0)
|
||||
|
||||
|
||||
class TestComputeItlStats:
|
||||
"""_compute_itl_stats() computes ITL summary statistics."""
|
||||
|
||||
def test_empty_list(self):
|
||||
stats = _compute_itl_stats([])
|
||||
assert stats["mean"] == 0.0
|
||||
assert stats["median"] == 0.0
|
||||
assert stats["p90"] == 0.0
|
||||
assert stats["p95"] == 0.0
|
||||
assert stats["p99"] == 0.0
|
||||
assert stats["std"] == 0.0
|
||||
|
||||
def test_single_value(self):
|
||||
stats = _compute_itl_stats([10.0])
|
||||
assert stats["mean"] == pytest.approx(10.0)
|
||||
assert stats["median"] == pytest.approx(10.0)
|
||||
assert stats["std"] == 0.0 # single value
|
||||
|
||||
def test_known_sequence(self):
|
||||
values = [10.0, 20.0, 30.0, 40.0, 50.0]
|
||||
stats = _compute_itl_stats(values)
|
||||
assert stats["mean"] == pytest.approx(30.0)
|
||||
assert stats["median"] == pytest.approx(30.0)
|
||||
assert stats["p90"] == pytest.approx(46.0)
|
||||
assert stats["p95"] == pytest.approx(48.0)
|
||||
assert stats["p99"] == pytest.approx(49.6)
|
||||
assert stats["std"] > 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Streaming tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestStreamTelemetry:
|
||||
"""InstrumentedEngine.stream() records telemetry with ITL."""
|
||||
|
||||
def test_stream_creates_telemetry_record(self):
|
||||
bus = EventBus()
|
||||
engine = _mock_engine_with_stream(["a", "b", "c"])
|
||||
ie = InstrumentedEngine(engine, bus)
|
||||
|
||||
records = []
|
||||
bus.subscribe(
|
||||
EventType.TELEMETRY_RECORD,
|
||||
lambda e: records.append(e.data["record"]),
|
||||
)
|
||||
|
||||
async def run():
|
||||
tokens = []
|
||||
async for tok in ie.stream(
|
||||
[Message(role=Role.USER, content="hi")], model="m"
|
||||
):
|
||||
tokens.append(tok)
|
||||
return tokens
|
||||
|
||||
tokens = asyncio.run(run())
|
||||
assert tokens == ["a", "b", "c"]
|
||||
assert len(records) == 1
|
||||
rec = records[0]
|
||||
assert rec.is_streaming is True
|
||||
assert rec.completion_tokens == 3
|
||||
|
||||
def test_stream_computes_itl(self):
|
||||
bus = EventBus()
|
||||
engine = _mock_engine_with_stream(["a", "b", "c", "d", "e"])
|
||||
ie = InstrumentedEngine(engine, bus)
|
||||
|
||||
records = []
|
||||
bus.subscribe(
|
||||
EventType.TELEMETRY_RECORD,
|
||||
lambda e: records.append(e.data["record"]),
|
||||
)
|
||||
|
||||
async def run():
|
||||
async for _ in ie.stream(
|
||||
[Message(role=Role.USER, content="hi")], model="m"
|
||||
):
|
||||
pass
|
||||
|
||||
asyncio.run(run())
|
||||
rec = records[0]
|
||||
# 5 tokens → 4 ITL deltas
|
||||
assert rec.mean_itl_ms >= 0
|
||||
assert rec.median_itl_ms >= 0
|
||||
assert rec.p90_itl_ms >= 0
|
||||
assert rec.p95_itl_ms >= 0
|
||||
assert rec.p99_itl_ms >= 0
|
||||
|
||||
def test_stream_with_energy_monitor(self):
|
||||
bus = EventBus()
|
||||
engine = _mock_engine_with_stream(["x", "y"])
|
||||
monitor = _mock_energy_monitor(energy_joules=5.0, power_watts=100.0)
|
||||
ie = InstrumentedEngine(engine, bus, energy_monitor=monitor)
|
||||
|
||||
records = []
|
||||
bus.subscribe(
|
||||
EventType.TELEMETRY_RECORD,
|
||||
lambda e: records.append(e.data["record"]),
|
||||
)
|
||||
|
||||
async def run():
|
||||
async for _ in ie.stream(
|
||||
[Message(role=Role.USER, content="hi")], model="m"
|
||||
):
|
||||
pass
|
||||
|
||||
asyncio.run(run())
|
||||
rec = records[0]
|
||||
assert rec.energy_joules == 5.0
|
||||
assert rec.energy_per_output_token_joules == pytest.approx(5.0 / 2)
|
||||
|
||||
def test_stream_empty_tokens(self):
|
||||
bus = EventBus()
|
||||
engine = _mock_engine_with_stream([])
|
||||
ie = InstrumentedEngine(engine, bus)
|
||||
|
||||
records = []
|
||||
bus.subscribe(
|
||||
EventType.TELEMETRY_RECORD,
|
||||
lambda e: records.append(e.data["record"]),
|
||||
)
|
||||
|
||||
async def run():
|
||||
async for _ in ie.stream(
|
||||
[Message(role=Role.USER, content="hi")], model="m"
|
||||
):
|
||||
pass
|
||||
|
||||
asyncio.run(run())
|
||||
rec = records[0]
|
||||
assert rec.completion_tokens == 0
|
||||
assert rec.mean_itl_ms == 0.0
|
||||
assert rec.ttft == 0.0
|
||||
|
||||
def test_stream_single_token(self):
|
||||
bus = EventBus()
|
||||
engine = _mock_engine_with_stream(["only"])
|
||||
ie = InstrumentedEngine(engine, bus)
|
||||
|
||||
records = []
|
||||
bus.subscribe(
|
||||
EventType.TELEMETRY_RECORD,
|
||||
lambda e: records.append(e.data["record"]),
|
||||
)
|
||||
|
||||
async def run():
|
||||
async for _ in ie.stream(
|
||||
[Message(role=Role.USER, content="hi")], model="m"
|
||||
):
|
||||
pass
|
||||
|
||||
asyncio.run(run())
|
||||
rec = records[0]
|
||||
assert rec.completion_tokens == 1
|
||||
# No ITL deltas with single token
|
||||
assert rec.mean_itl_ms == 0.0
|
||||
assert rec.std_itl_ms == 0.0
|
||||
|
||||
|
||||
class TestGenerateMeanItlApproximation:
|
||||
"""generate() computes mean_itl_ms from decode_latency/completion_tokens."""
|
||||
|
||||
def test_mean_itl_computed(self):
|
||||
bus = EventBus()
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
|
||||
def _slow_generate(*args, **kwargs):
|
||||
time.sleep(0.05) # ensure latency > ttft
|
||||
return {
|
||||
"content": "hi",
|
||||
"usage": {
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 20,
|
||||
"total_tokens": 30,
|
||||
},
|
||||
"model": "m",
|
||||
"ttft": 0.01,
|
||||
}
|
||||
|
||||
engine.generate.side_effect = _slow_generate
|
||||
ie = InstrumentedEngine(engine, bus)
|
||||
|
||||
records = []
|
||||
bus.subscribe(
|
||||
EventType.TELEMETRY_RECORD,
|
||||
lambda e: records.append(e.data["record"]),
|
||||
)
|
||||
|
||||
ie.generate([Message(role=Role.USER, content="hi")], model="m")
|
||||
rec = records[0]
|
||||
# decode_latency > 0 because latency > ttft
|
||||
assert rec.decode_latency_seconds > 0
|
||||
# mean_itl_ms = (decode_latency / completion_tokens) * 1000
|
||||
expected = (rec.decode_latency_seconds / 20) * 1000
|
||||
assert rec.mean_itl_ms == pytest.approx(expected)
|
||||
assert rec.is_streaming is False
|
||||
|
||||
def test_no_ttft_no_itl(self):
|
||||
bus = EventBus()
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
engine.generate.return_value = {
|
||||
"content": "hi",
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
|
||||
"model": "m",
|
||||
"ttft": 0.0,
|
||||
}
|
||||
ie = InstrumentedEngine(engine, bus)
|
||||
|
||||
records = []
|
||||
bus.subscribe(
|
||||
EventType.TELEMETRY_RECORD,
|
||||
lambda e: records.append(e.data["record"]),
|
||||
)
|
||||
|
||||
ie.generate([Message(role=Role.USER, content="hi")], model="m")
|
||||
rec = records[0]
|
||||
assert rec.mean_itl_ms == 0.0 # no decode_latency → no ITL
|
||||
|
||||
|
||||
class TestItlStorage:
|
||||
"""ITL fields are stored and queryable."""
|
||||
|
||||
def test_store_and_query_itl(self, tmp_path):
|
||||
store = TelemetryStore(tmp_path / "test.db")
|
||||
store.record(TelemetryRecord(
|
||||
timestamp=time.time(),
|
||||
model_id="m1",
|
||||
engine="mock",
|
||||
mean_itl_ms=15.0,
|
||||
median_itl_ms=14.0,
|
||||
p90_itl_ms=20.0,
|
||||
p95_itl_ms=25.0,
|
||||
p99_itl_ms=30.0,
|
||||
std_itl_ms=5.0,
|
||||
is_streaming=True,
|
||||
))
|
||||
|
||||
agg = TelemetryAggregator(tmp_path / "test.db")
|
||||
stats = agg.per_model_stats()
|
||||
assert len(stats) == 1
|
||||
assert stats[0].avg_mean_itl_ms == pytest.approx(15.0)
|
||||
assert stats[0].avg_median_itl_ms == pytest.approx(14.0)
|
||||
assert stats[0].avg_p95_itl_ms == pytest.approx(25.0)
|
||||
agg.close()
|
||||
store.close()
|
||||
@@ -0,0 +1,231 @@
|
||||
"""Tests for Tier 2.1 — phase energy split: decode_latency, prefill/decode energy."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.types import Message, Role, TelemetryRecord
|
||||
from openjarvis.telemetry.aggregator import TelemetryAggregator
|
||||
from openjarvis.telemetry.instrumented_engine import InstrumentedEngine
|
||||
from openjarvis.telemetry.store import TelemetryStore
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _mock_engine(ttft=0.1):
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
engine.generate.return_value = {
|
||||
"content": "hello world",
|
||||
"usage": {
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 50,
|
||||
"total_tokens": 60,
|
||||
},
|
||||
"model": "test-model",
|
||||
"ttft": ttft,
|
||||
}
|
||||
return engine
|
||||
|
||||
|
||||
def _mock_energy_monitor(energy_joules=10.0, power_watts=200.0):
|
||||
monitor = MagicMock()
|
||||
sample = MagicMock()
|
||||
sample.energy_joules = energy_joules
|
||||
sample.mean_power_watts = power_watts
|
||||
sample.peak_power_watts = power_watts
|
||||
sample.mean_utilization_pct = 80.0
|
||||
sample.peak_utilization_pct = 95.0
|
||||
sample.mean_memory_used_gb = 16.0
|
||||
sample.peak_memory_used_gb = 20.0
|
||||
sample.mean_temperature_c = 65.0
|
||||
sample.peak_temperature_c = 72.0
|
||||
sample.duration_seconds = 0.5
|
||||
sample.num_snapshots = 10
|
||||
sample.energy_method = "hw_counter"
|
||||
sample.vendor = "nvidia"
|
||||
sample.cpu_energy_joules = 0.0
|
||||
sample.gpu_energy_joules = energy_joules
|
||||
sample.dram_energy_joules = 0.0
|
||||
|
||||
@contextmanager
|
||||
def _sample():
|
||||
yield sample
|
||||
|
||||
monitor.sample = _sample
|
||||
return monitor
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDecodeLatency:
|
||||
"""decode_latency = latency - ttft when ttft > 0."""
|
||||
|
||||
def test_decode_latency_computed(self):
|
||||
bus = EventBus()
|
||||
engine = _mock_engine(ttft=0.1)
|
||||
ie = InstrumentedEngine(engine, bus)
|
||||
|
||||
records = []
|
||||
bus.subscribe(
|
||||
EventType.TELEMETRY_RECORD,
|
||||
lambda e: records.append(e.data["record"]),
|
||||
)
|
||||
|
||||
ie.generate([Message(role=Role.USER, content="hi")], model="m")
|
||||
rec = records[0]
|
||||
# decode_latency = latency - ttft
|
||||
assert rec.decode_latency_seconds == pytest.approx(
|
||||
rec.latency_seconds - 0.1
|
||||
)
|
||||
|
||||
def test_decode_latency_zero_when_no_ttft(self):
|
||||
bus = EventBus()
|
||||
engine = _mock_engine(ttft=0.0)
|
||||
ie = InstrumentedEngine(engine, bus)
|
||||
|
||||
records = []
|
||||
bus.subscribe(
|
||||
EventType.TELEMETRY_RECORD,
|
||||
lambda e: records.append(e.data["record"]),
|
||||
)
|
||||
|
||||
ie.generate([Message(role=Role.USER, content="hi")], model="m")
|
||||
rec = records[0]
|
||||
assert rec.decode_latency_seconds == 0.0
|
||||
|
||||
|
||||
class TestPhaseEnergySplit:
|
||||
"""prefill_energy + decode_energy ≈ total energy."""
|
||||
|
||||
def test_energy_split_proportional(self):
|
||||
bus = EventBus()
|
||||
engine = _mock_engine(ttft=0.1)
|
||||
monitor = _mock_energy_monitor(energy_joules=10.0)
|
||||
ie = InstrumentedEngine(engine, bus, energy_monitor=monitor)
|
||||
|
||||
records = []
|
||||
bus.subscribe(
|
||||
EventType.TELEMETRY_RECORD,
|
||||
lambda e: records.append(e.data["record"]),
|
||||
)
|
||||
|
||||
ie.generate([Message(role=Role.USER, content="hi")], model="m")
|
||||
rec = records[0]
|
||||
|
||||
# Sum should equal total energy
|
||||
assert rec.prefill_energy_joules + rec.decode_energy_joules == pytest.approx(
|
||||
rec.energy_joules
|
||||
)
|
||||
|
||||
# Prefill fraction should be proportional to ttft/latency
|
||||
expected_prefill_frac = 0.1 / rec.latency_seconds
|
||||
actual_prefill_frac = rec.prefill_energy_joules / rec.energy_joules
|
||||
assert actual_prefill_frac == pytest.approx(expected_prefill_frac)
|
||||
|
||||
def test_no_energy_no_split(self):
|
||||
bus = EventBus()
|
||||
engine = _mock_engine(ttft=0.1)
|
||||
ie = InstrumentedEngine(engine, bus) # no energy monitor
|
||||
|
||||
records = []
|
||||
bus.subscribe(
|
||||
EventType.TELEMETRY_RECORD,
|
||||
lambda e: records.append(e.data["record"]),
|
||||
)
|
||||
|
||||
ie.generate([Message(role=Role.USER, content="hi")], model="m")
|
||||
rec = records[0]
|
||||
assert rec.prefill_energy_joules == 0.0
|
||||
assert rec.decode_energy_joules == 0.0
|
||||
|
||||
def test_no_ttft_no_split(self):
|
||||
bus = EventBus()
|
||||
engine = _mock_engine(ttft=0.0)
|
||||
monitor = _mock_energy_monitor(energy_joules=10.0)
|
||||
ie = InstrumentedEngine(engine, bus, energy_monitor=monitor)
|
||||
|
||||
records = []
|
||||
bus.subscribe(
|
||||
EventType.TELEMETRY_RECORD,
|
||||
lambda e: records.append(e.data["record"]),
|
||||
)
|
||||
|
||||
ie.generate([Message(role=Role.USER, content="hi")], model="m")
|
||||
rec = records[0]
|
||||
# No ttft → no prefill_latency → no split
|
||||
assert rec.prefill_energy_joules == 0.0
|
||||
assert rec.decode_energy_joules == 0.0
|
||||
|
||||
def test_latency_equals_ttft_decode_energy_zero(self):
|
||||
"""When latency == ttft, all energy is prefill."""
|
||||
bus = EventBus()
|
||||
# We'll use a ttft that's close to the measured latency
|
||||
engine = _mock_engine(ttft=0.001)
|
||||
monitor = _mock_energy_monitor(energy_joules=5.0)
|
||||
ie = InstrumentedEngine(engine, bus, energy_monitor=monitor)
|
||||
|
||||
records = []
|
||||
bus.subscribe(
|
||||
EventType.TELEMETRY_RECORD,
|
||||
lambda e: records.append(e.data["record"]),
|
||||
)
|
||||
|
||||
ie.generate([Message(role=Role.USER, content="hi")], model="m")
|
||||
rec = records[0]
|
||||
# prefill + decode should still sum to total
|
||||
assert rec.prefill_energy_joules + rec.decode_energy_joules == pytest.approx(
|
||||
rec.energy_joules
|
||||
)
|
||||
|
||||
|
||||
class TestPhaseEnergyInTelemetryDict:
|
||||
"""Phase energy fields appear in result['_telemetry']."""
|
||||
|
||||
def test_telemetry_dict_contains_phase_energy(self):
|
||||
bus = EventBus()
|
||||
engine = _mock_engine(ttft=0.1)
|
||||
monitor = _mock_energy_monitor(energy_joules=10.0)
|
||||
ie = InstrumentedEngine(engine, bus, energy_monitor=monitor)
|
||||
|
||||
result = ie.generate([Message(role=Role.USER, content="hi")], model="m")
|
||||
t = result["_telemetry"]
|
||||
assert "prefill_energy_joules" in t
|
||||
assert "decode_energy_joules" in t
|
||||
assert "decode_latency_seconds" in t
|
||||
assert t["prefill_energy_joules"] + t["decode_energy_joules"] == pytest.approx(
|
||||
t["energy_joules"] if t["energy_joules"] > 0 else 0.0
|
||||
)
|
||||
|
||||
|
||||
class TestPhaseEnergyStorage:
|
||||
"""Phase energy fields are stored and queryable."""
|
||||
|
||||
def test_store_and_aggregate(self, tmp_path):
|
||||
store = TelemetryStore(tmp_path / "test.db")
|
||||
store.record(TelemetryRecord(
|
||||
timestamp=time.time(),
|
||||
model_id="m1",
|
||||
engine="mock",
|
||||
energy_joules=10.0,
|
||||
prefill_energy_joules=3.0,
|
||||
decode_energy_joules=7.0,
|
||||
))
|
||||
|
||||
agg = TelemetryAggregator(tmp_path / "test.db")
|
||||
stats = agg.per_model_stats()
|
||||
assert len(stats) == 1
|
||||
assert stats[0].total_prefill_energy_joules == pytest.approx(3.0)
|
||||
assert stats[0].total_decode_energy_joules == pytest.approx(7.0)
|
||||
agg.close()
|
||||
store.close()
|
||||
@@ -0,0 +1,155 @@
|
||||
"""Tests for SteadyStateConfig, SteadyStateDetector, and SteadyStateResult."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from openjarvis.telemetry.steady_state import (
|
||||
SteadyStateConfig,
|
||||
SteadyStateDetector,
|
||||
SteadyStateResult,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: SteadyStateConfig
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSteadyStateConfig:
|
||||
def test_defaults(self):
|
||||
cfg = SteadyStateConfig()
|
||||
assert cfg.warmup_samples == 5
|
||||
assert cfg.window_size == 5
|
||||
assert cfg.cv_threshold == 0.05
|
||||
assert cfg.min_steady_samples == 3
|
||||
assert cfg.metric == "throughput"
|
||||
|
||||
def test_custom_values(self):
|
||||
cfg = SteadyStateConfig(warmup_samples=10, cv_threshold=0.1)
|
||||
assert cfg.warmup_samples == 10
|
||||
assert cfg.cv_threshold == 0.1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: SteadyStateResult
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSteadyStateResult:
|
||||
def test_default_fields(self):
|
||||
r = SteadyStateResult()
|
||||
assert r.total_samples == 0
|
||||
assert r.warmup_samples == 0
|
||||
assert r.steady_state_samples == 0
|
||||
assert r.steady_state_reached is False
|
||||
assert r.warmup_throughputs == []
|
||||
assert r.warmup_energies == []
|
||||
assert r.steady_throughputs == []
|
||||
assert r.steady_energies == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: SteadyStateDetector
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSteadyStateDetector:
|
||||
def test_constant_throughput_reaches_steady_state(self):
|
||||
"""Constant values should produce CV=0, reaching steady state quickly."""
|
||||
cfg = SteadyStateConfig(warmup_samples=3, window_size=3, min_steady_samples=2)
|
||||
detector = SteadyStateDetector(cfg)
|
||||
|
||||
reached = False
|
||||
for _ in range(20):
|
||||
reached = detector.record(throughput=100.0, energy=10.0)
|
||||
if reached:
|
||||
break
|
||||
assert reached is True
|
||||
assert detector.result.steady_state_reached is True
|
||||
|
||||
def test_erratic_values_no_steady_state(self):
|
||||
"""Highly variable values should not reach steady state."""
|
||||
cfg = SteadyStateConfig(
|
||||
warmup_samples=2, window_size=3, cv_threshold=0.01, min_steady_samples=3,
|
||||
)
|
||||
detector = SteadyStateDetector(cfg)
|
||||
|
||||
# Alternate between wildly different values
|
||||
values = [10.0, 1000.0, 10.0, 1000.0, 10.0, 1000.0, 10.0, 1000.0, 10.0, 1000.0]
|
||||
for v in values:
|
||||
detector.record(throughput=v)
|
||||
|
||||
assert detector.result.steady_state_reached is False
|
||||
|
||||
def test_warmup_boundary(self):
|
||||
"""First N samples are always warmup regardless of stability."""
|
||||
cfg = SteadyStateConfig(warmup_samples=5, window_size=3, min_steady_samples=1)
|
||||
detector = SteadyStateDetector(cfg)
|
||||
|
||||
# Feed 5 constant values (all warmup)
|
||||
for _ in range(5):
|
||||
result = detector.record(throughput=100.0)
|
||||
assert result is False # still in warmup
|
||||
|
||||
r = detector.result
|
||||
assert r.warmup_samples == 5
|
||||
assert r.steady_state_samples == 0
|
||||
|
||||
def test_cv_calculation_correctness(self):
|
||||
"""Verify CV-based detection with known values."""
|
||||
cfg = SteadyStateConfig(
|
||||
warmup_samples=2, window_size=3, cv_threshold=0.05, min_steady_samples=1,
|
||||
)
|
||||
detector = SteadyStateDetector(cfg)
|
||||
|
||||
# 2 warmup
|
||||
detector.record(throughput=50.0)
|
||||
detector.record(throughput=60.0)
|
||||
|
||||
# Post-warmup: values with low CV (100, 101, 100 -> CV ~ 0.006)
|
||||
detector.record(throughput=100.0)
|
||||
detector.record(throughput=101.0)
|
||||
result = detector.record(throughput=100.0)
|
||||
|
||||
assert result is True
|
||||
assert detector.result.steady_state_reached is True
|
||||
|
||||
def test_reset_clears_state(self):
|
||||
"""After reset, detector starts fresh."""
|
||||
cfg = SteadyStateConfig(warmup_samples=2, window_size=2, min_steady_samples=1)
|
||||
detector = SteadyStateDetector(cfg)
|
||||
|
||||
# Get to steady state
|
||||
for _ in range(10):
|
||||
detector.record(throughput=100.0)
|
||||
assert detector.result.steady_state_reached is True
|
||||
|
||||
# Reset
|
||||
detector.reset()
|
||||
r = detector.result
|
||||
assert r.total_samples == 0
|
||||
assert r.steady_state_reached is False
|
||||
assert r.warmup_throughputs == []
|
||||
assert r.steady_throughputs == []
|
||||
|
||||
def test_result_partitions_warmup_and_steady(self):
|
||||
"""Result should correctly partition samples into warmup and steady."""
|
||||
cfg = SteadyStateConfig(warmup_samples=3, window_size=2, min_steady_samples=1)
|
||||
detector = SteadyStateDetector(cfg)
|
||||
|
||||
# 3 warmup + 4 steady
|
||||
for i in range(7):
|
||||
detector.record(throughput=float(100 + i), energy=float(10 + i))
|
||||
|
||||
r = detector.result
|
||||
assert r.total_samples == 7
|
||||
assert r.warmup_samples == 3
|
||||
assert r.steady_state_samples == 4
|
||||
assert len(r.warmup_throughputs) == 3
|
||||
assert len(r.warmup_energies) == 3
|
||||
assert len(r.steady_throughputs) == 4
|
||||
assert len(r.steady_energies) == 4
|
||||
|
||||
def test_default_config_when_none(self):
|
||||
"""Passing None config should use defaults."""
|
||||
detector = SteadyStateDetector(None)
|
||||
assert detector._config.warmup_samples == 5
|
||||
assert detector._config.window_size == 5
|
||||
@@ -171,6 +171,15 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/fb/76/641ae371508676492379f16e2fa48f4e2c11741bd63c48be4b12a6b09cba/aiosignal-1.4.0-py3-none-any.whl", hash = "sha256:053243f8b92b990551949e63930a839ff0cf0b0ebbe0597b0f3fb19e1a0fe82e", size = 7490, upload-time = "2025-07-03T22:54:42.156Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "amdsmi"
|
||||
version = "7.0.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/45/c1/330da195623ec7d9f699be2dbec98df364b1def9b48aa169f1abe369804f/amdsmi-7.0.2.tar.gz", hash = "sha256:3e622e48c630b889045a6f57387455cdf082066348718172dd8af6d275baf8f2", size = 61577, upload-time = "2025-10-11T05:17:44.898Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/2b/cf/bacc741d9926d662e76fd694f5bc63dd4c2471e32cedfb1c3cca7e47aa3c/amdsmi-7.0.2-py3-none-any.whl", hash = "sha256:db5aa757f8ed82dfd799c4d39e2828542678dc4e485b0ab7fabe5f398fec5652", size = 64366, upload-time = "2025-10-11T05:17:44.047Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "annotated-doc"
|
||||
version = "0.0.4"
|
||||
@@ -2185,6 +2194,60 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/32/28/79f0f8de97cce916d5ae88a7bee1ad724855e83e6019c0b4d5b3fabc80f3/mkdocstrings_python-2.0.3-py3-none-any.whl", hash = "sha256:0b83513478bdfd803ff05aa43e9b1fca9dd22bcd9471f09ca6257f009bc5ee12", size = 104779, upload-time = "2026-02-20T10:38:34.517Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "mlx"
|
||||
version = "0.30.6"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "mlx-metal", marker = "sys_platform == 'darwin'" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a0/2e/016527cf1012a68bb25f1ba3a73914f87807a7fee58d7a54fa69adcd2f55/mlx-0.30.6-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:6c4df52aebfac40563259c04fca4a0c4d05b2061e09cdaad24e4233baa560b4f", size = 573214, upload-time = "2026-02-06T03:45:00.344Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/a4/8f/600c6bed6eb6574e4a9d15e7a20a2ec903c2c5b54e2fd782c592a00ff933/mlx-0.30.6-cp310-cp310-macosx_15_0_arm64.whl", hash = "sha256:0df8715b5cb84b6b6314aa868302873a0a94e63e6d195bc9858b8c58c79aa5a4", size = 573213, upload-time = "2026-02-06T03:45:02.208Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/11/f7/d15af26c639c3d6000b6478fc0d54a7a528d71e79255190a0abc42f31608/mlx-0.30.6-cp310-cp310-macosx_26_0_arm64.whl", hash = "sha256:7b4742ec2b748d2406c884e364fcd6f89d7f2b3f834f7b65c4c07acfa139cae8", size = 573254, upload-time = "2026-02-06T03:45:03.575Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/93/81/21d745beeda53ee29e9c027d806f1e1cac983e8ddb3d6b18d44a1b30a11b/mlx-0.30.6-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:e721d29c4250ada3cba7a5ad43d358b42401600e792c378ed6b52c9d692aaba8", size = 573359, upload-time = "2026-02-06T03:45:08.41Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/05/08/826286458df5ea91efc380d71fd8058ee7338207c6b547204f2758e168d8/mlx-0.30.6-cp311-cp311-macosx_15_0_arm64.whl", hash = "sha256:23f55c1c160a38ab350f4f7ce3ab10c490df39800ad35c4821c3ef5fa89ec24e", size = 573359, upload-time = "2026-02-06T03:45:09.688Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/56/aa/3fc9ac795934182e680a0cbeb99202838e4548139cfd580015dcfbfb7ee8/mlx-0.30.6-cp311-cp311-macosx_26_0_arm64.whl", hash = "sha256:37c37571f8c1567c2b7e4871237b92a2b321fb8157d6426373be946c03e49ebd", size = 573406, upload-time = "2026-02-06T03:45:11.383Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/85/fe/85acff870a9949494fd505b22c34d63eb127442f5f8751a159d3a78f7ef6/mlx-0.30.6-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:47d20016cb5733d06c1d017412a31983dbe3237cf70942760430188922ffc1ba", size = 573484, upload-time = "2026-02-06T03:45:15.88Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/e1/14/5546082ee37118b33afb6300d8e07d03efea2dbba838d514d9465f87489b/mlx-0.30.6-cp312-cp312-macosx_15_0_arm64.whl", hash = "sha256:6b8c133df2d6a2ed173d2b7bb50d7032a13be84e1792b7d79171ad8f50a8c0ea", size = 573486, upload-time = "2026-02-06T03:45:17.506Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ef/b5/ae04666a7b8bda74e2c6903756710103e283ea6fa4edd2c92449ad4547d6/mlx-0.30.6-cp312-cp312-macosx_26_0_arm64.whl", hash = "sha256:31eabb5d1da4ac7b16f2042fdb046b993cdf0f32bc3312e0af469232bb67720b", size = 573509, upload-time = "2026-02-06T03:45:18.68Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ae/5b/e460e144a34d5529e010056cccf50b538d56ed001473bc6b246018fd58cb/mlx-0.30.6-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:ed86f8bffc174c2f259ca589ea25464c96cf69d1bb457074a2bf2ef53737e54f", size = 573515, upload-time = "2026-02-06T03:45:23.405Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/60/25/69833fefb9a3fef30b56792b1bcd022496c4fea83e45411d289b77ef7546/mlx-0.30.6-cp313-cp313-macosx_15_0_arm64.whl", hash = "sha256:c52294958269e20f300639a17c1900ca8fc737d859ddda737f9811e94bd040e5", size = 573516, upload-time = "2026-02-06T03:45:24.618Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/9c/6a/7e7fbeebc5cb51b6a5eba96b263a6298707bcbdc059f4b0b73e088bc3dea/mlx-0.30.6-cp313-cp313-macosx_26_0_arm64.whl", hash = "sha256:b5b6636f7c49a4d86d8ec82643b972f45a144a7a9f3a967b27b2e6e22cf71e6a", size = 573592, upload-time = "2026-02-06T03:45:25.928Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/60/23/361dc7a5797634e4d7e9bdd6564c6b28f9b1246672632def2f91bf066b18/mlx-0.30.6-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:78804a89dcff4a838f7c2da72392fe87a523e95122a3c840e53df019122aad45", size = 575028, upload-time = "2026-02-06T03:45:31.549Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/a8/69/1854484d414171586814dfbe8def95f75c4ea2c7341ba13ba8ee675f7c62/mlx-0.30.6-cp314-cp314-macosx_15_0_arm64.whl", hash = "sha256:ec13584ab069665cc7ad34a05494d9291cd623aef6ae96be48875fc87cfc25d6", size = 575026, upload-time = "2026-02-06T03:45:33.072Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/6b/b8/3adbc441924209a7e4c568308b2a0b54bd09aee6a68db5bae85304791e54/mlx-0.30.6-cp314-cp314-macosx_26_0_arm64.whl", hash = "sha256:b2c5e8a090a753ef99a1380a4d059c983083f36198864f6df9faaf1223d083df", size = 575041, upload-time = "2026-02-06T03:45:34.814Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "mlx-lm"
|
||||
version = "0.30.7"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "jinja2", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
|
||||
{ name = "mlx", marker = "sys_platform == 'darwin'" },
|
||||
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
|
||||
{ name = "numpy", version = "2.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and sys_platform != 'emscripten' and sys_platform != 'win32'" },
|
||||
{ name = "protobuf", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
|
||||
{ name = "pyyaml", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
|
||||
{ name = "sentencepiece", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
|
||||
{ name = "transformers", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/66/0d/56542e2ae13ec6f542d3977d7cff89a205d4f6c5122e0ce23f33265f61c9/mlx_lm-0.30.7.tar.gz", hash = "sha256:e5f31ac58d9f2381f28e1ba639ff903e64f7cff1bdc245c0bc97f72264be329c", size = 275764, upload-time = "2026-02-12T18:41:11.86Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/1e/17/a41c798a3d9cbdc47f39c6db5bba4c2cd199203ead26bf911cb03b644070/mlx_lm-0.30.7-py3-none-any.whl", hash = "sha256:17442a4bf01c4c2d3bca1e647712fe44f19890c3f1eadc8589d389e57b44b9bf", size = 386591, upload-time = "2026-02-12T18:41:10.236Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "mlx-metal"
|
||||
version = "0.30.6"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/f3/85/44406b521f920248fad621334d4dc15e77660a494edf890e7cbee33bf38d/mlx_metal-0.30.6-py3-none-macosx_14_0_arm64.whl", hash = "sha256:ea6d0c973def9a5b4f652cc77036237db3f88c9d0af63701d76b5fddde99b820", size = 38437818, upload-time = "2026-02-06T03:44:56.19Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d0/cb/10a516995f7d0c154b0d7e633c54b51e96977a86a355105b6474cfcbe0d0/mlx_metal-0.30.6-py3-none-macosx_15_0_arm64.whl", hash = "sha256:0f8cb94634d07e06a372d6ad9a090f38a18bab1ff19a140aede60eacf707bb94", size = 38433701, upload-time = "2026-02-06T03:44:59.678Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/4c/7d/70cb272f7373c334709f210ed8420511fc9d64d05a7a646c0b3b94c29c04/mlx_metal-0.30.6-py3-none-macosx_26_0_arm64.whl", hash = "sha256:d761ae26304f2c4b454eeea7f612a56919d9e5e57dbb1dc0788f8e34aa6f41c2", size = 47718448, upload-time = "2026-02-06T03:45:03.133Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "more-itertools"
|
||||
version = "10.8.0"
|
||||
@@ -2813,6 +2876,17 @@ docs = [
|
||||
{ name = "mkdocs-material" },
|
||||
{ name = "mkdocstrings", extra = ["python"] },
|
||||
]
|
||||
energy-all = [
|
||||
{ name = "amdsmi" },
|
||||
{ name = "pynvml" },
|
||||
{ name = "zeus-ml" },
|
||||
]
|
||||
energy-amd = [
|
||||
{ name = "amdsmi" },
|
||||
]
|
||||
energy-apple = [
|
||||
{ name = "zeus-ml" },
|
||||
]
|
||||
gpu-metrics = [
|
||||
{ name = "pynvml" },
|
||||
]
|
||||
@@ -2826,6 +2900,9 @@ inference-google = [
|
||||
inference-litellm = [
|
||||
{ name = "litellm" },
|
||||
]
|
||||
inference-mlx = [
|
||||
{ name = "mlx-lm", marker = "sys_platform == 'darwin'" },
|
||||
]
|
||||
memory-bm25 = [
|
||||
{ name = "rank-bm25" },
|
||||
]
|
||||
@@ -2863,6 +2940,8 @@ tools-search = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "amdsmi", marker = "extra == 'energy-all'", specifier = ">=6.1" },
|
||||
{ name = "amdsmi", marker = "extra == 'energy-amd'", specifier = ">=6.1" },
|
||||
{ name = "anthropic", marker = "extra == 'inference-cloud'", specifier = ">=0.30" },
|
||||
{ name = "click", specifier = ">=8" },
|
||||
{ name = "colbert-ai", marker = "extra == 'memory-colbert'", specifier = ">=0.2" },
|
||||
@@ -2876,12 +2955,14 @@ requires-dist = [
|
||||
{ name = "mkdocs", marker = "extra == 'docs'", specifier = ">=1.6" },
|
||||
{ name = "mkdocs-material", marker = "extra == 'docs'", specifier = ">=9.5" },
|
||||
{ name = "mkdocstrings", extras = ["python"], marker = "extra == 'docs'", specifier = ">=0.25" },
|
||||
{ name = "mlx-lm", marker = "sys_platform == 'darwin' and extra == 'inference-mlx'", specifier = ">=0.19" },
|
||||
{ name = "numpy", marker = "extra == 'memory-faiss'", specifier = ">=1.24" },
|
||||
{ name = "openai", marker = "extra == 'inference-cloud'", specifier = ">=1.30" },
|
||||
{ name = "openhands-sdk", marker = "python_full_version >= '3.12' and extra == 'openhands'", specifier = ">=1.0" },
|
||||
{ name = "pdfplumber", marker = "extra == 'memory-pdf'", specifier = ">=0.10" },
|
||||
{ name = "pydantic", marker = "extra == 'server'", specifier = ">=2.0" },
|
||||
{ name = "pynvml", specifier = ">=13.0.1" },
|
||||
{ name = "pynvml", marker = "extra == 'energy-all'", specifier = ">=12.0" },
|
||||
{ name = "pynvml", marker = "extra == 'gpu-metrics'", specifier = ">=12.0" },
|
||||
{ name = "pytest", marker = "extra == 'dev'", specifier = ">=8" },
|
||||
{ name = "pytest-asyncio", marker = "extra == 'dev'", specifier = ">=0.24" },
|
||||
@@ -2899,8 +2980,10 @@ requires-dist = [
|
||||
{ name = "torch", marker = "extra == 'orchestrator-training'", specifier = ">=2.0" },
|
||||
{ name = "transformers", marker = "extra == 'orchestrator-training'", specifier = ">=4.40" },
|
||||
{ name = "uvicorn", marker = "extra == 'server'", specifier = ">=0.30" },
|
||||
{ name = "zeus-ml", extras = ["apple"], marker = "extra == 'energy-all'" },
|
||||
{ name = "zeus-ml", extras = ["apple"], marker = "extra == 'energy-apple'" },
|
||||
]
|
||||
provides-extras = ["dev", "inference-ollama", "inference-vllm", "inference-llamacpp", "inference-cloud", "inference-google", "inference-litellm", "tools-search", "memory-faiss", "memory-colbert", "memory-pdf", "memory-bm25", "server", "agents", "openhands", "claude-code", "gpu-metrics", "learning", "orchestrator-training", "channel-telegram", "channel-discord", "channel-slack", "channel-webhook", "channel-email", "channel-whatsapp", "channel-signal", "channel-google-chat", "channel-irc", "channel-webchat", "channel-teams", "channel-matrix", "channel-mattermost", "channel-feishu", "channel-bluebubbles", "channel-whatsapp-baileys", "scheduler", "docs"]
|
||||
provides-extras = ["dev", "inference-ollama", "inference-vllm", "inference-llamacpp", "inference-mlx", "inference-cloud", "inference-google", "inference-litellm", "tools-search", "memory-faiss", "memory-colbert", "memory-pdf", "memory-bm25", "server", "agents", "openhands", "claude-code", "gpu-metrics", "energy-amd", "energy-apple", "energy-all", "learning", "orchestrator-training", "channel-telegram", "channel-discord", "channel-slack", "channel-webhook", "channel-email", "channel-whatsapp", "channel-signal", "channel-google-chat", "channel-irc", "channel-webchat", "channel-teams", "channel-matrix", "channel-mattermost", "channel-feishu", "channel-bluebubbles", "channel-whatsapp-baileys", "scheduler", "docs"]
|
||||
|
||||
[[package]]
|
||||
name = "opentelemetry-api"
|
||||
@@ -4819,6 +4902,35 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/cc/21/7e925890636791386e81b52878134f114d63072e79fffe14cdcc5e7a5e6a/sentence_transformers-5.2.2-py3-none-any.whl", hash = "sha256:280ac54bffb84c110726b4d8848ba7b7c60813b9034547f8aea6e9a345cd1c23", size = 494106, upload-time = "2026-01-27T11:11:00.983Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sentencepiece"
|
||||
version = "0.2.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/15/15/2e7a025fc62d764b151ae6d0f2a92f8081755ebe8d4a64099accc6f77ba6/sentencepiece-0.2.1.tar.gz", hash = "sha256:8138cec27c2f2282f4a34d9a016e3374cd40e5c6e9cb335063db66a0a3b71fad", size = 3228515, upload-time = "2025-08-12T07:00:51.718Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/af/31/5b7cccb307b485db1a2372d6d2980b0a65d067f8be5ca943a103b4acd5b3/sentencepiece-0.2.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:e10fa50bdbaa5e2445dbd387979980d391760faf0ec99a09bd7780ff37eaec44", size = 1942557, upload-time = "2025-08-12T06:59:12.379Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/1f/41/0ac923a8e685ad290c5afc8ae55c5844977b8d75076fcc04302b9a324274/sentencepiece-0.2.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:2f27ae6deea72efdb6f361750c92f6c21fd0ad087445082770cc34015213c526", size = 1325384, upload-time = "2025-08-12T06:59:14.334Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/fc/ef/3751555d67daf9003384978f169d31c775cb5c7baf28633caaf1eb2b2b4d/sentencepiece-0.2.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:60937c959e6f44159fdd9f56fbdd302501f96114a5ba436829496d5f32d8de3f", size = 1253317, upload-time = "2025-08-12T06:59:16.247Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d8/15/46afbab00733d81788b64be430ca1b93011bb9388527958e26cc31832de5/sentencepiece-0.2.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:6356d0986b8b8dc351b943150fcd81a1c6e6e4d439772e8584c64230e58ca987", size = 1942560, upload-time = "2025-08-12T06:59:25.82Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/fa/79/7c01b8ef98a0567e9d84a4e7a910f8e7074fcbf398a5cd76f93f4b9316f9/sentencepiece-0.2.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:8f8ba89a3acb3dc1ae90f65ec1894b0b9596fdb98ab003ff38e058f898b39bc7", size = 1325385, upload-time = "2025-08-12T06:59:27.722Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/bb/88/2b41e07bd24f33dcf2f18ec3b74247aa4af3526bad8907b8727ea3caba03/sentencepiece-0.2.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:02593eca45440ef39247cee8c47322a34bdcc1d8ae83ad28ba5a899a2cf8d79a", size = 1253319, upload-time = "2025-08-12T06:59:29.306Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/4a/be/32ce495aa1d0e0c323dcb1ba87096037358edee539cac5baf8755a6bd396/sentencepiece-0.2.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:57cae326c8727de58c85977b175af132a7138d84c764635d7e71bbee7e774133", size = 1943152, upload-time = "2025-08-12T06:59:40.048Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/88/7e/ff23008899a58678e98c6ff592bf4d368eee5a71af96d0df6b38a039dd4f/sentencepiece-0.2.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:56dd39a3c4d6493db3cdca7e8cc68c6b633f0d4195495cbadfcf5af8a22d05a6", size = 1325651, upload-time = "2025-08-12T06:59:41.536Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/19/84/42eb3ce4796777a1b5d3699dfd4dca85113e68b637f194a6c8d786f16a04/sentencepiece-0.2.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d9381351182ff9888cc80e41c632e7e274b106f450de33d67a9e8f6043da6f76", size = 1253645, upload-time = "2025-08-12T06:59:42.903Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ba/4a/85fbe1706d4d04a7e826b53f327c4b80f849cf1c7b7c5e31a20a97d8f28b/sentencepiece-0.2.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:dcd8161eee7b41aae57ded06272905dbd680a0a04b91edd0f64790c796b2f706", size = 1943150, upload-time = "2025-08-12T06:59:53.588Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/c2/83/4cfb393e287509fc2155480b9d184706ef8d9fa8cbf5505d02a5792bf220/sentencepiece-0.2.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c6c8f42949f419ff8c7e9960dbadcfbc982d7b5efc2f6748210d3dd53a7de062", size = 1325651, upload-time = "2025-08-12T06:59:55.073Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/8d/de/5a007fb53b1ab0aafc69d11a5a3dd72a289d5a3e78dcf2c3a3d9b14ffe93/sentencepiece-0.2.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:097f3394e99456e9e4efba1737c3749d7e23563dd1588ce71a3d007f25475fff", size = 1253641, upload-time = "2025-08-12T06:59:56.562Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/4a/b6/08fe2ce819e02ccb0296f4843e3f195764ce9829cbda61b7513f29b95718/sentencepiece-0.2.1-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:8dd4b477a7b069648d19363aad0cab9bad2f4e83b2d179be668efa672500dc94", size = 1946052, upload-time = "2025-08-12T07:00:08.136Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ab/d9/1ea0e740591ff4c6fc2b6eb1d7510d02f3fb885093f19b2f3abd1363b402/sentencepiece-0.2.1-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:0c0f672da370cc490e4c59d89e12289778310a0e71d176c541e4834759e1ae07", size = 1327408, upload-time = "2025-08-12T07:00:09.572Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/99/7e/1fb26e8a21613f6200e1ab88824d5d203714162cf2883248b517deb500b7/sentencepiece-0.2.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:ad8493bea8432dae8d6830365352350f3b4144415a1d09c4c8cb8d30cf3b6c3c", size = 1254857, upload-time = "2025-08-12T07:00:11.021Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/24/9c/89eb8b2052f720a612478baf11c8227dcf1dc28cd4ea4c0c19506b5af2a2/sentencepiece-0.2.1-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:5d0350b686c320068702116276cfb26c066dc7e65cfef173980b11bb4d606719", size = 1943147, upload-time = "2025-08-12T07:00:21.809Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/82/0b/a1432bc87f97c2ace36386ca23e8bd3b91fb40581b5e6148d24b24186419/sentencepiece-0.2.1-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:c7f54a31cde6fa5cb030370566f68152a742f433f8d2be458463d06c208aef33", size = 1325624, upload-time = "2025-08-12T07:00:23.289Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ea/99/bbe054ebb5a5039457c590e0a4156ed073fb0fe9ce4f7523404dd5b37463/sentencepiece-0.2.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c83b85ab2d6576607f31df77ff86f28182be4a8de6d175d2c33ca609925f5da1", size = 1253670, upload-time = "2025-08-12T07:00:24.69Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/a1/11/5b414b9fae6255b5fb1e22e2ed3dc3a72d3a694e5703910e640ac78346bb/sentencepiece-0.2.1-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:a19adcec27c524cb7069a1c741060add95f942d1cbf7ad0d104dffa0a7d28a2b", size = 1946081, upload-time = "2025-08-12T07:00:36.97Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/77/eb/7a5682bb25824db8545f8e5662e7f3e32d72a508fdce086029d89695106b/sentencepiece-0.2.1-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:e37e4b4c4a11662b5db521def4e44d4d30ae69a1743241412a93ae40fdcab4bb", size = 1327406, upload-time = "2025-08-12T07:00:38.669Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/03/b0/811dae8fb9f2784e138785d481469788f2e0d0c109c5737372454415f55f/sentencepiece-0.2.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:477c81505db072b3ab627e7eab972ea1025331bd3a92bacbf798df2b75ea86ec", size = 1254846, upload-time = "2025-08-12T07:00:40.611Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "setuptools"
|
||||
version = "82.0.0"
|
||||
@@ -5201,6 +5313,18 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/f6/56/6113c23ff46c00aae423333eb58b3e60bdfe9179d542781955a5e1514cb3/triton-3.6.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:46bd1c1af4b6704e554cad2eeb3b0a6513a980d470ccfa63189737340c7746a7", size = 188397994, upload-time = "2026-01-20T16:01:14.236Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "typeguard"
|
||||
version = "4.5.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "typing-extensions" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/2b/e8/66e25efcc18542d58706ce4e50415710593721aae26e794ab1dec34fb66f/typeguard-4.5.1.tar.gz", hash = "sha256:f6f8ecbbc819c9bc749983cc67c02391e16a9b43b8b27f15dc70ed7c4a007274", size = 80121, upload-time = "2026-02-19T16:09:03.392Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/91/88/b55b3117287a8540b76dbdd87733808d4d01c8067a3b339408c250bb3600/typeguard-4.5.1-py3-none-any.whl", hash = "sha256:44d2bf329d49a244110a090b55f5f91aa82d9a9834ebfd30bcc73651e4a8cc40", size = 36745, upload-time = "2026-02-19T16:09:01.6Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "typer"
|
||||
version = "0.24.0"
|
||||
@@ -5249,6 +5373,20 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tyro"
|
||||
version = "1.0.8"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "docstring-parser" },
|
||||
{ name = "typeguard" },
|
||||
{ name = "typing-extensions" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/9e/89/f4d76ba5ac4ca627c3a0413fdc8d90f202929c45d9c25deec29c9a86a0ae/tyro-1.0.8.tar.gz", hash = "sha256:c22bf238ce029b8cc262759129fccfd968300c3761f00ce186570af3891e14bb", size = 460626, upload-time = "2026-02-25T00:18:45.388Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/61/7c/9c5a4ab9e24f1428fd5ab2728b7f68878f54d549a11fa431ead76d3d127b/tyro-1.0.8-py3-none-any.whl", hash = "sha256:abb6a4054f20616d9ca5181ba0c6ae0e35cf8f35ae0c78d3d4dbd1f678e0643d", size = 181935, upload-time = "2026-02-25T00:18:43.689Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tzdata"
|
||||
version = "2025.3"
|
||||
@@ -5880,6 +6018,30 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/73/ae/b48f95715333080afb75a4504487cbe142cae1268afc482d06692d605ae6/yarl-1.22.0-py3-none-any.whl", hash = "sha256:1380560bdba02b6b6c90de54133c81c9f2a453dee9912fe58c1dcced1edb7cff", size = 46814, upload-time = "2025-10-06T14:12:53.872Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "zeus-ml"
|
||||
version = "0.11.0.post1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "amdsmi" },
|
||||
{ name = "httpx" },
|
||||
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
|
||||
{ name = "numpy", version = "2.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
|
||||
{ name = "nvidia-ml-py" },
|
||||
{ name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
|
||||
{ name = "pandas", version = "3.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
|
||||
{ name = "pydantic" },
|
||||
{ name = "python-dateutil" },
|
||||
{ name = "rich" },
|
||||
{ name = "scikit-learn", version = "1.7.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
|
||||
{ name = "scikit-learn", version = "1.8.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
|
||||
{ name = "tyro" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/c7/26/e7e70121272fffc4c9ce3ec232e457d33d02b490eefe1b7a2b022bf296aa/zeus_ml-0.11.0.post1.tar.gz", hash = "sha256:0bc061b6c34edcfc2e86b3de81000157e30227fe11e7d0c683f6f5cf7bc3be22", size = 178924, upload-time = "2025-02-03T01:17:07.93Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/37/f4/d787c837308a26d7db69154761ae8b7e7f0a0703952cd57dc40fd13471c8/zeus_ml-0.11.0.post1-py3-none-any.whl", hash = "sha256:d0554ff3a4b8a27d818f4ea8688daa500be5f10322fb5824365aef24c1682e43", size = 227090, upload-time = "2025-02-03T01:17:04.018Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "zipp"
|
||||
version = "3.23.0"
|
||||
|
||||
Reference in New Issue
Block a user