Repository navigation
[Bugfix] Preserve object storage URIs during model resolution - #5036
hsliuustc0106 merged 3 commits into
Conversation
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
There was a problem hiding this comment.
Pull request overview
Fixes incorrect early “resolution” of object-storage model URIs in omni_snapshot_download() by preserving s3://, gs://, and az:// URIs until stage-specific ModelConfig initialization, allowing vLLM’s Run:AI Model Streamer path to materialize only config/tokenizer files and keep streamed weight loading intact.
Changes:
- Add a small allowlist of object-storage URI schemes.
- Short-circuit
omni_snapshot_download()to return object-storage URIs unchanged (avoiding the Hugging Face pre-download path).
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # ModelConfig. vLLM then materializes only config/tokenizer files locally | ||
| # and keeps the URI in model_weights for Run:AI streaming. Treating the URI | ||
| # as a Hugging Face repo here either fails validation or downloads through | ||
| # the wrong backend before the stage processes are created. | ||
| if model_id.lower().startswith(_OBJECT_STORAGE_SCHEMES): |
There was a problem hiding this comment.
I see the CI does not support the S3 loading, so I did not add any test here
There was a problem hiding this comment.
Perhaps we could use a mock approach instead of the real S3 environment, and add the following verification points?
Unit test: When omni_snapshot_download receives an object storage URI (e.g., s3://bucket/model, gs://..., az://...), the return value should be the original URI, and the Hugging Face pre-download path should not be invoked. Use mock assertions to verify that the HF download function is not called.
Unit test: When a local path or relative path is passed in, the existing behavior should be preserved. If the expected behavior is to use HF's local resolution/download, then assert that the HF path is called.
There was a problem hiding this comment.
Thanks for the comments – these are great suggestions. I'll address them and add the unit tests today.
There was a problem hiding this comment.
Perhaps we could use a mock approach instead of the real S3 environment, and add the following verification points? Unit test: When omni_snapshot_download receives an object storage URI (e.g., s3://bucket/model, gs://..., az://...), the return value should be the original URI, and the Hugging Face pre-download path should not be invoked. Use mock assertions to verify that the HF download function is not called.
Unit test: When a local path or relative path is passed in, the existing behavior should be preserved. If the expected behavior is to use HF's local resolution/download, then assert that the HF path is called.
Fixed
|
|
||
| logger = init_logger(__name__) | ||
|
|
||
| _OBJECT_STORAGE_SCHEMES = ("s3://", "gs://", "az://") |
There was a problem hiding this comment.
vLLM's Run:AI streamer only recognizes s3:// and gs:// (SUPPORTED_SCHEMES in vllm/transformers_utils/runai_utils.py). An az:// URI gets preserved here but is_runai_obj_uri() returns False downstream, so vLLM falls back to treating it as an HF repo id — the exact failure this PR fixes. Drop az:// unless you've confirmed streaming works for it.
| # and keeps the URI in model_weights for Run:AI streaming. Treating the URI | ||
| # as a Hugging Face repo here either fails validation or downloads through | ||
| # the wrong backend before the stage processes are created. | ||
| if model_id.lower().startswith(_OBJECT_STORAGE_SCHEMES): |
There was a problem hiding this comment.
Reuse from vllm.transformers_utils.runai_utils import is_runai_obj_uri here instead of the local _OBJECT_STORAGE_SCHEMES + .lower().startswith(...). Same check, and it stays in sync if vLLM adds a scheme.
| # as a Hugging Face repo here either fails validation or downloads through | ||
| # the wrong backend before the stage processes are created. | ||
| if model_id.lower().startswith(_OBJECT_STORAGE_SCHEMES): | ||
| return model_id |
There was a problem hiding this comment.
Add a test for this branch, e.g. asserting omni_snapshot_download("s3://bucket/model") returns the URI unchanged, so a future refactor doesn't silently re-download.
| @pytest.mark.parametrize( | ||
| "model_uri", | ||
| [ | ||
| "s3://bucket/model", | ||
| "gs://bucket/model", | ||
| ], | ||
| ) |
There was a problem hiding this comment.
Added az://bucket/model to the parametrized object-storage cases.
fc1fc74 to
3538437
Compare
| import huggingface_hub | ||
| from vllm.logger import init_logger | ||
| from vllm.transformers_utils.runai_utils import is_runai_obj_uri | ||
| from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError |
There was a problem hiding this comment.
Not guarding this. vLLM is a hard dependency of vLLM-Omni and vllm.transformers_utils.runai_utils (with is_runai_obj_uri) has been part of it for many releases — every vLLM version vLLM-Omni supports already provides this symbol, so an unconditional import cannot raise ImportError in practice, and a local fallback would risk silently diverging from the streamer's actual scheme list.
| @pytest.mark.parametrize( | ||
| "model_uri", | ||
| [ | ||
| "s3://bucket/model", | ||
| "gs://bucket/model", | ||
| ], | ||
| ) |
There was a problem hiding this comment.
Added az://bucket/model to the parametrized object-storage cases; see also the reply on the companion thread.
3538437 to
afc2913
Compare
|
Hi @lishunyang12 @yenuo26 – I think this PR is almost ready. Could you please take a look when you have a moment? |
3a3478e to
c397654
Compare
|
Hi @lishunyang12 @yenuo26 – I think this PR is almost ready. Could you please take a look when you have a moment? |
c397654 to
ddb3319
Compare
|
Hi @lishunyang12 @yenuo26 – I think this PR is almost ready. Could you please take a look when you have a moment? |
ddb3319 to
93d7f85
Compare
|
Please resolve conflicts. Thanks |
|
fix precommits |
|
Hi, I already fix it, please take a look. |
linyueqian
left a comment
There was a problem hiding this comment.
Reviewed at 31ac9c7ad against main@c98517ce7.
The added branch itself is correct and I see no regression risk: nothing that previously worked can newly match is_runai_obj_uri, and os.path.exists still runs first. Delegating to vLLM's own scheme list also settles the earlier az:// question, since on a vLLM build without az:// the URI just takes the pre-PR path.
The issue is that it is not sufficient. vllm-omni serve s3://... still crashes in the parent process before any stage is created, and that applies to Qwen3-TTS too, so #2408 is not fixed on the default path. Trace is inline on the new branch.
I have not run this against real object storage and I do not see evidence in the PR that anyone has. A single end to end run with an actual S3 URI would be worth having before merge.
| # and keeps the URI in model_weights for Run:AI streaming. Treating the URI | ||
| # as a Hugging Face repo here either fails validation or downloads through | ||
| # the wrong backend before the stage processes are created. | ||
| if is_runai_obj_uri(model_id): |
There was a problem hiding this comment.
[High] This unblocks omni_snapshot_download, but the same failure returns a few frames later, so vllm-omni serve s3://bucket/... still crashes in the parent process. Qwen3-TTS is affected as well, so #2408 is not fixed on the default path.
Path on this head:
OmniBase.__init__(omni_base.py:184) ->AsyncOmniEngine.__init__(async_omni_engine.py:245) ->_resolve_stage_configs(async_omni_engine.py:1215) ->load_and_resolve_stage_configs.- With no
--deploy-configand no--stage-configs-path, that takes theelif stage_configs_path is None:branch and callsconfig_path = resolve_model_config_path(model)at entrypoints/utils.py:681, unconditionally and beforeload_stage_configs_from_model. - Inside
resolve_model_config_path, onlyget_configis guarded (utils.py:288-291). The next callget_diffusion_model_index(model)at utils.py:293 is not, and it hands the raw URI toget_hf_file_to_dict->_try_download_from_hf_hub->hf_hub_download. - vLLM's
_try_download_from_hf_hubcatchesOfflineModeIsEnabled, the four not-found errors, andHfHubHTTPError.HFValidationErroris not in that list, so it propagates. Confirmed on huggingface_hub 1.27.0:hf_hub_download("s3://bucket/model", "model_index.json")raisesRepo id must be in the form 'repo_name' or 'namespace/repo_name': 's3://bucket/model', the same message this PR removes.
Passing --deploy-config takes the branch at utils.py:664 and skips resolve_model_config_path entirely, so that configuration does get further. Is that how the fix was verified?
Two things still stand in the way even on that route:
get_hf_config(config_factory.py:120) can never return a config for an object URI, because vLLM'sget_confighas no object-storage handling.try_infer_model_typetherefore falls through to the substring match at config_factory.py:207, which is unanchored and scans the whole URI including the bucket name.s3://qwen3-tts-models/Qwen3-Omni-30B-A3B-Instructresolves toqwen3_tts, andbagel,lance,moss_tts_realtimeare all registered keys (pipeline_registry.py:134, 137, 171), so ordinary bucket names can select the wrong pipeline. Separately,qwen3_omni_moeis a resolver (pipeline_registry.py:129) whose wrapper returnsNonewhenhf_config is None(stage_config.py:39-40), so Qwen3-Omni cannot resolve from object storage at all.- qwen3_tts_code2wav.py:548 and qwen3_tts_talker.py:1010 load the
speech_tokenizer/weights viaDefaultModelLoader.Source(model_or_path=self.model_path, ...), andself.model_pathismodel_config.model(code2wav.py:67, talker.py:288), which after the Run:AI pull points at the localmodel_streamer/<hash>dir. vLLM pulls that dir withallow_pattern=["*.model", "*.py", "*.json"], so those auxiliary weights are never there.
Smallest fix for the blocker: materialize the config files once for is_runai_obj_uri models, for example ObjectStorageModel(url).pull_files(url, allow_pattern=["*.json"]), and use that local dir for get_hf_config / try_infer_model_type / resolve_model_config_path while still passing the URI as each stage's model. That also removes the dependence on the bucket name. If that is too large for this PR, it would at least help to guard resolve_model_config_path so the failure names object storage rather than an HF repo id, and to narrow the "any vLLM-Omni pipeline" claim in the description to what is actually covered.
There was a problem hiding this comment.
Thanks for the thorough trace — you were right on every point. Addressed in 8f458e9:
Parent-process crash: added _materialize_object_storage_configs in config_factory.py. For is_runai_obj_uri models it pulls ["*.model", "*.py", "*.json"] once via ObjectStorageModel into the deterministic model_streamer/<sha256(url)[:8]> directory (same dir the stage workers later reuse in maybe_pull_model_tokenizer_for_runai). All config reads now go through it:
resolve_model_config_path—get_config, theconfig.jsonfallback, andget_diffusion_model_indexall receive the local dir, sohf_hub_download("s3://...", "model_index.json")no longer happens and noHFValidationError.StageConfigFactory.get_hf_config— materializes beforeget_config, sohf_configis now a real config for object URIs.qwen3_omni_moe's resolver no longer receivesNone, and Qwen3-Omni resolves from object storage.- The original URI is untouched everywhere stages receive it, so
model_config.model_weights/streaming behavior is unchanged.
Bucket-name hijack: both name-match fallbacks (_try_infer_model_type and _try_resolve_omni_model_type) now scan only the basename. Your example s3://qwen3-tts-models/Qwen3-Omni-30B-A3B-Instruct now resolves via its real config.json to qwen3_omni_moe, and even in fallback mode the bucket cannot select a pipeline. Regression tests included for both.
Remaining gap acknowledged: the speech_tokenizer/ subfolder weights in qwen3_tts_code2wav/talker are still not obtainable from object storage — vLLM's pull pattern never fetches weight files, and DefaultModelLoader has no object-storage-aware subfolder path (a raw URI goes to download_weights_from_hf). I cannot verify a streamer-side subfolder pull without real object storage, so I documented this as the known remaining limitation in the PR description rather than shipping an untested fix. Proposing it as a follow-up; happy to do it here instead if you're OK with the verification constraint.
| mocker.patch.object(omni_base.os.path, "exists", return_value=False) | ||
| hf_download = mocker.patch.object(omni_base, "download_weights_from_hf_specific") |
There was a problem hiding this comment.
[Low] This test reaches the network. file_or_path_exists is not stubbed, so omni_snapshot_download("org/model") gets to omni_base.py:93 and does a live lookup against huggingface.co for modular_model_index.json. Online that returns False and the test passes, but with HF_HUB_OFFLINE=1 it raises OfflineModeIsEnabled, and an HF 5xx raises HfHubHTTPError. Neither is in the except (GatedRepoError, RepositoryNotFoundError) at omni_base.py:95, so the test errors rather than fails. The download_backend fixture below already stubs this at lines 113-117.
Adding mocker.patch.object(omni_base, "file_or_path_exists", return_value=False) makes it hermetic.
[Nit] Line 69 patches omni_base.os.path, which is the stdlib os.path module, so os.path.exists returns False process wide for the duration of the test. os.path.exists("org/model") is already False here, so the line can just be dropped.
There was a problem hiding this comment.
Fixed — file_or_path_exists is now stubbed via mocker.patch.object(omni_base, "file_or_path_exists", return_value=False) (matching the download_backend fixture below), and the omni_base.os.path patch is dropped since "org/model" never exists on disk. The test no longer reaches the network under either setting.
| "s3://bucket/model", | ||
| "gs://bucket/model", | ||
| ], |
There was a problem hiding this comment.
[Nit] az:// is missing here, though the description and vLLM's SUPPORTED_SCHEMES both include it. Copilot raised this twice already.
Worth a line in the description too: az:// only entered SUPPORTED_SCHEMES between vLLM 0.16.0 and 0.18.0, so the az:// claim holds on 0.18 and later.
There was a problem hiding this comment.
Fixed — az://bucket/model is now a parametrized case in test_omni_snapshot_download_preserves_object_storage_uri. The PR description also notes the version window: az:// entered vLLM's SUPPORTED_SCHEMES between 0.16.0 and 0.18.0, so the claim holds on 0.18 and later.
31ac9c7 to
8f458e9
Compare
|
Rebased onto latest New commit 8f458e9 addresses the outstanding review:
Review comment threads above each have a point-by-point reply. |
Fixed — the branch has been rebased onto latest Buildkite (previous head |
Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
Use the vLLM is_runai_obj_uri helper when resolving model paths so supported object-storage schemes stay aligned with the downstream streamer. Add mock-based coverage for object-storage URIs, existing local paths, and relative Hugging Face repository IDs. Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
Address the remaining "serve s3://..." parent-process crash from the review: resolve_model_config_path handed raw object-storage URIs to huggingface_hub helpers via get_config/get_hf_file_to_dict, which reject them with HFValidationError, and get_hf_config could never succeed, so pipeline selection fell back to scanning the whole URI (bucket name included) for a registered pipeline key. For is_runai_obj_uri models, pull the lightweight config/code/tokenizer files once into vLLM's deterministic model_streamer/<hash> directory and route every config read (get_config, config.json/model_index.json fallbacks, diffusers _class_name lookup) through it. Stage workers later reuse the same directory for their own Run:AI pull. The original URI is unchanged everywhere stages receive it, so streaming is unaffected. Also anchor both name-match fallbacks (_try_infer_model_type and _try_resolve_omni_model_type) to the URI/path basename so non-model segments can no longer select an unrelated pipeline, e.g. a bucket "qwen3-tts-models" holding a Qwen3-Omni checkpoint. Mock-based tests cover passthrough, one-pull-per-URI, config-driven pipeline selection under a deceptive bucket name, basename-only name matching, and end-to-end resolve_model_config_path with an s3 URI. The s3/gs/az parametrization now includes az://, and the HF-repo-id test no longer reaches the network or patches the stdlib os.path module. Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
8f458e9 to
785f10a
Compare
…roject#5036) Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
Purpose
Fix #2408.
Object storage model URIs must remain unchanged until vLLM constructs the stage-specific
ModelConfig.Currently,
omni_snapshot_download()sends every non-local model identifier through the Hugging Face pre-download path. This is incorrect for object storage URIs such as:These URIs are handled by vLLM's Run:AI Model Streamer integration. During ModelConfig initialization, vLLM materializes the small configuration and tokenizer files into a local model_streamer/ cache directory and
preserves the original object storage URI in model_weights for streamed weight loading.
Resolving the model before stage initialization can either:
If the local cache path is unavailable or incomplete in the receiving stage process, Hugging Face treats the absolute path as a repository ID and raises:
Repo id must be in the form 'repo_name' or 'namespace/repo_name'
This issue is not specific to Qwen3-TTS. It can affect any vLLM-Omni pipeline using Run:AI Model Streamer and an object storage model URI.
What this change does
omni_snapshot_download()preserves the URI unchanged for every scheme in vLLM's Run:AISUPPORTED_SCHEMES(s3://,gs://;az://since it entered that list between vLLM 0.16.0 and 0.18.0), detected viais_runai_obj_uriso the set stays aligned with the streamer.Parent-process resolution no longer crashes on the URI. For
is_runai_obj_urimodels, the lightweight config/code/tokenizer files (*.json,*.py,*.model) are pulled once into vLLM's deterministicmodel_streamer/<hash>directory, and every config read goes through it:StageConfigFactory.get_hf_config/try_infer_model_type(config read,config.jsonfallback,model_index.jsonfallback) — so pipeline keys resolved from the real checkpoint config, and resolver-based pipelines such asqwen3_omni_moereceive a realhf_config.resolve_model_config_path— no moreHFValidationErrorfromhf_hub_download("s3://...", "model_index.json").The stage workers' own Run:AI pull reuses the same directory, and each stage still receives the original URI as its
model, so streamed weight loading is unaffected.Both name-match fallbacks (
_try_infer_model_type,_try_resolve_omni_model_type) now scan only the URI/path basename. Non-model segments such as the bucket name can no longer select an unrelated pipeline (a bucketqwen3-tts-modelsholding a Qwen3-Omni checkpoint previously resolved toqwen3_tts).Testing: mock-based unit tests cover URI passthrough in
omni_snapshot_download(incl.az://), local-path and HF-repo-id behavior without network access, one-pull-per-URI materialization, config-driven pipeline selection under a deceptive bucket name, basename-only name matching, andresolve_model_config_pathon ans3://URI without handing it to HF helpers. No run against real object storage has been performed yet.Known remaining gap (follow-up): the stage-side auxiliary weights under
speech_tokenizer/for Qwen3-TTS (qwen3_tts_code2wav.py,qwen3_tts_talker.py) load viaDefaultModelLoaderwithsubfolder="speech_tokenizer"from the locally materialized directory, which vLLM populates with only*.model/*.py/*.json— those weight files are therefore absent for object-storage URIs. Making the streamer fetch that subfolder needs verification against real object storage and is left to a follow-up PR.vLLM Version:
0.18.0 for the original issue reproduction.
vLLM-Omni Commit:
9e1a069
BEFORE SUBMITTING: read CONTRIBUTING.md and run the precheck-pr skill with the code agent for a self-check against project conventions.
(anything written below this line will be removed by GitHub Actions)