Spaces:
Running
Running
| """Start the same server locally for notebooks and jobs.""" | |
| import json | |
| import os | |
| import socket | |
| import subprocess | |
| import sys | |
| import time | |
| from contextlib import contextmanager | |
| from pathlib import Path | |
| import requests | |
| def local_server( | |
| snapshot, | |
| sessions=16, | |
| web=False, | |
| *, | |
| source_root=None, | |
| cache_dir=None, | |
| local_source=False, | |
| ): | |
| with socket.socket() as sock: | |
| sock.bind(("127.0.0.1", 0)) | |
| port = sock.getsockname()[1] | |
| env = { | |
| **os.environ, | |
| "NAYANA_MAX_SESSIONS": str(sessions), | |
| "ENABLE_WEB_INTERFACE": str(web).lower(), | |
| } | |
| env.pop("NAYANA_CORPUS_MANIFEST", None) | |
| env.pop("NAYANA_SNAPSHOT", None) | |
| for key in ("NAYANA_SOURCE_ROOT", "NAYANA_LOCAL_SOURCE"): | |
| env.pop(key, None) | |
| if str(snapshot).startswith("hf://"): | |
| env["NAYANA_CORPUS_MANIFEST"] = str(snapshot) | |
| else: | |
| path = Path(snapshot).resolve() | |
| manifest = path / "manifest.json" if path.is_dir() else path | |
| storage = json.loads(manifest.read_text()).get("storage") | |
| env[ | |
| "NAYANA_CORPUS_MANIFEST" | |
| if storage == "bucket-parquet" | |
| else "NAYANA_SNAPSHOT" | |
| ] = str(path) | |
| if source_root is not None: | |
| env["NAYANA_SOURCE_ROOT"] = str(Path(source_root).resolve()) | |
| if cache_dir is not None: | |
| env["NAYANA_CACHE_DIR"] = str(Path(cache_dir).resolve()) | |
| if local_source: | |
| env["NAYANA_LOCAL_SOURCE"] = "true" | |
| process = subprocess.Popen( | |
| [ | |
| sys.executable, | |
| "-m", | |
| "uvicorn", | |
| "nayana_ocr.server.app:create_server", | |
| "--factory", | |
| "--host", | |
| "127.0.0.1", | |
| "--port", | |
| str(port), | |
| "--ws", | |
| "websockets", | |
| "--log-level", | |
| "warning", | |
| ], | |
| env=env, | |
| ) | |
| url = f"http://127.0.0.1:{port}" | |
| try: | |
| deadline = time.monotonic() + 90 | |
| while time.monotonic() < deadline: | |
| if process.poll() is not None: | |
| raise RuntimeError(f"Server exited with code {process.returncode}") | |
| try: | |
| if requests.get(f"{url}/healthz", timeout=2).ok: | |
| break | |
| except requests.RequestException: | |
| pass | |
| time.sleep(0.25) | |
| else: | |
| raise TimeoutError("Server did not become healthy in 90 seconds") | |
| yield url | |
| finally: | |
| process.terminate() | |
| try: | |
| process.wait(timeout=10) | |
| except subprocess.TimeoutExpired: | |
| process.kill() | |
| process.wait() | |