feat: select CUDA device consistently across inference backends

[be-student]

Summary

Add Separator(cuda_device_index=...) and --cuda_device_index so a multi-GPU user can select the same visible device for PyTorch and ONNX Runtime. configure_cuda(..., device_index=...) also supports an explicit index.

Fixes #225.

The constructor argument is appended to preserve existing positional calls. With no index specified, current automatic behavior remains unchanged. Explicit values are validated before initialization; unavailable CUDA and out-of-range indices raise ValueError rather than silently choosing another backend. Indices are relative to CUDA_VISIBLE_DEVICES.

The change also handles two downstream requirements of indexed devices:

• MDX compares the provider name when ONNX configuration includes device options, avoiding a false fallback warning.
• CUDA cache cleanup recognizes indexed devices and enters the selected device context before clearing the cache.

The README documents CLI/API usage and validation behavior.

Validation

• New tests on unchanged code: 10 failed, 1 passed.
• Focused API/CLI/runtime tests after implementation: 42 passed, 1 skipped.
• Final pytest tests/unit -q: 566 passed, 4 skipped, 13 warnings on Python 3.12 with poetry install -E cpu.
• git diff --check: passed.

Tests cover index propagation to both backends, invalid values, range checks, unavailable CUDA, unchanged defaults, positional compatibility, CLI forwarding, MDX provider options/fallback reporting, and cache cleanup for indexed CUDA devices. These device-selection checks mock CUDA availability and ONNX sessions; real multi-GPU inference was not performed.

AI assistance: Codex implemented this change and ran the listed validation.

Summary by CodeRabbit

• New Features

  • Added support for selecting a specific CUDA device through the CLI or Python API.
  • CUDA device selection now applies consistently to PyTorch and ONNX Runtime.
  • Added validation with clear errors for unavailable or out-of-range devices.
  • Documented device indexing, including behavior with CUDA_VISIBLE_DEVICES.

• Bug Fixes

  • Improved GPU cache clearing for the active CUDA device.
  • Preserved ONNX Runtime provider configuration when validating selected devices.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Team

Run ID: e262f94e-372d-40ed-8203-b7d053fb4bdc

📥 Commits

Reviewing files that changed from the base of the PR and between bf1164a and 2b04211.

📒 Files selected for processing (8)

• README.md
• audio_separator/separator/architectures/mdx_separator.py
• audio_separator/separator/common_separator.py
• audio_separator/separator/separator.py
• audio_separator/utils/cli.py
• tests/unit/test_cli.py
• tests/unit/test_cuda_device_selection.py
• tests/unit/test_separator_api_compatibility.py

Included review availability: Your plan provides up to 4 included reviews per hour; 3 remain after this review.

---

Walkthrough

Changes

CUDA device selection

Layer / File(s)	Summary
CUDA selection API and configuration audio_separator/separator/separator.py, tests/unit/test_cuda_device_selection.py	Separator validates cuda_device_index, applies it to PyTorch and ONNX Runtime, and raises ValueError for invalid, unavailable, or out-of-range devices.
Runtime provider and cache handling audio_separator/separator/architectures/mdx_separator.py, audio_separator/separator/common_separator.py, tests/unit/test_cuda_device_selection.py	ONNX provider tuples retain device options, and GPU cache cleanup runs in the selected CUDA device context.
CLI, documentation, and API compatibility audio_separator/utils/cli.py, README.md, tests/unit/test_cli.py, tests/unit/test_separator_api_compatibility.py	The CLI and Python documentation expose the new option. CLI defaults and constructor signature tests include it.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Merge Risk: ⚪ Minimal · up to 2b042

This change adds validated CUDA device selection across the API and CLI while preserving default behavior and positional compatibility. Current coverage indicates no actionable merge-blocking risk remains.

Suggested reviewers: ntamotsu, beveradb

Sequence Diagram(s)

sequenceDiagram
  participant CLI
  participant Separator
  participant PyTorch
  participant ONNXRuntime
  CLI->>Separator: Pass cuda_device_index
  Separator->>Separator: Validate device index
  Separator->>PyTorch: Select cuda:index
  Separator->>ONNXRuntime: Set CUDA provider device_id

Loading

Poem

A rabbit picked a CUDA lane
And checked each index twice
PyTorch hopped to cuda:n
ONNX kept its device slice
The cache cleared clean and nice

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 21 functions across 7 files. (1 skipped: …	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: consistent CUDA device selection across PyTorch and ONNX Runtime inference backends.
Linked Issues check	✅ Passed	The PR satisfies issue [#225] by adding an optional device index to configure_cuda and extending the selected index consistently to PyTorch and ONNX Runtime. It also preserves default behavior and val…
Out of Scope Changes check	✅ Passed	The changes remain within the linked objective. The README updates, CLI support, MDX provider handling, CUDA cache cleanup, and tests directly support indexed CUDA selection and compatibility.

Full details: Docstring Coverage

Explanation

Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 21 functions across 7 files. (1 skipped: 1 unsupported.)

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}