Spaces:
Running on RTX PRO 6000
Running on RTX PRO 6000
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| """Process wiring: build the store, bind the servicers, serve.""" | |
| from __future__ import annotations | |
| import asyncio | |
| import logging | |
| import signal | |
| from dataclasses import dataclass | |
| import grpc | |
| from .auth import BearerAuthInterceptor | |
| from .catalog.loader import CatalogLoader | |
| from .health import add_health_service, enter_graceful_shutdown | |
| from .hf.client import HFClient | |
| from .registration import planned_bindings, register | |
| from .services.capabilities import CapabilitiesService | |
| from .services.filefolder import FileFolderService | |
| from .services.fileobject import FileObjectService | |
| from .services.metadata_service import MetadataService | |
| from .settings import Settings | |
| from .store import AssetStore | |
| logger = logging.getLogger(__name__) | |
| class Service: | |
| """A wired-up, not-yet-serving service. Split out so tests can drive it over | |
| an in-process channel without a port.""" | |
| settings: Settings | |
| hf: HFClient | |
| store: AssetStore | |
| loader: CatalogLoader | |
| server: grpc.aio.Server | |
| health: object | |
| async def aclose(self) -> None: | |
| await self.loader.stop() | |
| await self.hf.aclose() | |
| async def build(settings: Settings, *, hf: HFClient | None = None) -> Service: | |
| """Load the catalog and wire everything up. | |
| The catalog is the only thing loaded here. No HuggingFace asset data is | |
| fetched — folders load on first request (plan §5.2) — so startup does not | |
| depend on the Hub being reachable. | |
| """ | |
| hf = hf or HFClient( | |
| endpoint=settings.hf_endpoint, | |
| token=settings.resolve_hf_token(), | |
| timeout=settings.http_timeout, | |
| max_connections=settings.http_max_connections, | |
| max_retries=settings.http_max_retries, | |
| backoff_base=settings.http_backoff_base, | |
| ) | |
| loader = CatalogLoader(settings, hf) | |
| catalog = await loader.load() | |
| store = AssetStore(catalog, hf, settings) | |
| interceptors = [] | |
| if not settings.allow_anonymous: | |
| interceptors.append(BearerAuthInterceptor(settings.resolve_auth_token())) | |
| else: | |
| logger.warning("inbound auth is DISABLED (HF_ADAPTER_ALLOW_ANONYMOUS=true)") | |
| server = grpc.aio.server( | |
| interceptors=interceptors, | |
| maximum_concurrent_rpcs=settings.max_concurrent_rpcs, | |
| ) | |
| # One plan drives both what gets bound and what ListServices reports. | |
| bindings = planned_bindings( | |
| advertise_filefolder=settings.advertise_filefolder, | |
| advertise_v1beta=settings.advertise_v1beta, | |
| ) | |
| capabilities = CapabilitiesService(store, settings, bindings) | |
| register( | |
| server, | |
| bindings=bindings, | |
| capabilities=capabilities, | |
| fileobject=FileObjectService(store, settings), | |
| metadata=MetadataService(store, settings), | |
| filefolder=FileFolderService(store, settings) if settings.advertise_filefolder else None, | |
| ) | |
| logger.info( | |
| "advertising: %s", | |
| ", ".join(f"{name} ({'+'.join(v)})" for name, v in bindings.items()), | |
| ) | |
| # Registered last, and only once everything above succeeded: the catalog is | |
| # loaded and every servicer is bound, which is exactly what SERVING claims. | |
| health = await add_health_service(server, bindings) | |
| await loader.start_refresh(store.swap_catalog) | |
| return Service(settings=settings, hf=hf, store=store, loader=loader, server=server, health=health) | |
| async def serve(settings: Settings | None = None) -> None: | |
| settings = settings or Settings() | |
| logging.basicConfig( | |
| level=settings.log_level.upper(), | |
| format="%(asctime)s %(levelname)s %(name)s %(message)s", | |
| ) | |
| service = await build(settings) | |
| address = f"{settings.grpc_host}:{settings.grpc_port}" | |
| service.server.add_insecure_port(address) | |
| await service.server.start() | |
| logger.info("listening on %s", address) | |
| stop = asyncio.Event() | |
| loop = asyncio.get_running_loop() | |
| for sig in (signal.SIGINT, signal.SIGTERM): | |
| loop.add_signal_handler(sig, stop.set) | |
| try: | |
| await stop.wait() | |
| finally: | |
| # NOT_SERVING first, so Kubernetes pulls the pod from the Service | |
| # endpoints before we start refusing connections. | |
| logger.info("draining...") | |
| await enter_graceful_shutdown(service.health) | |
| await service.server.stop(grace=10) | |
| await service.aclose() | |