Add SIGMOSFilterStage for audio quality assessment - #1577
Conversation
Greptile SummaryThis PR adds Confidence Score: 4/5Safe to merge after addressing the non-atomic model download, which can leave a corrupt cache file that fails on subsequent runs without re-downloading. Two P2 findings remain: the non-atomic/non-streaming download (where a partial write passes the size guard and causes an opaque ORT error) and the undeclared nemo_curator/stages/audio/filtering/sigmos.py (download logic); pyproject.toml (missing requests in audio_common) Important Files Changed
Sequence DiagramsequenceDiagram
participant C as Caller
participant S as SIGMOSFilterStage
participant D as _download_model
participant B as build_sigmos_model
participant O as OnnxRuntime
C->>S: setup(worker_metadata)
S->>S: ensure_cudnn_loaded()
S->>S: _initialize_model()
S->>D: _resolve_model_path()
D->>D: Check cache / download from GitHub
D-->>S: model_path
S->>B: build_sigmos_model(force_cpu, device_id, model_path)
B->>O: InferenceSession(CUDA or CPU)
O-->>B: session
B-->>S: SigMOS instance
S-->>C: ready
C->>S: process(AudioTask)
S->>S: _get_audio_numpy_sr()
S->>O: model.run(audio, sr)
O-->>S: MOS scores dict
S->>S: _check_thresholds(scores)
alt passes all thresholds
S-->>C: AudioTask with sigmos_* fields
else fails any threshold
S-->>C: []
end
Reviews (39): Last reviewed commit: "Adding Microsoft license to LICENSE file" | Re-trigger Greptile |
| _cudnn_loaded = True | ||
| except OSError: | ||
| logger.warning("Failed to load %s", cudnn_so, exc_info=True) | ||
| return False | ||
|
|
There was a problem hiding this comment.
Hardcoded cuDNN version libcudnn.so.9 will fail on other cuDNN versions
The function hardcodes libcudnn.so.9, which is specific to cuDNN 9.x. This will fail silently (returning False) on systems running cuDNN 8.x (e.g., libcudnn.so.8) or any future cuDNN 10+ version, even when cuDNN is correctly installed.
Consider searching for any available libcudnn.so.* in the directory:
import glob
cudnn_candidates = sorted(glob.glob(os.path.join(cudnn_lib_dir, "libcudnn.so.*")), reverse=True)
if not cudnn_candidates:
logger.warning("No libcudnn.so.* found in %s", cudnn_lib_dir)
return False
cudnn_so = cudnn_candidates[0] # use newest available| # Model cache for storing initialized models by GPU ID | ||
| _MODEL_CACHE = {} |
There was a problem hiding this comment.
Global _MODEL_CACHE is not thread-safe
The module-level _MODEL_CACHE dictionary is shared across all threads. In a multi-worker environment (e.g., Ray actors on the same process, or Python threading), concurrent SIGMOSPipeline initializations can race on the if cache_key not in _MODEL_CACHE check and build the same model multiple times, or read a partially-constructed model.
Consider protecting it with a lock:
import threading
_MODEL_CACHE = {}
_MODEL_CACHE_LOCK = threading.Lock()And wrapping the check-and-insert:
with _MODEL_CACHE_LOCK:
if cache_key not in _MODEL_CACHE:
_MODEL_CACHE[cache_key] = build_sigmos_model(...)
self.model = _MODEL_CACHE[cache_key]|
|
||
|
|
||
| @dataclass | ||
| class SIGMOSConfig: |
There was a problem hiding this comment.
Discussed offline. We do not need a config class per stage.
| options.inter_op_num_threads = 1 | ||
| options.intra_op_num_threads = 1 | ||
|
|
||
| use_gpu = True if is_gpu_support_available() and not force_cpu else False |
There was a problem hiding this comment.
This can be handled by the stage instead.
| def test_defaults(self) -> None: | ||
| cfg = SIGMOSConfig() | ||
| assert cfg.noise_threshold == 4.0 | ||
| assert cfg.ovrl_threshold == 3.5 | ||
| assert cfg.sig_threshold is None | ||
| assert cfg.col_threshold is None | ||
| assert cfg.disc_threshold is None | ||
| assert cfg.loud_threshold is None | ||
| assert cfg.reverb_threshold is None | ||
|
|
||
| def test_from_dict(self) -> None: | ||
| cfg = SIGMOSConfig.from_dict({"noise_threshold": 3.0, "sig_threshold": 2.5}) | ||
| assert cfg.noise_threshold == 3.0 | ||
| assert cfg.sig_threshold == 2.5 | ||
| assert cfg.ovrl_threshold == 3.5 | ||
|
|
||
| def test_from_dict_none(self) -> None: | ||
| cfg = SIGMOSConfig.from_dict(None) | ||
| assert cfg.noise_threshold == 4.0 | ||
|
|
||
| def test_from_dict_ignores_unknown(self) -> None: | ||
| cfg = SIGMOSConfig.from_dict({"unknown": 99, "noise_threshold": 1.0}) | ||
| assert cfg.noise_threshold == 1.0 |
| ) | ||
| assert cfg.get_active_thresholds() == {} | ||
|
|
||
| def test_get(self) -> None: |
| def test_stage_properties(self) -> None: | ||
| stage = SIGMOSFilterStage() | ||
| assert stage.name == "SIGMOSFilter" | ||
| assert stage.inputs() == (["data"], []) | ||
| _, output_keys = stage.outputs() | ||
| for key in ["sigmos_noise", "sigmos_ovrl", "sigmos_sig", "sigmos_col", | ||
| "sigmos_disc", "sigmos_loud", "sigmos_reverb"]: | ||
| assert key in output_keys | ||
|
|
||
| def test_config_overrides_params(self) -> None: | ||
| cfg = SIGMOSConfig(noise_threshold=2.0, ovrl_threshold=None, sig_threshold=3.0) | ||
| stage = SIGMOSFilterStage(config=cfg) | ||
| assert stage.noise_threshold == 2.0 | ||
| assert stage.ovrl_threshold is None | ||
| assert stage.sig_threshold == 3.0 |
| assert len(result.data) == 0 | ||
|
|
||
| @patch("nemo_curator.stages.audio.filtering.sigmos.SIGMOSFilterStage._initialize_model") | ||
| def test_preserves_task_metadata(self, mock_init) -> None: |
…_audio_mos, simplified pipeline
c02ad9b to
b76b6ad
Compare
…ction, remove shebang and __main__, rewrite SIGMOS tests to match stage API
b76b6ad to
ffd5888
Compare
sarahyurick
left a comment
There was a problem hiding this comment.
Kicking off the tests and Ruff formatter.
| # --------------------------------------------------------------------------- | ||
|
|
||
| TEST_DATA_DIR = Path( | ||
| "/lustre/fsw/portfolios/maxine/users/shbhawsar/debug/nemo_curator" |
There was a problem hiding this comment.
We can't have this in the codebase.
There was a problem hiding this comment.
This code will not be the part of pr its for internal testing, will remove this
|
/ok to test 3d7a714 |
|
/ok to test 672dfe0 |
|
/ok to test a1f3c82 |
|
/ok to test 3747312 |
|
/ok to test e089988 |
There was a problem hiding this comment.
Following up from our call, this can be deleted.
There was a problem hiding this comment.
Following up from our call, this can be deleted.
| Use .with_(resources=Resources(gpus=X)) to configure GPU allocation. | ||
| """ | ||
|
|
||
| model_path: str = "model/model-sigmos_1697718653_41d092e8-epo-200.onnx" |
There was a problem hiding this comment.
| model_path: str = "model/model-sigmos_1697718653_41d092e8-epo-200.onnx" | |
| model_path: str |
| "sigmos_reverb", | ||
| ] | ||
|
|
||
| def setup(self, _: WorkerMetadata | None = None) -> None: |
There was a problem hiding this comment.
The setup_on_node function can verify whether model_path exists, and if not, raise an error along with a message about where it can be downloaded from.
There was a problem hiding this comment.
Done. setup_on_node now auto-downloads the ONNX model from Microsoft's SIG-Challenge repo into ~/.cache/nemo_curator/sigmos_model/. If download fails it raises RuntimeError with the source URL and expected local path so the user can download manually. Users can also skip auto-download entirely by passing model_path= directly.
|
/ok to test b3228af |
Signed-off-by: shbhawsar <shbhawsar@nvidia.com>
| @@ -0,0 +1,121 @@ | |||
| import logging | |||
| @@ -0,0 +1,33 @@ | |||
| # What is SIGMOS? | |||
There was a problem hiding this comment.
Done please review
| import onnxruntime as ort | ||
| import scipy | ||
|
|
||
| logger = logging.getLogger(__name__) |
There was a problem hiding this comment.
Our repo uses loguru's logger throughout for logging, maybe that's a modification that can be made here.
| current_dir = os.path.dirname(os.path.abspath(__file__)) | ||
| sys.path.append(os.path.join(current_dir, "third_party")) | ||
| from sigmos.sigmos import build_sigmos_model |
There was a problem hiding this comment.
I wouldn't recommend sys path hacks to account for imports. If import errors are common we can fix the top level import to be more stable
| return audio, int(sample_rate) | ||
|
|
||
| path = item.get("audio_filepath") | ||
| if path and os.path.isfile(path): |
There was a problem hiding this comment.
not blocking for this release since it's a pattern across audio but we should make this more cloud friendly.
| return None | ||
|
|
||
| try: | ||
| score_data = self._predict_audio_mos(audio_np, sample_rate, model_path=self._resolve_model_path()) |
There was a problem hiding this comment.
Each predict call will instantiate a new class object and load the cached model? Wouldn't it be better to ensure the model is loaded once per actor and predict is called throughout?
| logger.debug( | ||
| "nvidia-cudnn-cu12 is not installed; " | ||
| "cuDNN must be available on the system LD_LIBRARY_PATH for GPU inference." | ||
| ) |
There was a problem hiding this comment.
Should this be a warning or a debug call? (I'm not sure asking for curiosity)
There was a problem hiding this comment.
Kept as debug because the pip package is only one of several ways to provide cuDNN — system installs are common. The actual GPU fallback is already surfaced as a warning in the ONNX session setup.
| ld_path = os.environ.get("LD_LIBRARY_PATH", "") | ||
| if cudnn_lib_dir not in ld_path: | ||
| os.environ["LD_LIBRARY_PATH"] = cudnn_lib_dir + (":" + ld_path if ld_path else "") | ||
|
|
||
| # Eagerly load cuDNN shared libraries into the process address space. | ||
| # Setting LD_LIBRARY_PATH alone is not enough once the process has started | ||
| # because the dynamic linker caches its search paths at startup. | ||
| # ONNX Runtime's CUDA provider uses dlopen() for sub-libraries like | ||
| # libcudnn_adv.so.9, so we must pre-load all of them. | ||
| import glob | ||
|
|
||
| cudnn_libs = sorted(glob.glob(os.path.join(cudnn_lib_dir, "libcudnn*.so*"))) | ||
| if not cudnn_libs: | ||
| logger.warning("No libcudnn*.so* files found in %s", cudnn_lib_dir) | ||
| return False | ||
|
|
||
| # Load the main library first (other libs depend on it), then the rest. | ||
| # The main library matches "libcudnn.so.<version>" (no underscore after "libcudnn"). | ||
| cudnn_libs.sort(key=lambda p: (not os.path.basename(p).startswith("libcudnn.so."), p)) | ||
|
|
||
| for lib_path in cudnn_libs: | ||
| try: | ||
| ctypes.cdll.LoadLibrary(lib_path) | ||
| logger.debug("Pre-loaded %s", lib_path) | ||
| except OSError: # noqa: PERF203 | ||
| logger.warning("Failed to load %s", lib_path, exc_info=True) | ||
|
|
||
| _cudnn_loaded = True | ||
|
|
||
| return True |
There was a problem hiding this comment.
This sort of logic can be fragile in a lot of deployment setups, we should open an issue to track if we can improve this behavior.
| ) | ||
|
|
||
|
|
||
| class TestSIGMOSFilterStage: |
There was a problem hiding this comment.
Is it possible to have a test that uses the real model?
If not will the benchmark handle both CPU and GPU fallbacks for this stage?
There was a problem hiding this comment.
The ONNX model will be removed from the repo (users download from source), so a real-model unit test would need download infrastructure in CI, adding fragility and latency to every test run. Unit tests use mocks to cover all stage logic (thresholds, nested segments, error paths); real-model end-to-end validation will be added to the benchmark pipeline, same pattern as FLEURS/ASR. The ONNX layer already handles CPU/GPU fallback transparently tries CUDAExecutionProvider first, falls back to CPUExecutionProvider if unavailable.
- Merge upstream/main, resolve __init__.py conflicts keeping SIGMOS + UTMOS + SpeakerSep - Add MIT (Microsoft) + NVIDIA license headers to third-party sigmos files - Switch third-party sigmos.py from logging to loguru for consistency - Remove sys.path hack in sigmos_pipeline.py, use direct relative import - Refactor SIGMOSFilterStage: load model once in setup(), call self._model.run() directly - Update tests to mock _initialize_model and use _make_mock_model helper - Fix Ruff lint (__all__ sorting) Signed-off-by: shbhawsar <shbhawsar@nvidia.com>
|
/ok to test 75809cc |
- Remove ONNX model files from repo (third_party/sigmos/ and model/) - Add auto-download from Microsoft's SIG-Challenge GitHub repository - Model cached at ~/.cache/nemo_curator/sigmos_model/ by default - Users can override with model_path= or model_dir= - Add setup_on_node() for multi-node pre-download - Add file validation (missing/empty check with clear error message) - Uses requests.get() matching NSFW filter pattern Signed-off-by: shbhawsar <shbhawsar@nvidia.com>
|
git history commit 5faad7a (added), 75809cc (removed). Models no longer in latest HEAD of PR, but 50 MB still in git history. Two identical copies of the SIGMOS ONNX model were checked into git:
That's 50 MB of binary permanently in git history. Every Follow-up needed:
|
| import librosa | ||
| import numpy as np | ||
| import onnxruntime as ort | ||
| import scipy |
There was a problem hiding this comment.
The third-party sigmos.py (Microsoft SigMOS) does import librosa and import scipy at module level (lines 7-8 of sigmos_filter_module/third_party/sigmos/sigmos.py). Neither librosa nor scipy is declared in pyproject.toml — not in base deps, not in any audio extras.
Impact:
from nemo_curator.stages.audio.filtering import SIGMOSFilterStagecrashes withModuleNotFoundError: No module named 'librosa'- Since
filtering/__init__.pyimportsSIGMOSFilterStageat module level, the entirenemo_curator.stages.audio.filteringpackage is broken without librosa installed - This also breaks
AudioDataFilterStage(this PR) since it imports fromfiltering
Fix: add librosa and scipy to the audio_cpu and/or audio_cuda12 extras in pyproject.toml.
This I think will come up when running unit tests in this pr ci/cd as well.
|
|
||
| def teardown(self) -> None: | ||
| self._model = None | ||
| torch.cuda.empty_cache() |
There was a problem hiding this comment.
SIGMOSFilterStage.teardown() in nemo_curator/stages/audio/filtering/sigmos.py calls torch.cuda.empty_cache() unconditionally:
def teardown(self) -> None:
self._model = None
torch.cuda.empty_cache() # crashes if CUDA not availableThere was a problem hiding this comment.
The unconditional torch.cuda.empty_cache() in teardown() now affects 4 of 6 new stages:
| Stage | PR | teardown() guarded? |
|---|---|---|
| SIGMOSFilterStage | #1577 | ❌ No |
| VADSegmentationStage | #1578 | ❌ No |
| UTMOSFilterStage | #1639 | ❌ No |
| SpeakerSeparationStage | #1579 | ❌ No |
| BandFilterStage | #1576 | ✅ Yes |
| MonoConversionStage | #1575 | N/A (no GPU) |
BandFilterStage (#1576) already has the correct pattern:
if torch.cuda.is_available():
torch.cuda.empty_cache()
|
|
||
| from .third_party.sigmos.sigmos import build_sigmos_model | ||
|
|
||
| _MODEL_CACHE: dict[str, Any] = {} |
There was a problem hiding this comment.
_SIGMOSPipeline and predict_audio_mos() with a module-level _MODEL_CACHE dict, but SIGMOSFilterStage never imports or uses any of it. The stage has its own _initialize_model() that directly calls build_sigmos_model().
This file seems like dead code. Remove it, or if model caching across workers was intended, wire it into the stage.
Please let me know if this file is used somewhere. Also add a readme edit somewhere relevant so that it is clear to users what exactly is being done by each added stage/file.
There was a problem hiding this comment.
Thanks for bringing this up. Removed this file
|
|
||
| def teardown(self) -> None: | ||
| self._model = None | ||
| torch.cuda.empty_cache() |
There was a problem hiding this comment.
The unconditional torch.cuda.empty_cache() in teardown() now affects 4 of 6 new stages:
| Stage | PR | teardown() guarded? |
|---|---|---|
| SIGMOSFilterStage | #1577 | ❌ No |
| VADSegmentationStage | #1578 | ❌ No |
| UTMOSFilterStage | #1639 | ❌ No |
| SpeakerSeparationStage | #1579 | ❌ No |
| BandFilterStage | #1576 | ✅ Yes |
| MonoConversionStage | #1575 | N/A (no GPU) |
BandFilterStage (#1576) already has the correct pattern:
if torch.cuda.is_available():
torch.cuda.empty_cache()
|
/ok to test d0936c4 |
- Guard torch.cuda.empty_cache() with torch.cuda.is_available() in SIGMOSFilterStage.teardown() to prevent crash on CPU-only machines - Add librosa and scipy to audio_common extras in pyproject.toml (required by third-party sigmos.py module-level imports) - Remove unused sigmos_pipeline.py (dead code, never imported by stage) - Regenerate uv.lock
Signed-off-by: Sarah Yurick <53962159+sarahyurick@users.noreply.github.com>
|
/ok to test f297f01 |
|
/ok to test 9aa21a3 |
Signed-off-by: shubhamNvidia <shbhawsar@nvidia.com>
|
/ok to test 961e8d0 |
Adds SIGMOSFilterStage which filters audio based on SIGMOS (Signal-based Mean Opinion Score) quality metrics using an ONNX model. Predicts 7 quality dimensions (NOISE, OVRL, SIG, COL, DISC, LOUD, REVERB) on a 0-5 scale. Each threshold is independently configurable and can be set to None to disable. Items failing any active threshold are filtered out. Includes SIGMOSConfig with get_active_thresholds(), the sigmos_filter_module pipeline, and unit tests with mocked prediction function.