Function _init_api

checkpoint_engine/api.py:46–108

(ps: ParameterServer)

def _init_api(ps: ParameterServer) -> Any:
    app = fastapi.FastAPI()

    class RegisterRequest(BaseModel):
        files: list[str]

    class UpdateRequest(BaseModel):
        ranks: list[int] = []
        update_url: str | None = None
        inference_group_ranks: list[int] = []
        timeout: float = 300.0
        uds: str | None = None

    def wrap_exception(func: Callable[[], None]) -> Response:
        try:
            func()
        except Exception as e:  # noqa: BLE001
            logger.exception(f"wrap exception {func} failed")
            return JSONResponse(content=str(e), status_code=500)
        return Response(status_code=200)

    @app.post("/v1/checkpoints/{checkpoint_name}/files")
    async def register_files(checkpoint_name: str, req: RegisterRequest, raw: Request) -> Response:
        return wrap_exception(lambda: ps.register_checkpoint(checkpoint_name, files=req.files))

    @app.delete("/v1/checkpoints/{checkpoint_name}")
    async def unregister_checkpoint(checkpoint_name: str) -> Response:
        return wrap_exception(lambda: ps.unregister_checkpoint(checkpoint_name))

    @app.get("/v1/healthz")
    async def healthz() -> Response:
        return Response(status_code=200)

    @app.post("/v1/checkpoints/{checkpoint_name}/gather-metas")
    async def gather_metas(checkpoint_name: str) -> Response:
        return wrap_exception(lambda: ps.gather_metas(checkpoint_name))

    @app.get("/v1/metas")
    async def get_metas() -> dict[int, MemoryBufferMetaList]:
        try:
            return ps.get_metas()
        except Exception as e:
            logger.exception("get_metas failed")
            raise HTTPException(status_code=500, detail=str(e)) from e

    @app.post("/v1/metas")
    async def load_metas(metas: dict[int, MemoryBufferMetaList]) -> Response:
        return wrap_exception(lambda: ps.load_metas(metas))

    @app.post("/v1/checkpoints/{checkpoint_name}/update")
    async def update(checkpoint_name: str, req: UpdateRequest) -> Response:
        def update_func(socket_paths: list[tuple[str, str]]):
            if req.update_url is None:
                return
            if req.inference_group_ranks:
                socket_paths = [socket_paths[i] for i in req.inference_group_ranks]
            request_inference_to_update(
                req.update_url, dict(socket_paths), timeout=req.timeout, uds=req.uds