Repository: mosecorg/mosec (language: python)

Issue:
[FEATURE] Streaming output
Triton supports returning the response in stream by `triton_client.async_stream_infer`. Hope mosec can have a similar function, which is useful for language models


Discussion / hints:
@VoVAllen  this would be necessary to support streaming response for large language models as they usually require a longer response time, which wears down users' patience. In this case, generating the response in the stream (e.g., word by word) would significantly improve users' experience. Btw, ChatGPT is exactly acting like this.
```
curl 'https://chat.openai.com/backend-api/conversation' \
  -H 'authority: chat.openai.com' \
  -H 'accept: text/event-stream' \
  -H 'accept-language: zh-CN,zh;q=0.9' \
  -H 'content-type: application/json' \
  -H 'origin: https://chat.openai.com' \
  -H 'referer: https://chat.openai.com/chat' \
  -H 'sec-ch-ua: "Google Chrome";v="105", "Not)A;Brand";v="8", "Chromium";v="105"' \
  -H 'sec-ch-ua-mobile: ?0' \
  -H 'sec-ch-ua-platform: "Linux"' \
  -H 'sec-fetch-dest: empty' \
  -H 'sec-fetch-mode: cors' \
  -H 'sec-fetch-site: same-origin' \
  -H 'user-agent: Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/105.0.0.0 Safari/537.36' \
  --data-raw '{"action":"next","messages":[{"id":"1a5f62f1-9c1a-4485-8f74-df2fd3723589","role":"user","content":{"content_type":"text","parts":["envd is here https://github.com/tensorchord/envd\n"]}}],"conversation_id":"409004f7-0e7e-4b66-a264-09bd909d12db","parent_message_id":"6db6f583-d9d0-430a-a929-d3a4be84e75a","model":"text-davinci-002-render-sha"}' \
  --compressed
```

This is the ChatGPT request

It is a `text/event-stream'` req


I will work on this one.

So far, [SSE](https://en.wikipedia.org/wiki/Server-sent_events) should be enough for this feature.

/assign
I think supporting websocket like MOSS did would be a more elegant solution than using SSE. Just let mosec spin up multiple ws is enough for streaming output.


Pros of ws:
- Disentangle with the inference endpoint. More flexibility, such as mixing streaming output with non-streaming output.
- More control of the inference process. For example, user can use custom protocol to interrupt the streaming inference process if needed. 
- `ws.write_msg` is more intuitive than writing generator with `yield xxx`, and much easier to understand
> I think supporting websocket like MOSS did would be a more elegant solution than using SSE. Just let mosec spin up multiple ws is enough for streaming output.

Do you mean we don't need to add any features? Because WebSocket is more like user's choice. Nothing prevents users from adding another communication tunnel like WebSocket or gRPC or something else.
Mosec can add ws ability so that ws can be on the same endpoint (or same port) instead of a new port. And mosec can provide wrapper in client SDK to enable user to use streaming without knowing websockets. Otherwise, I think it's hard to support both streaming and non-streaming scenario in the same function.
This is also "make the easy things easy, and the hard things possible". 
> Mosec can add ws ability so that ws can be on the same endpoint (or same port) instead of a new port. And mosec can provide wrapper in client SDK to enable user to use streaming without knowing websockets. Otherwise, I think it's hard to support both streaming and non-streaming scenario in the same function. This is also "make the easy things easy, and the hard things possible".

But how?

- Do you mean add WebSocket to the rust web layer?
- If we accept streaming input, how can we know which part should be dynamically batched? If we know, then it might be possible to support both streaming and non-streaming simultaneously (maybe in different functions).
- We don't provide SDK because we don't know which language to support, also HTTP is a general one so we don't need to write an HTTP client SDK.


Fix patch:
diff --git a/Cargo.lock b/Cargo.lock
index 764a753c..ca8d7eb2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -41,12 +41,72 @@ dependencies = [
  "futures-core",
 ]
 
+[[package]]
+name = "async-trait"
+version = "0.1.68"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b9ccdd8f2a161be9bd5c023df56f1b2a0bd1d83872ae53b71a84a12c9bf6e842"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.13",
+]
+
 [[package]]
 name = "autocfg"
 version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa"
 
+[[package]]
+name = "axum"
+version = "0.6.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "349f8ccfd9221ee7d1f3d4b33e1f8319b3a81ed8f61f2ea40b37b859794b4491"
+dependencies = [
+ "async-trait",
+ "axum-core",
+ "bitflags",
+ "bytes",
+ "futures-util",
+ "http",
+ "http-body",
+ "hyper",
+ "itoa",
+ "matchit",
+ "memchr",
+ "mime",
+ "percent-encoding",
+ "pin-project-lite",
+ "rustversion",
+ "serde",
+ "serde_json",
+ "serde_path_to_error",
+ "serde_urlencoded",
+ "sync_wrapper",
+ "tokio",
+ "tower",
+ "tower-layer",
+ "tower-service",
+]
+
+[[package]]
+name = "axum-core"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b2f958c80c248b34b9a877a643811be8dbca03ca5ba827f2b63baf3a81e5fc4e"
+dependencies = [
+ "async-trait",
+ "bytes",
+ "futures-util",
+ "http",
+ "http-body",
+ "mime",
+ "rustversion",
+ "tower-layer",
+ "tower-service",
+]
+
 [[package]]
 name = "bitflags"
 version = "1.3.2"
@@ -114,6 +174,15 @@ version = "1.0.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
 
+[[package]]
+name = "form_urlencoded"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a9c384f161156f5260c24a097c56119f9be8c798586aecc13afbcbe7b7e26bf8"
+dependencies = [
+ "percent-encoding",
+]
+
 [[package]]
 name = "futures-channel"
 version = "0.3.27"
@@ -250,12 +319,24 @@ dependencies = [
  "cfg-if",
 ]
 
+[[package]]
+name = "matchit"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b87248edafb776e59e6ee64a79086f65890d3510f2c656c000bf2a7e8a0aea40"
+
 [[package]]
 name = "memchr"
 version = "2.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2dffe52ecf27772e601905b7522cb4ef790d2cc203488bbd0e2fe85fcb74566d"
 
+[[package]]
+name = "mime"
+version = "0.3.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
+
 [[package]]
 name = "mio"
 version = "0.8.6"
@@ -274,6 +355,7 @@ version = "0.6.4"
 dependencies = [
  "argh",
  "async-channel",
+ "axum",
  "bytes",
  "derive_more",
  "hyper",
@@ -348,6 +430,32 @@ dependencies = [
  "windows-sys",
 ]
 
+[[package]]
+name = "percent-encoding"
+version = "2.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "478c572c3d73181ff3c2539045f6eb99e5491218eae919370993b890cdbdd98e"
+
+[[package]]
+name = "pin-project"
+version = "1.0.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ad29a609b6bcd67fee905812e544992d216af9d755757c05ed2d0e15a74c6ecc"
+dependencies = [
+ "pin-project-internal",
+]
+
+[[package]]
+name = "pin-project-internal"
+version = "1.0.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "069bdb1e05adc7a8990dce9cc75370895fbe4e3d58b9b73bf1aee56359344a55"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 1.0.109",
+]
+
 [[package]]
 name = "pin-project-lite"
 version = "0.2.9"
@@ -362,9 +470,9 @@ checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184"
 
 [[package]]
 name = "proc-macro2"
-version = "1.0.52"
+version = "1.0.55"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1d0e1ae9e836cc3beddd63db0df682593d7e2d3d891ae8c9083d2113e1744224"
+checksum = "1d0dd4be24fcdcfeaa12a432d588dc59bbad6cad3510c67e74a2b6b2fc950564"
 dependencies = [
  "unicode-ident",
 ]
@@ -417,6 +525,12 @@ dependencies = [
  "semver",
 ]
 
+[[package]]
+name = "rustversion"
+version = "1.0.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4f3208ce4d8448b3f3e7d168a73f5e0c43a61e32930de3bceeccedb388b6bf06"
+
 [[package]]
 name = "ryu"
 version = "1.0.13"
@@ -452,6 +566,27 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "serde_path_to_error"
+version = "0.1.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f7f05c1d5476066defcdfacce1f52fc3cae3af1d3089727100c02ae92e5abbe0"
+dependencies = [
+ "serde",
+]
+
+[[package]]
+name = "serde_urlencoded"
+version = "0.7.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd"
+dependencies = [
+ "form_urlencoded",
+ "itoa",
+ "ryu",
+ "serde",
+]
+
 [[package]]
 name = "sharded-slab"
 version = "0.1.4"
@@ -499,15 +634,21 @@ dependencies = [
 
 [[package]]
 name = "syn"
-version = "2.0.0"
+version = "2.0.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4cff13bb1732bccfe3b246f3fdb09edfd51c01d6f5299b7ccd9457c2e4e37774"
+checksum = "4c9da457c5285ac1f936ebd076af6dac17a61cfe7826f2076b4d015cf47bc8ec"
 dependencies = [
  "proc-macro2",
  "quote",
  "unicode-ident",
 ]
 
+[[package]]
+name = "sync_wrapper"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2047c6ded9c721764247e62cd3b03c09ffc529b2ba5b10ec482ae507a4a70160"
+
 [[package]]
 name = "thiserror"
 version = "1.0.40"
@@ -525,7 +666,7 @@ checksum = "f9456a42c5b0d803c8cd86e73dd7cc9edd429499f37a3550d286d5e86720569f"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.0",
+ "syn 2.0.13",
 ]
 
 [[package]]
@@ -597,6 +738,28 @@ dependencies = [
  "syn 1.0.109",
 ]
 
+[[package]]
+name = "tower"
+version = "0.4.13"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b8fa9be0de6cf49e536ce1851f987bd21a43b771b09473c3549a6c853db37c1c"
+dependencies = [
+ "futures-core",
+ "futures-util",
+ "pin-project",
+ "pin-project-lite",
+ "tokio",
+ "tower-layer",
+ "tower-service",
+ "tracing",
+]
+
+[[package]]
+name = "tower-layer"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c20c8dbed6283a09604c3e69b4b7eeb54e298b8a600d4d5ecb5ad39de609f1d0"
+
 [[package]]
 name = "tower-service"
 version = "0.3.2"
@@ -610,6 +773,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8ce8c33a8d48bd45d624a6e523445fd21ec13d3653cd51f681abf67418f54eb8"
 dependencies = [
  "cfg-if",
+ "log",
  "pin-project-lite",
  "tracing-attributes",
  "tracing-core",
diff --git a/Cargo.toml b/Cargo.toml
index f1d6f8c5..12b15e9b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -23,3 +23,4 @@ async-channel = { version = "1" }
 once_cell = "1.17"
 prometheus = "0.13"
 argh = "0.1"
+axum = "0.6.12"
diff --git a/src/main.rs b/src/main.rs
index 7e1bfecb..dd37b28c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -21,9 +21,12 @@ mod tasks;
 
 use std::net::SocketAddr;
 
+use axum::{
+    routing::{get, post},
+    Router,
+};
 use bytes::Bytes;
-use hyper::service::{make_service_fn, service_fn};
-use hyper::{body::to_bytes, header::HeaderValue, Body, Method, Request, Response, StatusCode};
+use hyper::{body::to_bytes, header::HeaderValue, Body, Request, Response, StatusCode};
 use prometheus::{Encoder, TextEncoder};
 use tokio::signal::unix::{signal, SignalKind};
 use tracing::info;
@@ -38,7 +41,6 @@ use crate::tasks::{TaskCode, TaskManager};
 
 const SERVER_INFO: &str = concat!(env!("CARGO_PKG_NAME"), "/", env!("CARGO_PKG_VERSION"));
 const RESPONSE_DEFAULT: &[u8] = b"MOSEC service";
-const RESPONSE_NOT_FOUND: &[u8] = b"not found";
 const RESPONSE_EMPTY: &[u8] = b"no data provided";
 const RESPONSE_SHUTDOWN: &[u8] = b"gracefully shutting down";
 
@@ -124,18 +126,6 @@ fn build_response(status: StatusCode, content: Bytes) -> Response<Body> {
         .unwrap()
 }
 
-async fn service_func(req: Request<Body>) -> Result<Response<Body>, hyper::Error> {
-    match (req.method(), req.uri().path()) {
-        (&Method::GET, "/") => Ok(index(req).await),
-        (&Method::GET, "/metrics") => Ok(metrics(req).await),
-        (&Method::POST, "/inference") => Ok(inference(req).await),
-        _ => Ok(build_response(
-            StatusCode::NOT_FOUND,
-            Bytes::from(RESPONSE_NOT_FOUND),
-        )),
-    }
-}
-
 async fn shutdown_signal() {
     let mut interrupt = signal(SignalKind::interrupt()).unwrap();
     let mut terminate = signal(SignalKind::terminate()).unwrap();
@@ -163,14 +153,18 @@ async fn run(opts: &Opts) {
     let barrier = coordinator.run();
     barrier.wait().await;
 
-    let service = make_service_fn(|_| async { Ok::<_, hyper::Error>(service_fn(service_func)) });
+    let app = Router::new()
+        .route("/", get(index))
+        .route("/metrics", get(metrics))
+        .route("/inference", post(inference));
+
     let addr: SocketAddr = format!("{}:{}", opts.address, opts.port).parse().unwrap();
-    let server = hyper::Server::bind(&addr).serve(service);
     info!(?addr, "http service is running");
-    let graceful = server.with_graceful_shutdown(shutdown_signal());
-    if let Err(err) = graceful.await {
-        tracing::error!(%err, "shutdown service error");
-    }
+    axum::Server::bind(&addr)
+        .serve(app.into_make_service())
+        .with_graceful_shutdown(shutdown_signal())
+        .await
+        .unwrap();
 }
 
 fn main() {


Test patch:
diff --git a/tests/bad_req.py b/tests/bad_req.py
index 48ea03a8..01b1ec7e 100644
--- a/tests/bad_req.py
+++ b/tests/bad_req.py
@@ -61,7 +61,7 @@ def main():
                 print("[~]", data)
                 count += 1
 
-    print(f">> {count}/{REQ_NUM} requests recevied before disconnection")
+    print(f">> {count}/{REQ_NUM} requests received before disconnection")
 
     # re-try to check if the service is still alive
     resp = httpx.post(URL, json={"time": 0})