sleepy
4ea36783d6
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.
321 lines
11 KiB
Python
321 lines
11 KiB
Python
"""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()
|