Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion funasr/auto/auto_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -527,7 +527,8 @@ def __init__(self, **kwargs):
Args:
model (str): Model name (hub alias or full ID) or local path.
device (str): Device for inference. "cuda:0", "cpu", "mps", "npu:0".
Falls back to CPU if specified device is unavailable.
Logs a warning and falls back to CPU if unavailable (default: "cuda").
Explicit CPU use or ngpu=0 does not emit a fallback warning.
vad_model (str, optional): VAD model for long audio segmentation.
Enables processing of any-length audio.
vad_kwargs (dict, optional): VAD config, e.g. {"device": "cpu"}.
Expand Down Expand Up @@ -673,6 +674,14 @@ def build_model(**kwargs):
or (device.startswith("npu") and not is_npu_available())
or kwargs.get("ngpu", 1) == 0
):
if kwargs.get("ngpu", 1) != 0:
logging.warning(
"AutoModel: requested device '%s' is unavailable for model '%s'; "
"falling back to CPU. Check your PyTorch build and accelerator "
"availability, or set device='cpu' to use CPU explicitly.",
device,
kwargs["model"],
)
device = "cpu"
kwargs["batch_size"] = 1
kwargs["device"] = device
Expand Down
124 changes: 124 additions & 0 deletions tests/test_submodel_device.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,12 @@
"""

import copy
import logging
import unittest
from unittest.mock import patch

import numpy as np
import pytest
import torch

from funasr.auto.auto_model import AutoModel
Expand Down Expand Up @@ -153,5 +155,127 @@ def test_unavailable_explicit_submodel_devices_fall_back_to_cpu(self):
self.assertEqual(configs[role]["device"], "cuda:0")


@pytest.fixture
def fallback_model(monkeypatch):
from funasr.auto.auto_model import _current_blas_threads, _limit_blas_threads

original_threads = torch.get_num_threads()
original_blas = _current_blas_threads()
for role in ("asr", "vad", "punc", "spk"):
monkeypatch.setitem(tables.model_classes, f"device-test-{role}", RecordingModel)

def reject_download(**kwargs):
raise AssertionError("device tests must not access model hubs")

monkeypatch.setattr("funasr.auto.auto_model.download_model", reject_download)
with patch("funasr.auto.auto_model.ClusterBackend", create=True):
def build(**kwargs):
return AutoModel(
model="device-test-asr",
model_conf={},
ncpu=1,
disable_update=True,
disable_pbar=True,
**kwargs,
)

yield build
if original_blas is not None:
_limit_blas_threads(original_blas)
torch.set_num_threads(original_threads)


_ACCELERATOR_PROBES = [
("cuda:0", "torch.cuda.is_available"),
("xpu:0", "torch.xpu.is_available"),
("mps", "torch.backends.mps.is_available"),
("npu:0", "funasr.auto.auto_model.is_npu_available"),
]


@pytest.mark.parametrize("device,probe", _ACCELERATOR_PROBES)
def test_unavailable_accelerator_warns_and_preserves_cpu_fallback(
fallback_model, caplog, device, probe
):
with patch(probe, return_value=False), caplog.at_level(logging.WARNING):
model = fallback_model(device=device, batch_size=8)
assert model.kwargs["device"] == model.model.placement == "cpu"
assert model.kwargs["batch_size"] == 1
assert len(caplog.records) == 1
assert caplog.records[0].levelno == logging.WARNING
message = caplog.records[0].getMessage()
assert device in message
assert "device-test-asr" in message
assert "falling back to CPU" in message
assert "PyTorch" in message


@pytest.mark.parametrize("device,probe", _ACCELERATOR_PROBES)
def test_available_accelerator_does_not_warn(fallback_model, caplog, device, probe):
with patch(probe, return_value=True), caplog.at_level(logging.WARNING):
model = fallback_model(device=device, batch_size=8)
assert model.kwargs["device"] == model.model.placement == device
assert model.kwargs["batch_size"] == 8
assert not caplog.records


@pytest.mark.parametrize("device,probe", _ACCELERATOR_PROBES)
@pytest.mark.parametrize("available", [False, True])
def test_ngpu_zero_is_intentional_cpu_use(
fallback_model, caplog, device, probe, available
):
with patch(probe, return_value=available), caplog.at_level(logging.WARNING):
model = fallback_model(device=device, ngpu=0, batch_size=8)
assert model.kwargs["device"] == model.model.placement == "cpu"
assert model.kwargs["batch_size"] == 1
assert not caplog.records


def test_explicit_cpu_does_not_warn(fallback_model, caplog):
with caplog.at_level(logging.WARNING):
model = fallback_model(device="cpu", batch_size=8)
assert model.kwargs["device"] == model.model.placement == "cpu"
assert model.kwargs["batch_size"] == 8
assert not caplog.records


def test_default_cuda_fallback_warns_once_with_inherited_submodels(fallback_model, caplog):
configs = {
f"{role}_{kind}": f"device-test-{role}" if kind == "model" else {"model_conf": {}}
for role in ("vad", "punc", "spk")
for kind in ("model", "kwargs")
}
original = copy.deepcopy(configs)
with patch("torch.cuda.is_available", return_value=False), caplog.at_level(logging.WARNING):
model = fallback_model(**configs)
assert model.kwargs["device"] == "cpu"
for role in ("vad", "punc", "spk"):
assert getattr(model, f"{role}_kwargs")["device"] == "cpu"
assert getattr(model, f"{role}_model").placement == "cpu"
assert configs == original
assert len(caplog.records) == 1
assert "cuda" in caplog.records[0].getMessage()


def test_explicit_submodel_fallback_identifies_the_submodel(fallback_model, caplog):
config = {"model_conf": {}, "device": "cuda:0"}
original = copy.deepcopy(config)
with patch("torch.cuda.is_available", return_value=False), caplog.at_level(logging.WARNING):
model = fallback_model(
device="cpu", punc_model="device-test-punc", punc_kwargs=config
)
assert model.punc_kwargs["device"] == model.punc_model.placement == "cpu"
assert config == original
assert len(caplog.records) == 1
assert "device-test-punc" in caplog.records[0].getMessage()


def test_fallback_respects_application_logging_level(fallback_model, caplog):
with patch("torch.cuda.is_available", return_value=False), caplog.at_level(logging.ERROR):
model = fallback_model(device="cuda:0")
assert model.kwargs["device"] == "cpu"
assert not caplog.records


if __name__ == "__main__":
unittest.main()
Loading