Compare commits
1 Commits
fix/749
...
feat/674-a
| Author | SHA1 | Date | |
|---|---|---|---|
|
12b34f6928 |
219
agent/atlas_provider.py
Normal file
219
agent/atlas_provider.py
Normal file
@@ -0,0 +1,219 @@
|
||||
"""Atlas Inference Engine provider integration.
|
||||
|
||||
Atlas is a Rust+CUDA LLM inference engine that is 3x faster than vLLM.
|
||||
It exposes an OpenAI-compatible API at http://localhost:8888/v1.
|
||||
|
||||
This module provides:
|
||||
- Atlas provider configuration and validation
|
||||
- Health check for Atlas server
|
||||
- Model discovery via Atlas API
|
||||
- Benchmark comparison utilities
|
||||
|
||||
Usage:
|
||||
from agent.atlas_provider import AtlasProvider
|
||||
atlas = AtlasProvider()
|
||||
if atlas.is_available():
|
||||
models = atlas.list_models()
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default Atlas configuration
|
||||
ATLAS_DEFAULT_BASE_URL = os.getenv("ATLAS_BASE_URL", "http://localhost:8888/v1")
|
||||
ATLAS_DEFAULT_PORT = int(os.getenv("ATLAS_PORT", "8888"))
|
||||
|
||||
# Known Atlas-compatible models (as of alpha-2.8)
|
||||
ATLAS_SUPPORTED_MODELS = [
|
||||
"Sehyo/Qwen3.5-35B-A3B-NVFP4",
|
||||
"Sehyo/Qwen3.5-122B-A10B-NVFP4",
|
||||
"Sehyo/Qwen3-Next-80B-A3B-NVFP4",
|
||||
"Sehyo/Qwen3-Coder-Next-FP8",
|
||||
"Sehyo/Qwen3-VL-30B-NVFP4",
|
||||
"Sehyo/Gemma-4-26B-NVFP4",
|
||||
"Sehyo/Nemotron-3-Nano-30B-NVFP4",
|
||||
"Sehyo/Mistral-Small-4-119B-NVFP4",
|
||||
]
|
||||
|
||||
|
||||
class AtlasProvider:
|
||||
"""Atlas Inference Engine provider.
|
||||
|
||||
Wraps the Atlas OpenAI-compatible API with health checks,
|
||||
model discovery, and configuration validation.
|
||||
"""
|
||||
|
||||
def __init__(self, base_url: str = ""):
|
||||
self.base_url = (base_url or ATLAS_DEFAULT_BASE_URL).rstrip("/")
|
||||
self._api_url = self.base_url
|
||||
if not self._api_url.endswith("/v1"):
|
||||
self._api_url += "/v1"
|
||||
|
||||
def is_available(self) -> bool:
|
||||
"""Check if Atlas server is running and responding."""
|
||||
try:
|
||||
import urllib.request
|
||||
req = urllib.request.Request(f"{self._api_url}/models", method="GET")
|
||||
with urllib.request.urlopen(req, timeout=5) as resp:
|
||||
return resp.status == 200
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def list_models(self) -> List[Dict[str, Any]]:
|
||||
"""List models available on the Atlas server."""
|
||||
try:
|
||||
import urllib.request
|
||||
req = urllib.request.Request(f"{self._api_url}/models", method="GET")
|
||||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||||
data = json.loads(resp.read())
|
||||
return data.get("data", [])
|
||||
except Exception as exc:
|
||||
logger.warning("Atlas model list failed: %s", exc)
|
||||
return []
|
||||
|
||||
def health_check(self) -> Dict[str, Any]:
|
||||
"""Comprehensive health check of the Atlas server."""
|
||||
result = {
|
||||
"available": False,
|
||||
"base_url": self.base_url,
|
||||
"models": [],
|
||||
"model_count": 0,
|
||||
"latency_ms": 0,
|
||||
"error": None,
|
||||
}
|
||||
|
||||
t0 = time.monotonic()
|
||||
try:
|
||||
import urllib.request
|
||||
req = urllib.request.Request(f"{self._api_url}/models", method="GET")
|
||||
with urllib.request.urlopen(req, timeout=5) as resp:
|
||||
result["latency_ms"] = int((time.monotonic() - t0) * 1000)
|
||||
if resp.status == 200:
|
||||
data = json.loads(resp.read())
|
||||
models = data.get("data", [])
|
||||
result["available"] = True
|
||||
result["models"] = [m.get("id", "") for m in models]
|
||||
result["model_count"] = len(models)
|
||||
except Exception as exc:
|
||||
result["latency_ms"] = int((time.monotonic() - t0) * 1000)
|
||||
result["error"] = str(exc)
|
||||
|
||||
return result
|
||||
|
||||
def get_provider_config(self) -> Dict[str, Any]:
|
||||
"""Return a provider config dict suitable for hermes config.yaml."""
|
||||
return {
|
||||
"name": "atlas",
|
||||
"base_url": self._api_url,
|
||||
"api_mode": "openai",
|
||||
"description": "Atlas Inference Engine (Rust+CUDA, 3x faster than vLLM)",
|
||||
}
|
||||
|
||||
def benchmark_inference(
|
||||
self,
|
||||
prompt: str = "Explain the theory of relativity in three sentences.",
|
||||
model: str = "",
|
||||
num_tokens: int = 100,
|
||||
) -> Dict[str, Any]:
|
||||
"""Run a quick inference benchmark against Atlas.
|
||||
|
||||
Returns timing metrics for comparison with vLLM or other backends.
|
||||
"""
|
||||
result = {
|
||||
"provider": "atlas",
|
||||
"model": model or "unknown",
|
||||
"prompt_tokens": 0,
|
||||
"completion_tokens": 0,
|
||||
"total_time_ms": 0,
|
||||
"tokens_per_second": 0.0,
|
||||
"time_to_first_token_ms": 0,
|
||||
"error": None,
|
||||
}
|
||||
|
||||
try:
|
||||
import urllib.request
|
||||
|
||||
messages = [{"role": "user", "content": prompt}]
|
||||
body = {
|
||||
"model": model or "",
|
||||
"messages": messages,
|
||||
"max_tokens": num_tokens,
|
||||
"stream": False,
|
||||
}
|
||||
|
||||
t0 = time.monotonic()
|
||||
req = urllib.request.Request(
|
||||
f"{self._api_url}/chat/completions",
|
||||
data=json.dumps(body).encode(),
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=60) as resp:
|
||||
elapsed = time.monotonic() - t0
|
||||
data = json.loads(resp.read())
|
||||
|
||||
usage = data.get("usage", {})
|
||||
result["prompt_tokens"] = usage.get("prompt_tokens", 0)
|
||||
result["completion_tokens"] = usage.get("completion_tokens", 0)
|
||||
result["total_time_ms"] = int(elapsed * 1000)
|
||||
if elapsed > 0 and result["completion_tokens"] > 0:
|
||||
result["tokens_per_second"] = round(
|
||||
result["completion_tokens"] / elapsed, 1
|
||||
)
|
||||
|
||||
except Exception as exc:
|
||||
result["error"] = str(exc)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def get_atlas_config_hint() -> str:
|
||||
"""Return a config.yaml snippet for adding Atlas as a provider."""
|
||||
return """# Atlas Inference Engine configuration
|
||||
# Add to config.yaml under providers:
|
||||
|
||||
providers:
|
||||
atlas:
|
||||
base_url: http://localhost:8888/v1
|
||||
api_mode: openai
|
||||
# No API key needed for local Atlas
|
||||
|
||||
# Then set model:
|
||||
model:
|
||||
default: atlas/<model-name>
|
||||
provider: atlas
|
||||
|
||||
# Or use as fallback:
|
||||
fallback_model:
|
||||
provider: atlas
|
||||
model: Sehyo/Qwen3.5-35B-A3B-NVFP4
|
||||
"""
|
||||
|
||||
|
||||
def get_atlas_docker_command(
|
||||
model: str = "Sehyo/Qwen3.5-35B-A3B-NVFP4",
|
||||
port: int = 8888,
|
||||
speculative: bool = True,
|
||||
max_seq_len: int = 131072,
|
||||
max_batch_size: int = 1,
|
||||
) -> str:
|
||||
"""Return the docker run command for Atlas."""
|
||||
cmd = (
|
||||
"docker run -d --gpus all --ipc=host "
|
||||
f"-p {port}:8888 "
|
||||
"-v ~/.cache/huggingface:/root/.cache/huggingface "
|
||||
"avarok/atlas-gb10:alpha-2.8 serve "
|
||||
f"{model} "
|
||||
)
|
||||
if speculative:
|
||||
cmd += "--speculative --scheduling-policy slai "
|
||||
cmd += f"--max-seq-len {max_seq_len} --max-batch-size {max_batch_size} "
|
||||
cmd += "--max-prefill-tokens 0"
|
||||
return cmd
|
||||
@@ -43,6 +43,13 @@ model:
|
||||
# Set OLLAMA_API_KEY in .env — automatically picked up when base_url
|
||||
# points to ollama.com.
|
||||
#
|
||||
# Atlas Inference Engine (Rust+CUDA, 3x faster than vLLM):
|
||||
# provider: "atlas"
|
||||
# base_url: "http://localhost:8888/v1"
|
||||
# Start with: docker run -d --gpus all --ipc=host -p 8888:8888
|
||||
# avarok/atlas-gb10:alpha-2.8 serve <model> --speculative
|
||||
# See: agent/atlas_provider.py for full config.
|
||||
#
|
||||
# Can also be overridden with --provider flag or HERMES_INFERENCE_PROVIDER env var.
|
||||
provider: "auto"
|
||||
|
||||
|
||||
@@ -924,6 +924,7 @@ def resolve_provider(
|
||||
# Local server aliases — route through the generic custom provider
|
||||
"lmstudio": "custom", "lm-studio": "custom", "lm_studio": "custom",
|
||||
"ollama": "custom", "vllm": "custom", "llamacpp": "custom",
|
||||
"atlas": "custom",
|
||||
"llama.cpp": "custom", "llama-cpp": "custom",
|
||||
}
|
||||
normalized = _PROVIDER_ALIASES.get(normalized, normalized)
|
||||
|
||||
@@ -1,77 +0,0 @@
|
||||
"""Tests for batch tool execution (#749)."""
|
||||
|
||||
import pytest
|
||||
from tools.batch_executor import (
|
||||
classify_tool_call,
|
||||
classify_batch,
|
||||
)
|
||||
|
||||
|
||||
class TestClassifyToolCall:
|
||||
def test_read_file_is_parallel(self):
|
||||
assert classify_tool_call("read_file") == "parallel"
|
||||
|
||||
def test_search_files_is_parallel(self):
|
||||
assert classify_tool_call("search_files") == "parallel"
|
||||
|
||||
def test_write_file_is_sequential(self):
|
||||
assert classify_tool_call("write_file") == "sequential"
|
||||
|
||||
def test_terminal_is_sequential(self):
|
||||
assert classify_tool_call("terminal") == "sequential"
|
||||
|
||||
def test_execute_code_is_sequential(self):
|
||||
assert classify_tool_call("execute_code") == "sequential"
|
||||
|
||||
def test_cronjob_list_is_parallel(self):
|
||||
assert classify_tool_call("cronjob", {"action": "list"}) == "parallel"
|
||||
|
||||
def test_cronjob_create_is_sequential(self):
|
||||
assert classify_tool_call("cronjob", {"action": "create"}) == "sequential"
|
||||
|
||||
def test_fact_store_search_is_parallel(self):
|
||||
assert classify_tool_call("fact_store", {"action": "search"}) == "parallel"
|
||||
|
||||
def test_fact_store_add_is_sequential(self):
|
||||
assert classify_tool_call("fact_store", {"action": "add"}) == "sequential"
|
||||
|
||||
def test_unknown_tool_is_sequential(self):
|
||||
assert classify_tool_call("unknown_tool") == "sequential"
|
||||
|
||||
|
||||
class TestClassifyBatch:
|
||||
def test_splits_correctly(self):
|
||||
calls = [
|
||||
{"name": "read_file", "args": {"path": "a"}},
|
||||
{"name": "write_file", "args": {"path": "b"}},
|
||||
{"name": "search_files", "args": {"pattern": "c"}},
|
||||
{"name": "terminal", "args": {"command": "d"}},
|
||||
]
|
||||
parallel, sequential = classify_batch(calls)
|
||||
assert len(parallel) == 2
|
||||
assert len(sequential) == 2
|
||||
assert parallel[0]["name"] == "read_file"
|
||||
assert sequential[0]["name"] == "write_file"
|
||||
|
||||
def test_all_parallel(self):
|
||||
calls = [
|
||||
{"name": "read_file", "args": {}},
|
||||
{"name": "search_files", "args": {}},
|
||||
]
|
||||
parallel, sequential = classify_batch(calls)
|
||||
assert len(parallel) == 2
|
||||
assert len(sequential) == 0
|
||||
|
||||
def test_all_sequential(self):
|
||||
calls = [
|
||||
{"name": "write_file", "args": {}},
|
||||
{"name": "terminal", "args": {}},
|
||||
]
|
||||
parallel, sequential = classify_batch(calls)
|
||||
assert len(parallel) == 0
|
||||
assert len(sequential) == 2
|
||||
|
||||
def test_empty(self):
|
||||
parallel, sequential = classify_batch([])
|
||||
assert len(parallel) == 0
|
||||
assert len(sequential) == 0
|
||||
@@ -1,250 +0,0 @@
|
||||
"""
|
||||
Batch tool execution with parallel safety checks (#749).
|
||||
|
||||
Classifies tool calls as parallel-safe or sequential, then executes
|
||||
parallel-safe calls concurrently while keeping destructive ops serialized.
|
||||
|
||||
Safety classification:
|
||||
- PARALLEL-SAFE: read_file, search_files, browser_snapshot, session_search,
|
||||
fact_store (search/probe/list), skill_view
|
||||
- SEQUENTIAL: write_file, patch, terminal, execute_code, browser_click,
|
||||
browser_type, browser_navigate, cronjob (create/update/delete),
|
||||
memory (add/update/remove), skill_manage
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Tools that only read state — safe to parallelize
|
||||
PARALLEL_SAFE_TOOLS = frozenset([
|
||||
"read_file",
|
||||
"search_files",
|
||||
"browser_snapshot",
|
||||
"browser_get_images",
|
||||
"browser_back",
|
||||
"browser_vision",
|
||||
"browser_console",
|
||||
"session_search",
|
||||
"fact_store", # search/probe/list are read-only; add/update are not
|
||||
"skill_view",
|
||||
"skills_list",
|
||||
"cronjob", # list is read-only; create/update/run are not (filtered below)
|
||||
"clarify", # asking questions is safe
|
||||
"memory", # probe/search/list are read-only
|
||||
"vision_analyze",
|
||||
])
|
||||
|
||||
# Tools that modify state — must be serialized
|
||||
SEQUENTIAL_TOOLS = frozenset([
|
||||
"write_file",
|
||||
"patch",
|
||||
"terminal",
|
||||
"execute_code",
|
||||
"browser_click",
|
||||
"browser_type",
|
||||
"browser_press",
|
||||
"browser_scroll",
|
||||
"browser_navigate",
|
||||
"cronjob", # create/update/run/pause/resume/remove
|
||||
"memory", # add/update/remove
|
||||
"skill_manage",
|
||||
"todo",
|
||||
"text_to_speech",
|
||||
"image_generate",
|
||||
"delegate_task",
|
||||
"clarify", # clarify with choices needs user input
|
||||
"process",
|
||||
])
|
||||
|
||||
# Cronjob sub-actions that are read-only
|
||||
_CRON_READ_ONLY = frozenset(["list"])
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchResult:
|
||||
"""Result of a batch tool execution."""
|
||||
results: List[Dict[str, Any]] = field(default_factory=list)
|
||||
parallel_count: int = 0
|
||||
sequential_count: int = 0
|
||||
elapsed_ms: float = 0
|
||||
|
||||
|
||||
def classify_tool_call(tool_name: str, tool_args: Optional[Dict] = None) -> str:
|
||||
"""Classify a tool call as 'parallel' or 'sequential'.
|
||||
|
||||
Returns 'parallel' or 'sequential'.
|
||||
"""
|
||||
# Special cases based on sub-action
|
||||
if tool_name == "cronjob":
|
||||
action = (tool_args or {}).get("action", "")
|
||||
if action in _CRON_READ_ONLY:
|
||||
return "parallel"
|
||||
return "sequential"
|
||||
|
||||
if tool_name == "fact_store":
|
||||
action = (tool_args or {}).get("action", "")
|
||||
if action in ("search", "probe", "list", "related", "reason", "contradict"):
|
||||
return "parallel"
|
||||
return "sequential"
|
||||
|
||||
if tool_name == "memory":
|
||||
action = (tool_args or {}).get("action", "")
|
||||
if action in ("probe", "search", "list"):
|
||||
return "parallel"
|
||||
return "sequential"
|
||||
|
||||
# Check sequential first (more restrictive)
|
||||
if tool_name in SEQUENTIAL_TOOLS:
|
||||
return "sequential"
|
||||
|
||||
if tool_name in PARALLEL_SAFE_TOOLS:
|
||||
return "parallel"
|
||||
|
||||
# Unknown tools default to sequential (safe)
|
||||
return "sequential"
|
||||
|
||||
|
||||
def classify_batch(tool_calls: List[Dict]) -> Tuple[List[Dict], List[Dict]]:
|
||||
"""Split a list of tool calls into parallel-safe and sequential groups.
|
||||
|
||||
Args:
|
||||
tool_calls: List of dicts with 'name' and 'args' keys
|
||||
|
||||
Returns:
|
||||
(parallel_calls, sequential_calls)
|
||||
"""
|
||||
parallel = []
|
||||
sequential = []
|
||||
|
||||
for call in tool_calls:
|
||||
name = call.get("name", "")
|
||||
args = call.get("args", {})
|
||||
classification = classify_tool_call(name, args)
|
||||
|
||||
if classification == "parallel":
|
||||
parallel.append(call)
|
||||
else:
|
||||
sequential.append(call)
|
||||
|
||||
return parallel, sequential
|
||||
|
||||
|
||||
async def execute_parallel(
|
||||
tool_calls: List[Dict],
|
||||
executor: Callable,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Execute parallel-safe tool calls concurrently.
|
||||
|
||||
Args:
|
||||
tool_calls: List of tool call dicts
|
||||
executor: Async callable(tool_name, tool_args) -> result
|
||||
|
||||
Returns:
|
||||
List of results in same order as input
|
||||
"""
|
||||
tasks = []
|
||||
for call in tool_calls:
|
||||
task = asyncio.create_task(
|
||||
executor(call["name"], call.get("args", {})),
|
||||
name=f"tool:{call['name']}"
|
||||
)
|
||||
tasks.append((call, task))
|
||||
|
||||
results = []
|
||||
for call, task in tasks:
|
||||
try:
|
||||
result = await task
|
||||
results.append({
|
||||
"tool_name": call["name"],
|
||||
"result": result,
|
||||
"parallel": True,
|
||||
"error": None,
|
||||
})
|
||||
except Exception as e:
|
||||
logger.error("Parallel tool '%s' failed: %s", call["name"], e)
|
||||
results.append({
|
||||
"tool_name": call["name"],
|
||||
"result": None,
|
||||
"parallel": True,
|
||||
"error": str(e),
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
|
||||
async def execute_sequential(
|
||||
tool_calls: List[Dict],
|
||||
executor: Callable,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Execute sequential tool calls one at a time."""
|
||||
results = []
|
||||
for call in tool_calls:
|
||||
try:
|
||||
result = await executor(call["name"], call.get("args", {}))
|
||||
results.append({
|
||||
"tool_name": call["name"],
|
||||
"result": result,
|
||||
"parallel": False,
|
||||
"error": None,
|
||||
})
|
||||
except Exception as e:
|
||||
logger.error("Sequential tool '%s' failed: %s", call["name"], e)
|
||||
results.append({
|
||||
"tool_name": call["name"],
|
||||
"result": None,
|
||||
"parallel": False,
|
||||
"error": str(e),
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
|
||||
async def execute_batch(
|
||||
tool_calls: List[Dict],
|
||||
executor: Callable,
|
||||
) -> BatchResult:
|
||||
"""Execute a batch of tool calls with parallel safety checks.
|
||||
|
||||
1. Classify each call as parallel-safe or sequential
|
||||
2. Execute all parallel-safe calls concurrently
|
||||
3. Execute sequential calls one at a time
|
||||
4. Merge results in original order
|
||||
|
||||
Args:
|
||||
tool_calls: List of dicts with 'name' and 'args' keys
|
||||
executor: Async callable(tool_name, tool_args) -> result
|
||||
|
||||
Returns:
|
||||
BatchResult with all results and timing
|
||||
"""
|
||||
start = time.monotonic()
|
||||
|
||||
parallel_calls, sequential_calls = classify_batch(tool_calls)
|
||||
|
||||
# Execute parallel-safe calls concurrently
|
||||
parallel_results = []
|
||||
if parallel_calls:
|
||||
parallel_results = await execute_parallel(parallel_calls, executor)
|
||||
|
||||
# Execute sequential calls in order
|
||||
sequential_results = []
|
||||
if sequential_calls:
|
||||
sequential_results = await execute_sequential(sequential_calls, executor)
|
||||
|
||||
# Merge results — parallel first, then sequential (order preserved within groups)
|
||||
all_results = parallel_results + sequential_results
|
||||
|
||||
elapsed = (time.monotonic() - start) * 1000
|
||||
|
||||
return BatchResult(
|
||||
results=all_results,
|
||||
parallel_count=len(parallel_calls),
|
||||
sequential_count=len(sequential_calls),
|
||||
elapsed_ms=elapsed,
|
||||
)
|
||||
Reference in New Issue
Block a user