golem-vm-provider 0.1.26__py3-none-any.whl → 0.1.28__py3-none-any.whl

provider/discovery/advertiser.py

@@ -1,90 +1,71 @@
 import aiohttp
 import asyncio
 import logging
-import psutil
-from datetime import datetime
-from typing import Dict, Optional
+from abc import ABC, abstractmethod
+from typing import Optional
 
 from ..config import settings
 from ..utils.retry import async_retry
+from .resource_tracker import ResourceTracker
 
 logger = logging.getLogger(__name__)
 
-class ResourceMonitor:
-    """Monitor system resources."""
-    
-    @staticmethod
-    def get_cpu_count() -> int:
-        """Get number of CPU cores."""
-        return psutil.cpu_count()
-
-    @staticmethod
-    def get_memory_gb() -> int:
-        """Get available memory in GB."""
-        return psutil.virtual_memory().available // (1024 ** 3)
-
-    @staticmethod
-    def get_storage_gb() -> int:
-        """Get available storage in GB."""
-        return psutil.disk_usage("/").free // (1024 ** 3)
-
-    @staticmethod
-    def get_cpu_percent() -> float:
-        """Get CPU usage percentage."""
-        return psutil.cpu_percent(interval=1)
-
-    @staticmethod
-    def get_memory_percent() -> float:
-        """Get memory usage percentage."""
-        return psutil.virtual_memory().percent
-
-    @staticmethod
-    def get_storage_percent() -> float:
-        """Get storage usage percentage."""
-        return psutil.disk_usage("/").percent
-
-class ResourceAdvertiser:
-    """Advertise available resources to discovery service."""
+class Advertiser(ABC):
+    """Abstract base class for advertisers."""
+
+    @abstractmethod
+    async def initialize(self):
+        """Initialize the advertiser."""
+        pass
+
+    @abstractmethod
+    async def start_loop(self):
+        """Start the advertising loop."""
+        pass
+
+    @abstractmethod
+    async def stop(self):
+        """Stop the advertising loop."""
+        pass
+
+    @abstractmethod
+    async def post_advertisement(self):
+        """Post a single advertisement."""
+        pass
+
+class DiscoveryServerAdvertiser(Advertiser):
+    """Advertise available resources to a discovery service."""
     
     def __init__(
         self,
         resource_tracker: 'ResourceTracker',
         discovery_url: Optional[str] = None,
         provider_id: Optional[str] = None,
-        update_interval: Optional[int] = None
     ):
         self.resource_tracker = resource_tracker
         self.discovery_url = discovery_url or settings.DISCOVERY_URL
         self.provider_id = provider_id or settings.PROVIDER_ID
-        self.update_interval = update_interval or settings.ADVERTISEMENT_INTERVAL
         self.session: Optional[aiohttp.ClientSession] = None
         self._stop_event = asyncio.Event()
 
-    async def start(self):
-        """Start advertising resources."""
+    async def initialize(self):
+        """Initialize the advertiser."""
         self.session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=10))
-        # Register for resource updates
-        self.resource_tracker.on_update(self._post_advertisement)
-        
-        # Test discovery service connection with retries
+        self.resource_tracker.on_update(
+            lambda: asyncio.create_task(self.post_advertisement())
+        )
         try:
             await self._check_discovery_health()
         except Exception as e:
             logger.warning(f"Could not connect to discovery service after retries, continuing without advertising: {e}")
             return
-            
+
+    async def start_loop(self):
+        """Start advertising resources in a loop."""
         try:
             while not self._stop_event.is_set():
-                try:
-                    await self._post_advertisement()
-                except aiohttp.ClientError as e:
-                    logger.error(f"Network error posting advertisement: {e}")
-                    await asyncio.sleep(min(60, self.update_interval))
-                except Exception as e:
-                    logger.error(f"Failed to post advertisement: {e}")
-                    await asyncio.sleep(min(60, self.update_interval))
-                else:
-                    await asyncio.sleep(self.update_interval)
+                await self.post_advertisement()
+                await asyncio.sleep(settings.ADVERTISEMENT_INTERVAL)
         finally:
             await self.stop()
 
@@ -106,19 +87,17 @@ class ResourceAdvertiser:
                 raise Exception(f"Discovery service health check failed: {response.status}")
 
     @async_retry(retries=3, delay=1.0, backoff=2.0, exceptions=(aiohttp.ClientError, asyncio.TimeoutError))
-    async def _post_advertisement(self):
+    async def post_advertisement(self):
         """Post resource advertisement to discovery service."""
         if not self.session:
             raise RuntimeError("Session not initialized")
 
         resources = self.resource_tracker.get_available_resources()
         
-        # Don't advertise if resources are too low
         if not self.resource_tracker._meets_minimum_requirements(resources):
             logger.warning("Resources too low, skipping advertisement")
             return
 
-        # Get public IP with retries
         try:
             ip_address = await self._get_public_ip()
         except Exception as e:
@@ -130,7 +109,7 @@ class ResourceAdvertiser:
                 f"{self.discovery_url}/api/v1/advertisements",
                 headers={
                     "X-Provider-ID": self.provider_id,
-                    "X-Provider-Signature": "signature",  # TODO: Implement signing
+                    "X-Provider-Signature": "signature",
                     "Content-Type": "application/json"
                 },
                 json={
@@ -138,7 +117,7 @@ class ResourceAdvertiser:
                     "country": settings.PROVIDER_COUNTRY,
                     "resources": resources
                 },
-                timeout=aiohttp.ClientTimeout(total=5)  # 5 second timeout for advertisement
+                timeout=aiohttp.ClientTimeout(total=5)
             ) as response:
                 if not response.ok:
                     error_text = await response.text()
@@ -159,7 +138,6 @@ class ResourceAdvertiser:
         if not self.session:
             raise RuntimeError("Session not initialized")
 
-        # Try multiple IP services in case one fails
         services = [
             "https://api.ipify.org",
             "https://ifconfig.me/ip",

provider/discovery/golem_base_advertiser.py

@@ -2,6 +2,7 @@ import asyncio
 from typing import Optional
 
 from golem_base_sdk import GolemBaseClient, GolemBaseCreate, GolemBaseUpdate, GolemBaseDelete, Annotation
+from .advertiser import Advertiser
 from .golem_base_utils import get_provider_entity_keys
 from ..config import settings
 from ..utils.logging import setup_logger
@@ -9,7 +10,7 @@ from ..utils.logging import setup_logger
 logger = setup_logger(__name__)
 
 
-class GolemBaseAdvertiser:
+class GolemBaseAdvertiser(Advertiser):
     """Advertise available resources to the Golem Base network."""
 
     def __init__(self, resource_tracker: "ResourceTracker"):
@@ -89,11 +90,16 @@ class GolemBaseAdvertiser:
                 current_annotations = {ann.key: ann.value for ann in metadata.numeric_annotations}
                 current_annotations.update({ann.key: ann.value for ann in metadata.string_annotations})
 
-                if (current_annotations.get("golem_cpu") == resources["cpu"] and
-                    current_annotations.get("golem_memory") == resources["memory"] and
-                    current_annotations.get("golem_storage") == resources["storage"] and
-                    current_annotations.get("golem_ip_address") == ip_address and
-                    current_annotations.get("golem_provider_name") == settings.PROVIDER_NAME):
+                # Full comparison of all annotations
+                expected_annotations = {ann.key: ann.value for ann in string_annotations}
+                expected_annotations.update({ann.key: ann.value for ann in numeric_annotations})
+
+                # Debugging logs to compare annotations
+                logger.info(f"IP address from settings: {ip_address}")
+                logger.info(f"Current on-chain annotations: {current_annotations}")
+                logger.info(f"Expected annotations based on current config: {expected_annotations}")
+
+                if sorted(current_annotations.items()) == sorted(expected_annotations.items()):
                     logger.info("Advertisement is up-to-date. Waiting for expiration.")
                 else:
                     logger.info("Advertisement is outdated. Updating.")

provider/discovery/resource_monitor.py

@@ -0,0 +1,34 @@
+import psutil
+
+class ResourceMonitor:
+    """Monitor system resources."""
+    
+    @staticmethod
+    def get_cpu_count() -> int:
+        """Get number of CPU cores."""
+        return psutil.cpu_count()
+
+    @staticmethod
+    def get_memory_gb() -> int:
+        """Get available memory in GB."""
+        return psutil.virtual_memory().available // (1024 ** 3)
+
+    @staticmethod
+    def get_storage_gb() -> int:
+        """Get available storage in GB."""
+        return psutil.disk_usage("/").free // (1024 ** 3)
+
+    @staticmethod
+    def get_cpu_percent() -> float:
+        """Get CPU usage percentage."""
+        return psutil.cpu_percent(interval=1)
+
+    @staticmethod
+    def get_memory_percent() -> float:
+        """Get memory usage percentage."""
+        return psutil.virtual_memory().percent
+
+    @staticmethod
+    def get_storage_percent() -> float:
+        """Get storage usage percentage."""
+        return psutil.disk_usage("/").percent

provider/discovery/resource_tracker.py

@@ -11,7 +11,7 @@ class ResourceTracker:
 
     def __init__(self):
         """Initialize resource tracker."""
-        from .advertiser import ResourceMonitor
+        from .resource_monitor import ResourceMonitor
         self.total_resources = {
             "cpu": ResourceMonitor.get_cpu_count(),
             "memory": ResourceMonitor.get_memory_gb(),

provider/discovery/service.py

@@ -0,0 +1,24 @@
+import asyncio
+from typing import Optional
+
+from .advertiser import Advertiser
+from ..config import settings
+
+class AdvertisementService:
+    """Service for managing the advertisement lifecycle."""
+
+    def __init__(self, advertiser: Advertiser):
+        self.advertiser = advertiser
+        self._task: Optional[asyncio.Task] = None
+
+    async def start(self):
+        """Initialize and start the advertiser."""
+        await self.advertiser.initialize()
+        self._task = asyncio.create_task(self.advertiser.start_loop())
+
+    async def stop(self):
+        """Stop the advertiser."""
+        if self._task:
+            self._task.cancel()
+            await self._task
+        await self.advertiser.stop()

provider/main.py

@@ -1,6 +1,7 @@
 from .api import routes
 import asyncio
 import os
+import socket
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from typing import Optional
@@ -9,15 +10,35 @@ from .config import settings
 from .utils.logging import setup_logger, PROCESS, SUCCESS
 from .utils.ascii_art import startup_animation
 from .discovery.resource_tracker import ResourceTracker
-from .discovery.advertiser import ResourceAdvertiser
-from .discovery.golem_base_advertiser import GolemBaseAdvertiser
-from .vm.multipass import MultipassProvider
-from .vm.port_manager import PortManager
-from .security.faucet import FaucetClient
-
-logger = setup_logger(__name__)
+from .discovery.advertiser import DiscoveryServerAdvertiser
+from .container import Container
+from .service import ProviderService
+from .utils.logging import setup_logger
 
 app = FastAPI(title="VM on Golem Provider")
+container = Container()
+container.config.from_pydantic(settings)
+app.container = container
+container.wire(modules=[".api.routes"])
+
+from .vm.models import VMNotFoundError
+from fastapi import Request
+from fastapi.responses import JSONResponse
+
+@app.exception_handler(VMNotFoundError)
+async def vm_not_found_exception_handler(request: Request, exc: VMNotFoundError):
+    return JSONResponse(
+        status_code=404,
+        content={"message": str(exc)},
+    )
+
+@app.exception_handler(Exception)
+async def generic_exception_handler(request: Request, exc: Exception):
+    return JSONResponse(
+        status_code=500,
+        content={"message": "An unexpected error occurred"},
+    )
+
 # Add CORS middleware
 app.add_middleware(
     CORSMiddleware,
@@ -27,156 +48,18 @@ app.add_middleware(
     allow_headers=["*"],  # Allows all headers
 )
 
-
-async def setup_provider() -> None:
-    """Setup and initialize the provider components."""
-    try:
-        # Create resource tracker first
-        logger.process("🔄 Initializing resource tracker...")
-        resource_tracker = ResourceTracker()
-        app.state.resource_tracker = resource_tracker
-
-        # Create provider with resource tracker and temporary port manager
-        logger.process("🔄 Initializing VM provider...")
-        provider = MultipassProvider(
-            resource_tracker, port_manager=None)  # Will be set later
-
-        try:
-            # Initialize provider (without port operations)
-            await asyncio.wait_for(provider.initialize(), timeout=30)
-
-            # Store provider reference
-            app.state.provider = provider
-            app.state.proxy_manager = provider.proxy_manager
-
-            # Initialize port manager first to verify all ports
-            logger.process("🔄 Initializing port manager...")
-            port_manager = PortManager(
-                start_port=settings.PORT_RANGE_START,
-                end_port=settings.PORT_RANGE_END,
-                discovery_port=settings.PORT,
-                skip_verification=settings.SKIP_PORT_VERIFICATION
-            )
-
-            if not await port_manager.initialize():
-                raise RuntimeError("Port verification failed")
-
-            # Store port manager references
-            app.state.port_manager = port_manager
-            provider.port_manager = port_manager
-            app.state.proxy_manager.port_manager = port_manager
-
-            # Now restore proxy configurations using only verified ports
-            logger.process("🔄 Restoring proxy configurations...")
-            await app.state.proxy_manager._load_state()
-
-        except asyncio.TimeoutError:
-            logger.error("Provider initialization timed out")
-            raise
-        except Exception as e:
-            logger.error(f"Failed to initialize provider: {e}")
-            raise
-
-        # Create advertiser
-        logger.process("🔄 Initializing resource advertiser...")
-        if settings.DISCOVERY_DRIVER == "golem-base":
-            advertiser = GolemBaseAdvertiser(
-                resource_tracker=resource_tracker
-            )
-            await advertiser.initialize()
-        else:
-            advertiser = ResourceAdvertiser(
-                resource_tracker=resource_tracker,
-                discovery_url=settings.DISCOVERY_URL,
-                provider_id=settings.PROVIDER_ID
-            )
-        app.state.advertiser = advertiser
-
-        logger.success(
-            "✨ Provider setup complete and ready to accept requests")
-    except Exception as e:
-        logger.error(f"Failed to setup provider: {e}")
-        # Attempt cleanup of any initialized components
-        await cleanup_provider()
-        raise
-
-
-async def cleanup_provider() -> None:
-    """Cleanup provider components."""
-    cleanup_errors = []
-
-    # Stop advertiser
-    if hasattr(app.state, "advertiser"):
-        try:
-            await app.state.advertiser.stop()
-            if hasattr(app.state, "advertiser_task"):
-                app.state.advertiser_task.cancel()
-                try:
-                    await app.state.advertiser_task
-                except asyncio.CancelledError:
-                    pass
-        except Exception as e:
-            cleanup_errors.append(f"Failed to stop advertiser: {e}")
-
-    # Cleanup proxy manager first to stop all proxy servers
-    if hasattr(app.state, "proxy_manager"):
-        try:
-            await asyncio.wait_for(app.state.proxy_manager.cleanup(), timeout=30)
-        except asyncio.TimeoutError:
-            cleanup_errors.append("Proxy manager cleanup timed out")
-        except Exception as e:
-            cleanup_errors.append(f"Failed to cleanup proxy manager: {e}")
-
-    # Cleanup provider
-    if hasattr(app.state, "provider"):
-        try:
-            await asyncio.wait_for(app.state.provider.cleanup(), timeout=30)
-        except asyncio.TimeoutError:
-            cleanup_errors.append("Provider cleanup timed out")
-        except Exception as e:
-            cleanup_errors.append(f"Failed to cleanup provider: {e}")
-
-    if cleanup_errors:
-        error_msg = "\n".join(cleanup_errors)
-        logger.error(f"Errors during cleanup:\n{error_msg}")
-    else:
-        logger.success("✨ Provider cleanup complete")
-
-
 @app.on_event("startup")
 async def startup_event():
     """Handle application startup."""
-    try:
-        # Display startup animation
-        await startup_animation()
- 
-        # Initialize provider
-        await setup_provider()
-
-        # Check wallet balance and request funds if needed
-        faucet_client = FaucetClient(
-            faucet_url=settings.FAUCET_URL,
-            captcha_url=settings.CAPTCHA_URL,
-            captcha_api_key=settings.CAPTCHA_API_KEY,
-        )
-        await faucet_client.get_funds(settings.PROVIDER_ID)
-
-        # Post initial advertisement and start advertising loop
-        if isinstance(app.state.advertiser, GolemBaseAdvertiser):
-            await app.state.advertiser.post_advertisement()
-            app.state.advertiser_task = asyncio.create_task(app.state.advertiser.start_loop())
-
-    except Exception as e:
-        logger.error(f"Startup failed: {e}")
-        # Ensure proper cleanup
-        await cleanup_provider()
-        raise
+    provider_service = container.provider_service()
+    await provider_service.setup(app)
 
 
 @app.on_event("shutdown")
 async def shutdown_event():
     """Handle application shutdown."""
-    await cleanup_provider()
+    provider_service = container.provider_service()
+    await provider_service.cleanup()
 
 # Import routes after app creation to avoid circular imports
 app.include_router(routes.router, prefix="/api/v1")
@@ -225,6 +108,9 @@ async def verify_provider_port(port: int) -> bool:
         return False
 
 
+# The get_local_ip function has been removed as this logic is now handled in config.py
+
+
 import typer
 
 cli = typer.Typer()
@@ -232,21 +118,36 @@ cli = typer.Typer()
 @cli.command()
 def start(no_verify_port: bool = typer.Option(False, "--no-verify-port", help="Skip provider port verification.")):
     """Start the provider server."""
+    run_server(dev_mode=False, no_verify_port=no_verify_port)
+
+@cli.command()
+def dev(no_verify_port: bool = typer.Option(True, "--no-verify-port", help="Skip provider port verification.")):
+    """Start the provider server in development mode."""
+    run_server(dev_mode=True, no_verify_port=no_verify_port)
+
+def run_server(dev_mode: bool, no_verify_port: bool):
+    """Helper to run the uvicorn server."""
     import sys
     from pathlib import Path
     from dotenv import load_dotenv
     import uvicorn
     from .utils.logging import setup_logger
+    
+    # Load appropriate .env file
+    env_file = ".env.dev" if dev_mode else ".env"
+    env_path = Path(__file__).parent.parent / env_file
+    load_dotenv(dotenv_path=env_path)
+    
+    # The logic for setting the public IP in dev mode is now handled in config.py
+    # The following lines are no longer needed and have been removed.
+
+    # Import settings after loading env
     from .config import settings
 
     # Configure logging with debug mode
-    logger = setup_logger(__name__, debug=True)
+    logger = setup_logger(__name__, debug=dev_mode)
 
     try:
-        # Load environment variables from .env file
-        env_path = Path(__file__).parent.parent / '.env'
-        load_dotenv(dotenv_path=env_path)
-
         # Log environment variables
         logger.info("Environment variables:")
         for key, value in os.environ.items():
@@ -275,7 +176,7 @@ def start(no_verify_port: bool = typer.Option(False, "--no-verify-port", help="S
             host=settings.HOST,
             port=settings.PORT,
             reload=settings.DEBUG,
-            log_level="info" if not settings.DEBUG else "debug",
+            log_level="debug" if dev_mode else "info",
             log_config=log_config,
             timeout_keep_alive=60,  # Increase keep-alive timeout
             limit_concurrency=100,  # Limit concurrent connections

provider/service.py

@@ -0,0 +1,67 @@
+import asyncio
+from fastapi import FastAPI
+
+from .utils.logging import setup_logger
+from .vm.service import VMService
+from .discovery.service import AdvertisementService
+
+logger = setup_logger(__name__)
+
+
+class ProviderService:
+    """Service for managing the provider's lifecycle."""
+
+    def __init__(self, vm_service: VMService, advertisement_service: AdvertisementService, port_manager):
+        self.vm_service = vm_service
+        self.advertisement_service = advertisement_service
+        self.port_manager = port_manager
+
+    async def setup(self, app: FastAPI):
+        """Setup and initialize the provider components."""
+        from .config import settings
+        from .utils.ascii_art import startup_animation
+        from .security.faucet import FaucetClient
+
+        try:
+            # Display startup animation
+            await startup_animation()
+
+            logger.process("🔄 Initializing provider...")
+
+            # Setup directories
+            self._setup_directories()
+
+            # Initialize services
+            await self.port_manager.initialize()
+            await self.vm_service.provider.initialize()
+            await self.advertisement_service.start()
+
+            # Check wallet balance and request funds if needed
+            faucet_client = FaucetClient(
+                faucet_url=settings.FAUCET_URL,
+                captcha_url=settings.CAPTCHA_URL,
+                captcha_api_key=settings.CAPTCHA_API_KEY,
+            )
+            await faucet_client.get_funds(settings.PROVIDER_ID)
+
+            logger.success("✨ Provider setup complete")
+        except Exception as e:
+            logger.error(f"Startup failed: {e}")
+            await self.cleanup()
+            raise
+
+    async def cleanup(self):
+        """Cleanup provider components."""
+        logger.process("🔄 Cleaning up provider...")
+        await self.advertisement_service.stop()
+        await self.vm_service.provider.cleanup()
+        logger.success("✨ Provider cleanup complete")
+
+    def _setup_directories(self):
+        """Create necessary directories for the provider."""
+        from .config import settings
+        from pathlib import Path
+        
+        Path(settings.VM_DATA_DIR).mkdir(parents=True, exist_ok=True)
+        Path(settings.SSH_KEY_DIR).mkdir(parents=True, exist_ok=True)
+        Path(settings.CLOUD_INIT_DIR).mkdir(parents=True, exist_ok=True)

provider/utils/logging.py

@@ -39,42 +39,26 @@ def setup_logger(name: Optional[str] = None, debug: bool = False) -> logging.Log
         Configured logger instance
     """
     logger = logging.getLogger(name or __name__)
-    logger.handlers = []  # Clear existing handlers
-    
-    # Fancy handler for important logs
-    fancy_handler = colorlog.StreamHandler(sys.stdout)
-    fancy_formatter = colorlog.ColoredFormatter(
-        "%(log_color)s[%(asctime)s] %(message)s",
+    if logger.handlers:
+        return logger  # Already configured
+
+    handler = colorlog.StreamHandler(sys.stdout)
+    formatter = colorlog.ColoredFormatter(
+        "%(log_color)s[%(asctime)s] %(levelname)s: %(message)s",
         datefmt="%Y-%m-%d %H:%M:%S",
-        reset=True,
         log_colors={
+            'DEBUG': 'cyan',
             'INFO': 'green',
             'PROCESS': 'yellow',
             'WARNING': 'yellow',
             'SUCCESS': 'green,bold',
             'ERROR': 'red',
             'CRITICAL': 'red,bold',
-        },
-        secondary_log_colors={},
-        style='%'
+        }
     )
-    fancy_handler.setFormatter(fancy_formatter)
-    fancy_handler.addFilter(lambda record: record.levelno in [INFO, PROCESS, SUCCESS, WARNING, ERROR, CRITICAL])
-    logger.addHandler(fancy_handler)
-    
-    if debug:
-        # Debug handler for detailed logs
-        debug_handler = logging.StreamHandler(sys.stdout)
-        debug_formatter = logging.Formatter(
-            '%(asctime)s - %(name)s - %(levelname)s - %(message)s',
-            datefmt="%Y-%m-%d %H:%M:%S"
-        )
-        debug_handler.setFormatter(debug_formatter)
-        debug_handler.addFilter(lambda record: record.levelno == DEBUG)
-        logger.addHandler(debug_handler)
-        logger.setLevel(logging.DEBUG)
-    else:
-        logger.setLevel(logging.INFO)
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+    logger.setLevel(logging.DEBUG if debug else logging.INFO)
     
     return logger