/ haystack / testing / test_utils.py
test_utils.py
 1  # SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
 2  #
 3  # SPDX-License-Identifier: Apache-2.0
 4  
 5  import os
 6  import random
 7  
 8  from haystack import logging
 9  
10  logger = logging.getLogger(__name__)
11  
12  
13  def set_all_seeds(seed: int, deterministic_cudnn: bool = False) -> None:
14      """
15      Setting multiple seeds to make runs reproducible.
16  
17      Important: Enabling `deterministic_cudnn` gives you full reproducibility with CUDA,
18      but might slow down your training (see https://pytorch.org/docs/stable/notes/randomness.html#cudnn) !
19  
20      :param seed:number to use as seed
21      :param deterministic_cudnn: Enable for full reproducibility when using CUDA. Caution: might slow down training.
22      """
23      random.seed(seed)
24      os.environ["PYTHONHASHSEED"] = str(seed)
25  
26      try:
27          import torch
28  
29          torch.manual_seed(seed)
30          torch.cuda.manual_seed_all(seed)
31          if deterministic_cudnn:
32              torch.backends.cudnn.deterministic = True
33              torch.backends.cudnn.benchmark = False
34  
35      except (ImportError, ModuleNotFoundError) as exc:
36          logger.info("Could not set PyTorch seed because torch is not installed. Exception: {exception}", exception=exc)