[RFC] [Misc] Convert untyped OmniStageTask dict to TypedDict

[NickCao]

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

Currently the task messages passed between stages are untyped dict, requiring the receiving stage to validate all incoming messages manually. Thus I propose we gradually type them as TypedDict, and eventually as pydantic tagged union to improve the maintainability of these messages.

Test Plan

TypedDicts are just regular dicts at runtime, no additional testing is required, apart from testing if the stages still work.

Test Result

N/A

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
• The test results. Please paste the results comparison before and after, or the e2e results.
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model. Please run mkdocs serve to sync the documentation editions to ./docs.
• (Optional) Release notes update. If your change is user-facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 7f2d3b0440

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[lishunyang12]

Left a couple comments. Main concern is that OmniStageTaskGenerate tries to cover both seed tasks and connector notifications, which have almost entirely different field sets. The most important construction sites (stage-0 seeds in omni.py and async_omni.py) are still raw dicts, which kind of defeats the purpose of the typing effort.

[NickCao]

The most important construction sites (stage-0 seeds in omni.py and async_omni.py) are still raw dicts, which kind of defeats the purpose of the typing effort.

Converted these to OmniStageTaskGenerate, whether to further factor them into a new task type can happen in a followup PR.

[lishunyang12]

Seed-site conversions look good. Two things to fix:

1. sampling_params is declared twice in OmniStageTaskGenerate — second definition shadows the first. Pick one type (probably NotRequired[Any]).
2. connector_metadata is still dynamically assigned in adapter.py but missing from the TypedDict.

[lishunyang12]

See reviews

[NickCao]

See reviews

Already resolved in the latest push?

[NickCao]

Rebased on current main.

[Gaohan123]

@natureofnature PTAL

[hsliuustc0106]

fix conflicts

[NickCao]

No longer relevant after #1908