Repository: deepset-ai/haystack

Issue:
Better model security with new PyTorch version
When loading PyTorch models from the modelhub, arbitrary code can be executed. See [here.](https://huggingface.co/docs/hub/security-pickle#pickle-scanning)

Seems like PyTorch already solved this:
- Issue on the PyTorch Repo: https://github.com/pytorch/pytorch/issues/52596
- By default, this is not available in torch==1.13. TORCH_FORCE_WEIGHTS_ONLY_LOAD needs to be set to True which allows global override to safe only model loading via.

Test
Test haystack tests with new flags

**Solution:**
- Bump up PyTorch version to 1.13

**Behaviour**
- Secure by default TORCH_FORCE_WEIGHTS_ONLY_LOAD always set to true when Haystack loads any models




Discussion:
Something that I have been struggling for a couple of days: https://github.com/python-poetry/poetry/issues/6409
@mayankjobanputra I also tried to get this working. This seams to be a bug on pytorchs side as mentioned in https://github.com/pytorch/pytorch/pull/89924. It will probably get fixes in the [1.13.1](https://github.com/pytorch/pytorch/issues/89855) release which will be released in few days. I'll take another look when they fixed the marker problem in pytorch.
@LLukas22 thanks a lot :) 

I am also tracking this PyTorch issue and looking forward to upgrading PyTorch version in Haystack!
https://github.com/pytorch/pytorch/issues/52181#issuecomment-1335560667

Discussion is still going on about enabling TORCH_FORCE_WEIGHTS_ONLY_LOAD by default.

Resolving patch:
diff --git a/haystack/__init__.py b/haystack/__init__.py
--- a/haystack/__init__.py
+++ b/haystack/__init__.py
@@ -20,6 +20,8 @@
 from haystack.schema import Document, Answer, Label, MultiLabel, Span, EvaluationResult
 from haystack.nodes.base import BaseComponent
 from haystack.pipelines.base import Pipeline
+from haystack.environment import set_pytorch_secure_model_loading
 
 
 pd.options.display.max_colwidth = 80
+set_pytorch_secure_model_loading()
diff --git a/haystack/environment.py b/haystack/environment.py
--- a/haystack/environment.py
+++ b/haystack/environment.py
@@ -1,3 +1,4 @@
+import logging
 import os
 import platform
 import sys
@@ -17,6 +18,18 @@
 
 env_meta_data: Dict[str, Any] = {}
 
+logger = logging.getLogger(__name__)
+
+
+def set_pytorch_secure_model_loading(flag_val="1"):
+    # To load secure only model pytorch requires value of
+    # TORCH_FORCE_WEIGHTS_ONLY_LOAD to be ["1", "y", "yes", "true"]
+    os_flag_val = os.getenv("TORCH_FORCE_WEIGHTS_ONLY_LOAD")
+    if os_flag_val is None:
+        os.environ["TORCH_FORCE_WEIGHTS_ONLY_LOAD"] = flag_val
+    else:
+        logger.info("TORCH_FORCE_WEIGHTS_ONLY_LOAD is already set to %s, Haystack will use the same.", os_flag_val)
+
 
 def get_or_create_env_meta_data() -> Dict[str, Any]:
     """


Test patch:
diff --git a/test/others/test_utils.py b/test/others/test_utils.py
--- a/test/others/test_utils.py
+++ b/test/others/test_utils.py
@@ -1,4 +1,6 @@
+import importlib
 import logging
+import os
 from random import random
 from typing import List
 
@@ -12,6 +14,7 @@
 from ..conftest import fail_at_version, haystack_version
 
 from haystack.errors import OpenAIRateLimitError
+from haystack.environment import set_pytorch_secure_model_loading
 from haystack.schema import Answer, Document, Span, Label
 from haystack.utils.deepsetcloud import DeepsetCloud, DeepsetCloudExperiments
 from haystack.utils.labels import aggregate_labels
@@ -1245,6 +1248,17 @@ def greet2(name: str):
     assert greet2("John") == "Hello John"
 
 
+def test_secure_model_loading(monkeypatch, caplog):
+    caplog.set_level(logging.INFO)
+    monkeypatch.setenv("TORCH_FORCE_WEIGHTS_ONLY_LOAD", "0")
+
+    # now testing if just importing haystack is enough to enable secure loading of pytorch models
+    import haystack
+
+    importlib.reload(haystack)
+    assert "already set to" in caplog.text
+
+
 class TestAggregateLabels:
     @pytest.fixture
     def standard_labels(self) -> List[Label]: