Skip to content

Commit 470fbe9

Browse files
alando46andreaskoepfolliestanley
authored
Inference Documentation (#3119)
This is a mostly done (although not totally complete) PR with a technical overview of the inference architecture. I'm looking forward to high level feedback (general layout, flow of documentation) or specific suggestions (I'm sure I made some errors or missed some details.) I will try to wrap up the final section soon. See related discussion on the issue: #1473 (comment) --------- Co-authored-by: Andreas Koepf <andreas.koepf@provisio.com> Co-authored-by: Oliver Stanley <olivergestanley@gmail.com>
1 parent d381aa1 commit 470fbe9

3 files changed

Lines changed: 170 additions & 0 deletions

File tree

‎docs/docs/architecture/README.md‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
# Architecture
2+
3+
## Inference
4+
5+
The Inference architecture is comprised of several core components: a text, or
6+
frontend client, a FastAPI webserver, a database with several tables, Reddis
7+
used for queueing, and distributed gpu workers.
8+
9+
A more detailed overview can be viewed [here](inference.md).

‎docs/docs/architecture/inference.md‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
# Inference
2+
3+
### From the perspective of the Text Client
4+
5+
```mermaid
6+
---
7+
title: Overview of Text Client Workflow
8+
---
9+
graph LR;
10+
subgraph OA Server
11+
/auth/callback/debug
12+
id0["/chats"]
13+
id1["/chats/{chat_id}/messages"]
14+
id2["{backend_url}/chats/{chat_id}/messages/{message_id}/events"]
15+
end
16+
17+
subgraph text-client
18+
id3["SSEClient(response)"]
19+
end
20+
21+
text-client == 1. GET ==> /auth/callback/debug;
22+
/auth/callback/debug -- bearer_token --> text-client;
23+
24+
text-client == 2. POST ==> id0["/chats"]
25+
id0["/chats"] -- chat_id --> text-client;
26+
27+
text-client == 3. requests prompt ==> user;
28+
user == submits prompt ==> text-client;
29+
30+
text-client == 4. POST ==> id1["/chats/{chat_id}/messages"];
31+
id1["/chats/{chat_id}/messages"] -- message_id --> text-client;
32+
33+
text-client == 5. GET ==> id2["{backend_url}/chats/{chat_id}/messages/{message_id}/events"];
34+
id2["{backend_url}/chats/{chat_id}/messages/{message_id}/events"] -- events --> id3["SSEClient(response)"];
35+
36+
linkStyle 0 stroke-width:2px,fill:none,stroke:green;
37+
linkStyle 1 stroke-width:2px,fill:none,stroke:green;
38+
linkStyle 2 stroke-width:2px,fill:none,stroke:purple;
39+
linkStyle 3 stroke-width:2px,fill:none,stroke:purple;
40+
linkStyle 4 stroke-width:2px,fill:none,stroke:blue;
41+
linkStyle 5 stroke-width:2px,fill:none,stroke:blue;
42+
linkStyle 6 stroke-width:2px,fill:none,stroke:grey;
43+
linkStyle 7 stroke-width:2px,fill:none,stroke:grey;
44+
```
45+
46+
For development, a basic REPL client is used as a chat interface, built around
47+
[Typer](https://typer.tiangolo.com/). The bulk of the logic is in
48+
`inference.text-client.text_client_utils.DebugClient`. The basic steps are as
49+
follows:
50+
51+
1. The debug client first authenticates with a GET request to
52+
`/auth/callback/debug` which returns a bearer token, which is then set in
53+
future request headers.
54+
55+
2. After authenticating, the `DebugClient` creates a "chat" by posting to
56+
`/chats`, which returns a `chat_id` for session.
57+
58+
3. The script then collects the user prompt.
59+
60+
4. The `DebugClient` posts to the endpoint `/chats/{chat_id}/messages`. Included
61+
in this request is the message content, a `parent_id:str` (ID of assistant's
62+
response to prior message, if there is one), the `model_config_name:str`, and
63+
inference `sampling_parameters:dict`. In the response, the server will return
64+
a `message_id:str`.
65+
66+
5. The client will use this id to make a GET request to
67+
`{backend_url}/chats/{chat_id}/messages/{message_id}/events`. Critically, in
68+
this `get()` method of requests, the `stream=True` is passed, meaning that
69+
the response content will not be immediately downloaded but rather streamed.
70+
Additionally, the actual GET request must have
71+
`"Accept": "text/event-stream"` in the headers to let the server know we are
72+
awaiting an event stream.
73+
74+
6. The response is then used to instantiate an SSEClient, which - via its
75+
events() method, returns an iterable that can be used to print out inference
76+
results, one token at a time.
77+
78+
7. After exhausting the events iterable (ie inference is complete), the user is
79+
prompted for a new message.
80+
81+
### From the perspective of the OA Inference Server
82+
83+
The inference server is built around [FastAPI](https://fastapi.tiangolo.com/).
84+
85+
1. When the client posts to `/chats`, the UserChatRepository - an interface
86+
between application logic and chats message tables - creates a chat in the
87+
chat table, which assigns the needed chat id and is returned in the response
88+
to the client.
89+
90+
2. Next, the client POSTs to `/chats/{chat_id}/prompter_message`, which uses the
91+
`UserChatRepository` to:
92+
93+
1. Check to see if the message content included in the POST request is longer
94+
than the configured settings.
95+
2. Using the database session and `user_id` which was configured when this
96+
instance of `UserChatRepository` was instantiated, we lookup in the `chat`
97+
table the corresponding entry for this `chat_id` and `user_id`.
98+
3. Check to see if the total number of assistant + prompter messages in the
99+
`chat` is greater than the configured settings.
100+
4. If this is a new chat (i.e. no `parent_id` is set in the `chat`), and the
101+
chat table contains no title for chat of that id, updates the chat table
102+
with the initial user prompt as title.
103+
5. Creates a message in the message table for the user prompt.
104+
105+
3. The client then POSTs to `/chats/{chat_id}/assistant_message`.
106+
107+
1. First we load the model config for the `model_config_name` specified by
108+
the client's request.
109+
2. Then, we use the `UserChatRepository` to post to create a message in the
110+
`message` table, with a pending state. Note, we also will update the state
111+
for any other currently pending messages in the chat to
112+
`inference.MessageState.cancelled`.
113+
3. After updating the `message` table, we create a RedisQueue for this
114+
specific message and enque the message.
115+
4. Finally, we return an `inference.MessageRead` (a Pydantic model) to the
116+
client. This is the object contains the needed `message_id`.
117+
118+
4. Once the client has the `message_id`, it will make a GET request to
119+
`/chats/{chat_id}/messages/{message_id}/events`. Upon receiving this request,
120+
the server retrieves the message from the message table and verifies that the
121+
message isn't finished and that the message role is assistant.
122+
123+
After this, we get the Redis queue for this specific message id. From here, we
124+
poll the queue, using the `dequeue()` method. By default this method will block
125+
for up to 1 second, however if a message is added to the queue in that interval,
126+
it will automatically be returned. Otherwise, `dequeue()` will return `None`.
127+
128+
If `None` is returned and this is the first time we have polled the queue since
129+
its creation, we will yield a `chat_schema.PendingResponseEvent`. The loop will
130+
continue, with 1 second blocking (by default), until a message item is returned.
131+
A message item is a tuple where the first element is a string of the message id,
132+
and the second element is a response. After some checks on the message response
133+
type (for example if there is a safety intervention or an internal error from
134+
the worker), we yield a `chat_schema.MessageResponseEvent`. The message polling
135+
and yielding is bundled in an `EventSourceResponse` object.
136+
137+
The EventSourceResponse is Server Sent Events (SSE) plugin for
138+
Starlette/FastAPI, which takes content and returns it as part of a HTTP event
139+
stream. You can check out the EventSourceResposethe library's
140+
[source code](https://github.com/sysid/sse-starlette/tree/master) for more
141+
details about how this works.
142+
143+
### From the perspective of the OA Worker
144+
145+
**This section is not yet written**. If you are interested in helping write it
146+
please get in touch on GitHub or Discord.

‎docs/sidebars.js‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,21 @@ const sidebars = {
3232
},
3333
],
3434
},
35+
{
36+
type: "category",
37+
label: "Architecture",
38+
link: {
39+
type: "doc",
40+
id: "architecture/README",
41+
},
42+
items: [
43+
"architecture/inference",
44+
{
45+
type: "doc",
46+
id: "architecture/inference",
47+
},
48+
],
49+
},
3550
{
3651
type: "category",
3752
label: "Plugins",

0 commit comments

Comments
 (0)