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:
@@ -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()
|
||||
Reference in New Issue
Block a user