worker.py

import asyncio
import json
import shutil
import subprocess
from pathlib import Path

from fastapi import FastAPI, HTTPException, Request
from huggingface_hub import HfApi

app = FastAPI(title="Altamira Worker")

CACHE_DIR = Path("/tmp/altamira-cache")
WORKSPACE_DIR = Path("/tmp/altamira-workspace")
STAGE_DIRS = [CACHE_DIR / "stage_0", CACHE_DIR / "stage_1", CACHE_DIR / "stage_2"]

api = HfApi()


def _ensure_dirs():
    for d in [CACHE_DIR, WORKSPACE_DIR] + STAGE_DIRS:
        d.mkdir(parents=True, exist_ok=True)


def _rotate_stages():
    for i in range(len(STAGE_DIRS) - 1, 0, -1):
        src = STAGE_DIRS[i - 1]
        dst = STAGE_DIRS[i]
        if src.exists():
            if dst.exists():
                shutil.rmtree(dst)
            shutil.copytree(src, dst, dirs_exist_ok=True)


def _cleanup_stale(keep_last: int = 2):
    for d in STAGE_DIRS[:-keep_last] if keep_last < len(STAGE_DIRS) else []:
        if d.exists():
            shutil.rmtree(d)


@app.post("/sync")
async def sync_repo(repo_id: str, local_path: str = "/tmp/altamira-checkout"):
    _ensure_dirs()
    try:
        subprocess.run(
            ["git", "clone", f"https://huggingface.co/spaces/{repo_id}", local_path],
            capture_output=True, text=True, check=False
        )
        if not Path(local_path).exists():
            subprocess.run(
                ["git", "init", local_path],
                capture_output=True, check=True,
            )
        _rotate_stages()
        shutil.copytree(local_path, STAGE_DIRS[0], dirs_exist_ok=True)
        _cleanup_stale()
        return {"status": "synced", "repo": repo_id, "stage": str(STAGE_DIRS[0])}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/rotate")
async def rotate_cache():
    _ensure_dirs()
    _rotate_stages()
    _cleanup_stale()
    return {"status": "rotated", "stages": [str(d) for d in STAGE_DIRS]}


@app.post("/webhook/results")
async def receive_results(request: Request):
    try:
        payload = await request.json()
        token = request.query_params.get("token")
        if not token:
            raise HTTPException(status_code=401, detail="Verification key required")
        
        # Log diagnostic payload to local file for Orchestrator to pull
        log_file = Path("/tmp/altamira-results.json")
        with open(log_file, "a") as f:
            f.write(json.dumps(payload) + "\n")
            
        return {"status": "received", "payload_size": len(str(payload))}
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Invalid payload: {str(e)}")


@app.get("/health")
async def health():
    return {"status": "healthy", "worker": "altamira"}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)