diff --git a/src/bentoml/_internal/frameworks/transformers.py b/src/bentoml/_internal/frameworks/transformers.py index 34b137fc5a9..773e1861d2a 100644 --- a/src/bentoml/_internal/frameworks/transformers.py +++ b/src/bentoml/_internal/frameworks/transformers.py @@ -609,7 +609,9 @@ def make_default_signatures(pretrained_cls: t.Any) -> ModelSignaturesType: ) return {} - return {k: default_config for k in infer_fn} + # NOTE: filter out methods that are not available on the given class, since + # the set of generation methods varies between transformers versions. + return {k: default_config for k in infer_fn if hasattr(pretrained_cls, k)} def import_model( @@ -1208,7 +1210,8 @@ def __init__(self): "cuda" if available_gpus not in ("", "-1") else "cpu" ) ) - torch.set_default_tensor_type("torch.cuda.FloatTensor") + if available_gpus not in ("", "-1"): + torch.set_default_tensor_type("torch.cuda.FloatTensor") elif "tf" == bento_model.info.metadata["_framework"]: with tf.device( "/device:CPU:0" @@ -1239,6 +1242,15 @@ def __init__(self): self.predict_fns: dict[str, t.Callable[..., t.Any]] = {} for method_name in bento_model.info.signatures: + if not hasattr(self.model, method_name): + # model was saved with a transformers version that had + # methods the installed version no longer provides. + logger.warning( + "Saved signature method '%s' is not available on %s and will be skipped.", + method_name, + type(self.model).__name__, + ) + continue self.predict_fns[method_name] = getattr(self.model, method_name) def add_runnable_method(method_name: str, options: ModelSignature): diff --git a/tests/integration/frameworks/test_transformers_unit.py b/tests/integration/frameworks/test_transformers_unit.py index 5963abc5477..43117bdef81 100644 --- a/tests/integration/frameworks/test_transformers_unit.py +++ b/tests/integration/frameworks/test_transformers_unit.py @@ -293,6 +293,69 @@ def test_custom_pipeline(pair_classification_pipeline: PairClassificationPipelin runner.destroy() +def test_pretrained_runnable_init_without_gpu(monkeypatch: pytest.MonkeyPatch): + """ + Instantiating a runnable for a torch pretrained model on a machine without an + assigned GPU must not set the default tensor type to CUDA. + See https://github.com/bentoml/BentoML/issues/4376 + """ + import torch + + if torch.cuda.is_available(): + pytest.skip("requires a machine without CUDA") + + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + + config = transformers.BertConfig( + vocab_size=99, + hidden_size=32, + num_hidden_layers=2, + num_attention_heads=4, + intermediate_size=37, + ) + model = transformers.BertForSequenceClassification(config).eval() + bento_model = bentoml.transformers.save_model("tiny_pretrained_pt", model) + + runnable = bento_model.to_runnable()() + + assert next(runnable.model.parameters()).device.type == "cpu" + # default tensor type must remain a CPU type + assert torch.tensor([1.0]).device.type == "cpu" + + +def test_pretrained_runnable_skips_unavailable_signature_methods( + caplog: pytest.LogCaptureFixture, +): + """ + Models saved with an older transformers version may have recorded signature + methods (e.g. 'greedy_search') that no longer exist on the loaded model. + The runnable should skip them with a warning instead of crashing. + """ + config = transformers.BertConfig( + vocab_size=99, + hidden_size=32, + num_hidden_layers=2, + num_attention_heads=4, + intermediate_size=37, + ) + model = transformers.BertForSequenceClassification(config).eval() + bento_model = bentoml.transformers.save_model( + "tiny_pretrained_pt_stale", + model, + signatures={ + "__call__": {"batchable": False}, + "greedy_search": {"batchable": False}, + }, + ) + + with caplog.at_level(logging.WARNING): + runnable = bento_model.to_runnable()() + + assert "greedy_search" in caplog.text + assert "__call__" in runnable.predict_fns + assert "greedy_search" not in runnable.predict_fns + + def test_import_model_with_synced_version(): revision = "3956d303d3cddf0708ff20660c1ea5f6ec30e434" bento_model = bentoml.transformers.import_model(