feat(nemotron-parse): richer inference metrics, engine_kwargs passthrough, fanout stage spec#2054
Conversation
Greptile SummaryThis PR enriches
Confidence Score: 5/5Safe to merge; all previously flagged issues are resolved and the new code is well-tested. The previously flagged problems (duplicate-kwarg TypeError, unreachable fallback return, and overly broad engine-startup retry markers) are all addressed. The engine_kwargs merge uses a plain dict with **-spread so Python's normal keyword-binding rules prevent duplicates. The unreachable path now raises RuntimeError. The retry classification walks the full exception chain and fast-fails on OOM markers. New tests cover metric emission, fanout spec, empty-output guard, kwargs forwarding, and the wrapped vLLM v1 startup path. The benchmark script hardcodes the metric key prefix as a string literal tied to the stage's default name; if the stage is ever renamed or instantiated with a different name, throughput figures silently become zero. Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant P as process()
participant IV as _infer_vllm()
participant VM as _vllm_metrics_from_outputs()
participant LM as _log_metrics()
P->>P: load images, log image_load_time
P->>IV: valid_images
loop up to 3 attempts
IV->>IV: llm.generate(prompts)
alt success
IV-->>P: texts, raw_outputs, vllm_retries
else RuntimeError
IV->>IV: _reset_vllm() + increment vllm_retries
end
end
P->>VM: raw_outputs + timing + page counts
VM-->>P: metrics dict (tokens, chars, truncations, retries)
P->>LM: merged metrics dict
LM->>LM: update _custom_metrics (consumed by StageTimer)
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant P as process()
participant IV as _infer_vllm()
participant VM as _vllm_metrics_from_outputs()
participant LM as _log_metrics()
P->>P: load images, log image_load_time
P->>IV: valid_images
loop up to 3 attempts
IV->>IV: llm.generate(prompts)
alt success
IV-->>P: texts, raw_outputs, vllm_retries
else RuntimeError
IV->>IV: _reset_vllm() + increment vllm_retries
end
end
P->>VM: raw_outputs + timing + page counts
VM-->>P: metrics dict (tokens, chars, truncations, retries)
P->>LM: merged metrics dict
LM->>LM: update _custom_metrics (consumed by StageTimer)
Reviews (5): Last reviewed commit: "Merge remote-tracking branch 'upstream/m..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| _ENGINE_STARTUP_FAILURE_MARKERS = ( | ||
| "EADDRINUSE", | ||
| "address already in use", | ||
| "Engine core initialization failed", | ||
| "EngineCore failed to start", | ||
| ) | ||
|
|
||
|
|
||
| def _is_engine_startup_failure(exc: BaseException) -> bool: | ||
| """Return True if ``exc`` looks like a retryable vLLM engine-startup failure.""" | ||
| message = str(exc) | ||
| return any(marker in message for marker in _ENGINE_STARTUP_FAILURE_MARKERS) |
There was a problem hiding this comment.
"Engine core initialization failed" and "EngineCore failed to start" are the generic wrapper messages that vLLM emits when the EngineCore subprocess dies for any reason — including OOM during model loading, an invalid config, or CUDA errors. Matching these strings means that a legitimate non-retryable error (e.g., torch.cuda.OutOfMemoryError wrapped into one of these messages) will now be retried up to max_port_retries times, sleeping up to ~5 s per attempt before the real error finally surfaces. The four existing EADDRINUSE-specific markers are precise; the two new ones are not. A tighter heuristic (e.g., also requiring "EADDRINUSE" somewhere in the exception chain, or inspecting __cause__) would keep retries limited to actual port collisions.
There was a problem hiding this comment.
Agree with this one.
There was a problem hiding this comment.
Maybe something like this?
# Errors that should NEVER be retried (non-transient)
_NON_RETRYABLE_MARKERS = (
"out of memory",
"CUDA out of memory",
"cudaErrorMemoryAllocation",
)
# Substrings that indicate a (usually transient) vLLM engine-startup failure
# caused by a MASTER_PORT collision.
_ENGINE_STARTUP_FAILURE_MARKERS = (
"EADDRINUSE",
"address already in use",
"Engine core initialization failed",
"EngineCore failed to start",
)
and
def _is_engine_startup_failure(exc: BaseException) -> bool:
"""Return True if ``exc`` looks like a retryable vLLM engine-startup failure."""
message = str(exc)
if any(marker in message for marker in _NON_RETRYABLE_MARKERS):
return False # OOM / config errors — fail fast, don't retry
return any(marker in message for marker in _ENGINE_STARTUP_FAILURE_MARKERS)
There was a problem hiding this comment.
Yes, this is a clean improvement. The blocklist-first pattern is easy to reason about and directly addresses the concern — OOM errors fail fast instead of burning retry budget.
A couple of thoughts:
- The OOM markers may be incomplete. Other non-transient failures (e.g. invalid model config,
RuntimeError: CUDA error: device-side assert triggered) would still slip through to the_ENGINE_STARTUP_FAILURE_MARKERScheck and get retried. The blocklist approach only works well if it's kept up to date. A small comment acknowledging this would set expectations:
| _ENGINE_STARTUP_FAILURE_MARKERS = ( | |
| "EADDRINUSE", | |
| "address already in use", | |
| "Engine core initialization failed", | |
| "EngineCore failed to start", | |
| ) | |
| def _is_engine_startup_failure(exc: BaseException) -> bool: | |
| """Return True if ``exc`` looks like a retryable vLLM engine-startup failure.""" | |
| message = str(exc) | |
| return any(marker in message for marker in _ENGINE_STARTUP_FAILURE_MARKERS) | |
| # Errors that should NEVER be retried (non-transient). | |
| # Add new entries here if vLLM introduces other non-retryable startup errors. | |
| _NON_RETRYABLE_MARKERS = ( | |
| "out of memory", | |
| "CUDA out of memory", | |
| "cudaErrorMemoryAllocation", | |
| ) |
str(exc)coverage: if vLLM constructs the wrapper message as"Engine core initialization failed"without embedding the original OOM text, the blocklist check onmessagewould miss it. It may be worth also walking__cause__/__context__:
def _is_engine_startup_failure(exc: BaseException) -> bool:
"""Return True if ``exc`` looks like a retryable vLLM engine-startup failure."""
# Collect message text from the full exception chain
parts = []
e: BaseException | None = exc
while e is not None:
parts.append(str(e))
e = e.__cause__ or e.__context__
message = " ".join(parts)
if any(marker in message for marker in _NON_RETRYABLE_MARKERS):
return False
return any(marker in message for marker in _ENGINE_STARTUP_FAILURE_MARKERS)This is optional — if vLLM reliably propagates the full traceback text into the outer exception string, the simpler version is fine. But the chain walk makes it more defensive.
Overall the proposed structure is a solid fix.
Tip: You can customize Greptile's behavior for this repo with .greptile/rules.md and .greptile/config.json.
| else: | ||
| texts = [output.outputs[0].text for output in outputs] | ||
| return texts, outputs, vllm_retries | ||
| return [], [], vllm_retries |
There was a problem hiding this comment.
The final return [], [], vllm_retries is never reached: on the last retry attempt the except branch unconditionally raises, and every successful attempt returns from the else block. If this line were ever reached (e.g., after a future refactor changes the loop structure), it would silently return empty valid_outputs, causing the downstream process() to write empty text_content for every page rather than propagating an error. Consider replacing it with raise RuntimeError("unreachable") or removing it entirely, consistent with the pragma comment in create_vllm_llm.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
praateekmahajan
changed the title
|
@claude review |
| texts = [output.outputs[0].text for output in outputs] | ||
| return texts, outputs, vllm_retries |
There was a problem hiding this comment.
_vllm_metrics_from_outputs helper correctly skips requests where req_out.outputs is empty, but the text extraction list-comprehension right below it in _infer_vllm accesses output.outputs[0].text unconditionally. If vLLM ever returns a RequestOutput with an empty completions list (e.g., during a cancelled or filtered request), the metrics path records the page as an empty output while the text path crashes with IndexError, propagating as an unhandled task failure. Using the same guard here makes the two paths consistent.
| texts = [output.outputs[0].text for output in outputs] | |
| return texts, outputs, vllm_retries | |
| texts = [output.outputs[0].text if output.outputs else "" for output in outputs] | |
| return texts, outputs, vllm_retries |
…fanout stage spec
- Add Dynamo server inference mode to NemotronParseInferenceStage: routes requests
to a running Dynamo HTTP server instead of launching an in-process vLLM engine
- Add engine_kwargs passthrough to NemotronParseInferenceStage and create_vllm_llm
so callers can tune gpu_memory_utilization, max_num_batched_tokens, etc.
- Mark PDFPartitioningStage as a fanout stage via ray_stage_spec to enable
parallelism across PDF pages
- Improve vLLM port-collision retry: handle wrapped vLLM v1 startup failures
("Engine core initialization failed") in addition to raw EADDRINUSE, and add
random jitter so concurrent workers de-stagger instead of re-colliding
- Update PDF parsing benchmark script for new inference stage API
- Add tests covering all the above
Signed-off-by: Praateek <praateekm@gmail.com>
Signed-off-by: Praateek <praateekm@gmail.com>
praateekmahajan
force-pushed
the
praateek/nemotron-parse-fixes
branch
from
e8e32a8 to
f579e3d
Compare
Signed-off-by: Praateek <praateekm@gmail.com>
…arse-fixes Signed-off-by: Praateek <praateekm@gmail.com>
3 participants
Summary
Richer per-task metrics in
NemotronParseInferenceStage(vLLM backend): now emits token counts (prompt + output), output character count, inference wall time, truncated/empty output counts, and retry count via_log_metricssoTaskPerfUtilscan aggregate them across the pipeline run.engine_kwargspassthrough:NemotronParseInferenceStageandcreate_vllm_llmnow accept arbitrary extra engine kwargs (e.g.gpu_memory_utilization,max_num_batched_tokens) forwarded tovllm.LLM, eliminating the need to subclass for tuning.Fanout stage spec on
PDFPartitioningStage: marked as a fanout stage inray_stage_specto unlock page-level parallelism under the Ray executor.vLLM v1 port-collision retry in
create_vllm_llm: also retries on"Engine core initialization failed"— in vLLM v1 the distributed-store bind happens inside theEngineCoresubprocess so the rawEADDRINUSEnever reaches the caller. Added random jitter to de-stagger concurrent workers.Benchmark script (
nemotron_parse_pdf_benchmark.py): switched from manual stage-timing loops toTaskPerfUtils.aggregate_task_metrics, reporting pages/s and output-tokens/s derived from the new stage metrics.Test plan
tests/stages/interleaved/pdf/nemotron_parse/test_stages.py— metrics emission,engine_kwargsforwarding, fanout spectests/utils/test_vllm_utils.py— wrapped vLLM v1 startup retry, extra kwargs forwarding🤖 Generated with Claude Code