Backend: Added dynamic infrastructure management (DigitalOcean)

backend/main.py

@@ -1,11 +1,17 @@
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse, Response
+import asyncio
+import logging
 import os
 import json
 from pathlib import Path
 from dotenv import load_dotenv
 import sentry_sdk
+from services.orchestrator_service import orchestrator_service
+from services.infrastructure_service import infrastructure_service
+from services.config import settings
+from worker import AubmWorker
 
 
 def _load_app_version() -> str:
@@ -21,6 +27,9 @@ def _load_app_version() -> str:
 load_dotenv()
 FRONTEND_DIST = Path(__file__).resolve().parent.parent / "frontend" / "dist"
 APP_VERSION = _load_app_version()
+logger = logging.getLogger("aubm.api")
+embedded_worker: AubmWorker | None = None
+embedded_worker_task: asyncio.Task | None = None
 
 # Sentry Initialization
 SENTRY_DSN = os.getenv("SENTRY_DSN")
@@ -48,6 +57,54 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+
+def _log_embedded_worker_result(task: asyncio.Task) -> None:
+    if task.cancelled():
+        return
+
+    exc = task.exception()
+    if exc:
+        logger.error(
+            "Embedded worker stopped unexpectedly",
+            exc_info=(type(exc), exc, exc.__traceback__),
+        )
+
+
+@app.on_event("startup")
+async def start_embedded_worker() -> None:
+    global embedded_worker, embedded_worker_task
+
+    if settings.TASK_EXECUTION_MODE != "queue" or not settings.TASK_QUEUE_EMBEDDED_WORKER:
+        return
+
+    if embedded_worker_task and not embedded_worker_task.done():
+        return
+
+    embedded_worker = AubmWorker()
+    embedded_worker_task = asyncio.create_task(embedded_worker.start())
+    embedded_worker_task.add_done_callback(_log_embedded_worker_result)
+    logger.info("Embedded task worker started: %s", embedded_worker.worker_id)
+
+
+@app.on_event("shutdown")
+async def stop_embedded_worker() -> None:
+    global embedded_worker, embedded_worker_task
+
+    if not embedded_worker or not embedded_worker_task:
+        return
+
+    embedded_worker.stop()
+    try:
+        await asyncio.wait_for(embedded_worker_task, timeout=10)
+        await embedded_worker.heartbeat("stopping")
+    except asyncio.TimeoutError:
+        embedded_worker_task.cancel()
+        logger.warning("Embedded task worker did not stop before timeout")
+    finally:
+        embedded_worker = None
+        embedded_worker_task = None
+
+
 @app.get("/")
 async def root():
     index_path = FRONTEND_DIST / "index.html"
@@ -81,6 +138,20 @@ async def runtime_config():
         media_type="application/javascript",
     )
 
+@app.get("/{path:path}", include_in_schema=False)
+async def serve_frontend(path: str):
+    if not FRONTEND_DIST.exists():
+        return await root()
+
+    requested_path = FRONTEND_DIST / path
+    if requested_path.is_file():
+        return FileResponse(requested_path)
+
+    return Response(
+        content=f"window.__AUBM_CONFIG__ = {json.dumps(config)};",
+        media_type="application/javascript",
+    )
+
 @app.get("/{path:path}", include_in_schema=False)
 async def serve_frontend(path: str):
     if not FRONTEND_DIST.exists():
@@ -96,7 +167,32 @@ async def serve_frontend(path: str):
 
     return await root()
 
+# --- Infrastructure Management ---
+
+@app.post("/infrastructure/nodes/provision")
+async def provision_node(name: str = "aubm-inference-node", size: str = "s-4vcpu-8gb-amd"):
+    """Creates a new inference node on DigitalOcean."""
+    node = await infrastructure_service.create_inference_node(name, size)
+    if not node:
+        raise HTTPException(status_code=500, detail="Failed to initiate node provisioning.")
+    return node
+
+@app.get("/infrastructure/nodes/{droplet_id}/ip")
+async def get_node_ip(droplet_id: int):
+    """Wait and return the public IP of a node."""
+    ip = await infrastructure_service.wait_for_ip(droplet_id)
+    if not ip:
+        raise HTTPException(status_code=404, detail="IP not assigned or timed out.")
+    return {"ip": ip}
+
+@app.delete("/infrastructure/nodes/{droplet_id}")
+async def terminate_node(droplet_id: int):
+    """Destroy an inference node."""
+    success = await infrastructure_service.terminate_node(droplet_id)
+    if not success:
+        raise HTTPException(status_code=500, detail="Failed to terminate node.")
+    return {"status": "termination_requested"}
+
 if __name__ == "__main__":
     import uvicorn
-    from services.config import settings
-    uvicorn.run("main:app", host="0.0.0.0", port=settings.PORT, reload=True)
+    uvicorn.run(app, host="0.0.0.0", port=int(settings.PORT))

backend/services/config.py

@@ -16,6 +16,10 @@ class Settings(BaseSettings):
     AMD_API_KEY: Optional[str] = None
     TAVILY_API_KEY: Optional[str] = None
     
+    # Infrastructure (DigitalOcean)
+    DO_API_TOKEN: Optional[str] = None
+    DO_REGION: str = "nyc3"
+    
     # App Config
     TASK_QUEUE_EMBEDDED_WORKER: bool = True
     TASK_QUEUE_HEARTBEAT_ENABLED: bool = True

backend/services/infrastructure_service.py

@@ -0,0 +1,97 @@
+import httpx
+import logging
+import asyncio
+from typing import Optional, Dict, Any
+from .config import settings
+
+logger = logging.getLogger("infrastructure")
+
+class InfrastructureService:
+    """
+    Manages on-the-fly compute resources on DigitalOcean for AI inference.
+    """
+    API_URL = "https://api.digitalocean.com/v2"
+
+    def __init__(self):
+        self.headers = {
+            "Authorization": f"Bearer {settings.DO_API_TOKEN}",
+            "Content-Type": "application/json"
+        }
+
+    async def create_inference_node(self, name: str, size: str = "s-4vcpu-8gb-amd") -> Optional[Dict[str, Any]]:
+        """
+        Provision a new AMD-based droplet with Ollama pre-installed.
+        Default size is a capable AMD-based node.
+        """
+        if not settings.DO_API_TOKEN:
+            logger.error("DO_API_TOKEN not configured.")
+            return None
+
+        # Cloud-init script to setup the inference environment
+        user_data = """#cloud-config
+runcmd:
+  - curl -fsSL https://get.docker.com | sh
+  - docker run -d -v ollama:/root/.ollama -p 11434:11434 --name ollama -e OLLAMA_HOST=0.0.0.0 ollama/ollama
+  - sleep 10
+  - docker exec ollama ollama pull llama3
+"""
+
+        payload = {
+            "name": name,
+            "region": settings.DO_REGION,
+            "size": size,
+            "image": "ubuntu-22-04-x64",
+            "user_data": user_data,
+            "tags": ["aubm-worker", "inference-node"]
+        }
+
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.post(f"{self.API_URL}/droplets", headers=self.headers, json=payload)
+                response.raise_for_status()
+                data = response.json()
+                droplet_id = data["droplet"]["id"]
+                logger.info(f"Inference node creation initiated: {name} (ID: {droplet_id})")
+                return data["droplet"]
+            except Exception as e:
+                logger.error(f"Failed to create droplet: {e}")
+                return None
+
+    async def wait_for_ip(self, droplet_id: int, timeout: int = 300) -> Optional[str]:
+        """
+        Polls the API until the droplet has a public IP assigned.
+        """
+        start_time = asyncio.get_event_loop().time()
+        async with httpx.AsyncClient() as client:
+            while (asyncio.get_event_loop().time() - start_time) < timeout:
+                try:
+                    response = await client.get(f"{self.API_URL}/droplets/{droplet_id}", headers=self.headers)
+                    response.raise_for_status()
+                    droplet = response.json()["droplet"]
+                    
+                    networks = droplet.get("networks", {}).get("v4", [])
+                    for nw in networks:
+                        if nw["type"] == "public":
+                            return nw["ip_address"]
+                            
+                except Exception as e:
+                    logger.warning(f"Error polling droplet {droplet_id}: {e}")
+                
+                await asyncio.sleep(10)
+        return None
+
+    async def terminate_node(self, droplet_id: int):
+        """
+        Destroy the inference node to stop billing.
+        """
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.delete(f"{self.API_URL}/droplets/{droplet_id}", headers=self.headers)
+                response.raise_for_status()
+                logger.info(f"Inference node {droplet_id} termination requested.")
+                return True
+            except Exception as e:
+                logger.error(f"Failed to terminate droplet {droplet_id}: {e}")
+                return False
+
+infrastructure_service = InfrastructureService()