Spaces:
Runtime error
Runtime error
File size: 3,446 Bytes
e4f9cbe dbed4d4 b708786 e4f9cbe 55dc3dd e4f9cbe 0dc939d e4f9cbe 54369d2 e4f9cbe 8796ec1 b708786 dbed4d4 55dc3dd dbed4d4 54369d2 0dc939d 54369d2 0dc939d dbed4d4 54369d2 ba8f8ef 54369d2 ba8f8ef 54369d2 bb6ad73 7b46386 e4f9cbe |
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 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
"""Serves the Lilac server."""
import logging
import os
import shutil
from typing import Any
from fastapi import APIRouter, FastAPI
from fastapi.responses import ORJSONResponse
from fastapi.routing import APIRoute
from fastapi.staticfiles import StaticFiles
from huggingface_hub import snapshot_download
from . import router_concept, router_data_loader, router_dataset, router_signal, router_tasks
from .config import CONFIG, data_path
from .router_utils import RouteErrorHandler
from .tasks import task_manager
from .utils import get_dataset_output_dir, list_datasets
DIST_PATH = os.path.abspath(os.path.join('web', 'blueprint', 'build'))
tags_metadata: list[dict[str, Any]] = [{
'name': 'datasets',
'description': 'API for querying a dataset.',
}, {
'name': 'concepts',
'description': 'API for managing concepts.',
}, {
'name': 'data_loaders',
'description': 'API for loading data.',
}, {
'name': 'signals',
'description': 'API for managing signals.',
}]
def custom_generate_unique_id(route: APIRoute) -> str:
"""Generate the name for the API endpoint."""
return route.name
app = FastAPI(
default_response_class=ORJSONResponse,
generate_unique_id_function=custom_generate_unique_id,
openapi_tags=tags_metadata)
v1_router = APIRouter(route_class=RouteErrorHandler)
v1_router.include_router(router_dataset.router, prefix='/datasets', tags=['datasets'])
v1_router.include_router(router_concept.router, prefix='/concepts', tags=['concepts'])
v1_router.include_router(router_data_loader.router, prefix='/data_loaders', tags=['data_loaders'])
v1_router.include_router(router_signal.router, prefix='/signals', tags=['signals'])
v1_router.include_router(router_tasks.router, prefix='/tasks', tags=['tasks'])
app.include_router(v1_router, prefix='/api/v1')
# Serve static files in production mode.
app.mount('/', StaticFiles(directory=os.path.join(DIST_PATH), html=True, check_dir=False))
@app.on_event('startup')
def startup() -> None:
"""Download dataset files from the HF space that is uploaded before building the image."""
# Setup.
repo_id = 'lilacai/nikhil_staging'
# Download the huggingface space data. This includes code and datasets, so we move the datasets
# alone to the data directory.
spaces_download_dir = os.path.join(data_path(), '.hf-spaces', repo_id)
snapshot_download(
repo_id=repo_id,
repo_type='space',
local_dir=spaces_download_dir,
local_dir_use_symlinks=False,
token=CONFIG['HF_ACCESS_TOKEN'])
for dataset in list_datasets(spaces_download_dir):
dataset_output_dir = get_dataset_output_dir(data_path(), dataset.namespace,
dataset.dataset_name)
shutil.rmtree(dataset_output_dir)
print('~~~~moving', os.path.join(spaces_download_dir, dataset.namespace, dataset.dataset_name),
'to', dataset_output_dir)
shutil.move(
os.path.join(spaces_download_dir, dataset.namespace, dataset.dataset_name),
dataset_output_dir)
@app.on_event('shutdown')
async def shutdown_event() -> None:
"""Kill the task manager when FastAPI shuts down."""
await task_manager().stop()
class GetTasksFilter(logging.Filter):
"""Task filter for /tasks."""
def filter(self, record: logging.LogRecord) -> bool:
"""Filters out /api/v1/tasks/ from the logs."""
return record.getMessage().find('/api/v1/tasks/') == -1
logging.getLogger('uvicorn.access').addFilter(GetTasksFilter())
|