Skip to content

Commit 50cb3b0

Browse files
olliestanleyshahules786
olliestanley
and
shahules786
authored
Add safety server to inference (LAION-AI#2449)
Close LAION-AI#2447. --------- Co-authored-by: Shahules786 <Shahules786@gmail.com>
1 parent b9be654 commit 50cb3b0

12 files changed

Lines changed: 238 additions & 2 deletions

File tree

‎.github/workflows/release.yaml‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,8 +75,21 @@ jobs:
7575
context: .
7676
dockerfile: docker/inference/Dockerfile.worker-standalone
7777
build-args: ""
78+
build-inference-safety:
79+
uses: ./.github/workflows/docker-build.yaml
80+
needs: pre-commit
81+
with:
82+
image-name: oasst-inference-safety
83+
context: .
84+
dockerfile: docker/inference/Dockerfile.safety
85+
build-args: ""
7886
deploy-to-node:
79-
needs: [build-backend, build-web, build-bot, build-inference-server]
87+
needs:
88+
- build-backend
89+
- build-web
90+
- build-bot
91+
- build-inference-server
92+
- build-inference-safety
8093
uses: ./.github/workflows/deploy-to-node.yaml
8194
secrets: inherit
8295
with:

‎docker-compose.yaml‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -256,6 +256,18 @@ services:
256256
replicas: 1
257257
profiles: ["inference"]
258258

259+
inference-safety:
260+
build:
261+
dockerfile: docker/inference/Dockerfile.safety
262+
context: .
263+
image: oasst-inference-safety:dev
264+
environment:
265+
PORT: 8002
266+
volumes:
267+
- "./oasst-shared:/opt/inference/lib/oasst-shared"
268+
- "./inference/safety:/opt/inference/safety"
269+
profiles: ["inference"]
270+
259271
prometheus:
260272
image: prom/prometheus
261273
container_name: prometheus

‎docker/inference/Dockerfile.safety‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
# syntax=docker/dockerfile:1
2+
3+
ARG MODULE="inference"
4+
ARG SERVICE="safety"
5+
6+
ARG APP_USER="${MODULE}-${SERVICE}"
7+
ARG APP_RELATIVE_PATH="${MODULE}/${SERVICE}"
8+
9+
10+
FROM python:3.10-slim as build
11+
ARG APP_RELATIVE_PATH
12+
13+
WORKDIR /build
14+
15+
RUN apt-get update && apt-get install -y --no-install-recommends \
16+
git \
17+
&& rm -rf /var/lib/apt/lists/*
18+
19+
COPY ./${APP_RELATIVE_PATH}/requirements.txt requirements.txt
20+
21+
RUN --mount=type=cache,target=/var/cache/pip \
22+
pip install \
23+
--cache-dir=/var/cache/pip \
24+
--target=lib \
25+
-r requirements.txt
26+
27+
28+
29+
FROM python:3.10-slim as base-env
30+
ARG APP_USER
31+
ARG APP_RELATIVE_PATH
32+
ARG MODULE
33+
ARG SERVICE
34+
35+
ENV APP_BASE="/opt/${MODULE}"
36+
ENV APP_ROOT="${APP_BASE}/${SERVICE}"
37+
ENV APP_LIBS="/var/opt/${APP_RELATIVE_PATH}/lib"
38+
ENV SHARED_LIBS_BASE="${APP_BASE}/lib"
39+
40+
ENV PATH="${PATH}:${APP_LIBS}/bin"
41+
ENV PYTHONPATH="${PYTHONPATH}:${APP_LIBS}"
42+
ENV PORT="8002"
43+
44+
45+
RUN adduser \
46+
--disabled-password \
47+
"${APP_USER}"
48+
49+
USER ${APP_USER}
50+
51+
WORKDIR ${APP_ROOT}
52+
53+
54+
COPY --chown="${APP_USER}:${APP_USER}" --from=build /build/lib ${APP_LIBS}
55+
COPY --chown="${APP_USER}:${APP_USER}" ./${APP_RELATIVE_PATH}/*.py .
56+
57+
58+
CMD python3 __main__.py
59+
60+
FROM base-env as prod
61+
ARG APP_USER
62+
63+
64+
COPY --chown="${APP_USER}:${APP_USER}" ./oasst-shared /tmp/lib/oasst-shared
65+
RUN --mount=type=cache,target=/var/cache/pip,from=base-env \
66+
pip install \
67+
--cache-dir=/var/cache/pip \
68+
--target="${APP_LIBS}" \
69+
/tmp/lib/oasst-shared
70+
71+
COPY --chown="${APP_USER}:${APP_USER}" ./inference/safety/safety_main.sh /entrypoint.sh
72+
73+
CMD ["/entrypoint.sh"]

‎inference/safety/README.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
# OpenAssistant Inference Safety Server
2+
3+
Basic FastAPI server to serve safety models using
4+
[Blade2Blade](https://github.com/LAION-AI/blade2blade/).

‎inference/safety/main.py‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
# A FastAPI server to run the safety pipeline
2+
3+
import asyncio
4+
from concurrent.futures import ThreadPoolExecutor
5+
6+
import fastapi
7+
import uvicorn
8+
from blade2blade import Blade2Blade
9+
from loguru import logger
10+
from oasst_shared.schemas import inference
11+
from settings import settings
12+
13+
app = fastapi.FastAPI()
14+
15+
16+
@app.middleware("http")
17+
async def log_exceptions(request: fastapi.Request, call_next):
18+
try:
19+
response = await call_next(request)
20+
except Exception:
21+
logger.exception("Exception in request")
22+
raise
23+
return response
24+
25+
26+
pipeline_loaded: bool = False
27+
pipeline: Blade2Blade
28+
executor = ThreadPoolExecutor()
29+
30+
31+
@app.on_event("startup")
32+
async def load_pipeline():
33+
global pipeline_loaded, pipeline
34+
pipeline = Blade2Blade(settings.safety_model_name)
35+
# warmup
36+
input = "|prompter|Hey,how are you?|endoftext|"
37+
_ = pipeline.predict(input)
38+
pipeline_loaded = True
39+
logger.info("Pipeline loaded")
40+
41+
42+
async def async_predict(pipeline: Blade2Blade, inputs: str):
43+
return await asyncio.get_event_loop().run_in_executor(executor, pipeline.predict, inputs)
44+
45+
46+
@app.post("/safety", response_model=inference.SafetyResponse)
47+
async def safety(request: inference.SafetyRequest):
48+
global pipeline_loaded, pipeline
49+
while not pipeline_loaded:
50+
await asyncio.sleep(1)
51+
outputs = await async_predict(pipeline, request.inputs)
52+
return inference.SafetyResponse(outputs=outputs)
53+
54+
55+
@app.get("/health")
56+
async def health():
57+
if not pipeline_loaded:
58+
raise fastapi.HTTPException(status_code=503, detail="Server not fully loaded")
59+
return {"status": "ok"}
60+
61+
62+
if __name__ == "__main__":
63+
uvicorn.run(app, host="0.0.0.0", port=8008)

‎inference/safety/requirements.txt‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
git+https://github.com/LAION-AI/blade2blade@8fd43bcbc5ff35fd59663c77ef08b3ec6c239dd4#egg=blade2blade
2+
fastapi
3+
loguru
4+
pydantic
5+
uvicorn

‎inference/safety/safety_main.sh‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
#!/bin/bash
2+
3+
uvicorn main:app --host 0.0.0.0 --port "$PORT"

‎inference/safety/settings.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
import pydantic
2+
3+
4+
class Settings(pydantic.BaseSettings):
5+
safety_model_name: str = "shahules786/blade2blade-t5-base"
6+
7+
8+
settings = Settings()

‎inference/worker/README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# OpenAssitant Workers
1+
# OpenAssistant Workers
22

33
## Running the worker
44

‎inference/worker/settings.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ class Settings(pydantic.BaseSettings):
55
backend_url: str = "ws://localhost:8000"
66
model_config_name: str = "distilgpt2"
77
inference_server_url: str = "http://localhost:8001"
8+
safety_server_url: str = "http://localhost:8002"
89
api_key: str = "0000"
910

1011
oa_protocol_version: str = "v2"
@@ -23,5 +24,7 @@ class Settings(pydantic.BaseSettings):
2324
basic_auth_username: str | None = None
2425
basic_auth_password: str | None = None
2526

27+
enable_safety: bool = False
28+
2629

2730
settings = Settings()

0 commit comments

Comments
 (0)