Test(rollout): add unit tests for vllm_rollout helpers and mocked generate paths - #25
Conversation
There was a problem hiding this comment.
Code Review
This pull request transitions the base Docker image to vllm/vllm-openai, updates the build process with necessary CUDA development headers and dependencies, and introduces a build-time smoke test. It also adds a comprehensive suite of unit tests and shared stubs for the vllm_rollout module. Feedback was provided regarding the inclusion of IPython in the production image, which may be unnecessary and increases the image size.
| RUN pip install --no-deps "sglang==${SGLANG_VERSION}" sglang-router==0.3.2 && \ | ||
| pip install IPython |
There was a problem hiding this comment.
The installation of IPython in the production Docker image seems unnecessary and increases the image size. Unless it is explicitly required for a runtime feature (e.g., interactive debugging within the container, which is usually discouraged for production), it should be removed to keep the image lean.
RUN pip install --no-deps "sglang==${SGLANG_VERSION}" sglang-router==0.3.2
There was a problem hiding this comment.
No changes were made to the Dockerfile.
Summary
Adds unit tests for
slime.rollout.vllm_rollout, covering token preprocessing, vLLM/inference/v1/generaterequest/response handling, and core async paths (generate,generate_and_rm,abort) via mocks.tests/unit/rollout/test_vllm_rollout.py— 55 test cases marked with@pytest.mark.unittests/unit/rollout/conftest.py— stubs for optional deps (vllm_router,PIL,transformers,aiohttp,pylatexenc) so tests collect on minimal dev machinestests/test_vllm_generate_endpoint.py— still used for real vLLM + GPU integration checksCoverage
Pure helpers
_coerce_flat_int_token_ids,_prepare_prompt_ids,_base_dataset_prompt_idsget_model_url_vllm_meta_from_generate_choice,_decode_vllm_routed_experts,_apply_vllm_routed_experts(including disabled / invalid-shape branches)_inference_generate_tokens_and_logprobs,_align_engine_tokens_and_logprobs_build_inference_sampling_params,_mm_render_response_to_generate_body_router_worker_urls,_resume_vllm_workersMocked async paths
generate(): text rollout, partial-continuation budget, zero-budget truncation, consistent-hashing headers, multimodal render→generate, routed-experts assignmentgenerate_and_rm(): skip completed samples, abort handling, custom generate hookgenerate_and_rm_group(): per-samplesession_idassignmentabort(): worker pause + resumeOut of scope (existing integration tests)
generate_rollout_async/generate_rollout/eval_rollout*tests/test_vllm_generate_endpoint.py)Test plan
cd vime pytest tests/unit/rollout/test_vllm_rollout.py -m unit -q