refactor(cli): break down main.py into modular CLI components
Extracted main.py (556 lines) into focused modules: - cli/parser.py: Argument parsing (151 lines) - cli/main_runner.py: Main application logic (320 lines) - cli/test_runner.py: Test mode runner (81 lines) - cli/tool_server.py: Tool server runner (69 lines) - utils/network.py: Network utilities (IP detection) main.py is now 99 lines (down from 556). All 35 tests pass. Note: main_runner.py at 320 lines is slightly over 300 limit, will address in subsequent refactoring.
This commit is contained in:
@@ -10,223 +10,63 @@ import sys
|
|||||||
import multiprocessing as mp
|
import multiprocessing as mp
|
||||||
|
|
||||||
# CRITICAL: Set spawn method BEFORE any other imports on macOS
|
# CRITICAL: Set spawn method BEFORE any other imports on macOS
|
||||||
# This prevents fork-related issues with Metal GPU
|
|
||||||
if sys.platform == "darwin":
|
if sys.platform == "darwin":
|
||||||
try:
|
try:
|
||||||
mp.set_start_method("spawn", force=True)
|
mp.set_start_method("spawn", force=True)
|
||||||
except RuntimeError:
|
except RuntimeError:
|
||||||
pass # Already set
|
pass
|
||||||
|
|
||||||
import argparse
|
|
||||||
import asyncio
|
import asyncio
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
# Add src to path - resolve for Windows compatibility
|
# Add src to path
|
||||||
src_path = Path(__file__).parent.resolve() / "src"
|
src_path = Path(__file__).parent.resolve() / "src"
|
||||||
sys.path.insert(0, str(src_path))
|
sys.path.insert(0, str(src_path))
|
||||||
|
|
||||||
# Also add parent dir for Windows import issues
|
|
||||||
if str(Path(__file__).parent.resolve()) not in sys.path:
|
if str(Path(__file__).parent.resolve()) not in sys.path:
|
||||||
sys.path.insert(0, str(Path(__file__).parent.resolve()))
|
sys.path.insert(0, str(Path(__file__).parent.resolve()))
|
||||||
|
|
||||||
# These imports must come AFTER setting spawn method on macOS
|
from cli.parser import parse_args
|
||||||
from hardware.detector import detect_hardware
|
from cli.tool_server import run_tool_server
|
||||||
from models.selector import select_optimal_model
|
from utils.network import get_local_ip
|
||||||
from models.downloader import download_model_for_config
|
|
||||||
from swarm import SwarmManager
|
|
||||||
from api import create_server
|
|
||||||
from api.routes import set_federated_swarm
|
|
||||||
from mcp_server import create_mcp_server
|
|
||||||
from interactive import (
|
|
||||||
interactive_model_selection,
|
|
||||||
show_startup_summary,
|
|
||||||
show_runtime_menu,
|
|
||||||
custom_configuration,
|
|
||||||
)
|
|
||||||
from network import create_discovery_service, FederatedSwarm
|
|
||||||
from tools.executor import ToolExecutor, set_tool_executor
|
|
||||||
from utils.logging_config import setup_logging
|
from utils.logging_config import setup_logging
|
||||||
|
from hardware.detector import detect_hardware
|
||||||
|
from interactive import print_hardware_info
|
||||||
|
|
||||||
# Set up logging (DEBUG level for development)
|
# Set up logging
|
||||||
setup_logging()
|
setup_logging()
|
||||||
|
|
||||||
|
|
||||||
async def setup_swarm(model_config, hardware):
|
def handle_detect_mode(hardware) -> int:
|
||||||
"""Download model and initialize swarm."""
|
"""Handle --detect mode."""
|
||||||
# Download model
|
print_hardware_info(hardware)
|
||||||
print("\n⬇️ Downloading model...")
|
print("\n✅ Detection complete")
|
||||||
try:
|
return 0
|
||||||
model_path = download_model_for_config(model_config)
|
|
||||||
print(f"✓ Model ready at: {model_path}")
|
|
||||||
except Exception as e:
|
def handle_tool_server_mode(args, hardware) -> int:
|
||||||
print(f"\n❌ Error downloading model: {e}", file=sys.stderr)
|
"""Handle --tool-server mode."""
|
||||||
return None
|
print("\n🔧 Starting Tool Execution Server...")
|
||||||
|
host = args.host if args.host else get_local_ip()
|
||||||
|
|
||||||
# Initialize swarm
|
|
||||||
print("\n🚀 Initializing swarm...")
|
|
||||||
try:
|
try:
|
||||||
swarm = SwarmManager(
|
asyncio.run(run_tool_server(host, args.tool_port))
|
||||||
model_config=model_config,
|
return 0
|
||||||
hardware=hardware,
|
except KeyboardInterrupt:
|
||||||
consensus_strategy="similarity"
|
print("\n\nTool server stopped")
|
||||||
)
|
return 0
|
||||||
|
|
||||||
success = await swarm.initialize(str(model_path))
|
|
||||||
if not success:
|
|
||||||
print("❌ Failed to initialize swarm")
|
|
||||||
return None
|
|
||||||
|
|
||||||
return swarm
|
|
||||||
except Exception as e:
|
|
||||||
print(f"\n❌ Error initializing swarm: {e}", file=sys.stderr)
|
|
||||||
return None
|
|
||||||
|
|
||||||
|
|
||||||
|
async def run_main_mode(args, hardware) -> int:
|
||||||
def get_local_ip():
|
"""Run the main application mode."""
|
||||||
"""Get the local network IP address (private networks only)."""
|
from cli.main_runner import MainRunner
|
||||||
import socket
|
|
||||||
try:
|
|
||||||
# Create a socket and connect to a public DNS server
|
|
||||||
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
|
||||||
s.settimeout(2)
|
|
||||||
# Try to connect to Google's DNS - this doesn't actually send data
|
|
||||||
s.connect(("8.8.8.8", 80))
|
|
||||||
ip = s.getsockname()[0]
|
|
||||||
s.close()
|
|
||||||
|
|
||||||
# Check if it's a private IP (only 192.168.x.x for this network)
|
|
||||||
is_private = (
|
|
||||||
ip.startswith('192.168.')
|
|
||||||
)
|
|
||||||
|
|
||||||
if is_private:
|
|
||||||
print(f" 📡 Detected local IP: {ip}")
|
|
||||||
return ip
|
|
||||||
else:
|
|
||||||
# If not private, return localhost for safety
|
|
||||||
print(f" ⚠️ IP {ip} is not a private network, binding to localhost")
|
|
||||||
return "127.0.0.1"
|
|
||||||
except Exception as e:
|
|
||||||
print(f" ⚠️ Could not detect local IP: {e}, using localhost")
|
|
||||||
return "127.0.0.1"
|
|
||||||
|
|
||||||
def main():
|
|
||||||
parser = argparse.ArgumentParser(
|
|
||||||
description="Local Swarm - AI-powered coding LLM swarm",
|
|
||||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
||||||
epilog="""
|
|
||||||
Examples:
|
|
||||||
python main.py # Interactive setup and start
|
|
||||||
python main.py --auto # Auto-detect and start without menu
|
|
||||||
python main.py --detect # Show hardware detection only
|
|
||||||
python main.py --model qwen:3b:q4 # Use specific model (skip menu)
|
|
||||||
python main.py --port 17615 # Use custom port (default: 17615)
|
|
||||||
python main.py --host 192.168.1.5 # Bind to specific IP
|
|
||||||
python main.py --instances 4 # Force number of instances
|
|
||||||
python main.py --download-only # Download model only
|
|
||||||
python main.py --test # Test with sample prompt
|
|
||||||
python main.py --mcp # Enable MCP server
|
|
||||||
python main.py --federation # Enable federation with other instances
|
|
||||||
python main.py --federation --peer 192.168.1.10:17615 # Manual peer
|
|
||||||
"""
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
runner = MainRunner(hardware, args)
|
||||||
"--auto",
|
return await runner.run()
|
||||||
action="store_true",
|
|
||||||
help="Auto-detect best configuration without interactive menu"
|
|
||||||
)
|
def main() -> int:
|
||||||
parser.add_argument(
|
"""Main entry point."""
|
||||||
"--detect",
|
args = parse_args()
|
||||||
action="store_true",
|
|
||||||
help="Show hardware detection and exit"
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--model",
|
|
||||||
type=str,
|
|
||||||
help="Model to use (format: name:size:quant, e.g., qwen:3b:q4)"
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--port",
|
|
||||||
type=int,
|
|
||||||
default=17615,
|
|
||||||
help="Port to run the API server on (default: 17615)"
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--instances",
|
|
||||||
type=int,
|
|
||||||
help="Force number of instances (overrides auto-calculation)"
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--download-only",
|
|
||||||
action="store_true",
|
|
||||||
help="Download models only, don't start server"
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--test",
|
|
||||||
action="store_true",
|
|
||||||
help="Test with a sample prompt"
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--mcp",
|
|
||||||
action="store_true",
|
|
||||||
help="Enable MCP server alongside HTTP API"
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--config",
|
|
||||||
type=str,
|
|
||||||
default="config.yaml",
|
|
||||||
help="Path to config file"
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--host",
|
|
||||||
type=str,
|
|
||||||
default=None,
|
|
||||||
help="Host IP to bind to (default: auto-detect)"
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--federation",
|
|
||||||
action="store_true",
|
|
||||||
help="Enable federation with other Local Swarm instances on the network"
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--peer",
|
|
||||||
action="append",
|
|
||||||
dest="peers",
|
|
||||||
help="Manually add a peer (format: host:port, can be used multiple times)"
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--tool-server",
|
|
||||||
action="store_true",
|
|
||||||
help="Run as dedicated tool execution server (executes read/write/bash tools)"
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--tool-port",
|
|
||||||
type=int,
|
|
||||||
default=17616,
|
|
||||||
help="Port for tool execution server (default: 17616)"
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--tool-host",
|
|
||||||
type=str,
|
|
||||||
default=None,
|
|
||||||
nargs='?',
|
|
||||||
const='', # When --tool-host is used without a value, use empty string
|
|
||||||
help="URL of tool execution server. Use without value for auto-detected local IP (http://<local-ip>:17616), or provide explicit URL."
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--use-opencode-tools",
|
|
||||||
action="store_true",
|
|
||||||
help="Use opencode's tool definitions (adds ~27k tokens to context). Default: use local tool server (saves tokens)"
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--version",
|
|
||||||
action="version",
|
|
||||||
version="%(prog)s 0.1.0"
|
|
||||||
)
|
|
||||||
|
|
||||||
args = parser.parse_args()
|
|
||||||
|
|
||||||
# Detect hardware first
|
# Detect hardware first
|
||||||
print("\n🔍 Detecting hardware...")
|
print("\n🔍 Detecting hardware...")
|
||||||
@@ -234,323 +74,26 @@ Examples:
|
|||||||
hardware = detect_hardware()
|
hardware = detect_hardware()
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"\n❌ Error detecting hardware: {e}", file=sys.stderr)
|
print(f"\n❌ Error detecting hardware: {e}", file=sys.stderr)
|
||||||
sys.exit(1)
|
return 1
|
||||||
|
|
||||||
|
# Handle detect mode
|
||||||
if args.detect:
|
if args.detect:
|
||||||
# Just show hardware info
|
return handle_detect_mode(hardware)
|
||||||
from interactive import print_hardware_info
|
|
||||||
print_hardware_info(hardware)
|
|
||||||
print("\n✅ Detection complete")
|
|
||||||
return
|
|
||||||
|
|
||||||
# Tool server mode - run minimal tool-only server
|
# Handle tool server mode
|
||||||
if args.tool_server:
|
if args.tool_server:
|
||||||
print("\n🔧 Starting Tool Execution Server...")
|
return handle_tool_server_mode(args, hardware)
|
||||||
from fastapi import FastAPI
|
|
||||||
import uvicorn
|
|
||||||
|
|
||||||
# Initialize local tool executor
|
|
||||||
tool_executor = ToolExecutor(tool_host_url=None)
|
|
||||||
set_tool_executor(tool_executor)
|
|
||||||
|
|
||||||
app = FastAPI(title="Local Swarm Tool Server")
|
|
||||||
|
|
||||||
@app.post("/v1/tools/execute")
|
|
||||||
async def execute_tool(request: dict):
|
|
||||||
tool_name = request.get("tool", "")
|
|
||||||
tool_args = request.get("arguments", {})
|
|
||||||
result = await tool_executor.execute(tool_name, tool_args)
|
|
||||||
return {"result": result}
|
|
||||||
|
|
||||||
@app.get("/health")
|
|
||||||
async def health():
|
|
||||||
return {"status": "healthy", "mode": "tool-server"}
|
|
||||||
|
|
||||||
host = args.host if args.host else get_local_ip()
|
|
||||||
tool_port = args.tool_port
|
|
||||||
print(f"🔗 Tool server running at http://{host}:{tool_port}")
|
|
||||||
print(f" Endpoints:")
|
|
||||||
print(f" - POST /v1/tools/execute")
|
|
||||||
print(f" - GET /health")
|
|
||||||
print(f"\n✅ Tool server ready!")
|
|
||||||
|
|
||||||
uvicorn.run(app, host=host, port=tool_port)
|
|
||||||
return
|
|
||||||
|
|
||||||
# Determine model configuration
|
# Run main mode
|
||||||
config = None
|
try:
|
||||||
|
return asyncio.run(run_main_mode(args, hardware))
|
||||||
if args.model or args.instances or args.auto:
|
except KeyboardInterrupt:
|
||||||
# Use command-line arguments or auto-detect
|
print("\n\nReceived stop signal")
|
||||||
print("\n📊 Calculating optimal configuration...")
|
return 0
|
||||||
try:
|
except Exception as e:
|
||||||
config = select_optimal_model(
|
print(f"\n❌ Error: {e}", file=sys.stderr)
|
||||||
hardware,
|
return 1
|
||||||
preferred_model=args.model,
|
|
||||||
force_instances=args.instances
|
|
||||||
)
|
|
||||||
|
|
||||||
if not config:
|
|
||||||
print("\n❌ No suitable model found for your hardware")
|
|
||||||
print(" Minimum requirement: 2 GB available memory")
|
|
||||||
sys.exit(1)
|
|
||||||
|
|
||||||
# Show brief summary
|
|
||||||
print(f"\n✓ Selected: {config.display_name}")
|
|
||||||
print(f" Instances: {config.instances}")
|
|
||||||
print(f" Memory: {config.total_memory_gb:.1f} GB")
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
print(f"\n❌ Error selecting model: {e}", file=sys.stderr)
|
|
||||||
sys.exit(1)
|
|
||||||
else:
|
|
||||||
# Interactive mode - show menu
|
|
||||||
config = interactive_model_selection(hardware)
|
|
||||||
|
|
||||||
if not config:
|
|
||||||
print("\n❌ No configuration selected")
|
|
||||||
sys.exit(1)
|
|
||||||
|
|
||||||
if args.download_only:
|
|
||||||
# Download model only
|
|
||||||
print("\n" + "=" * 70)
|
|
||||||
print("⬇️ Download Mode: Downloading model only")
|
|
||||||
print("=" * 70)
|
|
||||||
|
|
||||||
try:
|
|
||||||
model_path = download_model_for_config(config)
|
|
||||||
print(f"✓ Model downloaded to: {model_path}")
|
|
||||||
print("\n" + "=" * 70)
|
|
||||||
print("✅ Download complete")
|
|
||||||
print("=" * 70)
|
|
||||||
except Exception as e:
|
|
||||||
print(f"\n❌ Download failed: {e}", file=sys.stderr)
|
|
||||||
sys.exit(1)
|
|
||||||
|
|
||||||
elif args.test:
|
|
||||||
# Test mode with sample prompt
|
|
||||||
print("\n" + "=" * 70)
|
|
||||||
print("🧪 Test Mode: Running sample inference")
|
|
||||||
print("=" * 70)
|
|
||||||
|
|
||||||
async def test_inference():
|
|
||||||
show_startup_summary(hardware, config)
|
|
||||||
swarm = await setup_swarm(config, hardware)
|
|
||||||
if not swarm:
|
|
||||||
return False
|
|
||||||
|
|
||||||
try:
|
|
||||||
# Test prompt
|
|
||||||
prompt = "Write a Python function to calculate factorial:"
|
|
||||||
print(f"\nPrompt: {prompt}\n")
|
|
||||||
print("Generating responses...\n")
|
|
||||||
|
|
||||||
result = await swarm.generate(prompt, max_tokens=200)
|
|
||||||
|
|
||||||
print("\n" + "=" * 70)
|
|
||||||
print("SELECTED RESPONSE:")
|
|
||||||
print("=" * 70)
|
|
||||||
print(result.selected_response.text)
|
|
||||||
print("\n" + "=" * 70)
|
|
||||||
print(f"Strategy: {result.strategy}")
|
|
||||||
print(f"Confidence: {result.confidence:.2f}")
|
|
||||||
print(f"Latency: {result.selected_response.latency_ms:.1f}ms")
|
|
||||||
print(f"Tokens/sec: {result.selected_response.tokens_per_second:.1f}")
|
|
||||||
|
|
||||||
# Show all responses
|
|
||||||
print("\nAll responses received:")
|
|
||||||
for i, resp in enumerate(result.all_responses):
|
|
||||||
preview = resp.text[:60].replace('\n', ' ')
|
|
||||||
print(f" Worker {i}: {preview}... ({resp.latency_ms:.1f}ms)")
|
|
||||||
|
|
||||||
return True
|
|
||||||
finally:
|
|
||||||
await swarm.shutdown()
|
|
||||||
|
|
||||||
success = asyncio.run(test_inference())
|
|
||||||
|
|
||||||
if success:
|
|
||||||
print("\n" + "=" * 70)
|
|
||||||
print("✅ Test complete")
|
|
||||||
print("=" * 70)
|
|
||||||
else:
|
|
||||||
print("\n❌ Test failed")
|
|
||||||
sys.exit(1)
|
|
||||||
|
|
||||||
else:
|
|
||||||
# Full mode (download + start API server + optional MCP)
|
|
||||||
show_startup_summary(hardware, config)
|
|
||||||
|
|
||||||
async def run_server():
|
|
||||||
swarm = await setup_swarm(config, hardware)
|
|
||||||
if not swarm:
|
|
||||||
return False
|
|
||||||
|
|
||||||
# Initialize tool executor
|
|
||||||
if args.tool_host is not None:
|
|
||||||
# --tool-host was provided
|
|
||||||
if args.tool_host == "":
|
|
||||||
# --tool-host with no value - use local IP with default port
|
|
||||||
local_ip = get_local_ip()
|
|
||||||
tool_host_url = f"http://{local_ip}:17616"
|
|
||||||
print(f"\n🔧 Using remote tool host: {tool_host_url} (auto-detected local IP)")
|
|
||||||
else:
|
|
||||||
# --tool-host with explicit value
|
|
||||||
tool_host_url = args.tool_host
|
|
||||||
print(f"\n🔧 Using remote tool host: {tool_host_url}")
|
|
||||||
tool_executor = ToolExecutor(tool_host_url=tool_host_url)
|
|
||||||
set_tool_executor(tool_executor)
|
|
||||||
else:
|
|
||||||
# Local tool execution (default)
|
|
||||||
tool_executor = ToolExecutor(tool_host_url=None)
|
|
||||||
set_tool_executor(tool_executor)
|
|
||||||
|
|
||||||
# Update summary with runtime info
|
|
||||||
show_startup_summary(hardware, config, swarm)
|
|
||||||
|
|
||||||
# Initialize federation if enabled
|
|
||||||
discovery = None
|
|
||||||
federated_swarm = None
|
|
||||||
if args.federation:
|
|
||||||
print("\n🌐 Initializing federation...")
|
|
||||||
try:
|
|
||||||
# Use specified host for advertising if provided
|
|
||||||
advertise_ip = args.host if args.host else None
|
|
||||||
discovery = await create_discovery_service(args.port, advertise_ip=advertise_ip)
|
|
||||||
|
|
||||||
# Get swarm info for advertising
|
|
||||||
swarm_info = {
|
|
||||||
"version": "0.1.0",
|
|
||||||
"instances": config.instances,
|
|
||||||
"model_id": config.model_id,
|
|
||||||
"hardware_summary": f"{hardware.cpu_cores} CPU, {hardware.ram_gb:.1f}GB RAM"
|
|
||||||
}
|
|
||||||
|
|
||||||
await discovery.start_advertising(swarm_info)
|
|
||||||
await discovery.start_listening()
|
|
||||||
|
|
||||||
# Add manual peers if specified
|
|
||||||
if args.peers:
|
|
||||||
print(f" 📍 Adding {len(args.peers)} manual peer(s)...")
|
|
||||||
from network.discovery import PeerInfo
|
|
||||||
from datetime import datetime
|
|
||||||
for peer_str in args.peers:
|
|
||||||
try:
|
|
||||||
host, port = peer_str.rsplit(':', 1)
|
|
||||||
port = int(port)
|
|
||||||
peer = PeerInfo(
|
|
||||||
host=host,
|
|
||||||
port=port,
|
|
||||||
name=f"manual_{host}_{port}",
|
|
||||||
version="0.1.0",
|
|
||||||
instances=0,
|
|
||||||
model_id="unknown",
|
|
||||||
hardware_summary="manual",
|
|
||||||
last_seen=datetime.now()
|
|
||||||
)
|
|
||||||
discovery.peers[peer.name] = peer
|
|
||||||
print(f" ✓ Added peer: {host}:{port}")
|
|
||||||
except Exception as e:
|
|
||||||
print(f" ⚠️ Failed to add peer {peer_str}: {e}")
|
|
||||||
|
|
||||||
# Create federated swarm wrapper
|
|
||||||
federated_swarm = FederatedSwarm(swarm, discovery)
|
|
||||||
set_federated_swarm(federated_swarm)
|
|
||||||
|
|
||||||
# Start health check loop in background
|
|
||||||
asyncio.create_task(discovery.start_health_check_loop(interval_seconds=10))
|
|
||||||
|
|
||||||
print(f" ✓ Federation enabled")
|
|
||||||
print(f" ✓ Discovery active on port {discovery.discovery_port}")
|
|
||||||
print(f" ✓ Peer health checks every 10s")
|
|
||||||
except Exception as e:
|
|
||||||
print(f" ⚠️ Failed to initialize federation: {e}")
|
|
||||||
print(" Continuing without federation...")
|
|
||||||
|
|
||||||
mcp_server = None
|
|
||||||
try:
|
|
||||||
# Create and start API server
|
|
||||||
print("\n🌐 Starting HTTP API server...")
|
|
||||||
# Use provided host or auto-detect
|
|
||||||
if args.host:
|
|
||||||
host = args.host
|
|
||||||
print(f"🔗 Using specified host: {host}:{args.port}")
|
|
||||||
else:
|
|
||||||
# Use local network IP instead of 0.0.0.0 for security
|
|
||||||
host = get_local_ip()
|
|
||||||
print(f"🔗 Binding to {host}:{args.port}")
|
|
||||||
|
|
||||||
# Show tool mode being used
|
|
||||||
if args.use_opencode_tools:
|
|
||||||
print(f"🔧 Tool mode: opencode tools (~27k tokens, full capabilities)")
|
|
||||||
else:
|
|
||||||
print(f"🔧 Tool mode: local tool server (~125 tokens, saves tokens)")
|
|
||||||
|
|
||||||
server = create_server(swarm, host=host, port=args.port, use_opencode_tools=args.use_opencode_tools)
|
|
||||||
|
|
||||||
print(f"\n✅ Local Swarm is running!")
|
|
||||||
print(f" API: http://{host}:{args.port}/v1")
|
|
||||||
print(f" Health: http://{host}:{args.port}/health")
|
|
||||||
|
|
||||||
if args.federation and discovery:
|
|
||||||
peers = discovery.get_peers()
|
|
||||||
print(f"\n🌐 Federation: Enabled")
|
|
||||||
print(f" Discovery port: {discovery.discovery_port}")
|
|
||||||
if peers:
|
|
||||||
print(f" Peers discovered: {len(peers)}")
|
|
||||||
for peer in peers:
|
|
||||||
print(f" - {peer.name} ({peer.model_id})")
|
|
||||||
else:
|
|
||||||
print(f" Peers discovered: 0 (waiting for peers...)")
|
|
||||||
|
|
||||||
# Show tool server status
|
|
||||||
if args.tool_host is not None:
|
|
||||||
print(f"\n🔧 Tool Server: Remote")
|
|
||||||
if args.tool_host == "":
|
|
||||||
local_ip = get_local_ip()
|
|
||||||
print(f" URL: http://{local_ip}:17616 (auto-detected)")
|
|
||||||
else:
|
|
||||||
print(f" URL: {args.tool_host}")
|
|
||||||
print(f" Mode: Tools executed remotely on tool host")
|
|
||||||
else:
|
|
||||||
print(f"\n🔧 Tool Server: Local")
|
|
||||||
print(f" Mode: Tools executed on this machine")
|
|
||||||
|
|
||||||
if args.mcp:
|
|
||||||
# Start MCP server alongside HTTP API
|
|
||||||
print("\n🤖 Starting MCP server...")
|
|
||||||
mcp_server = await create_mcp_server(swarm)
|
|
||||||
print(" MCP server active (stdio)")
|
|
||||||
|
|
||||||
print(f"\n💡 Configure opencode to use:")
|
|
||||||
print(f' base_url: http://127.0.0.1:{args.port}/v1')
|
|
||||||
print(f' api_key: any (not used)')
|
|
||||||
print(f"\nPress Ctrl+C to stop...\n")
|
|
||||||
|
|
||||||
# Start HTTP server (this will block)
|
|
||||||
await server.start()
|
|
||||||
|
|
||||||
except KeyboardInterrupt:
|
|
||||||
print("\n\nReceived stop signal")
|
|
||||||
finally:
|
|
||||||
if federated_swarm:
|
|
||||||
await federated_swarm.close()
|
|
||||||
if discovery:
|
|
||||||
await discovery.stop()
|
|
||||||
await swarm.shutdown()
|
|
||||||
|
|
||||||
return True
|
|
||||||
|
|
||||||
try:
|
|
||||||
success = asyncio.run(run_server())
|
|
||||||
if success:
|
|
||||||
print("\n" + "=" * 70)
|
|
||||||
print("✅ Server stopped gracefully")
|
|
||||||
print("=" * 70)
|
|
||||||
except Exception as e:
|
|
||||||
print(f"\n❌ Error running server: {e}", file=sys.stderr)
|
|
||||||
sys.exit(1)
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
main()
|
sys.exit(main())
|
||||||
|
|||||||
@@ -0,0 +1,320 @@
|
|||||||
|
"""Main application runner for Local Swarm.
|
||||||
|
|
||||||
|
Handles the primary application modes: download-only, test, and full server mode.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import sys
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
from models.selector import select_optimal_model, ModelConfig
|
||||||
|
from models.downloader import download_model_for_config
|
||||||
|
from swarm import SwarmManager
|
||||||
|
from api import create_server
|
||||||
|
from api.routes import set_federated_swarm
|
||||||
|
from interactive import (
|
||||||
|
interactive_model_selection,
|
||||||
|
show_startup_summary,
|
||||||
|
show_runtime_menu,
|
||||||
|
)
|
||||||
|
from network import create_discovery_service, FederatedSwarm
|
||||||
|
from tools.executor import ToolExecutor, set_tool_executor
|
||||||
|
from utils.network import get_local_ip
|
||||||
|
|
||||||
|
|
||||||
|
class MainRunner:
|
||||||
|
"""Runs the main application logic."""
|
||||||
|
|
||||||
|
def __init__(self, hardware, args):
|
||||||
|
"""Initialize the main runner.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
hardware: Hardware profile
|
||||||
|
args: Parsed command line arguments
|
||||||
|
"""
|
||||||
|
self.hardware = hardware
|
||||||
|
self.args = args
|
||||||
|
self.config: Optional[ModelConfig] = None
|
||||||
|
self.swarm: Optional[SwarmManager] = None
|
||||||
|
self.discovery = None
|
||||||
|
self.federated_swarm = None
|
||||||
|
self.mcp_server = None
|
||||||
|
|
||||||
|
async def run(self) -> int:
|
||||||
|
"""Run the main application.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Exit code (0 for success, 1 for error)
|
||||||
|
"""
|
||||||
|
# Get configuration
|
||||||
|
self.config = self._get_configuration()
|
||||||
|
if not self.config:
|
||||||
|
return 1
|
||||||
|
|
||||||
|
# Handle download-only mode
|
||||||
|
if self.args.download_only:
|
||||||
|
return await self._run_download_mode()
|
||||||
|
|
||||||
|
# Handle test mode
|
||||||
|
if self.args.test:
|
||||||
|
return await self._run_test_mode()
|
||||||
|
|
||||||
|
# Run full server mode
|
||||||
|
return await self._run_server_mode()
|
||||||
|
|
||||||
|
def _get_configuration(self) -> Optional[ModelConfig]:
|
||||||
|
"""Get the model configuration."""
|
||||||
|
if self.args.model or self.args.instances or self.args.auto:
|
||||||
|
return self._get_auto_config()
|
||||||
|
else:
|
||||||
|
return interactive_model_selection(self.hardware)
|
||||||
|
|
||||||
|
def _get_auto_config(self) -> Optional[ModelConfig]:
|
||||||
|
"""Get auto-detected configuration."""
|
||||||
|
print("\n📊 Calculating optimal configuration...")
|
||||||
|
try:
|
||||||
|
config = select_optimal_model(
|
||||||
|
self.hardware,
|
||||||
|
preferred_model=self.args.model,
|
||||||
|
force_instances=self.args.instances
|
||||||
|
)
|
||||||
|
|
||||||
|
if not config:
|
||||||
|
print("\n❌ No suitable model found for your hardware")
|
||||||
|
print(" Minimum requirement: 2 GB available memory")
|
||||||
|
return None
|
||||||
|
|
||||||
|
print(f"\n✓ Selected: {config.display_name}")
|
||||||
|
print(f" Instances: {config.instances}")
|
||||||
|
print(f" Memory: {config.total_memory_gb:.1f} GB")
|
||||||
|
return config
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
print(f"\n❌ Error selecting model: {e}", file=sys.stderr)
|
||||||
|
return None
|
||||||
|
|
||||||
|
async def _run_download_mode(self) -> int:
|
||||||
|
"""Run download-only mode."""
|
||||||
|
print("\n" + "=" * 70)
|
||||||
|
print("⬇️ Download Mode: Downloading model only")
|
||||||
|
print("=" * 70)
|
||||||
|
|
||||||
|
try:
|
||||||
|
model_path = download_model_for_config(self.config)
|
||||||
|
print(f"✓ Model downloaded to: {model_path}")
|
||||||
|
print("\n" + "=" * 70)
|
||||||
|
print("✅ Download complete")
|
||||||
|
print("=" * 70)
|
||||||
|
return 0
|
||||||
|
except Exception as e:
|
||||||
|
print(f"\n❌ Download failed: {e}", file=sys.stderr)
|
||||||
|
return 1
|
||||||
|
|
||||||
|
async def _run_test_mode(self) -> int:
|
||||||
|
"""Run test mode with sample prompt."""
|
||||||
|
from cli.test_runner import run_test
|
||||||
|
return await run_test(self.hardware, self.config)
|
||||||
|
|
||||||
|
async def _run_server_mode(self) -> int:
|
||||||
|
"""Run full server mode."""
|
||||||
|
show_startup_summary(self.hardware, self.config)
|
||||||
|
|
||||||
|
# Setup swarm
|
||||||
|
if not await self._setup_swarm():
|
||||||
|
return 1
|
||||||
|
|
||||||
|
# Initialize tool executor
|
||||||
|
self._setup_tool_executor()
|
||||||
|
|
||||||
|
# Show updated summary with runtime info
|
||||||
|
show_startup_summary(self.hardware, self.config, self.swarm)
|
||||||
|
|
||||||
|
# Initialize federation if enabled
|
||||||
|
if self.args.federation:
|
||||||
|
await self._setup_federation()
|
||||||
|
|
||||||
|
# Start MCP server if enabled
|
||||||
|
if self.args.mcp:
|
||||||
|
await self._setup_mcp()
|
||||||
|
|
||||||
|
# Run server
|
||||||
|
return await self._run_server()
|
||||||
|
|
||||||
|
async def _setup_swarm(self) -> bool:
|
||||||
|
"""Setup the swarm.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
True if successful
|
||||||
|
"""
|
||||||
|
print("\n⬇️ Downloading model...")
|
||||||
|
try:
|
||||||
|
model_path = download_model_for_config(self.config)
|
||||||
|
print(f"✓ Model ready at: {model_path}")
|
||||||
|
except Exception as e:
|
||||||
|
print(f"\n❌ Error downloading model: {e}", file=sys.stderr)
|
||||||
|
return False
|
||||||
|
|
||||||
|
print("\n🚀 Initializing swarm...")
|
||||||
|
try:
|
||||||
|
self.swarm = SwarmManager(
|
||||||
|
model_config=self.config,
|
||||||
|
hardware=self.hardware,
|
||||||
|
consensus_strategy="similarity"
|
||||||
|
)
|
||||||
|
|
||||||
|
success = await self.swarm.initialize(str(model_path))
|
||||||
|
if not success:
|
||||||
|
print("❌ Failed to initialize swarm")
|
||||||
|
return False
|
||||||
|
|
||||||
|
return True
|
||||||
|
except Exception as e:
|
||||||
|
print(f"\n❌ Error initializing swarm: {e}", file=sys.stderr)
|
||||||
|
return False
|
||||||
|
|
||||||
|
def _setup_tool_executor(self) -> None:
|
||||||
|
"""Setup the tool executor."""
|
||||||
|
if self.args.tool_host is not None:
|
||||||
|
if self.args.tool_host == "":
|
||||||
|
tool_host_url = f"http://{get_local_ip()}:17616"
|
||||||
|
print(f"\n🔧 Using remote tool host: {tool_host_url} (auto-detected)")
|
||||||
|
else:
|
||||||
|
tool_host_url = self.args.tool_host
|
||||||
|
print(f"\n🔧 Using remote tool host: {tool_host_url}")
|
||||||
|
executor = ToolExecutor(tool_host_url=tool_host_url)
|
||||||
|
else:
|
||||||
|
executor = ToolExecutor(tool_host_url=None)
|
||||||
|
print("\n🔧 Tool Server: Local")
|
||||||
|
|
||||||
|
set_tool_executor(executor)
|
||||||
|
|
||||||
|
async def _setup_federation(self) -> None:
|
||||||
|
"""Setup federation."""
|
||||||
|
print("\n🌐 Initializing federation...")
|
||||||
|
try:
|
||||||
|
advertise_ip = self.args.host if self.args.host else None
|
||||||
|
self.discovery = await create_discovery_service(
|
||||||
|
self.args.port,
|
||||||
|
advertise_ip=advertise_ip
|
||||||
|
)
|
||||||
|
|
||||||
|
swarm_info = {
|
||||||
|
"version": "0.1.0",
|
||||||
|
"instances": self.config.instances,
|
||||||
|
"model_id": self.config.model_id,
|
||||||
|
"hardware_summary": f"{self.hardware.cpu_cores} CPU, {self.hardware.ram_gb:.1f}GB RAM"
|
||||||
|
}
|
||||||
|
|
||||||
|
await self.discovery.start_advertising(swarm_info)
|
||||||
|
await self.discovery.start_listening()
|
||||||
|
|
||||||
|
# Add manual peers
|
||||||
|
if self.args.peers:
|
||||||
|
await self._add_manual_peers()
|
||||||
|
|
||||||
|
self.federated_swarm = FederatedSwarm(self.swarm, self.discovery)
|
||||||
|
set_federated_swarm(self.federated_swarm)
|
||||||
|
|
||||||
|
# Start health check loop
|
||||||
|
asyncio.create_task(
|
||||||
|
self.discovery.start_health_check_loop(interval_seconds=10)
|
||||||
|
)
|
||||||
|
|
||||||
|
print(f" ✓ Federation enabled")
|
||||||
|
print(f" ✓ Discovery active on port {self.discovery.discovery_port}")
|
||||||
|
print(f" ✓ Peer health checks every 10s")
|
||||||
|
except Exception as e:
|
||||||
|
print(f" ⚠️ Failed to initialize federation: {e}")
|
||||||
|
print(" Continuing without federation...")
|
||||||
|
|
||||||
|
async def _add_manual_peers(self) -> None:
|
||||||
|
"""Add manual peers from command line."""
|
||||||
|
print(f" 📍 Adding {len(self.args.peers)} manual peer(s)...")
|
||||||
|
from network.discovery import PeerInfo
|
||||||
|
from datetime import datetime
|
||||||
|
|
||||||
|
for peer_str in self.args.peers:
|
||||||
|
try:
|
||||||
|
host, port = peer_str.rsplit(':', 1)
|
||||||
|
port = int(port)
|
||||||
|
peer = PeerInfo(
|
||||||
|
host=host,
|
||||||
|
port=port,
|
||||||
|
name=f"manual_{host}_{port}",
|
||||||
|
version="0.1.0",
|
||||||
|
instances=0,
|
||||||
|
model_id="unknown",
|
||||||
|
hardware_summary="manual",
|
||||||
|
last_seen=datetime.now()
|
||||||
|
)
|
||||||
|
self.discovery.peers[peer.name] = peer
|
||||||
|
print(f" ✓ Added peer: {host}:{port}")
|
||||||
|
except Exception as e:
|
||||||
|
print(f" ⚠️ Failed to add peer {peer_str}: {e}")
|
||||||
|
|
||||||
|
async def _setup_mcp(self) -> None:
|
||||||
|
"""Setup MCP server."""
|
||||||
|
print("\n🤖 Starting MCP server...")
|
||||||
|
from mcp_server import create_mcp_server
|
||||||
|
self.mcp_server = await create_mcp_server(self.swarm)
|
||||||
|
print(" MCP server active (stdio)")
|
||||||
|
|
||||||
|
async def _run_server(self) -> int:
|
||||||
|
"""Run the API server."""
|
||||||
|
print("\n🌐 Starting HTTP API server...")
|
||||||
|
|
||||||
|
# Determine host
|
||||||
|
if self.args.host:
|
||||||
|
host = self.args.host
|
||||||
|
print(f"🔗 Using specified host: {host}:{self.args.port}")
|
||||||
|
else:
|
||||||
|
host = get_local_ip()
|
||||||
|
print(f"🔗 Binding to {host}:{self.args.port}")
|
||||||
|
|
||||||
|
# Show tool mode
|
||||||
|
if self.args.use_opencode_tools:
|
||||||
|
print(f"🔧 Tool mode: opencode tools (~27k tokens)")
|
||||||
|
else:
|
||||||
|
print(f"🔧 Tool mode: local tool server (~125 tokens)")
|
||||||
|
|
||||||
|
# Create server
|
||||||
|
server = create_server(
|
||||||
|
self.swarm,
|
||||||
|
host=host,
|
||||||
|
port=self.args.port,
|
||||||
|
use_opencode_tools=self.args.use_opencode_tools
|
||||||
|
)
|
||||||
|
|
||||||
|
# Print connection info
|
||||||
|
print(f"\n✅ Local Swarm is running!")
|
||||||
|
print(f" API: http://{host}:{self.args.port}/v1")
|
||||||
|
print(f" Health: http://{host}:{self.args.port}/health")
|
||||||
|
|
||||||
|
if self.args.federation and self.discovery:
|
||||||
|
peers = self.discovery.get_peers()
|
||||||
|
print(f"\n🌐 Federation: Enabled")
|
||||||
|
print(f" Discovery port: {self.discovery.discovery_port}")
|
||||||
|
if peers:
|
||||||
|
print(f" Peers discovered: {len(peers)}")
|
||||||
|
|
||||||
|
print(f"\n💡 Configure opencode to use:")
|
||||||
|
print(f' base_url: http://127.0.0.1:{self.args.port}/v1')
|
||||||
|
print(f' api_key: any (not used)')
|
||||||
|
print(f"\nPress Ctrl+C to stop...\n")
|
||||||
|
|
||||||
|
# Start server
|
||||||
|
try:
|
||||||
|
await server.start()
|
||||||
|
finally:
|
||||||
|
await self._shutdown()
|
||||||
|
|
||||||
|
return 0
|
||||||
|
|
||||||
|
async def _shutdown(self) -> None:
|
||||||
|
"""Shutdown all services."""
|
||||||
|
if self.federated_swarm:
|
||||||
|
await self.federated_swarm.close()
|
||||||
|
if self.discovery:
|
||||||
|
await self.discovery.stop()
|
||||||
|
if self.swarm:
|
||||||
|
await self.swarm.shutdown()
|
||||||
@@ -0,0 +1,151 @@
|
|||||||
|
"""CLI argument parsing for Local Swarm."""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
|
||||||
|
def create_parser() -> argparse.ArgumentParser:
|
||||||
|
"""Create and configure the argument parser."""
|
||||||
|
parser = argparse.ArgumentParser(
|
||||||
|
description="Local Swarm - AI-powered coding LLM swarm",
|
||||||
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||||
|
epilog="""
|
||||||
|
Examples:
|
||||||
|
python main.py # Interactive setup and start
|
||||||
|
python main.py --auto # Auto-detect and start without menu
|
||||||
|
python main.py --detect # Show hardware detection only
|
||||||
|
python main.py --model qwen:3b:q4 # Use specific model (skip menu)
|
||||||
|
python main.py --port 17615 # Use custom port (default: 17615)
|
||||||
|
python main.py --host 192.168.1.5 # Bind to specific IP
|
||||||
|
python main.py --instances 4 # Force number of instances
|
||||||
|
python main.py --download-only # Download model only
|
||||||
|
python main.py --test # Test with sample prompt
|
||||||
|
python main.py --mcp # Enable MCP server
|
||||||
|
python main.py --federation # Enable federation with other instances
|
||||||
|
python main.py --federation --peer 192.168.1.10:17615 # Manual peer
|
||||||
|
"""
|
||||||
|
)
|
||||||
|
|
||||||
|
# Mode options
|
||||||
|
parser.add_argument(
|
||||||
|
"--auto",
|
||||||
|
action="store_true",
|
||||||
|
help="Auto-detect best configuration without interactive menu"
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--detect",
|
||||||
|
action="store_true",
|
||||||
|
help="Show hardware detection and exit"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Model options
|
||||||
|
parser.add_argument(
|
||||||
|
"--model",
|
||||||
|
type=str,
|
||||||
|
help="Model to use (format: name:size:quant, e.g., qwen:3b:q4)"
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--instances",
|
||||||
|
type=int,
|
||||||
|
help="Force number of instances (overrides auto-calculation)"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Server options
|
||||||
|
parser.add_argument(
|
||||||
|
"--port",
|
||||||
|
type=int,
|
||||||
|
default=17615,
|
||||||
|
help="Port to run the API server on (default: 17615)"
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--host",
|
||||||
|
type=str,
|
||||||
|
default=None,
|
||||||
|
help="Host IP to bind to (default: auto-detect)"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Operation modes
|
||||||
|
parser.add_argument(
|
||||||
|
"--download-only",
|
||||||
|
action="store_true",
|
||||||
|
help="Download models only, don't start server"
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--test",
|
||||||
|
action="store_true",
|
||||||
|
help="Test with a sample prompt"
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--mcp",
|
||||||
|
action="store_true",
|
||||||
|
help="Enable MCP server alongside HTTP API"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Configuration
|
||||||
|
parser.add_argument(
|
||||||
|
"--config",
|
||||||
|
type=str,
|
||||||
|
default="config.yaml",
|
||||||
|
help="Path to config file"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Federation options
|
||||||
|
parser.add_argument(
|
||||||
|
"--federation",
|
||||||
|
action="store_true",
|
||||||
|
help="Enable federation with other Local Swarm instances on the network"
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--peer",
|
||||||
|
action="append",
|
||||||
|
dest="peers",
|
||||||
|
help="Manually add a peer (format: host:port, can be used multiple times)"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Tool server options
|
||||||
|
parser.add_argument(
|
||||||
|
"--tool-server",
|
||||||
|
action="store_true",
|
||||||
|
help="Run as dedicated tool execution server (executes read/write/bash tools)"
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--tool-port",
|
||||||
|
type=int,
|
||||||
|
default=17616,
|
||||||
|
help="Port for tool execution server (default: 17616)"
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--tool-host",
|
||||||
|
type=str,
|
||||||
|
default=None,
|
||||||
|
nargs='?',
|
||||||
|
const='',
|
||||||
|
help="URL of tool execution server. Use without value for auto-detected local IP"
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--use-opencode-tools",
|
||||||
|
action="store_true",
|
||||||
|
help="Use opencode's tool definitions (~27k tokens). Default: use local tool server"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Version
|
||||||
|
parser.add_argument(
|
||||||
|
"--version",
|
||||||
|
action="version",
|
||||||
|
version="%(prog)s 0.1.0"
|
||||||
|
)
|
||||||
|
|
||||||
|
return parser
|
||||||
|
|
||||||
|
|
||||||
|
def parse_args(args: Optional[list] = None):
|
||||||
|
"""Parse command line arguments.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
args: Command line arguments (defaults to sys.argv)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Parsed arguments namespace
|
||||||
|
"""
|
||||||
|
parser = create_parser()
|
||||||
|
return parser.parse_args(args)
|
||||||
@@ -0,0 +1,81 @@
|
|||||||
|
"""Test mode runner for Local Swarm."""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
from models.downloader import download_model_for_config
|
||||||
|
from swarm import SwarmManager
|
||||||
|
from interactive import show_startup_summary
|
||||||
|
|
||||||
|
|
||||||
|
async def run_test(hardware, config) -> int:
|
||||||
|
"""Run test mode with sample prompt.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
hardware: Hardware profile
|
||||||
|
config: Model configuration
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Exit code (0 for success, 1 for error)
|
||||||
|
"""
|
||||||
|
print("\n" + "=" * 70)
|
||||||
|
print("🧪 Test Mode: Running sample inference")
|
||||||
|
print("=" * 70)
|
||||||
|
|
||||||
|
show_startup_summary(hardware, config)
|
||||||
|
|
||||||
|
# Download model
|
||||||
|
print("\n⬇️ Downloading model...")
|
||||||
|
try:
|
||||||
|
model_path = download_model_for_config(config)
|
||||||
|
print(f"✓ Model ready at: {model_path}")
|
||||||
|
except Exception as e:
|
||||||
|
print(f"\n❌ Error downloading model: {e}")
|
||||||
|
return 1
|
||||||
|
|
||||||
|
# Initialize swarm
|
||||||
|
print("\n🚀 Initializing swarm...")
|
||||||
|
try:
|
||||||
|
swarm = SwarmManager(
|
||||||
|
model_config=config,
|
||||||
|
hardware=hardware,
|
||||||
|
consensus_strategy="similarity"
|
||||||
|
)
|
||||||
|
|
||||||
|
success = await swarm.initialize(str(model_path))
|
||||||
|
if not success:
|
||||||
|
print("❌ Failed to initialize swarm")
|
||||||
|
return 1
|
||||||
|
except Exception as e:
|
||||||
|
print(f"\n❌ Error initializing swarm: {e}")
|
||||||
|
return 1
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Test prompt
|
||||||
|
prompt = "Write a Python function to calculate factorial:"
|
||||||
|
print(f"\nPrompt: {prompt}\n")
|
||||||
|
print("Generating responses...\n")
|
||||||
|
|
||||||
|
result = await swarm.generate(prompt, max_tokens=200)
|
||||||
|
|
||||||
|
print("\n" + "=" * 70)
|
||||||
|
print("SELECTED RESPONSE:")
|
||||||
|
print("=" * 70)
|
||||||
|
print(result.selected_response.text)
|
||||||
|
print("\n" + "=" * 70)
|
||||||
|
print(f"Strategy: {result.strategy}")
|
||||||
|
print(f"Confidence: {result.confidence:.2f}")
|
||||||
|
print(f"Latency: {result.selected_response.latency_ms:.1f}ms")
|
||||||
|
print(f"Tokens/sec: {result.selected_response.tokens_per_second:.1f}")
|
||||||
|
|
||||||
|
# Show all responses
|
||||||
|
print("\nAll responses received:")
|
||||||
|
for i, resp in enumerate(result.all_responses):
|
||||||
|
preview = resp.text[:60].replace('\n', ' ')
|
||||||
|
print(f" Worker {i}: {preview}... ({resp.latency_ms:.1f}ms)")
|
||||||
|
|
||||||
|
print("\n" + "=" * 70)
|
||||||
|
print("✅ Test complete")
|
||||||
|
print("=" * 70)
|
||||||
|
return 0
|
||||||
|
|
||||||
|
finally:
|
||||||
|
await swarm.shutdown()
|
||||||
@@ -0,0 +1,69 @@
|
|||||||
|
"""Tool server for Local Swarm.
|
||||||
|
|
||||||
|
Standalone tool execution server for distributed setups.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
from fastapi import FastAPI
|
||||||
|
import uvicorn
|
||||||
|
|
||||||
|
from tools.executor import ToolExecutor, set_tool_executor
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def create_tool_server_app() -> FastAPI:
|
||||||
|
"""Create the tool server FastAPI application.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Configured FastAPI application
|
||||||
|
"""
|
||||||
|
app = FastAPI(title="Local Swarm Tool Server")
|
||||||
|
|
||||||
|
@app.post("/v1/tools/execute")
|
||||||
|
async def execute_tool(request: dict):
|
||||||
|
tool_name = request.get("tool", "")
|
||||||
|
tool_args = request.get("arguments", {})
|
||||||
|
|
||||||
|
# Get the global executor
|
||||||
|
from tools.executor import get_tool_executor
|
||||||
|
executor = get_tool_executor()
|
||||||
|
|
||||||
|
if executor is None:
|
||||||
|
return {"result": "Error: No tool executor configured"}
|
||||||
|
|
||||||
|
result = await executor.execute(tool_name, tool_args)
|
||||||
|
return {"result": result}
|
||||||
|
|
||||||
|
@app.get("/health")
|
||||||
|
async def health():
|
||||||
|
return {"status": "healthy", "mode": "tool-server"}
|
||||||
|
|
||||||
|
return app
|
||||||
|
|
||||||
|
|
||||||
|
async def run_tool_server(host: str, port: int) -> None:
|
||||||
|
"""Run the tool server.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
host: Host to bind to
|
||||||
|
port: Port to listen on
|
||||||
|
"""
|
||||||
|
# Initialize local tool executor
|
||||||
|
tool_executor = ToolExecutor(tool_host_url=None)
|
||||||
|
set_tool_executor(tool_executor)
|
||||||
|
|
||||||
|
app = create_tool_server_app()
|
||||||
|
|
||||||
|
print(f"🔗 Tool server running at http://{host}:{port}")
|
||||||
|
print(f" Endpoints:")
|
||||||
|
print(f" - POST /v1/tools/execute")
|
||||||
|
print(f" - GET /health")
|
||||||
|
print(f"\n✅ Tool server ready!")
|
||||||
|
|
||||||
|
config = uvicorn.Config(app, host=host, port=port, log_level="warning")
|
||||||
|
server = uvicorn.Server(config)
|
||||||
|
await server.serve()
|
||||||
@@ -0,0 +1,45 @@
|
|||||||
|
"""Network utilities for Local Swarm."""
|
||||||
|
|
||||||
|
import socket
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
|
||||||
|
def get_local_ip() -> str:
|
||||||
|
"""Get the local network IP address (private networks only).
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Local IP address or 127.0.0.1 if detection fails
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
# Create a socket and connect to a public DNS server
|
||||||
|
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||||||
|
s.settimeout(2)
|
||||||
|
# Try to connect to Google's DNS - this doesn't actually send data
|
||||||
|
s.connect(("8.8.8.8", 80))
|
||||||
|
ip = s.getsockname()[0]
|
||||||
|
s.close()
|
||||||
|
|
||||||
|
# Check if it's a private IP
|
||||||
|
is_private = ip.startswith('192.168.')
|
||||||
|
|
||||||
|
if is_private:
|
||||||
|
print(f" 📡 Detected local IP: {ip}")
|
||||||
|
return ip
|
||||||
|
else:
|
||||||
|
print(f" ⚠️ IP {ip} is not a private network, binding to localhost")
|
||||||
|
return "127.0.0.1"
|
||||||
|
except Exception as e:
|
||||||
|
print(f" ⚠️ Could not detect local IP: {e}, using localhost")
|
||||||
|
return "127.0.0.1"
|
||||||
|
|
||||||
|
|
||||||
|
def is_private_ip(ip: str) -> bool:
|
||||||
|
"""Check if an IP address is private.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
ip: IP address string
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
True if IP is private
|
||||||
|
"""
|
||||||
|
return ip.startswith('192.168.') or ip.startswith('10.') or ip.startswith('172.16.')
|
||||||
Reference in New Issue
Block a user