Spaces:

PPAL-SongLab-UGA
/

watermelon-proxy

Paused

Add alternate models
38d3ceb 3 months ago
25.1 kB
import os
import time
import requests
from fastapi import FastAPI, UploadFile, File, Form, Query, Request, Body
from fastapi.responses import Response, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.concurrency import run_in_threadpool # <--- FIX: Added Threadpool
import uvicorn

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

HF_TOKEN = os.getenv("HF_TOKEN")

DEV_URL = "https://PPAL-SongLab-UGA-fruit-analyzer-dev.hf.space"
PROD_URL = "https://PPAL-SongLab-UGA-fruit-analyzer.hf.space"
HF_RATE_LIMIT_STATUS = 429
HF_TRANSIENT_STATUS = {429, 502, 503, 504}
PROCESS_PROXY_TIMEOUT_SECONDS = 40


def request_with_hf_backoff(method, url, *, max_retries=3, **kwargs):
    last_response = None
    for attempt in range(max_retries + 1):
        response = method(url, **kwargs)
        last_response = response
        if response.status_code not in HF_TRANSIENT_STATUS:
            return response

        retry_after = response.headers.get("retry-after")
        try:
            wait_s = float(retry_after) if retry_after else None
        except ValueError:
            wait_s = None
        if wait_s is None:
            wait_s = [2.0, 5.0, 10.0, 15.0][min(attempt, 3)] if response.status_code == HF_RATE_LIMIT_STATUS else [3.0, 8.0, 15.0, 20.0][min(attempt, 3)]
        if attempt < max_retries:
            time.sleep(min(max(wait_s, 0.5), 20.0))
    return last_response


def hf_error_message(exc):
    response = getattr(exc, "response", None)
    if response is not None and response.status_code == HF_RATE_LIMIT_STATUS:
        retry_after = response.headers.get("retry-after")
        suffix = f" Retry after about {retry_after} seconds." if retry_after else " Please wait a minute and try again."
        return f"Hugging Face is rate limiting this Space after several retry attempts.{suffix}"
    if response is not None and response.status_code in {502, 503, 504}:
        return "Hugging Face Space is temporarily unavailable or waking up after several retry attempts. Please wait a moment and try again."
    if isinstance(exc, requests.exceptions.Timeout):
        return "Hugging Face Space did not respond before the proxy timeout. Please try again."
    if isinstance(exc, requests.exceptions.ConnectionError):
        return "Could not connect to the Hugging Face Space. Please try again."
    return str(exc)


def is_retryable_hf_exception(exc):
    response = getattr(exc, "response", None)
    return bool(response is not None and response.status_code in HF_TRANSIENT_STATUS)


def proxy_error_payload(prefix, exc, **extra):
    return {
        "success": False,
        "message": f"{prefix}: {hf_error_message(exc)}",
        "retryable": is_retryable_hf_exception(exc),
        **extra,
    }

@app.get("/")
def read_root():
    return {"status": "Render Proxy is awake!"}

@app.get("/proxy_status")
async def proxy_status(username: str = Query("")):
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    status_url = f"{base_url}/queue_status"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    
    # FIX: Define the blocking request
    def fetch_status():
        return requests.get(status_url, headers=headers, timeout=5)
        
    try:
        # FIX: Run it in a background thread so the server doesn't freeze
        response = await run_in_threadpool(fetch_status)
        response.raise_for_status()
        return response.json()
    except Exception as e:
        return {"active_requests": 0, "max_concurrent": 2}


@app.post("/proxy_warmup")
async def proxy_warmup(payload: dict = Body(...)):
    """Warm the production backend only.
    This deliberately does not route devtest to the dev Space; it exists to
    keep the production user path responsive without burning analysis jobs.
    """
    target_url = f"{PROD_URL}/warmup"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_post():
        return request_with_hf_backoff(
            requests.post,
            target_url,
            max_retries=0,
            headers=headers,
            json=payload,
            timeout=PROCESS_PROXY_TIMEOUT_SECONDS,
        )

    try:
        response = await run_in_threadpool(make_post)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        return proxy_error_payload("Proxy Warmup Error (Hugging Face)", e)
    except Exception as e:
        return {"success": False, "message": f"Proxy Warmup Error (Internal): {str(e)}"}


@app.post("/proxy_process")
async def proxy_process(request: Request, file: UploadFile = File(...), password: str = Form(""), username: str = Form("")):
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/process_single"

    query_str = request.url.query
    if query_str: target_url += f"?{query_str}"

    file_bytes = await file.read()
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    files = {"file": (file.filename, file_bytes, file.content_type)}
    form = await request.form()
    data = {}
    for key, value in form.multi_items():
        if key == "file" or hasattr(value, "filename"):
            continue
        data[key] = str(value)
    data.setdefault("password", password)
    data.setdefault("username", username)

    # FIX: Define the blocking request
    def make_post():
        return request_with_hf_backoff(
            requests.post,
            target_url,
            max_retries=0,
            headers=headers,
            files=files,
            data=data,
            timeout=PROCESS_PROXY_TIMEOUT_SECONDS,
        )

    try:
        # FIX: Run it in a background thread so other users can check the status!
        response = await run_in_threadpool(make_post)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        return proxy_error_payload("Proxy Error (Hugging Face)", e)
    except Exception as e:
        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}"}


@app.post("/proxy_compatibility")
async def proxy_compatibility(request: Request, password: str = Form(""), username: str = Form("")):
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/compatibility_check"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    form = await request.form()
    files = []
    data = {}
    for key, value in form.multi_items():
        if hasattr(value, "filename"):
            files.append(("files", (value.filename, await value.read(), value.content_type)))
        else:
            data[key] = str(value)
    data.setdefault("password", password)
    data.setdefault("username", username)

    def make_post():
        return request_with_hf_backoff(
            requests.post,
            target_url,
            max_retries=0,
            headers=headers,
            files=files,
            data=data,
            timeout=PROCESS_PROXY_TIMEOUT_SECONDS,
        )

    try:
        response = await run_in_threadpool(make_post)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        return proxy_error_payload("Proxy Error (Hugging Face)", e, results=[])
    except Exception as e:
        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}", "results": []}


@app.get("/proxy_experts")
async def proxy_experts(username: str = Query(""), password: str = Query("")):
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/experts"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_get():
        return requests.get(target_url, headers=headers, params={"username": username, "password": password}, timeout=20)

    try:
        response = await run_in_threadpool(make_get)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        return proxy_error_payload("Proxy Error (Hugging Face)", e, experts=[])
    except Exception as e:
        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}", "experts": []}


@app.get("/proxy_training_model_options")
async def proxy_training_model_options(username: str = Query(""), password: str = Query("")):
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/training_model_options"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_get():
        return requests.get(target_url, headers=headers, params={"username": username, "password": password}, timeout=20)

    try:
        response = await run_in_threadpool(make_get)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        return proxy_error_payload("Proxy Error (Hugging Face)", e, options=[])
    except Exception as e:
        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}", "options": []}


@app.post("/proxy_finetune")
async def proxy_finetune(payload: dict = Body(...)):
    username = str(payload.get("username", ""))
    expert_id = str(payload.get("expert_id", "")).strip()
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/experts/{expert_id}/finetune"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_post():
        return request_with_hf_backoff(requests.post, target_url, max_retries=0, headers=headers, json=payload, timeout=60)

    try:
        response = await run_in_threadpool(make_post)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        return proxy_error_payload("Proxy Error (Hugging Face)", e)
    except Exception as e:
        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}"}


@app.get("/proxy_train_job")
async def proxy_train_job(job_id: str = Query(...), username: str = Query(""), password: str = Query("")):
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/train_jobs/status/{job_id}"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_get():
        return requests.get(target_url, headers=headers, params={"username": username, "password": password}, timeout=20)

    try:
        response = await run_in_threadpool(make_get)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        return proxy_error_payload("Proxy Error (Hugging Face)", e)
    except Exception as e:
        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}"}


@app.get("/proxy_train_jobs")
async def proxy_train_jobs(username: str = Query(""), password: str = Query("")):
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/train_jobs"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_get():
        return requests.get(target_url, headers=headers, params={"username": username, "password": password}, timeout=20)

    try:
        response = await run_in_threadpool(make_get)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        return proxy_error_payload("Proxy Error (Hugging Face)", e, jobs=[])
    except Exception as e:
        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}", "jobs": []}


DATASET_PROXY_TIMEOUT_SECONDS = 180


@app.api_route("/proxy_datasets{rest:path}", methods=["GET", "POST", "PUT", "DELETE"])
async def proxy_datasets(rest: str, request: Request):
    """Generic forwarder for all Labeling Studio dataset endpoints.
    Mirrors /proxy_datasets{rest} -> {base}/datasets{rest}, preserving the query
    string, body (multipart / JSON), and returning either JSON or raw binary
    (image/zip) responses. Routes dev vs prod by username (query, form, or JSON).
    """
    username = request.query_params.get("username", "")
    content_type = request.headers.get("content-type", "")

    files = []
    data = {}
    json_body = None
    if request.method in ("POST", "PUT", "DELETE"):
        if "multipart/form-data" in content_type:
            form = await request.form()
            for key, value in form.multi_items():
                if hasattr(value, "filename"):
                    files.append((key, (value.filename, await value.read(), value.content_type)))
                else:
                    data[key] = str(value)
            username = username or data.get("username", "")
        elif "application/json" in content_type:
            try:
                json_body = await request.json()
            except Exception:
                json_body = None
            if isinstance(json_body, dict):
                username = username or str(json_body.get("username", ""))

    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/datasets{rest}"
    if request.url.query:
        target_url += f"?{request.url.query}"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_request():
        return request_with_hf_backoff(
            getattr(requests, request.method.lower()),
            target_url,
            max_retries=0,
            headers=headers,
            files=files or None,
            data=data or None,
            json=json_body if json_body is not None else None,
            timeout=DATASET_PROXY_TIMEOUT_SECONDS,
        )

    try:
        response = await run_in_threadpool(make_request)
    except requests.exceptions.RequestException as e:
        return JSONResponse(status_code=502, content=proxy_error_payload("Proxy Error (Hugging Face)", e))
    except Exception as e:
        return JSONResponse(status_code=500, content={"success": False, "message": f"Proxy Error (Internal): {str(e)}"})

    resp_ct = response.headers.get("content-type", "application/octet-stream")
    passthrough_headers = {}
    disposition = response.headers.get("content-disposition")
    if disposition:
        passthrough_headers["content-disposition"] = disposition
    return Response(
        content=response.content,
        status_code=response.status_code,
        media_type=resp_ct,
        headers=passthrough_headers,
    )

@app.post("/proxy_batch_stage")
async def proxy_batch_stage(payload: dict = Body(...)):
    username = str(payload.get("username", ""))
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/batch_stage"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_post():
        return request_with_hf_backoff(
            requests.post,
            target_url,
            max_retries=0,
            headers=headers,
            json=payload,
            timeout=PROCESS_PROXY_TIMEOUT_SECONDS,
        )

    try:
        response = await run_in_threadpool(make_post)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        return proxy_error_payload("Proxy Error (Hugging Face)", e, rows=[])
    except Exception as e:
        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}", "rows": []}


@app.api_route("/proxy_process_jobs{rest:path}", methods=["GET", "POST"])
async def proxy_process_jobs(rest: str, request: Request):
    """Forward persistent processing-job uploads, polling, stop, and resume."""
    try:
        username = request.query_params.get("username", "")
        content_type = request.headers.get("content-type", "")
        files = []
        data = {}
        json_body = None
        if request.method == "POST":
            if "multipart/form-data" in content_type:
                form = await request.form()
                for key, value in form.multi_items():
                    if hasattr(value, "filename"):
                        files.append((key, (value.filename, await value.read(), value.content_type)))
                    else:
                        data[key] = str(value)
                username = username or data.get("username", "")
            elif "application/json" in content_type:
                json_body = await request.json()
                if isinstance(json_body, dict):
                    username = username or str(json_body.get("username", ""))
    except Exception as exc:
        return JSONResponse(status_code=400, content={"success": False, "message": f"Proxy Process Job Error (Request): {str(exc)}"})
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/process_jobs{rest}"
    if request.url.query:
        target_url += f"?{request.url.query}"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_request():
        return request_with_hf_backoff(
            getattr(requests, request.method.lower()),
            target_url,
            max_retries=2,
            headers=headers,
            files=files or None,
            data=data or None,
            json=json_body,
            timeout=DATASET_PROXY_TIMEOUT_SECONDS,
        )

    try:
        response = await run_in_threadpool(make_request)
        return Response(
            content=response.content,
            status_code=response.status_code,
            media_type=response.headers.get("content-type", "application/json"),
        )
    except requests.exceptions.RequestException as exc:
        return JSONResponse(status_code=502, content=proxy_error_payload("Proxy Error (Hugging Face)", exc))
    except Exception as exc:
        return JSONResponse(status_code=500, content={"success": False, "message": f"Proxy Process Job Error (Internal): {str(exc)}"})


@app.post("/proxy_flush_queue")
async def proxy_flush_queue(payload: dict = Body(...)):
    username = str(payload.get("username", ""))
    if username.strip().lower() != "devtest":
        return {"success": False, "message": "Queue flush is only available for devtest."}
    target_url = f"{DEV_URL}/flush_queue"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_post():
        return requests.post(target_url, headers=headers, json=payload, timeout=20)

    try:
        response = await run_in_threadpool(make_post)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        return proxy_error_payload("Proxy Error (Hugging Face)", e)
    except Exception as e:
        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}"}


@app.get("/proxy_preview/{session_id}/{row_id}/{preview_type}")
async def proxy_preview(session_id: str, row_id: str, preview_type: str, username: str = Query(""), password: str = Query(""), size: str = Query("thumb")):
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/preview/{session_id}/{row_id}/{preview_type}"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_get():
        return requests.get(
            target_url,
            headers=headers,
            params={"size": size, "username": username, "password": password},
            timeout=30,
        )

    response = await run_in_threadpool(make_get)
    if response.status_code >= 400:
        return Response(content=response.content, status_code=response.status_code, media_type=response.headers.get("content-type", "text/plain"))
    return Response(content=response.content, media_type=response.headers.get("content-type", "image/jpeg"))


@app.post("/proxy_preview_session_clear")
async def proxy_preview_session_clear(payload: dict = Body(...)):
    username = str(payload.get("username", ""))
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/preview_session/clear"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_post():
        return requests.post(target_url, headers=headers, json=payload, timeout=20)

    try:
        response = await run_in_threadpool(make_post)
        response.raise_for_status()
        return response.json()
    except Exception as e:
        return {"success": False, "message": f"Proxy Error: {str(e)}"}


@app.api_route("/proxy_adjustments{rest:path}", methods=["GET", "POST", "PUT"])
async def proxy_adjustments(rest: str, request: Request):
    """Forward preview-mask and traditional-feature adjustment requests."""
    username = request.query_params.get("username", "")
    content_type = request.headers.get("content-type", "")
    files = []
    data = {}
    json_body = None
    if request.method in {"POST", "PUT"}:
        if "multipart/form-data" in content_type:
            form = await request.form()
            for key, value in form.multi_items():
                if hasattr(value, "filename"):
                    files.append((key, (value.filename, await value.read(), value.content_type)))
                else:
                    data[key] = str(value)
            username = username or data.get("username", "")
        elif "application/json" in content_type:
            json_body = await request.json()
            if isinstance(json_body, dict):
                username = username or str(json_body.get("username", ""))
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/adjustments{rest}"
    if request.url.query:
        target_url += f"?{request.url.query}"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_request():
        return request_with_hf_backoff(
            getattr(requests, request.method.lower()),
            target_url,
            max_retries=0,
            headers=headers,
            files=files or None,
            data=data or None,
            json=json_body,
            timeout=DATASET_PROXY_TIMEOUT_SECONDS,
        )

    try:
        response = await run_in_threadpool(make_request)
        return Response(
            content=response.content,
            status_code=response.status_code,
            media_type=response.headers.get("content-type", "application/json"),
        )
    except requests.exceptions.RequestException as exc:
        return JSONResponse(status_code=502, content=proxy_error_payload("Proxy Error (Hugging Face)", exc))


SYNC_PROXY_TIMEOUT_SECONDS = 300


@app.api_route("/proxy_sync{rest:path}", methods=["GET", "POST"])
async def proxy_sync(rest: str, request: Request):
    """Forward admin sync requests to the dev backend only.
    Dev remains the controller for dev<->prod sync. The browser never receives
    production sync tokens, and this proxy never routes sync calls to prod based
    on username.
    """
    target_url = f"{DEV_URL}/sync{rest}"
    if request.url.query:
        target_url += f"?{request.url.query}"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    content_type = request.headers.get("content-type", "")
    files = []
    data = {}
    json_body = None
    if request.method == "POST":
        if "multipart/form-data" in content_type:
            form = await request.form()
            for key, value in form.multi_items():
                if hasattr(value, "filename"):
                    files.append((key, (value.filename, await value.read(), value.content_type)))
                else:
                    data[key] = str(value)
        elif "application/json" in content_type:
            try:
                json_body = await request.json()
            except Exception:
                json_body = None

    def make_request():
        return request_with_hf_backoff(
            getattr(requests, request.method.lower()),
            target_url,
            max_retries=0,
            headers=headers,
            files=files or None,
            data=data or None,
            json=json_body,
            timeout=SYNC_PROXY_TIMEOUT_SECONDS,
        )

    try:
        response = await run_in_threadpool(make_request)
    except requests.exceptions.RequestException as exc:
        return JSONResponse(status_code=502, content=proxy_error_payload("Proxy Error (Hugging Face)", exc))
    except Exception as exc:
        return JSONResponse(status_code=500, content={"success": False, "message": f"Proxy Error (Internal): {str(exc)}"})

    media_type = response.headers.get("content-type", "application/json")
    passthrough_headers = {}
    disposition = response.headers.get("content-disposition")
    if disposition:
        passthrough_headers["content-disposition"] = disposition
    return Response(content=response.content, status_code=response.status_code, media_type=media_type, headers=passthrough_headers)


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 7860)))