본문으로 건너뛰기

[sglang] [SGLang] 10초 걸리던 API를 1초로: JSON Pixel List의 NumPy 변환 최적화

PR 링크: sgl-project/sglang#41895 상태: Merged | 변경: +100 / -10

들어가며

VLA(Vision-Language-Action) 모델이나 Diffusion 기반의 로봇 제어 모델(예: FLUX 3 Action)을 서빙할 때, 클라이언트는 카메라 이미지를 서버로 전송해야 합니다. 이때 가장 직관적인 방법 중 하나는 이미지를 image.tolist()와 같이 중첩된 Python 리스트 형태로 JSON에 담아 보내는 것입니다.

하지만 이 방식에는 치명적인 성능 함정이 숨어 있습니다. 360x640 해상도의 카메라 3대 분량의 이미지는 약 207만 개의 Python 정수(int) 객체로 구성됩니다. SGLang과 같은 고성능 서빙 엔진은 내부적으로 여러 프로세스(Scheduler, Worker 등)가 협력하며, 이 과정에서 데이터는 IPC(Inter-Process Communication)를 거치게 됩니다. Python의 pickle은 이 거대한 객체 트리를 순회(walk)하며 직렬화하는데, 이 과정에서 엄청난 CPU 오버헤드가 발생합니다.

이번 PR은 이러한 '거대 JSON 리스트'가 유발하는 병목 현상을 엔트리포인트(Entrypoint)에서 NumPy 배열로 즉시 변환함으로써 해결한 사례입니다. 결과적으로 10.8초가 걸리던 요청 처리 시간이 1.48초로 단축되었습니다.

코드 분석: 핵심 변경 사항

1. protocol.py: 효율적인 NumPy 변환 로직 도입

가장 핵심적인 변화는 HTTP 요청을 받는 엔트리포인트에서 리스트를 NumPy 배열로 미리 변환하는 로직입니다. 단순히 np.array()를 호출하는 것이 아니라, 성능을 위해 C로 구현된 내장 함수를 영리하게 활용합니다.

[Before]

기존에는 리스트 형태의 이미지가 그대로 스케줄러로 전달되었습니다.

# 기존에는 별도의 변환 로직 없이 dict 구조만 확인
def _normalize_image_value(value: Any) -> Any:
    if not isinstance(value, dict):
        return value
    # ... base64 처리 로직만 존재 ...

[After]

새로운 로직은 리스트를 감지하면 즉시 _pixel_list_to_array를 호출합니다.

# python/sglang/multimodal_gen/runtime/entrypoints/action/protocol.py

def _uint8_pixels(rows: list) -> np.ndarray | None:
    """3-level lists of ints in [0, 255] -> uint8; ``bytes()`` runs in C."""
    try:
        height, width, channels = len(rows), len(rows[0]), len(rows[0][0])
        if any(len(row) != width for row in rows):
            return None
        # itertools.chain과 bytes()를 사용하여 C 레벨에서 빠르게 평탄화 및 변환
        flat = bytes(itertools.chain.from_iterable(itertools.chain.from_iterable(rows)))
    except (TypeError, ValueError):
        return None
    if len(flat) != height * width * channels:
        return None
    return np.frombuffer(flat, dtype=np.uint8).reshape(height, width, channels)

def _pixel_list_to_array(value: list) -> Any:
    """JSON pixel 리스트를 NumPy로 변환하여 IPC 비용 절감"""
    try:
        first = value[0][0][0]
    except (TypeError, IndexError, KeyError):
        return value
    
    # 정수형(uint8)인 경우 최적화된 _uint8_pixels 사용
    array = _uint8_pixels(value) if isinstance(first, int) else None
    if array is not None:
        return array
    
    try:
        array = np.asarray(value)
    except ValueError:
        return value
    return array if array.dtype.kind in "iuf" else value

분석:

  • _uint8_pixels 함수는 itertools.chain과 bytes()를 조합하여 중첩 리스트를 바이트 스트림으로 변환합니다. bytes(list_of_ints)는 Python 인터프리터 레벨이 아닌 C 레벨에서 루프를 돌기 때문에 매우 빠릅니다.
  • 이렇게 생성된 바이트 스트림을 np.frombuffer로 읽어오면 복사 비용을 최소화하면서 NumPy 배열을 만들 수 있습니다.

2. FLUX-3-Action.mdx: 사용자 가이드 및 모범 사례 제시

코드 최적화와 더불어, 사용자에게 더 효율적인 데이터 전송 방식을 권장하도록 문서를 수정했습니다.

[Before]

| `input.observation.images` | object | Camera name -> HWC RGB image, uint8 or float in `[0, 1]`. ... |

[After]

| `input.observation.images` | object | Camera name -> HWC RGB image: a base64 PNG or JPEG (`{"b64_json": ...}`), or a pixel array ... |

# ... 중략 ...
# PNG는 무손실이며, pixel list보다 훨씬 가볍다는 점을 명시
PNG is lossless. Pixel lists (`image.tolist()`) are accepted too, but ... encoding and parsing it adds about 0.3 s per request on top of the base64 PNG payload.

분석:

  • image.tolist()는 JSON 크기 자체가 너무 커지기 때문에(약 11MB), 네트워크 전송 및 JSON 파싱 단계에서도 오버헤드가 발생합니다. 이를 방지하기 위해 base64 PNG 사용을 권장하는 예제 코드를 추가했습니다.

왜 이게 좋은 최적화인가?

1. 지연 시간(Latency)의 드라마틱한 개선

테스트 결과에 따르면, RTX 5090 환경에서 다음과 같은 성능 향상이 확인되었습니다.

  • 기존: 10,787 ms (약 10.8초)
  • 개선 후: 1,479 ms (약 1.5초)
  • 개선율: 약 86% 감소

특히 preprocess_ms가 130ms에서 1.6ms로 줄어든 점이 인상적입니다. 이는 모델 연산 자체가 아닌, 데이터를 준비하고 프로세스 간에 전달하는 과정에서의 병목을 정확히 타격했음을 보여줍니다.

2. IPC 및 직렬화 비용 절감

Python의 다중 프로세스 아키텍처에서 거대한 중첩 리스트를 전달하면, 수신 측 프로세스는 이를 다시 파이썬 객체로 복원해야 합니다. 반면 NumPy 배열은 메모리 상에 연속적으로 배치된 바이너리 데이터이므로, pickle 시 훨씬 효율적으로 처리되며 공유 메모리 활용 가능성도 열어줍니다.

3. 하위 호환성 및 안정성 유지

단순히 리스트를 금지하는 것이 아니라, 엔트리포인트에서 투명하게(transparently) 변환해줌으로써 기존 클라이언트 코드를 깨뜨리지 않고 성능만 올렸습니다. 또한 test_json_pixel_lists_become_arrays_before_ipc 테스트 케이스를 통해 uint8, float, out-of-range 값 등에 대한 검증을 철저히 수행했습니다.

일반적인 교훈

  1. Python 객체 트리의 무서움: 수백만 개의 요소를 가진 Python 리스트는 그 자체로 메모리 관리와 직렬화의 재앙입니다. 대량의 수치 데이터는 가능한 한 빨리 NumPy나 PyTorch 텐서와 같은 바이너리 형식으로 변환해야 합니다.
  2. 엔트리포인트 최적화: 병목은 종종 모델 내부가 아니라, 데이터를 받아들이는 입구(Entrypoint)에 있습니다. 'Hops'가 발생하기 전에 데이터를 정규화하는 것이 중요합니다.
  3. C-API 활용: bytes()와 같이 C로 구현된 내장 함수를 적재적소에 활용하면 Python 코드만으로도 놀라운 성능 향상을 얻을 수 있습니다.

이 PR은 고성능 AI 서빙 엔진을 설계할 때 데이터 직렬화와 프로세스 간 통신이 얼마나 중요한지를 잘 보여주는 훌륭한 사례입니다.

참고 자료

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.

댓글

관련 포스트

PR Analysis 의 다른글