Migrate OvisImagePatchInputs to TensorSchema #22024

bbeckca · 2025-07-31T15:51:48Z

Purpose

This PR migrates OvisImagePatchInputs from a TypedDict-based definition to a structured TensorSchema model with runtime shape validation. This brings it in line with recent changes to Phi3VImagePixelInputs, and is part of a broader effort to improve input contract enforcement and debug-ability across multi-modal models.

Test Plan

Confirm validation works via standalone tests in tests/standalone_test/test_tensor_schema.py and rely on CI to check integration.

Test Result

(venv) benjibeck@Benjis-MacBook-Pro vllm % python3 -m pytest tests/utils_/test_tensor_schema.py -v --log-cli-level=DEBUG          
========================================================================== test session starts ==========================================================================
platform darwin -- Python 3.9.6, pytest-8.4.1, pluggy-1.6.0 -- /Users/benjibeck/Projects/vllm/venv/bin/python3
cachedir: .pytest_cache
rootdir: /Users/benjibeck/Projects/vllm
configfile: pyproject.toml
plugins: anyio-4.9.0
collected 19 items                                                                                                                                                      

tests/utils_/test_tensor_schema.py::test_tensor_schema_valid_tensor PASSED                                                                                        [  5%]
tests/utils_/test_tensor_schema.py::test_tensor_schema_optional_fields PASSED                                                                                     [ 10%]
tests/utils_/test_tensor_schema.py::test_tensor_schema_constant_dim_failure PASSED                                                                                [ 15%]
tests/utils_/test_tensor_schema.py::test_tensor_schema_invalid_types_in_list PASSED                                                                               [ 21%]
tests/utils_/test_tensor_schema.py::test_tensor_schema_rank_mismatch PASSED                                                                                       [ 26%]
tests/utils_/test_tensor_schema.py::test_tensor_schema_missing_required_field PASSED                                                                              [ 31%]
tests/utils_/test_tensor_schema.py::test_tensor_schema_symbolic_dim_mismatch PASSED                                                                               [ 36%]
tests/utils_/test_tensor_schema.py::test_tensor_schema_list_tensor_valid PASSED                                                                                   [ 42%]
tests/utils_/test_tensor_schema.py::test_tensor_schema_variable_patch_counts_valid PASSED                                                                         [ 47%]
tests/utils_/test_tensor_schema.py::test_tensor_schema_tuple_tensor_valid PASSED                                                                                  [ 52%]
tests/utils_/test_tensor_schema.py::test_tensor_schema_inconsistent_shapes_in_list PASSED                                                                         [ 57%]
tests/utils_/test_tensor_schema.py::test_tensor_schema_empty_list PASSED                                                                                          [ 63%]
tests/utils_/test_tensor_schema.py::test_tensor_schema_validation_disabled_skips_shape_check PASSED                                                               [ 68%]
tests/utils_/test_tensor_schema.py::test_tensor_schema_with_valid_resolve_binding_dims PASSED                                                                     [ 73%]
tests/utils_/test_tensor_schema.py::test_tensor_schema_with_invalid_resolve_binding_dims PASSED                                                                   [ 78%]
tests/utils_/test_tensor_schema.py::test_tensor_schema_with_list_of_symbolic_dim PASSED                                                                           [ 84%]
tests/utils_/test_tensor_schema.py::test_tensor_schema_with_list_of_symbolic_dim_mismatch_in_length PASSED                                                        [ 89%]
tests/utils_/test_tensor_schema.py::test_valid_tensor_schema_with_static_last_dim PASSED                                                                          [ 94%]
tests/utils_/test_tensor_schema.py::test_invalid_tensor_schema_with_static_last_dim PASSED                                                                        [100%]

github-actions · 2025-07-31T15:51:58Z

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

gemini-code-assist

Code Review

This pull request successfully migrates OvisImagePatchInputs from a TypedDict to a TensorSchema, which improves input validation and code clarity. However, a field rename introduces a breaking change because its usages elsewhere in the file were not updated. I've provided a critical comment with a suggestion to revert the rename to ensure the code remains functional.

vllm/model_executor/models/ovis.py

DarkLight1337

Thanks, LGTM now

bbeckca · 2025-08-22T15:07:47Z

Update: Working on failing MM test for:

[2025-08-22T04:45:50Z] (EngineCore_0 pid=8358) ERROR 08-21 21:45:50 [core.py:708]     image_input = self._parse_and_validate_image_input(**kwargs)
--
"/usr/local/lib/python3.12/dist-packages/vllm/model_executor/models/ovis.py", line 459, in _parse_and_validate_image_input
  | [2025-08-22T04:45:50Z] (EngineCore_0 pid=8358) ERROR 08-21 21:45:50 [core.py:708]     return OvisImagePatchInputs(
  | [2025-08-22T04:45:50Z] (EngineCore_0 pid=8358) ERROR 08-21 21:45:50 [core.py:708]            ^^^^^^^^^^^^^^^^^^^^^
  | [2025-08-22T04:45:50Z] (EngineCore_0 pid=8358) ERROR 08-21 21:45:50 [core.py:708]   File "/usr/local/lib/python3.12/dist-packages/vllm/utils/tensor_schema.py", line 67, in __init__
  | [2025-08-22T04:45:50Z] (EngineCore_0 pid=8358) ERROR 08-21 21:45:50 [core.py:708]     self.validate()
  | [2025-08-22T04:45:50Z] (EngineCore_0 pid=8358) ERROR 08-21 21:45:50 [core.py:708]   File "/usr/local/lib/python3.12/dist-packages/vllm/utils/tensor_schema.py", line 222, in validate
  | [2025-08-22T04:45:50Z] (EngineCore_0 pid=8358) ERROR 08-21 21:45:50 [core.py:708]     self._validate_tensor_shape_expected(
  | [2025-08-22T04:45:50Z] (EngineCore_0 pid=8358) ERROR 08-21 21:45:50 [core.py:708]   File "/usr/local/lib/python3.12/dist-packages/vllm/utils/tensor_schema.py", line 138, in _validate_tensor_shape_expected
  | [2025-08-22T04:45:50Z] (EngineCore_0 pid=8358) ERROR 08-21 21:45:50 [core.py:708]     raise ValueError(f"{field_name} has rank {len(actual_shape)} "
  | [2025-08-22T04:45:50Z] (EngineCore_0 pid=8358) ERROR 08-21 21:45:50 [core.py:708] ValueError: flat_data has rank 4 but expected 2
``

bbeckca · 2025-08-22T15:41:08Z

Update: Working on failing MM test for:

[2025-08-22T04:45:50Z] (EngineCore_0 pid=8358) ERROR 08-21 21:45:50 [core.py:708]     image_input = self._parse_and_validate_image_input(**kwargs)
--
"/usr/local/lib/python3.12/dist-packages/vllm/model_executor/models/ovis.py", line 459, in _parse_and_validate_image_input
  | [2025-08-22T04:45:50Z] (EngineCore_0 pid=8358) ERROR 08-21 21:45:50 [core.py:708]     return OvisImagePatchInputs(
  | [2025-08-22T04:45:50Z] (EngineCore_0 pid=8358) ERROR 08-21 21:45:50 [core.py:708]            ^^^^^^^^^^^^^^^^^^^^^
  | [2025-08-22T04:45:50Z] (EngineCore_0 pid=8358) ERROR 08-21 21:45:50 [core.py:708]   File "/usr/local/lib/python3.12/dist-packages/vllm/utils/tensor_schema.py", line 67, in __init__
  | [2025-08-22T04:45:50Z] (EngineCore_0 pid=8358) ERROR 08-21 21:45:50 [core.py:708]     self.validate()
  | [2025-08-22T04:45:50Z] (EngineCore_0 pid=8358) ERROR 08-21 21:45:50 [core.py:708]   File "/usr/local/lib/python3.12/dist-packages/vllm/utils/tensor_schema.py", line 222, in validate
  | [2025-08-22T04:45:50Z] (EngineCore_0 pid=8358) ERROR 08-21 21:45:50 [core.py:708]     self._validate_tensor_shape_expected(
  | [2025-08-22T04:45:50Z] (EngineCore_0 pid=8358) ERROR 08-21 21:45:50 [core.py:708]   File "/usr/local/lib/python3.12/dist-packages/vllm/utils/tensor_schema.py", line 138, in _validate_tensor_shape_expected
  | [2025-08-22T04:45:50Z] (EngineCore_0 pid=8358) ERROR 08-21 21:45:50 [core.py:708]     raise ValueError(f"{field_name} has rank {len(actual_shape)} "
  | [2025-08-22T04:45:50Z] (EngineCore_0 pid=8358) ERROR 08-21 21:45:50 [core.py:708] ValueError: flat_data has rank 4 but expected 2
``

@DarkLight1337 @Isotr0py Before trying to fix, I wanted to check whether expected format for OvisImagePatchInputs.flat_data should be 2D or 4D? Observed Fuyu using flat_data with 2 dims, but it's not clear to me whether Ovis should have the same. Trying to understand whether this issue is with the annotation or processor inputs.

class OvisImagePatchInputs(TypedDict):
    type: Literal["image_patches"]
    flat_data: torch.Tensor
    """
    Shape: 
    `(batch_size * num_patches, patch_size_x * patch_size_y * num_channels)`
    """

> /home/bbeckca/vllm/vllm/utils/tensor_schema.py(139)_validate_tensor_shape_expected()
-> raise ValueError(f"{field_name} has rank {len(actual_shape)} "
(Pdb) expected_shape
('batch_patches', 'patch_size')
(Pdb) actual_shape
torch.Size([30, 3, 448, 448])

mergify · 2025-09-01T18:20:24Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @bbeckca.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Signed-off-by: Benji Beck <benjibeck@meta.com>

gemini-code-assist bot reviewed Jul 31, 2025

View reviewed changes

vllm/model_executor/models/ovis.py Outdated Show resolved Hide resolved

bbeckca force-pushed the ovis branch from 1393744 to a5e9094 Compare August 21, 2025 15:49

DarkLight1337 approved these changes Aug 21, 2025

View reviewed changes

DarkLight1337 added the ready ONLY add when PR is ready to merge/full CI is needed label Aug 21, 2025

bbeckca mentioned this pull request Aug 28, 2025

Migrate Llama4ImagePatchInputs to TensorSchema #22021

Merged

bbeckca force-pushed the ovis branch 4 times, most recently from fd40f00 to 6d12850 Compare September 1, 2025 18:19

mergify bot added the needs-rebase label Sep 1, 2025

bbeckca added 3 commits September 1, 2025 11:24

Migrate OvisImagePatchInputs to TensorSchema

370992c

Signed-off-by: Benji Beck <benjibeck@meta.com>

Update type annotations in OvisImagePatchInputs

32592e5

Signed-off-by: Benji Beck <benjibeck@meta.com>

Fix shape mismatch for flat_data

8dd323e

Signed-off-by: Benji Beck <benjibeck@meta.com>

bbeckca force-pushed the ovis branch from 6d12850 to 8dd323e Compare September 1, 2025 18:25

mergify bot removed the needs-rebase label Sep 1, 2025

Merge branch 'main' into ovis

411a2bc

DarkLight1337 merged commit 1fa1d6a into vllm-project:main Sep 2, 2025
41 checks passed

eicherseiji pushed a commit to eicherseiji/vllm that referenced this pull request Sep 9, 2025

Migrate OvisImagePatchInputs to TensorSchema (vllm-project#22024)

7d0da1d

Signed-off-by: Benji Beck <benjibeck@meta.com>

FeiDaLI pushed a commit to FeiDaLI/vllm that referenced this pull request Sep 25, 2025

Migrate OvisImagePatchInputs to TensorSchema (vllm-project#22024)

fe85812

Signed-off-by: Benji Beck <benjibeck@meta.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

Migrate OvisImagePatchInputs to TensorSchema #22024

Migrate OvisImagePatchInputs to TensorSchema #22024

Uh oh!

bbeckca commented Jul 31, 2025 •

edited by github-actions bot

Loading

Uh oh!

github-actions bot commented Jul 31, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

DarkLight1337 left a comment

Uh oh!

bbeckca commented Aug 22, 2025

Uh oh!

bbeckca commented Aug 22, 2025

Uh oh!

mergify bot commented Sep 1, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Migrate OvisImagePatchInputs to TensorSchema #22024

Migrate OvisImagePatchInputs to TensorSchema #22024

Uh oh!

Conversation

bbeckca commented Jul 31, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Uh oh!

github-actions bot commented Jul 31, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

DarkLight1337 left a comment

Choose a reason for hiding this comment

Uh oh!

bbeckca commented Aug 22, 2025

Uh oh!

bbeckca commented Aug 22, 2025

Uh oh!

mergify bot commented Sep 1, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

bbeckca commented Jul 31, 2025 •

edited by github-actions bot

Loading