File size: 2,895 Bytes
e4f9cbe
 
 
 
dbed4d4
b708786
e4f9cbe
 
 
 
 
55dc3dd
e4f9cbe
 
55dc3dd
e4f9cbe
 
dbed4d4
e4f9cbe
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8796ec1
b708786
 
dbed4d4
55dc3dd
dbed4d4
 
 
 
 
 
 
 
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
"""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 data_path
from .router_utils import RouteErrorHandler
from .tasks import task_manager
from .utils import get_datasets_dir

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.
  tmpdir = os.path.join(data_path(), '.hf-spaces', repo_id)
  snapshot_download(repo_id=repo_id, repo_type='space', local_dir=tmpdir)

  shutil.move(get_datasets_dir(tmpdir), os.path.join(get_datasets_dir(data_path()), '..'))


@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())