From 6aa471ac5ea7168d0611ade34e8eb1954cbfeaf3 Mon Sep 17 00:00:00 2001
From: Stefano Fiorucci <44616784+anakin87@users.noreply.github.com>
Date: Mon, 25 Sep 2023 18:39:10 +0200
Subject: [PATCH] chore: make preview integration tests reproducible (#5871)

* relax extractive reader integration tests

* force reader to CPU

* ensure integration tests reproducibility

* move set_all_seeds to testing package
---
 haystack/preview/testing/test_utils.py        | 24 +++++++++++++++++++
 .../components/readers/test_extractive.py     |  4 ++--
 test/preview/conftest.py                      |  6 ++++-
 3 files changed, 31 insertions(+), 3 deletions(-)
 create mode 100644 haystack/preview/testing/test_utils.py

diff --git a/haystack/preview/testing/test_utils.py b/haystack/preview/testing/test_utils.py
new file mode 100644
index 0000000000..333685c8d3
--- /dev/null
+++ b/haystack/preview/testing/test_utils.py
@@ -0,0 +1,24 @@
+import os
+import random
+import numpy as np
+import torch
+
+
+def set_all_seeds(seed: int, deterministic_cudnn: bool = False) -> None:
+    """
+    Setting multiple seeds to make runs reproducible.
+
+    Important: Enabling `deterministic_cudnn` gives you full reproducibility with CUDA,
+    but might slow down your training (see https://pytorch.org/docs/stable/notes/randomness.html#cudnn) !
+
+    :param seed:number to use as seed
+    :param deterministic_cudnn: Enable for full reproducibility when using CUDA. Caution: might slow down training.
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    torch.cuda.manual_seed_all(seed)
+    if deterministic_cudnn:
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
diff --git a/test/preview/components/readers/test_extractive.py b/test/preview/components/readers/test_extractive.py
index ee0979c818..a6bbe7e2f5 100644
--- a/test/preview/components/readers/test_extractive.py
+++ b/test/preview/components/readers/test_extractive.py
@@ -244,7 +244,7 @@ def test_roberta():
     assert answers[1].data == "Angela Merkel"
     assert answers[1].probability == pytest.approx(0.857952892780304)
     assert answers[2].data is None
-    assert answers[2].probability == pytest.approx(0.0196738764278237)
+    assert answers[2].probability == pytest.approx(0.019673851661650588)
     # uncomment assertions below when there is batching in v2
     # assert answers[0][0].data == "Olaf Scholz"
     # assert answers[0][0].probability == pytest.approx(0.8614975214004517)
@@ -262,7 +262,7 @@ def test_roberta():
 
 @pytest.mark.integration
 def test_matches_hf_pipeline():
-    reader = ExtractiveReader("deepset/tinyroberta-squad2")
+    reader = ExtractiveReader("deepset/tinyroberta-squad2", device="cpu")
     reader.warm_up()
     answers = reader.run(example_queries[0], [[example_documents[0][0]]][0], top_k=20, no_answer=False)[
         "answers"
diff --git a/test/preview/conftest.py b/test/preview/conftest.py
index 2ff1c16a84..3a2f166120 100644
--- a/test/preview/conftest.py
+++ b/test/preview/conftest.py
@@ -1,7 +1,11 @@
 from pathlib import Path
-from unittest.mock import Mock, patch
+from unittest.mock import Mock
 import pytest
 
+from haystack.preview.testing.test_utils import set_all_seeds
+
+set_all_seeds(0)
+
 
 @pytest.fixture()
 def mock_tokenizer():