Upload production/worker.py

production/worker.py

@@ -0,0 +1,135 @@
+"""Background worker for ml-intern production system."""
+import asyncio
+import json
+import logging
+import os
+import signal
+import time
+from datetime import datetime, timedelta
+
+import aioredis
+import asyncpg
+
+logger = logging.getLogger("ml_intern.worker")
+REDIS_URL = os.environ.get("REDIS_URL", "redis://localhost:6379")
+DATABASE_URL = os.environ.get("DATABASE_URL", "postgresql://localhost/ml_intern")
+CLEANUP_INTERVAL_SECONDS = int(os.environ.get("CLEANUP_INTERVAL_SECONDS", "3600"))
+BUDGET_CHECK_INTERVAL_SECONDS = int(os.environ.get("BUDGET_CHECK_INTERVAL_SECONDS", "300"))
+
+class Worker:
+    def __init__(self):
+        self.redis: aioredis.Redis = None
+        self.db: asyncpg.Pool = None
+        self.running = True
+    
+    async def start(self):
+        self.redis = aioredis.from_url(REDIS_URL, decode_responses=True)
+        await self.redis.ping()
+        self.db = await asyncpg.create_pool(DATABASE_URL, min_size=2, max_size=5, command_timeout=60)
+        logger.info("Worker started")
+        tasks = [
+            asyncio.create_task(self._cleanup_expired_sessions()),
+            asyncio.create_task(self._check_budget_alerts()),
+            asyncio.create_task(self._aggregate_costs()),
+            asyncio.create_task(self._refresh_analytics()),
+        ]
+        await asyncio.gather(*tasks, return_exceptions=True)
+    
+    async def _cleanup_expired_sessions(self):
+        while self.running:
+            try:
+                expired = await self.db.fetch(
+                    "SELECT id, tenant_id FROM sessions WHERE expires_at < NOW() AND status = 'active' LIMIT 1000"
+                )
+                if expired:
+                    session_ids = [str(s["id"]) for s in expired]
+                    await self.db.execute(
+                        "INSERT INTO archived_sessions (id, tenant_id, data, archived_at) SELECT id, tenant_id, jsonb_build_object('metadata', metadata, 'context', context_json, 'spent_usd', spent_usd, 'budget_usd', budget_usd), NOW() FROM sessions WHERE id = ANY($1)",
+                        session_ids,
+                    )
+                    await self.db.execute("UPDATE sessions SET status = 'expired' WHERE id = ANY($1)", session_ids)
+                    pipe = self.redis.pipeline()
+                    for sid in session_ids:
+                        pipe.delete(f"ratelimit:*:{sid}")
+                        pipe.delete(f"session:{sid}:*")
+                    await pipe.execute()
+                    logger.info(f"Cleaned up {len(expired)} expired sessions")
+            except Exception as e:
+                logger.exception(f"Error in cleanup: {e}")
+            await asyncio.sleep(CLEANUP_INTERVAL_SECONDS)
+    
+    async def _check_budget_alerts(self):
+        while self.running:
+            try:
+                nearing_limit = await self.db.fetch(
+                    """SELECT id, tenant_id, budget_usd, spent_usd, (spent_usd / budget_usd) * 100 as pct_used
+                       FROM sessions WHERE status = 'active' AND spent_usd > budget_usd * 0.8 AND spent_usd < budget_usd * 0.95
+                       AND NOT (metadata ? 'budget_alert_80_sent')"""
+                )
+                for session in nearing_limit:
+                    session_id = str(session["id"])
+                    pct = session["pct_used"]
+                    await self.db.execute(
+                        "INSERT INTO budget_events (session_id, event_type, threshold_percent, spent_before, spent_after, budget) VALUES ($1, 'threshold_reached', $2, $3, $3, $4)",
+                        session["id"], 80, session["spent_usd"], session["budget_usd"],
+                    )
+                    await self.db.execute(
+                        "UPDATE sessions SET metadata = jsonb_set(metadata, '{budget_alert_80_sent}', 'true') WHERE id = $1",
+                        session["id"],
+                    )
+                    logger.warning(f"Budget alert: session {session_id} at {pct:.1f}% of budget")
+                
+                exceeded = await self.db.fetch(
+                    """SELECT id, tenant_id, budget_usd, spent_usd FROM sessions
+                       WHERE status = 'active' AND spent_usd >= budget_usd AND NOT (metadata ? 'budget_exceeded')"""
+                )
+                for session in exceeded:
+                    await self.db.execute(
+                        "INSERT INTO budget_events (session_id, event_type, threshold_percent, spent_before, spent_after, budget) VALUES ($1, 'exceeded', 100, $2, $2, $3)",
+                        session["id"], session["spent_usd"], session["budget_usd"],
+                    )
+                    await self.db.execute(
+                        "UPDATE sessions SET metadata = jsonb_set(metadata, '{budget_exceeded}', 'true') WHERE id = $1",
+                        session["id"],
+                    )
+                    logger.error(f"Budget exceeded: session {session['id']} spent ${session['spent_usd']} of ${session['budget_usd']}")
+            except Exception as e:
+                logger.exception(f"Error checking budgets: {e}")
+            await asyncio.sleep(BUDGET_CHECK_INTERVAL_SECONDS)
+    
+    async def _aggregate_costs(self):
+        while self.running:
+            try:
+                await self.db.execute("SELECT refresh_daily_costs()")
+                logger.info("Refreshed daily cost analytics")
+            except Exception as e:
+                logger.exception(f"Error aggregating costs: {e}")
+            await asyncio.sleep(3600)
+    
+    async def _refresh_analytics(self):
+        while self.running:
+            try:
+                top_models = await self.db.fetch(
+                    """SELECT provider, model, COUNT(*) as requests, SUM(total_tokens) as tokens, AVG(latency_ms)::int as avg_latency
+                       FROM requests WHERE created_at > NOW() - INTERVAL '24 hours'
+                       GROUP BY provider, model ORDER BY requests DESC LIMIT 10"""
+                )
+                logger.info("Top models (24h):")
+                for row in top_models:
+                    logger.info(f"  {row['provider']}/{row['model']}: {row['requests']} requests, {row['tokens']} tokens, {row['avg_latency']}ms avg latency")
+                active = await self.db.fetchval("SELECT COUNT(*) FROM sessions WHERE status = 'active'")
+                logger.info(f"Active sessions: {active}")
+            except Exception as e:
+                logger.exception(f"Error in analytics: {e}")
+            await asyncio.sleep(300)
+
+async def main():
+    worker = Worker()
+    loop = asyncio.get_event_loop()
+    for sig in (signal.SIGTERM, signal.SIGINT):
+        loop.add_signal_handler(sig, lambda: setattr(worker, 'running', False))
+    await worker.start()
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(name)s | %(message)s")
+    asyncio.run(main())