feat: Granary v2 pipeline — vLLM 0.19.1, throughput metadata export, standardized audio stages

[mohammadaaftabv]

Summary

• Upgrade vLLM to 0.19.1 with torch==2.10.0 for Qwen3.5 MoE architecture support
• Add throughput metadata export: per-stage inference timing via _log_metrics() in InferenceQwenOmniStage and PnCRestorationStage, with perf_summary.json aggregate output in ShardedManifestWriterStage
• Pin dependencies (vllm==0.19.1, torch==2.10.0+cu128, transformers==5.7.0) for stable ABI compatibility
• Fix config paths (/src/Curator -> /opt/Curator) for Docker image alignment
• Add pipeline architecture documentation

Changes

• nemo_curator/stages/audio/inference/qwen_omni.py — add inference timing + utterance count metrics
• nemo_curator/stages/audio/text_filtering/pnc_restoration.py — add inference timing + restoration count metrics
• nemo_curator/stages/audio/io/sharded_manifest_writer.py — add perf_summary.json + per-shard _perf.jsonl export
• tutorials/audio/qwen_omni_inprocess/main.py — log perf summary after pipeline completion
• pyproject.toml / uv.lock — dependency pins

Test plan

• Verified pipeline completes successfully on Kratos with 17,197 utterances (25.5 min)
• Docker image built and pushed with correct dependency versions
• Verify perf_summary.json output in next Kratos run

Made with Cursor

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[mohammadaaftabv]

Closing in favor of #1881 which targets main.

[greptile-apps]

Greptile Summary

This PR upgrades vLLM to 0.19.1 (torch 2.10.0) for Qwen3.5 MoE support and introduces a substantial set of new audio pipeline stages: ShardedManifestWriterStage, PnCRestorationStage, InferenceQwenASRStage, QwenTextLLM, plus throughput metadata export via _log_metrics() and perf_summary.json. The dependency constraint changes (removing transformers<5.0 and huggingface-hub<1.0 caps) are intentional and well-commented.

• P1: PnCRestorationStage.setup_on_node calls self._model.setup(), performing a full vLLM GPU allocation at node level. Unlike InferenceQwenOmniStage.setup_on_node which only prefetches weights, this leaks GPU memory in the node process — in distributed (Xenna) deployment the stage is serialized to workers which each re-allocate the engine independently, risking OOM before the pipeline starts.
• P2: InferenceQwenOmniStage.skip_me_key defaults to \"_skipme\" while PnCRestorationStage.skip_me_key defaults to \"_skip_me\" — these don't match, so the skip flag set by the omni stage is silently ignored by the PnC stage.

Confidence Score: 3/5

The P1 OOM risk in PnCRestorationStage.setup_on_node should be resolved before merging for distributed/multi-GPU deployments.

One P1 finding (GPU memory double-allocation in distributed setup_on_node) that could silently cause OOM at pipeline start in production Kratos runs; score pulled below the P1 ceiling of 4 due to combined risk with the skip_me_key inconsistency and the large surface area of 70 files.

nemo_curator/stages/audio/text_filtering/pnc_restoration.py (setup_on_node GPU allocation), nemo_curator/stages/audio/inference/qwen_omni.py (skip_me_key default mismatch)

Important Files Changed

Filename	Overview
nemo_curator/stages/audio/text_filtering/pnc_restoration.py	New PnC restoration stage — setup_on_node performs full GPU model load (P1 OOM risk in distributed), while the pattern from InferenceQwenOmniStage only prefetches weights at node level.
nemo_curator/stages/audio/inference/qwen_omni.py	Adds inference timing, language-aware prompt interpolation, and skipped_indices tracking; skip_me_key default differs from PnCRestorationStage causing silent skip-flag misalignment.
nemo_curator/stages/audio/io/sharded_manifest_writer.py	New stage — writes per-shard JSONL + perf stats; opens files per-task (no buffering), which is a throughput bottleneck for large pipelines.
nemo_curator/models/qwen_text_llm.py	New two-step completeness-check + PnC restoration LLM; clean design with thread-pool preprocessing and proper teardown.
tutorials/audio/qwen_omni_inprocess/main.py	Hydra entry point for Granary v2 pipeline; parallel model prefetch is well-implemented; auth token written only to parent-process env may miss spawned vLLM workers.
nemo_curator/stages/audio/inference/qwen_asr.py	New Qwen3-ASR inference stage; imports private _LANG_CODE_TO_NAME from sibling module rather than a shared utility location.
pyproject.toml	Pins vllm==0.19.1, torch==2.10.0, removes transformers<5.0 cap; transformers>=4.56.0 appears in both constraint-dependencies and override-dependencies (redundant).
nemo_curator/stages/audio/common.py	Adds ManifestReader, ManifestReaderStage, ManifestWriterStage, and read_jsonl_manifests helper; well-structured with fsspec for cloud path support.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[NemoTarredAudioReader] --> B[InitializeFieldsStage]
    B --> C[InferenceQwenOmniStage\nvLLM 0.19.1 sets _skipme]
    C --> D{followup_prompt?}
    D -- yes --> E[DisfluencyWerGuardStage]
    D -- no --> F[WhisperHallucinationStage omni]
    E --> F
    F --> G{asr_model_id?}
    G -- yes --> H[InferenceQwenASRStage]
    G -- no --> I[SelectBestPredictionStage]
    H --> H2[WhisperHallucinationStage asr]
    H2 --> I
    I --> J[FastTextLIDStage]
    J --> K[RegexSubstitutionStage]
    K --> L[AbbreviationConcatStage]
    L --> M{skip_pnc?}
    M -- no --> N[PnCRestorationStage\nsetup_on_node allocates GPU P1]
    N --> O[PnCContentGuardStage]
    O --> P{enable_itn?}
    M -- yes --> P
    P -- yes --> Q[ITNRestorationStage]
    P -- no --> R[ShardedManifestWriterStage\nper-shard JSONL + perf_summary.json]
    Q --> R

Loading

_{Reviews (1): Last reviewed commit: "docs: add Granary v2 pipeline architectu..." | Re-trigger Greptile}

[greptile-apps]

setup_on_node allocates GPU memory — diverges from intended node-level prefetch

setup_on_node calls self._model.setup(), which triggers a full vLLM engine allocation (GPU memory). In contrast, InferenceQwenOmniStage.setup_on_node explicitly only calls snapshot_download and states "no GPU allocation" in its docstring. In the Xenna executor, setup_on_node runs once per node process, and the stage is then serialized to per-worker actors — _model is lost on deserialization, so setup() re-allocates the full engine anyway. The node-level GPU allocation is therefore leaked: GPU memory is consumed and never freed in the node process, while each worker independently allocates again. On a node with multiple pipeline stages, this double allocation can cause OOM before the first batch is processed.

The fix is to align with the pattern used in InferenceQwenOmniStage: call only snapshot_download in setup_on_node and leave full model initialization to setup().

[greptile-apps]

skip_me_key default mismatches PnCRestorationStage

InferenceQwenOmniStage.skip_me_key defaults to "_skipme" (no second underscore), while PnCRestorationStage.skip_me_key defaults to "_skip_me". When InferenceQwenOmniStage marks a task as skipped, PnCRestorationStage reads the wrong key and gets an empty string, treating the task as eligible. In practice the downstream text_key is also empty for such tasks (preventing actual inference), but the skip-flag is silently missed for any stage that relies on it between the two. The defaults should be aligned to a single canonical field name across both stages.

[greptile-apps]

Importing private symbol from sibling module

_LANG_CODE_TO_NAME is a module-private dict defined in qwen_omni.py and imported here with its private-by-convention name. Sharing private symbols across module boundaries makes the dependency implicit and fragile — renaming or moving the dict in qwen_omni.py silently breaks qwen_asr.py. Consider moving _LANG_CODE_TO_NAME to a shared utility location (e.g., nemo_curator/stages/audio/common.py) and exporting it as a public constant.

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!

[greptile-apps]

Per-task file open/close is a throughput bottleneck

Both process() and _write_perf_line() open and close their respective files for every single AudioTask. On a 17K-utterance pipeline run, this is ~34K syscall pairs per shard file. For high-throughput pipelines this will be a significant bottleneck that can be avoided by buffering writes (e.g., holding an open file handle in stage state after setup_on_node, flushing in teardown). Since num_workers() already returns 1, thread-safety is not a concern.

[greptile-apps]

transformers>=4.56.0 appears in both constraint-dependencies and override-dependencies

After this PR, transformers>=4.56.0 is listed in both sections. While this probably does not break resolution today (the override supersedes the constraint), it is redundant and makes future version bumps error-prone — a maintainer may update only one entry, causing the two entries to diverge. Consider keeping the transformers pin only in override-dependencies and removing it from constraint-dependencies.