Gateway / src /services /space_client.py
Adidasa108's picture
Upload 3 files
24fe247 verified
Raw
History Blame Contribute Delete
3.71 kB
"""
HTTP client for calling remote HF Space microservices (Face Space + Object Space).
Handles cold-start timeouts (HF free tier sleeps after ~5 min inactivity),
retries on 503, and transparent fallback when a space is unavailable.
"""
import asyncio
import io
import os
from typing import Optional
import aiohttp
from src.core.config import FACE_SPACE_URL, OBJECT_SPACE_URL
SPACE_CALL_TIMEOUT = int(os.getenv("SPACE_CALL_TIMEOUT", "120")) # 2 min — covers cold-start wakeup
SPACE_CALL_RETRIES = int(os.getenv("SPACE_CALL_RETRIES", "2"))
_session: Optional[aiohttp.ClientSession] = None
async def get_session() -> aiohttp.ClientSession:
global _session
if _session is None or _session.closed:
timeout = aiohttp.ClientTimeout(total=SPACE_CALL_TIMEOUT)
_session = aiohttp.ClientSession(timeout=timeout)
return _session
async def close_session() -> None:
global _session
if _session and not _session.closed:
await _session.close()
_session = None
async def _post_image(url: str, image_bytes: bytes, params: dict) -> Optional[dict]:
"""POST image bytes as multipart form to a remote Space endpoint."""
session = await get_session()
for attempt in range(SPACE_CALL_RETRIES + 1):
try:
data = aiohttp.FormData()
data.add_field("file", io.BytesIO(image_bytes), filename="image.jpg", content_type="image/jpeg")
for k, v in params.items():
data.add_field(k, str(v))
async with session.post(url, data=data) as resp:
if resp.status == 200:
return await resp.json()
if resp.status == 503:
# Space is waking up — wait with exponential backoff and retry
wait = 15 * (attempt + 1)
await asyncio.sleep(wait)
continue
# Non-retryable error
return None
except asyncio.TimeoutError:
if attempt < SPACE_CALL_RETRIES:
await asyncio.sleep(5)
continue
except Exception:
return None
return None
async def embed_face(image_bytes: bytes, quality_gate: float = 0.35) -> Optional[list]:
"""
Call the Face Space to detect and embed all faces in the image.
Returns list of face dicts or None if the space is unavailable.
Falls back gracefully — caller should treat None as 'no faces found'.
"""
if not FACE_SPACE_URL:
return None
result = await _post_image(
f"{FACE_SPACE_URL}/embed/face",
image_bytes,
{"quality_gate": quality_gate},
)
return result.get("faces") if result else None
async def embed_object(image_bytes: bytes) -> Optional[list]:
"""
Call the Object Space to embed the full image + detected object crops.
Returns list of object dicts or None if the space is unavailable.
"""
if not OBJECT_SPACE_URL:
return None
result = await _post_image(
f"{OBJECT_SPACE_URL}/embed/object",
image_bytes,
{},
)
return result.get("objects") if result else None
async def ping_spaces() -> dict:
"""Health-check all spaces. Returns dict of space_name → bool."""
session = await get_session()
results = {}
for name, url in [("face", FACE_SPACE_URL), ("object", OBJECT_SPACE_URL)]:
if not url:
results[name] = None
continue
try:
async with session.get(f"{url}/health", timeout=aiohttp.ClientTimeout(total=5)) as resp:
results[name] = resp.status == 200
except Exception:
results[name] = False
return results