Repository: deepset-ai/haystack
Pull request: fix: Improve robustness of get_task HF pipeline invocations
### Related Issues

- fixes #4848 

### Proposed Changes:
Before this fix, the HF invocation layer was tested early which caused the process not to reach the correct invocation layer check if the HF servers were unresponsive. This PR changes the invocation order so HF is checked later and adds a timeout of 1s so it can't get stuck waiting for a response from the HF servers.
### How did you test it?
I added a unit test that checks that the timeout is handled correctly. There is also a new test checking that the huggingface checks are in the last five invocation layer checks.

### Notes for the reviewer
The timeout is set to 1s and can't be changed by the user. This could potentially lead to issues. I am also not sure about whether the test for the invocation layer order should be different because it does not make sure the methods are actually called in that order.

I have discussed these changes with @vblagoje who also created the original issue. I am not sure if it might make sense for him to review the PR.
### Checklist

- I have read the [contributors guidelines](https://github.com/deepset-ai/haystack/blob/main/CONTRIBUTING.md) and the [code of conduct](https://github.com/deepset-ai/haystack/blob/main/code_of_conduct.txt)
- I have updated the related issue with new insights and changes
- I added unit tests and updated the docstrings
- I've used one of the [conventional commit types](https://www.conventionalcommits.org/en/v1.0.0/) for my PR title: `fix:`, `feat:`, `build:`, `chore:`, `ci:`, `docs:`, `style:`, `refactor:`, `perf:`, `test:`.
- I documented my code
- I ran [pre-commit hooks](https://github.com/deepset-ai/haystack/blob/main/CONTRIBUTING.md#installation) and fixed any issue


Problem statement:
Improve robustness of get_task HF pipeline invocations
**Is your feature request related to a problem? Please describe.**
As a Haystack user, I noticed that when HF infra is temp down I can't use a newly added Cohere model. The issue has been traced to the `PromptModelInvocationLayer.supports` invocation order. Because HFLocalInvocationLayer is listed in `haystack/nodes/prompt/invocation_layer/__init__.py` before CohereInvocationLayer, HFLocalInvocationLayer gets registered via `__init_subclass__` mechanism before CohereInvocationLayer. Then when we invoke `HFLocalInvocationLayer.supports`, the get_task call blocks indefinitely, thus preventing the use of the Cohere-based PromptNode.  

**Describe the solution you'd like**
I would like Haystack PromptNode initialization to be as robust as possible and not dependent on HF infra outage, especially when I am not using one of the HF invocation layers for PromptNode. 

**Describe alternatives you've considered**
No resolution is possible atm aside from waiting for the outage to stop

**Additional context**
Add any other context or screenshots about the feature request here.

Hints from issue discussion:

Commit message: Merge branch 'main' into fix-invocation-layer-hf-timeout

Patch to review:
diff --git a/haystack/nodes/prompt/invocation_layer/__init__.py b/haystack/nodes/prompt/invocation_layer/__init__.py
index 72df2edfda..a50aafc082 100644
--- a/haystack/nodes/prompt/invocation_layer/__init__.py
+++ b/haystack/nodes/prompt/invocation_layer/__init__.py
@@ -3,10 +3,10 @@
 from haystack.nodes.prompt.invocation_layer.chatgpt import ChatGPTInvocationLayer
 from haystack.nodes.prompt.invocation_layer.azure_chatgpt import AzureChatGPTInvocationLayer
 from haystack.nodes.prompt.invocation_layer.handlers import TokenStreamingHandler, DefaultTokenStreamingHandler
-from haystack.nodes.prompt.invocation_layer.hugging_face import HFLocalInvocationLayer
-from haystack.nodes.prompt.invocation_layer.hugging_face_inference import HFInferenceEndpointInvocationLayer
 from haystack.nodes.prompt.invocation_layer.open_ai import OpenAIInvocationLayer
 from haystack.nodes.prompt.invocation_layer.anthropic_claude import AnthropicClaudeInvocationLayer
 from haystack.nodes.prompt.invocation_layer.cohere import CohereInvocationLayer
+from haystack.nodes.prompt.invocation_layer.hugging_face import HFLocalInvocationLayer
+from haystack.nodes.prompt.invocation_layer.hugging_face_inference import HFInferenceEndpointInvocationLayer
 from haystack.nodes.prompt.invocation_layer.sagemaker_hf_infer import SageMakerHFInferenceInvocationLayer
 from haystack.nodes.prompt.invocation_layer.sagemaker_hf_text_gen import SageMakerHFTextGenerationInvocationLayer
diff --git a/haystack/nodes/prompt/invocation_layer/hugging_face.py b/haystack/nodes/prompt/invocation_layer/hugging_face.py
index f7728b53fc..99d8d92ccf 100644
--- a/haystack/nodes/prompt/invocation_layer/hugging_face.py
+++ b/haystack/nodes/prompt/invocation_layer/hugging_face.py
@@ -21,7 +21,7 @@
         GenerationConfig,
         Pipeline,
     )
-    from transformers.pipelines import get_task
+    from huggingface_hub import model_info
     from haystack.modeling.utils import initialize_device_settings  # pylint: disable=ungrouped-imports
     from haystack.nodes.prompt.invocation_layer.handlers import HFTokenStreamingHandler
 
@@ -43,6 +43,15 @@ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwa
             stop_result = torch.isin(self.stop_words["input_ids"], input_ids[-1])
             return any(all(stop_word) for stop_word in stop_result)
 
+    def get_task(model: str, use_auth_token: Optional[Union[str, bool]] = None, timeout: float = 3.0) -> Optional[str]:
+        """
+        Simplified version of transformers.pipelines.get_task with support for timeouts
+        """
+        try:
+            return model_info(model, token=use_auth_token, timeout=timeout).pipeline_tag
+        except Exception as e:
+            raise RuntimeError(f"The task of {model} could not be checked because of the following error: {e}")
+
 
 class HFLocalInvocationLayer(PromptModelInvocationLayer):
     """
diff --git a/test/prompt/invocation_layer/test_hugging_face.py b/test/prompt/invocation_layer/test_hugging_face.py
index 5c30dd6b9c..80b00aa3fe 100644
--- a/test/prompt/invocation_layer/test_hugging_face.py
+++ b/test/prompt/invocation_layer/test_hugging_face.py
@@ -356,6 +356,10 @@ def test_streaming_stream_handler_param_in_constructor(mock_pipeline, mock_get_t
     assert hf_streamer.token_handler == dtsh
 
 
+def raise_runtime_error(*args, **kwargs):
+    raise RuntimeError()
+
+
 @pytest.mark.unit
 def test_supports(tmp_path):
     """
@@ -368,7 +372,9 @@ def test_supports(tmp_path):
         assert HFLocalInvocationLayer.supports("google/flan-t5-base")
         assert HFLocalInvocationLayer.supports("mosaicml/mpt-7b")
         assert HFLocalInvocationLayer.supports("CarperAI/stable-vicuna-13b-delta")
-        assert mock_get_task.call_count == 3
+        mock_get_task.side_effect = raise_runtime_error
+        assert not HFLocalInvocationLayer.supports("google/flan-t5-base")
+        assert mock_get_task.call_count == 4
 
     # some HF local model directory, let's use the one from test/prompt/invocation_layer
     assert HFLocalInvocationLayer.supports(str(tmp_path))
diff --git a/test/prompt/invocation_layer/test_invocation_layers.py b/test/prompt/invocation_layer/test_invocation_layers.py
new file mode 100644
index 0000000000..6004a9f7db
--- /dev/null
+++ b/test/prompt/invocation_layer/test_invocation_layers.py
@@ -0,0 +1,14 @@
+import pytest
+
+from haystack.nodes.prompt.prompt_model import PromptModelInvocationLayer
+from haystack.nodes.prompt.invocation_layer import HFLocalInvocationLayer, HFInferenceEndpointInvocationLayer
+
+
+@pytest.mark.unit
+def test_invocation_layer_order():
+    """
+    Checks that the huggingface invocation layer is checked late because it can timeout/be slow to respond.
+    """
+    last_invocation_layers = set(PromptModelInvocationLayer.invocation_layer_providers[-5:])
+    assert HFLocalInvocationLayer in last_invocation_layers
+    assert HFInferenceEndpointInvocationLayer in last_invocation_layers


Code review comments:
diff_hunk: @@ -43,6 +43,15 @@ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwa
             stop_result = torch.isin(self.stop_words["input_ids"], input_ids[-1])
             return any(all(stop_word) for stop_word in stop_result)
 
+    def get_task(model: str, use_auth_token: Optional[Union[str, bool]] = None, timeout: float = 3.0) -> Optional[str]:
+        """
+        Simplified version of transformers.pipelines.get_task with support for timeouts
+        """
+        try:
+            return model_info(model, token=use_auth_token, timeout=timeout).pipeline_tag
+        except Exception as e:
+            raise RuntimeError(f"The task of {model} could not be checked because of the following error: {e}")
path: haystack/nodes/prompt/invocation_layer/hugging_face.py
text: @user1:
```suggestion
            raise RuntimeError(f"The task of {model} could not be checked because of the following error: {e}") from e
```

I feel it can be slightly improved. (Read [this](https://www.pythontutorial.net/python-oop/python-raise-from/))

@author:
Thanks, I didn't know about this feature.

diff_hunk: @@ -368,7 +372,9 @@ def test_supports(tmp_path):
         assert HFLocalInvocationLayer.supports("google/flan-t5-base")
         assert HFLocalInvocationLayer.supports("mosaicml/mpt-7b")
         assert HFLocalInvocationLayer.supports("CarperAI/stable-vicuna-13b-delta")
-        assert mock_get_task.call_count == 3
+        mock_get_task.side_effect = raise_runtime_error
path: test/prompt/invocation_layer/test_hugging_face.py
text: @user1:
```suggestion
        mock_get_task.side_effect = RuntimeError
```

By doing so, you can avoid the definition of `raise_runtime_error`.

@author:
Ok, that seems more elegant.