diff --git a/src/wavespeed/serverless/__init__.py b/src/wavespeed/serverless/__init__.py index 407b500..ecac156 100644 --- a/src/wavespeed/serverless/__init__.py +++ b/src/wavespeed/serverless/__init__.py @@ -56,8 +56,8 @@ def _parse_args() -> argparse.Namespace: "--waverless_api_host", "--rp_api_host", type=str, - default="localhost", - help="API server host (default: localhost)", + default="0.0.0.0", + help="API server host (default: 0.0.0.0)", ) parser.add_argument( "--waverless_api_port", @@ -160,6 +160,17 @@ def handler(job): port=args.waverless_api_port, ) else: - # Standard worker mode + # Standard worker mode with health check server log.info("Starting serverless worker...") - run_worker(config) + from .modules.health import HealthServer + + health_server = HealthServer( + host=args.waverless_api_host, + port=args.waverless_api_port, + ) + health_server.start() + + try: + run_worker(config) + finally: + health_server.stop() diff --git a/src/wavespeed/serverless/modules/health.py b/src/wavespeed/serverless/modules/health.py new file mode 100644 index 0000000..b63e686 --- /dev/null +++ b/src/wavespeed/serverless/modules/health.py @@ -0,0 +1,59 @@ +"""Lightweight health check HTTP server for worker mode.""" + +import threading +from http.server import BaseHTTPRequestHandler, HTTPServer +from typing import Optional + +from .logger import log + + +class HealthHandler(BaseHTTPRequestHandler): + """Simple HTTP handler for health checks.""" + + def do_GET(self) -> None: + """Handle GET requests.""" + if self.path == "/health": + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(b'{"status": "ok"}') + else: + self.send_response(404) + self.end_headers() + + def log_message(self, format: str, *args) -> None: + """Suppress default logging.""" + pass + + +class HealthServer: + """Background health check server.""" + + def __init__(self, host: str = "0.0.0.0", port: int = 8000): + """Initialize the health server. + + Args: + host: Host to bind to. + port: Port to bind to. + """ + self.host = host + self.port = port + self._server: Optional[HTTPServer] = None + self._thread: Optional[threading.Thread] = None + + def start(self) -> None: + """Start the health server in a background thread.""" + try: + self._server = HTTPServer((self.host, self.port), HealthHandler) + self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) + self._thread.start() + log.info(f"Health server started at http://{self.host}:{self.port}/health") + except Exception as e: + log.error(f"Failed to start health server: {e}") + + def stop(self) -> None: + """Stop the health server.""" + if self._server: + self._server.shutdown() + self._server = None + log.debug("Health server stopped")