hindsight-api 0.3.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

hindsight_api/worker/main.py

@@ -0,0 +1,296 @@
+"""
+Command-line interface for Hindsight Worker.
+
+Run the worker with:
+    hindsight-worker
+
+Stop with Ctrl+C (graceful shutdown).
+"""
+
+import argparse
+import asyncio
+import atexit
+import logging
+import os
+import signal
+import socket
+import sys
+import warnings
+
+from ..config import get_config
+from ..engine.task_backend import SyncTaskBackend
+from .poller import WorkerPoller
+
+# Filter deprecation warnings from third-party libraries
+warnings.filterwarnings("ignore", message="websockets.legacy is deprecated")
+warnings.filterwarnings("ignore", message="websockets.server.WebSocketServerProtocol is deprecated")
+
+# Disable tokenizers parallelism to avoid warnings
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+logger = logging.getLogger(__name__)
+
+
+def create_worker_app(poller: WorkerPoller, memory):
+    """Create a minimal FastAPI app for worker metrics and health."""
+    from fastapi import FastAPI
+    from fastapi.responses import JSONResponse, Response
+    from prometheus_client import CONTENT_TYPE_LATEST, generate_latest
+
+    from ..metrics import create_metrics_collector, get_metrics_collector, initialize_metrics
+
+    app = FastAPI(
+        title="Hindsight Worker",
+        description="Worker process for distributed task execution",
+    )
+
+    # Initialize OpenTelemetry metrics
+    try:
+        prometheus_reader = initialize_metrics(service_name="hindsight-worker", service_version="1.0.0")
+        create_metrics_collector()
+        app.state.prometheus_reader = prometheus_reader
+        logger.info("Metrics initialized - available at /metrics endpoint")
+    except Exception as e:
+        logger.warning(f"Failed to initialize metrics: {e}. Metrics will be disabled.")
+        app.state.prometheus_reader = None
+
+    # Set up DB pool metrics if available
+    metrics_collector = get_metrics_collector()
+    if memory._pool is not None and hasattr(metrics_collector, "set_db_pool"):
+        metrics_collector.set_db_pool(memory._pool)
+        logger.info("DB pool metrics configured")
+
+    @app.get(
+        "/health",
+        summary="Health check endpoint",
+        description="Returns worker health status including database connectivity",
+        tags=["Monitoring"],
+    )
+    async def health_endpoint():
+        """Health check endpoint."""
+        health = await memory.health_check()
+        health["worker_id"] = poller.worker_id
+        health["is_shutdown"] = poller.is_shutdown
+        status_code = 200 if health.get("status") == "healthy" else 503
+        return JSONResponse(content=health, status_code=status_code)
+
+    @app.get(
+        "/metrics",
+        summary="Prometheus metrics endpoint",
+        description="Exports metrics in Prometheus format for scraping",
+        tags=["Monitoring"],
+    )
+    async def metrics_endpoint():
+        """Return Prometheus metrics."""
+        metrics_data = generate_latest()
+        return Response(content=metrics_data, media_type=CONTENT_TYPE_LATEST)
+
+    @app.get(
+        "/",
+        summary="Worker info",
+        description="Basic worker information",
+        tags=["Info"],
+    )
+    async def root():
+        """Return basic worker info."""
+        return {
+            "service": "hindsight-worker",
+            "worker_id": poller.worker_id,
+            "is_shutdown": poller.is_shutdown,
+        }
+
+    return app
+
+
+def main():
+    """Main entry point for the hindsight-worker CLI."""
+    # Load configuration from environment
+    config = get_config()
+
+    parser = argparse.ArgumentParser(
+        prog="hindsight-worker",
+        description="Hindsight Worker - distributed task processor",
+    )
+
+    # Worker options
+    parser.add_argument(
+        "--worker-id",
+        default=config.worker_id or socket.gethostname(),
+        help="Worker identifier (default: hostname, env: HINDSIGHT_API_WORKER_ID)",
+    )
+    parser.add_argument(
+        "--poll-interval",
+        type=int,
+        default=config.worker_poll_interval_ms,
+        help=f"Poll interval in milliseconds (default: {config.worker_poll_interval_ms}, env: HINDSIGHT_API_WORKER_POLL_INTERVAL_MS)",
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=config.worker_batch_size,
+        help=f"Tasks to claim per poll (default: {config.worker_batch_size}, env: HINDSIGHT_API_WORKER_BATCH_SIZE)",
+    )
+    parser.add_argument(
+        "--max-retries",
+        type=int,
+        default=config.worker_max_retries,
+        help=f"Max retries before marking failed (default: {config.worker_max_retries}, env: HINDSIGHT_API_WORKER_MAX_RETRIES)",
+    )
+
+    # HTTP server options
+    parser.add_argument(
+        "--http-port",
+        type=int,
+        default=config.worker_http_port,
+        help=f"HTTP port for metrics/health endpoints (default: {config.worker_http_port}, env: HINDSIGHT_API_WORKER_HTTP_PORT)",
+    )
+    parser.add_argument(
+        "--http-host",
+        default="0.0.0.0",
+        help="HTTP host to bind (default: 0.0.0.0)",
+    )
+
+    # Logging options
+    parser.add_argument(
+        "--log-level",
+        default=config.log_level,
+        choices=["critical", "error", "warning", "info", "debug", "trace"],
+        help=f"Log level (default: {config.log_level}, env: HINDSIGHT_API_LOG_LEVEL)",
+    )
+
+    args = parser.parse_args()
+
+    # Configure logging
+    config.configure_logging()
+
+    # Import MemoryEngine here to avoid circular imports
+    from .. import MemoryEngine
+
+    print(f"Starting Hindsight Worker: {args.worker_id}")
+    print(f"  Poll interval: {args.poll_interval}ms")
+    print(f"  Batch size: {args.batch_size}")
+    print(f"  Max retries: {args.max_retries}")
+    print(f"  HTTP server: {args.http_host}:{args.http_port}")
+    print()
+
+    # Global references for cleanup
+    memory = None
+    poller = None
+
+    async def run():
+        nonlocal memory, poller
+        import uvicorn
+
+        from ..extensions import TenantExtension, load_extension
+
+        # Initialize MemoryEngine
+        # Workers use SyncTaskBackend because they execute tasks directly,
+        # they don't need to store tasks (they poll from DB)
+        memory = MemoryEngine(
+            run_migrations=False,  # Workers don't run migrations
+            task_backend=SyncTaskBackend(),
+        )
+
+        await memory.initialize()
+
+        print(f"Database connected: {config.database_url}")
+
+        # Load tenant extension for dynamic schema discovery
+        tenant_extension = load_extension("TENANT", TenantExtension)
+
+        if tenant_extension:
+            print("Tenant extension loaded - schemas will be discovered dynamically on each poll")
+        else:
+            print("No tenant extension configured, using public schema only")
+
+        # Create a single poller that handles all schemas dynamically
+        poller = WorkerPoller(
+            pool=memory._pool,
+            worker_id=args.worker_id,
+            executor=memory.execute_task,
+            poll_interval_ms=args.poll_interval,
+            batch_size=args.batch_size,
+            max_retries=args.max_retries,
+            tenant_extension=tenant_extension,
+        )
+
+        # Create the HTTP app for metrics/health
+        app = create_worker_app(poller, memory)
+
+        # Setup signal handlers for graceful shutdown
+        shutdown_requested = asyncio.Event()
+
+        def signal_handler(signum, frame):
+            print(f"\nReceived signal {signum}, initiating graceful shutdown...")
+            shutdown_requested.set()
+
+        signal.signal(signal.SIGINT, signal_handler)
+        signal.signal(signal.SIGTERM, signal_handler)
+
+        # Create uvicorn config and server
+        uvicorn_config = uvicorn.Config(
+            app,
+            host=args.http_host,
+            port=args.http_port,
+            log_level="info",  # Reduce uvicorn noise
+            access_log=False,
+        )
+        server = uvicorn.Server(uvicorn_config)
+
+        # Run the poller and HTTP server concurrently
+        poller_task = asyncio.create_task(poller.run())
+        http_task = asyncio.create_task(server.serve())
+
+        print(f"Worker started. Metrics available at http://{args.http_host}:{args.http_port}/metrics")
+
+        # Wait for shutdown signal
+        await shutdown_requested.wait()
+
+        # Graceful shutdown
+        print("Shutting down HTTP server...")
+        server.should_exit = True
+
+        print("Waiting for poller to finish...")
+        await poller.shutdown_graceful(timeout=30.0)
+        poller_task.cancel()
+        try:
+            await poller_task
+        except asyncio.CancelledError:
+            pass
+
+        # Wait for HTTP server to finish
+        try:
+            await asyncio.wait_for(http_task, timeout=5.0)
+        except asyncio.TimeoutError:
+            http_task.cancel()
+            try:
+                await http_task
+            except asyncio.CancelledError:
+                pass
+
+        # Close memory engine
+        await memory.close()
+        print("Worker shutdown complete")
+
+    def cleanup():
+        """Synchronous cleanup for atexit."""
+        if memory is not None and memory._pg0 is not None:
+            try:
+                loop = asyncio.new_event_loop()
+                loop.run_until_complete(memory._pg0.stop())
+                loop.close()
+                print("\npg0 stopped.")
+            except Exception as e:
+                print(f"\nError stopping pg0: {e}")
+
+    atexit.register(cleanup)
+
+    try:
+        asyncio.run(run())
+    except KeyboardInterrupt:
+        print("\nWorker interrupted")
+        sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()