IO 프로세서 플러그인

IO 프로세서 플러그인 (IO Processor Plugins)

IO 프로세서 플러그인은 pooling 모델의 입력과 출력 전·후처리를 커스텀화할 수 있게 해주는 기능이에요. 아이디어는 이렇습니다. 사용자가 vLLM에 커스텀 입력을 주면, 그것을 하나 이상의 모델 프롬프트로 변환해 모델의 encode 메서드에 넣는 거예요. 이런 플러그인의 잠재적 쓰임새 하나는 멀티모달 데이터 생성에 vLLM을 쓰는 경우예요. 사용자가 이미지를 넣고 이미지를 출력으로 받는다고 상상해 보세요.

IO 프로세서 플러그인으로 추론할 때 프롬프트 타입은 플러그인이 정의하고, 최종 요청 출력도 마찬가지예요. vLLM은 입력/출력 데이터를 전혀 검증하지 않으므로, 올바른 데이터가 모델에 들어가고 사용자에게 반환되게 하는 건 전적으로 플러그인의 몫이에요. 현재 이 플러그인은 pooling 모델만 지원하고, LLM과 AsyncLLM의 encode 메서드 또는 온라인 서빙의 /pooling 엔드포인트로 트리거할 수 있어요.

IO 프로세서 플러그인 작성하기

IO 프로세서 플러그인은 IOProcessor 인터페이스를 구현해요.

IOProcessorInput = TypeVar("IOProcessorInput")
IOProcessorOutput = TypeVar("IOProcessorOutput")

class IOProcessor(ABC, Generic[IOProcessorInput, IOProcessorOutput]):
    """Abstract interface for pre/post-processing of engine I/O."""

    def __init__(self, vllm_config: VllmConfig, renderer: BaseRenderer):
        super().__init__()

        self.vllm_config = vllm_config

    def parse_data(self, data: object) -> IOProcessorInput:
        raise NotImplementedError

    def merge_sampling_params(
        self,
        params: SamplingParams | None = None,
    ) -> SamplingParams:
        return params or SamplingParams()

    def merge_pooling_params(
        self,
        params: PoolingParams | None = None,
    ) -> PoolingParams:
        return params or PoolingParams(task="plugin")

    @abstractmethod
    def pre_process(
        self,
        prompt: IOProcessorInput,
        request_id: str | None = None,
        **kwargs,
    ) -> PromptType | Sequence[PromptType]:
        raise NotImplementedError

    async def pre_process_async(
        self,
        prompt: IOProcessorInput,
        request_id: str | None = None,
        **kwargs,
    ) -> PromptType | Sequence[PromptType]:
        return self.pre_process(prompt, request_id, **kwargs)

    @abstractmethod
    def post_process(
        self,
        model_output: Sequence[PoolingRequestOutput],
        request_id: str | None = None,
        **kwargs,
    ) -> IOProcessorOutput:
        raise NotImplementedError

    async def post_process_async(
        self,
        model_output: AsyncGenerator[tuple[int, PoolingRequestOutput]],
        request_id: str | None = None,
        **kwargs,
    ) -> IOProcessorOutput:
        # We cannot guarantee outputs are returned in the same order they were
        # fed to vLLM.
        # Let's sort them by id before post_processing
        sorted_output = sorted(
            [(i, item) async for i, item in model_output], key=lambda output: output[0]
        )
        collected_output = [output[1] for output in sorted_output]
        return self.post_process(collected_output, request_id=request_id, **kwargs)
  • parse_data 메서드는 사용자 데이터를 검증하고 pre_process* 메서드가 기대하는 입력으로 변환하는 데 써요.
  • merge_sampling_params와 merge_pooling_params 메서드는 입력 SamplingParams·PoolingParams(있다면)를 기본값과 병합해요.
  • pre_process* 메서드는 검증된 플러그인 입력을 받아 일반 추론용 vLLM 모델 프롬프트를 생성해요.
  • post_process* 메서드는 PoolingRequestOutput 객체를 입력으로 받아 커스텀 플러그인 출력을 생성해요.

PrithviGeospatialMAE 모델로 geotiff 이미지를 생성하는 플러그인 구현 예시는 여기에서 볼 수 있어요. 온라인(examples/pooling/plugin/prithvi_geospatial_mae_online.py)과 오프라인(examples/pooling/plugin/prithvi_geospatial_mae_io_processor.py) 추론 예시도 참고하세요.

IO 프로세서 플러그인 사용하기

IO 프로세서 플러그인은 엔진 시작 시 로드되는데, 로드할 플러그인 이름을 지정하는 두 가지 방법이 있어요.

  1. vLLM의 EngineArgs를 통해: AsyncLLM을 초기화할 때 쓰는 EngineArgs에 io_processor_plugin 인자를 설정. 오프라인 모드에서는 LLM에 io_processor_plugin 인자를 넘기거나, 서빙 모드에서는 --io-processor-plugin 인자를 넘겨도 같은 결과를 얻을 수 있어요.
  2. 모델 HF 구성을 통해: 모델 구성(config.json)에 io_processor_plugin 필드 추가.

순서가 곧 우선순위를 결정해요. 즉 EngineArgs로 플러그인 이름을 설정하면 모델 HF 구성(config.json)에 지정된 플러그인 이름보다 우선해요.

더 알아보기 (Learn more)