[FastAPI] Server Sent Event(SSE)를 이용한 비동기 스트리밍 구현

이미지 생성 모델을 이용한 이미지 생성 및 이미지 처리 툴을 테스트하기 위해서 WebGUI가 필요했다. 이미지 생성은 오래 걸리면 10초가량도 소요될 수 있기 때문에, 태스크 처리 상황에 대한 클라이언트와 서버의 지속적인 통신이 필요하다. 나는 특히, 이미지 처리 상황을 ProgressBar로 보여주고 싶었다. HTTP/2의 주요 피쳐중의 하나인 SSE(Server Sent Event)를 통해서 이를 처리하는 방법을 공유하겠다.

 

사용한 서버쪽 프레임워크는 FastAPI다.

app = FastAPI()

async def progress_stream():
    for i in range(100):
        d = {
            "progress": i / 100,
        }
        yield f"data: {json.dumps(d)}\n\n"
        await asyncio.sleep(0.1)

@app.get("/mock_progress")
async def mock_progress():
    return StreamingResponse(progress_stream(), media_type="text/event-stream")

GPT가 만들어준…progress_stream()이라는 비동기 제너레이터 함수로 테스트를 진행했다. 서버를 열고, /mock_progress 로 가보면 0.1초마다 업데이트되는 progress 데이터가 쭉 찍히는 것을 확인할 수 있다.

 

서버쪽은 완성되었으니, 다음은 클라이언트쪽 javascript 코드다.

try {
  fetch(`/mock_progress`)
    .then((response) => {
      if (!response.ok) {
        throw new Error("Error:", response.statusText);
      }
      return response.body.getReader();
    })
    .then((reader) => {
      const decoder = new TextDecoder("utf-8");
      let buffer = "";
      function read() {
        reader
          .read()
          .then(({ done, value }) => {
            if (done) {
              return;
            }

            buffer += decoder.decode(value, { stream: true });

            let lines = buffer.split("\n");
            buffer = lines.pop(); // 마지막 줄은 다음 청크로 넘어갈 수 있음

            lines.forEach((line) => {
              if (line.startsWith("data:")) {
                const eventData = line.slice(5).trim();

                handleProgress(JSON.parse(eventData));
              }
            });

            // 다음 데이터 청크 읽기
            read();
          })
          .catch((err) => {
            console.error("SSE 연결 중 오류:", err);
          });
      }
      read();
    });
} catch (error) {
  console.error("Error:", error);
}

function handleProgress(data) {
    console.log(data.progress)
}

서버에서 보내주는 스트림이 끝날때까지 chunked data를 받아서 JSON Object로 만든다. UI에서 Progress나 상태를 업데이트하는 일은 handleProgress(data) 함수에서 처리하면 된다.

 

이를 응용해 만든 간단한 페이지를 예시로 남겨둔다.