feat: support server sent event (#333)

* init the sse_inference

Signed-off-by: Keming <kemingy94@gmail.com>

* impl

Signed-off-by: Keming <kemingy94@gmail.com>

* handle sse error

Signed-off-by: Keming <kemingy94@gmail.com>

* fix sse error event type

Signed-off-by: Keming <kemingy94@gmail.com>

* add sse metrics

Signed-off-by: Keming <kemingy94@gmail.com>

* pass the err msg to the client

Signed-off-by: Keming <kemingy94@gmail.com>

* add sse test

Signed-off-by: Keming <kemingy94@gmail.com>

* format rs import

Signed-off-by: Keming <kemingy94@gmail.com>

---------

Signed-off-by: Keming <kemingy94@gmail.com>

Cargo.toml

@@ -17,13 +17,14 @@ bytes = "1"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["local-time", "json"] }
 tokio = { version = "1", features = ["rt", "rt-multi-thread", "time", "macros", "sync", "signal", "io-util"] }
-derive_more = { version = "0.99", features = ["display", "error"] }
+derive_more = { version = "0.99", features = ["display", "error", "from"] }
 # MPMS that only one consumer sees each message & async
 async-channel = { version = "1" }
 once_cell = "1.18"
 prometheus-client = "0.21.1"
 argh = "0.1"
 axum = "0.6.18"
+async-stream = "0.3.5"
 utoipa = "3.3.0"
 serde_json = "1.0.96"
 serde = "1.0.163"

examples/server_side_event/client.py

@@ -0,0 +1,31 @@
+# Copyright 2023 MOSEC Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import httpx
+from httpx_sse import connect_sse
+
+with httpx.Client() as client:
+ with connect_sse(
+ client, "POST", "http://127.0.0.1:8000/sse_inference", json={"text": "mosec"}
+ ) as event_source:
+ for sse in event_source.iter_sse():
+ print(f"Event({sse.event}): {sse.data}")
+
+# error handling
+with httpx.Client() as client:
+ with connect_sse(
+ client, "POST", "http://127.0.0.1:8000/sse_inference", json={"error": "mosec"}
+ ) as event_source:
+ for sse in event_source.iter_sse():
+ print(f"Event({sse.event}): {sse.data}")

examples/server_side_event/server.py

@@ -0,0 +1,45 @@
+# Copyright 2023 MOSEC Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from mosec import Server, ValidationError, Worker, get_logger
+
+logger = get_logger()
+
+
+class Preprocess(Worker):
+ def forward(self, data):
+ text = data.get("text")
+ if text is None:
+ raise ValidationError("text is required")
+ return text
+
+
+class Inference(Worker):
+ def forward(self, data):
+ epoch = 5
+ for i in range(epoch):
+ for j in range(len(data)):
+ self.send_stream_event(
+ f"inference: ({i + 1}/{epoch}) {data[j]}", index=j
+ )
+
+ # this return value will be ignored
+ return data
+
+
+if __name__ == "__main__":
+ server = Server()
+ server.append_worker(Preprocess)
+ server.append_worker(Inference, max_batch_size=2)
+ server.run()

mosec/coordinator.py

@@ -15,9 +15,11 @@
 """The Coordinator is used to control the data flow between `Worker` and `Server`."""
 
 import os
+import queue
 import signal
 import socket
 import struct
+import threading
 import time
 import traceback
 from contextlib import contextmanager
@@ -106,8 +108,13 @@ def __init__(
  self.worker.worker_id = worker_id
  self.worker.max_batch_size = max_batch_size
  self.worker.stage = stage
+ self.stream: queue.SimpleQueue = queue.SimpleQueue()
+ self.semaphore: threading.Semaphore = threading.Semaphore()
+ self.worker._stream_queue = self.stream
+ self.worker._stream_semaphore = self.semaphore
  self.timeout = timeout
  self.name = f"<{stage_id}|{worker.__name__}|{worker_id}>"
+ self.current_ids: Sequence[bytes] = []
 
  self.protocol = Protocol(
  name=self.name,
@@ -132,6 +139,7 @@ def __init__(
 
  self.warmup()
  self.init_protocol()
+ threading.Thread(target=self.streaming, daemon=True).start()
  self.run()
 
  def init_protocol(self):
@@ -157,6 +165,19 @@ def init_protocol(self):
  self.protocol.addr,
  )
 
+ def streaming(self):
+ """Send stream data from the worker to the server through the socket."""
+ while not self.shutdown.is_set():
+ try:
+ text, index = self.stream.get(timeout=self.timeout)
+ # encode the text with UTF-8
+ payloads = (text.encode(),)
+ ids = (self.current_ids[index],)
+ self.protocol.send(HTTPStautsCode.STREAM_EVENT, ids, payloads)
+ self.semaphore.release()
+ except queue.Empty:
+ continue
+
  def warmup(self):
  """Warmup to allocate resources (useful for GPU workload)[Optional]."""
  need_warmup = self.worker.example or self.worker.multi_examples
@@ -263,6 +284,8 @@ def coordinate(self):
  while not self.shutdown.is_set():
  try:
  _, ids, payloads = protocol_recv()
+ # expose here to be used by stream event
+ self.current_ids = ids
  except socket.timeout:
  continue
  except (struct.error, OSError) as err:
@@ -299,10 +322,14 @@ def coordinate(self):
  payloads = ["inference internal error".encode()] * length
 
  try:
+ # pylint: disable=consider-using-with
+ self.semaphore.acquire(timeout=self.timeout)
  protocol_send(status, ids, payloads)
  except OSError as err:
  logger.error("%s failed to send to socket: %s", self.name, err)
  break
+ finally:
+ self.semaphore.release()
 
  self.protocol.close()
  time.sleep(CONN_CHECK_INTERVAL)

mosec/errors.py

@@ -29,8 +29,8 @@
 class MosecError(Exception):
  """Mosec basic exception."""
 
- code = HTTPStautsCode.INTERNAL_ERROR
- msg = "mosec error"
+ code: HTTPStautsCode = HTTPStautsCode.INTERNAL_ERROR
+ msg: str = "mosec error"
 
 
 class ClientError(MosecError):

mosec/protocol.py

@@ -38,6 +38,9 @@ class HTTPStautsCode(IntFlag):
  INTERNAL_ERROR = 8 # 500
  TIMEOUT_ERROR = 16 # 408
 
+ # special one, indicate that it's a SSE
+ STREAM_EVENT = 32768
+
 
 class Protocol:
  """IPC protocol.
@@ -130,9 +133,10 @@ def send(self, flag: int, ids: Sequence[bytes], payloads: Sequence[bytes]):
  self.socket.sendall(data.getbuffer())
  if logger.isEnabledFor(logging.DEBUG):
  logger.debug(
- "%s sent %d tasks with ids: %s",
+ "%s sent %d(%d) tasks with ids: %s",
  self.name,
  len(ids),
+ flag,
  struct.unpack("!" + "I" * len(ids), b"".join(ids)),
  )
 

mosec/worker.py

@@ -22,16 +22,22 @@
  4. data processing
 """
 
+from __future__ import annotations
+
 import abc
 import json
 import pickle
-from typing import Any, Dict, Sequence, Tuple
+from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Tuple
 
 from mosec.errors import DecodingError, EncodingError
 from mosec.utils import ParseTarget
 
 MOSEC_REF_TEMPLATE = "#/components/schemas/{name}"
 
+if TYPE_CHECKING:
+ from queue import SimpleQueue
+ from threading import Semaphore
+
 
 class Worker(abc.ABC):
  """MOSEC worker interface.
@@ -62,6 +68,8 @@ class Worker(abc.ABC):
  _worker_id: int = 0
  _stage: str = ""
  _max_batch_size: int = 1
+ _stream_queue: Optional[SimpleQueue] = None
+ _stream_semaphore: Optional[Semaphore] = None
 
  def __init__(self):
  """Initialize the worker.
@@ -85,6 +93,20 @@ def deserialize_ipc(self, data: bytes) -> Any:
  """
  return pickle.loads(data)
 
+ def send_stream_event(self, text: str, index: int = 0):
+ """Send a stream event to the client.
+
+ Args:
+ text: the text to be sent, needs to be UTF-8 compatible
+ index: the index of the stream event. For single request, this will always
+ be 0. For dynamic batch request, this should be the index of the
+ request in the batch.
+ """
+ if self._stream_queue is None or self._stream_semaphore is None:
+ raise RuntimeError("the worker stream or semaphore is not initialized")
+ self._stream_semaphore.acquire()
+ self._stream_queue.put((text, index))
+
  @property
  def stage(self) -> str:
  """Return the stage name."""
@@ -185,7 +207,7 @@ def forward(self, data: Any) -> Any:
  - for a single-stage worker, data will go through
  ``<deserialize> -> <forward> -> <serialize>``
 
- - for a multi-stage worker that is neithor `ingress` not `egress`, data
+ - for a multi-stage worker that is neither `ingress` not `egress`, data
  will go through ``<deserialize_ipc> -> <forward> -> <serialize_ipc>``
  """
  raise NotImplementedError

pyproject.toml

@@ -78,6 +78,7 @@ output-format = "colorized"
 reports = "no"
 score = "yes"
 max-args = 7
+max-attributes = 10
 
 [tool.pylint.messages_control]
 disable = []

requirements/dev.txt

@@ -12,3 +12,4 @@ pre-commit>=2.15.0
 msgpack>=1.0.5
 numpy>=1.24
 httpx==0.24.1
+httpx-sse==0.3.0

rustfmt.toml

@@ -1,2 +1,6 @@
 unstable_features = true
+reorder_imports = true
+format_strings = true
+imports_granularity = "Module"
 group_imports = "StdExternalCrate"
+reorder_impl_items = true