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:
Jon Saad-Falcon
2026-02-26 20:09:07 +00:00
co-authored by Claude Opus 4.6
parent 9f5e97eed9
commit 24972e3e52
78 changed files with 7762 additions and 205 deletions
+28 -13
View File
@@ -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 |
+27
View File
@@ -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"]
+28 -1
View File
@@ -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 |
+4
View File
@@ -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
+15
View File
@@ -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
View File
@@ -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__":
+3
View File
@@ -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
+258
View File
@@ -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
View File
@@ -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,
}
+17
View File
@@ -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)
+4
View File
@@ -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>
+2 -1
View File
@@ -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
View File
@@ -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,
+4
View File
@@ -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"]
+2
View File
@@ -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
+9 -1
View File
@@ -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
+132
View File
@@ -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"]
+12 -1
View File
@@ -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
+13 -3
View File
@@ -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:
+2
View File
@@ -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 -5
View File
@@ -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
View File
@@ -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
+329
View File
@@ -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()
+92
View File
@@ -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",
)
)
+84 -4
View File
@@ -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:
+52 -2
View File
@@ -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",
+2
View File
@@ -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)
+18
View File
@@ -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)
+1
View File
@@ -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
+1
View File
@@ -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,
}
+17
View File
@@ -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
View File
@@ -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()
+7 -1
View File
@@ -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,
+26 -6
View File
@@ -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)
+32 -1
View File
@@ -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",
]
+201 -9
View File
@@ -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,
*,
+130
View File
@@ -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",
]
+131
View File
@@ -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"]
+110
View File
@@ -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"]
+146
View File
@@ -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",
]
+248
View File
@@ -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"]
+190
View File
@@ -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"]
+10 -2
View File
@@ -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),
}
+275 -15
View File
@@ -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"]
+129
View File
@@ -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",
]
+66 -1
View File
@@ -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),
),
)
+117
View File
@@ -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")
+168
View File
@@ -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
+68
View File
@@ -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
+1 -1
View File
@@ -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(
+72
View File
@@ -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
View File
+205
View File
@@ -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()
+64 -7
View File
@@ -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
+4 -4
View File
@@ -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"
+7 -3
View File
@@ -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"
+83
View File
@@ -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)
+106
View File
@@ -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
+283
View File
@@ -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]
+191
View File
@@ -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()
+188
View File
@@ -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
+159
View File
@@ -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"
+216
View File
@@ -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"]
+362
View File
@@ -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
+248
View File
@@ -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
+932
View File
@@ -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"
+3
View File
@@ -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
+353
View File
@@ -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()
+231
View File
@@ -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()
+155
View File
@@ -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
Generated
+163 -1
View File
@@ -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"