From cfad4035b9d416cd1b7dff6ff0927486a92b2587 Mon Sep 17 00:00:00 2001 From: MushiSenpai <259025340+MushiSenpai@users.noreply.github.com> Date: Fri, 26 Jun 2026 19:51:51 +0800 Subject: [PATCH] Honor singular device= as an alias for devices= (fixes #907) Co-Authored-By: Claude Opus 4.8 --- FlagEmbedding/abc/inference/AbsEmbedder.py | 10 +++++ FlagEmbedding/abc/inference/AbsReranker.py | 10 +++++ tests/test_device_alias.py | 43 ++++++++++++++++++++++ 3 files changed, 63 insertions(+) create mode 100644 tests/test_device_alias.py diff --git a/FlagEmbedding/abc/inference/AbsEmbedder.py b/FlagEmbedding/abc/inference/AbsEmbedder.py index 32bf59ca0..a0b1308f7 100644 --- a/FlagEmbedding/abc/inference/AbsEmbedder.py +++ b/FlagEmbedding/abc/inference/AbsEmbedder.py @@ -69,6 +69,16 @@ def __init__( self.use_bf16 = use_bf16 self.query_instruction_for_retrieval = query_instruction_for_retrieval self.query_instruction_format = query_instruction_format + # `device` (singular) is a common alias for `devices`. Without this it is silently + # swallowed by **kwargs, leaving `devices=None` so the model spans ALL visible GPUs + # (out-of-memory). Honor it as `devices` when `devices` is not explicitly set. See #907. + if devices is None and "device" in kwargs: + devices = kwargs.pop("device") + logger.warning( + "`device` is not a documented argument for %s; interpreting it as " + "`devices=%r`. Pass `devices=` to silence this warning.", + type(self).__name__, devices, + ) self.target_devices = self.get_target_devices(devices) self.batch_size = batch_size diff --git a/FlagEmbedding/abc/inference/AbsReranker.py b/FlagEmbedding/abc/inference/AbsReranker.py index c7c6fec27..7c88aad50 100644 --- a/FlagEmbedding/abc/inference/AbsReranker.py +++ b/FlagEmbedding/abc/inference/AbsReranker.py @@ -66,6 +66,16 @@ def __init__( self.query_instruction_format = query_instruction_format self.passage_instruction_for_rerank = passage_instruction_for_rerank self.passage_instruction_format = passage_instruction_format + # `device` (singular) is a common alias for `devices`. Without this it is silently + # swallowed by **kwargs, leaving `devices=None` so the model spans ALL visible GPUs + # (out-of-memory). Honor it as `devices` when `devices` is not explicitly set. See #907. + if devices is None and "device" in kwargs: + devices = kwargs.pop("device") + logger.warning( + "`device` is not a documented argument for %s; interpreting it as " + "`devices=%r`. Pass `devices=` to silence this warning.", + type(self).__name__, devices, + ) self.target_devices = self.get_target_devices(devices) self.batch_size = batch_size diff --git a/tests/test_device_alias.py b/tests/test_device_alias.py new file mode 100644 index 000000000..f87972ca9 --- /dev/null +++ b/tests/test_device_alias.py @@ -0,0 +1,43 @@ +"""Regression test for #907. + +`device=` (singular) must be honored as an alias for `devices=` instead of being +silently swallowed by **kwargs — which left `devices=None`, so the model spanned ALL +visible GPUs and OOM'd. Model-free: only exercises argument handling in the abstract +base classes (no weights downloaded, no GPU required). +""" +import pytest + +from FlagEmbedding.abc.inference.AbsEmbedder import AbsEmbedder +from FlagEmbedding.abc.inference.AbsReranker import AbsReranker + + +class _DummyEmbedder(AbsEmbedder): + def encode_single_device(self, *args, **kwargs): # abstract stub + pass + + +class _DummyReranker(AbsReranker): + def compute_score_single_gpu(self, *args, **kwargs): # abstract stub + pass + + +@pytest.mark.parametrize("cls", [_DummyEmbedder, _DummyReranker]) +def test_singular_device_aliases_to_devices(cls): + # `device=` should be honored, not dropped into **kwargs. + obj = cls("dummy-model-name", device="cuda:5") + assert obj.target_devices == ["cuda:5"] + # ...and must NOT leak as a stray attribute or into self.kwargs. + assert not hasattr(obj, "device") + assert "device" not in obj.kwargs + + +@pytest.mark.parametrize("cls", [_DummyEmbedder, _DummyReranker]) +def test_plural_devices_still_works(cls): + obj = cls("dummy-model-name", devices="cuda:3") + assert obj.target_devices == ["cuda:3"] + + +@pytest.mark.parametrize("cls", [_DummyEmbedder, _DummyReranker]) +def test_explicit_devices_wins_over_device(cls): + obj = cls("dummy-model-name", devices="cuda:2", device="cuda:9") + assert obj.target_devices == ["cuda:2"]