Spaces:
Paused
Paused
File size: 2,608 Bytes
3ed7ad0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | """Controller entry point: `flock-controller`."""
from __future__ import annotations
import asyncio
import logging
import os
from flock.controller.api import Notifier, create_app
from flock.controller.jobs import FakeJobsClient, HubJobsClient, JobsClient
from flock.controller.reconcile import Settings, run_reconciler
from flock.controller.store import Store
log = logging.getLogger("flock.controller")
STATS_LOG_INTERVAL = 60.0
def _public_url(port: int) -> str:
if os.environ.get("FLOCK_URL"):
return os.environ["FLOCK_URL"]
if os.environ.get("SPACE_HOST"): # set by HF Spaces
return f"https://{os.environ['SPACE_HOST']}"
return f"http://localhost:{port}"
def parse_labels(raw: str) -> dict[str, str]:
"""'a=1,b=2' -> {'a': '1', 'b': '2'} (FLOCK_JOB_LABELS)."""
out = {}
for part in raw.split(","):
if part.strip():
k, _, v = part.partition("=")
out[k.strip()] = v.strip()
return out
def build_app():
port = int(os.environ.get("PORT", "7860"))
token = os.environ.get("FLOCK_TOKEN", "")
if not token:
raise SystemExit("FLOCK_TOKEN must be set")
store = Store(os.environ.get("FLOCK_DATA", "./data"))
jobs: JobsClient
if os.environ.get("FLOCK_FAKE_JOBS") == "1":
jobs = FakeJobsClient()
else:
jobs = HubJobsClient(token=os.environ.get("HF_TOKEN"))
settings = Settings(
public_url=_public_url(port),
token=token,
job_labels=parse_labels(os.environ.get("FLOCK_JOB_LABELS", "")),
)
interval = float(os.environ.get("FLOCK_RECONCILE_INTERVAL", "5"))
def stats() -> dict:
calls = getattr(jobs, "calls", None)
return {"hub_calls": dict(calls)} if calls is not None else {}
async def log_stats() -> None:
while True:
await asyncio.sleep(STATS_LOG_INTERVAL)
log.info("hub calls: %s", stats().get("hub_calls", {}))
async def background(notifier: Notifier) -> None:
stats_task = asyncio.create_task(log_stats())
try:
await run_reconciler(
store, jobs, settings, interval, wake=lambda t: notifier.wait("reconcile:", t)
)
finally:
stats_task.cancel()
return create_app(store, token, background=background, stats=stats), port
def main() -> None:
import uvicorn
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s")
app, port = build_app()
uvicorn.run(app, host="0.0.0.0", port=port, log_level="info")
if __name__ == "__main__":
main()
|