"""Servizio GPU di Neuronal (solo nel container su Hugging Face). Il container non ha una GPU: quando il motore deve calcolare, questo servizio chiede a uno Space ZeroGPU di avviare il modello per pochi secondi. Il modello si collega da solo al motore; qui si fa solo la richiesta. Se la GPU non arriva (coda, tempo del giorno finito) lo dice al motore, che continua sulla CPU. Variabili: NEURONAL_GPU_SPACE (utente/nome dello Space col modello), NEURONAL_GPU_KEY (la sua chiave), HF_TOKEN (di chi è il tempo di GPU; serve anche se lo Space è privato), SPACE_HOST (lo mette Hugging Face).""" import asyncio, json, os, re, secrets, sys, threading, time, urllib.error, urllib.request import websockets sys.stdout.reconfigure(encoding="utf-8", line_buffering=True) MOTORE = os.environ.get("NEURONAL_WS", "ws://127.0.0.1:8765/ws?role=voice").replace("role=voice", "role=gpu") SPACE = os.environ.get("NEURONAL_GPU_SPACE", "") CHIAVE = os.environ.get("NEURONAL_GPU_KEY", "") TOKEN = os.environ.get("HF_TOKEN", "") SECONDI = int(os.environ.get("NEURONAL_GPU_SECONDS", "40")) # durata di una sessione: più è corta, più a lungo basta il tempo del giorno MINIMO = 12 # sotto questa durata non vale la pena PUBBLICO = os.environ.get("NEURONAL_PUBLIC_WS") or ("wss://" + os.environ["SPACE_HOST"] + "/ws" if os.environ.get("SPACE_HOST") else "") BASE = "https://" + re.sub(r"[^a-z0-9]+", "-", SPACE.lower()).strip("-") + ".hf.space/" # Con il login attivo, chi non accede resta sulla CPU: la GPU veloce è per chi usa i propri minuti (così quelli dello Space non finiscono). # NEURONAL_GPU_ANONIMI=1 la dà anche a chi non accede, finché durano i minuti dello Space. ANONIMI = os.environ.get("NEURONAL_GPU_ANONIMI", "") == "1" or not os.environ.get("OAUTH_CLIENT_ID") INTEST = {"Content-Type": "application/json"} if TOKEN: INTEST["Authorization"] = "Bearer " + TOKEN funzione = None # numero della funzione "sessione" nello Space def apri(url, corpo=None, attesa=30, chiave=None): """chiave: quella del visitatore (login Hugging Face): il tempo di GPU è contato sul suo account e non sul tuo.""" intest = INTEST if not chiave else {"Content-Type": "application/json", "Authorization": "Bearer " + chiave} dati = json.dumps(corpo).encode() if corpo is not None else None return urllib.request.urlopen(urllib.request.Request(url, data=dati, headers=intest), timeout=attesa) def sveglia(attesa=240): """Uno Space che dorme (dopo giorni senza visite) si risveglia alla prima visita, ma ci mette un paio di minuti: lo si chiama e si aspetta che risponda. Non consuma tempo di GPU.""" fine = time.time() + attesa while time.time() < fine: try: with apri(BASE + "config", attesa=15) as r: if r.status == 200: return True except Exception: pass time.sleep(5) return False SVEGLIA = threading.Lock() def sveglia_in_fondo(): if SVEGLIA.acquire(blocking=False): def lavoro(): try: ok = sveglia() print("Servizio GPU: lo Space del modello " + ("è sveglio." if ok else "non risponde ancora.")) finally: SVEGLIA.release() threading.Thread(target=lavoro, daemon=True).start() def sessione(biglietto, secondi, utente=""): """Una sessione di GPU dall'inizio alla fine. Ritorna None se è andata, altrimenti il testo dell'errore.""" global funzione if funzione is None: with apri(BASE + "config") as r: deps = json.load(r)["dependencies"] funzione = next(d.get("id", i) for i, d in enumerate(deps) if d.get("api_name") == "sessione") coda = secrets.token_hex(8) corpo = {"data": [CHIAVE, f"{PUBBLICO}?role=brain&ticket={biglietto}", secondi], "fn_index": funzione, "session_hash": coda} apri(BASE + "gradio_api/queue/join", corpo, chiave=utente).close() with apri(BASE + "gradio_api/queue/data?session_hash=" + coda, attesa=secondi + 240, chiave=utente) as r: for riga in r: riga = riga.decode("utf-8", "replace").strip() if not riga.startswith("data:"): continue ev = json.loads(riga[5:]) if ev.get("msg") != "process_completed": continue out = ev.get("output") or {} if ev.get("success") is False or out.get("error"): return str(out.get("error") or "errore senza dettagli") testo = str((out.get("data") or [""])[0]) return None if testo.startswith("ok") else testo or "risposta vuota" return "lo Space ha chiuso senza rispondere" def chiedi(biglietto, utente=""): """Chiede la GPU e aspetta la fine della sessione. Ritorna None se è andata, altrimenti (motivo, secondi di pausa). Con la chiave del visitatore il tempo è il suo; senza, quello incluso nell'account dello Space, finché dura.""" chi = "il tuo" if utente else "di oggi" if not utente and not ANONIMI: return ("senza l'accesso con Hugging Face si resta sulla CPU", 600) # con il login il motore riprova subito try: errore = sessione(biglietto, SECONDI, utente) resta = re.search(r"vs\. (\d+)s left", errore or "") if resta: # il tempo del giorno sta finendo: una sessione più corta ci sta ancora? corta = int(int(resta.group(1)) / 1.5) - 9 # lo Space chiede 8 secondi in più e ZeroGPU li conta una volta e mezza if corta >= MINIMO: errore = sessione(biglietto, corta, utente) if not errore: return None if "quota" in errore.lower(): return (f"il tempo di GPU {chi} è finito", 900) return (errore[:140], 120) except urllib.error.HTTPError as e: if e.code >= 500 or e.code == 429: sveglia_in_fondo() if utente and e.code in (401, 403): return ("la chiave Hugging Face del visitatore non è accettata (accedi di nuovo)", 300) return (f"lo Space ha risposto {e.code}", 300 if e.code in (401, 403, 404) else 60) except Exception as e: sveglia_in_fondo() # probabilmente dorme: lo si sveglia mentre il motore va avanti sulla CPU return (f"lo Space non è raggiungibile ({type(e).__name__}): lo sveglio", 60) async def servi(): attesa, ultimo = 2, 0.0 while True: try: async with websockets.connect(MOTORE, max_size=None) as ws: print("Servizio GPU: collegato al motore, lo Space del modello è", SPACE) attesa = 2 sveglia_in_fondo() # se dormiva, intanto si risveglia async def tieni_sveglio(): # una visita ogni 6 ore: così non si addormenta while True: await asyncio.sleep(6 * 3600) sveglia_in_fondo() custode = asyncio.ensure_future(tieni_sveglio()) async def richiesta(biglietto, utente=""): t = time.time() esito = await asyncio.get_running_loop().run_in_executor(None, chiedi, biglietto, utente) if esito: print(f"Servizio GPU: niente GPU ({esito[0]}), il motore continua sulla CPU.") try: await ws.send(json.dumps({"cmd": "gpu_status", "ok": False, "message": esito[0], "retry": esito[1]})) except Exception: pass else: print(f"Servizio GPU: sessione di {time.time() - t:.0f} s conclusa.") async for m in ws: if isinstance(m, str) and '"gpu_need"' in m: ev = json.loads(m) asyncio.ensure_future(richiesta(ev["ticket"], ev.get("user") or "")) except Exception as e: try: custode.cancel() except NameError: pass if time.time() - ultimo > 60: print(f"Servizio GPU: aspetto il motore ({type(e).__name__}).") ultimo = time.time() await asyncio.sleep(attesa) attesa = min(10, attesa + 2) if not SPACE or not CHIAVE or not PUBBLICO: print("Servizio GPU spento: mancano NEURONAL_GPU_SPACE, NEURONAL_GPU_KEY o l'indirizzo pubblico. Neuronal usa la CPU.") sys.exit(2) asyncio.run(servi())