Repository: deepset-ai/haystack

Issue:
Add `devices` to FARMReader to manually select GPU, CPU etc 
**Is your feature request related to a problem? Please describe.**
Currently, we only have the arg `use_gpu` to configure if the FARMReader should use a GPU. If this is set to true and there's no GPU, we'll fallback to running on a CPU. This is convenient for many dev scenarios, but in production setups there are cases where we need more granular control (e.g. which cuda device to run on, or throwing an error if no GPU is available) 

**Describe the solution you'd like**
We can just use the same param `devices` as implemented [here](https://github.com/deepset-ai/haystack/blob/a846be99d165221790f70568651ecfcbf0c734b6/haystack/nodes/retriever/dense.py#L55) for DPR 

CC @oryx1729 



Discussion:

Resolving patch:
diff --git a/haystack/modeling/utils.py b/haystack/modeling/utils.py
--- a/haystack/modeling/utils.py
+++ b/haystack/modeling/utils.py
@@ -1,4 +1,4 @@
-from typing import Any, Iterator, Tuple, List
+from typing import Any, Iterator, Tuple, List, Optional
 
 import logging
 import os
@@ -49,38 +49,47 @@ def set_all_seeds(seed: int, deterministic_cudnn: bool = False) -> None:
 
 
 def initialize_device_settings(
-    use_cuda: bool, local_rank: int = -1, multi_gpu: bool = True
+    use_cuda: Optional[bool] = None,
+    local_rank: int = -1,
+    multi_gpu: bool = True,
+    devices: Optional[List[torch.device]] = None,
 ) -> Tuple[List[torch.device], int]:
     """
     Returns a list of available devices.
 
     :param use_cuda: Whether to make use of CUDA GPUs (if available).
-    :param local_rank: Ordinal of device to be used. If -1 and multi_gpu is True, all devices will be used.
+    :param local_rank: Ordinal of device to be used. If -1 and `multi_gpu` is True, all devices will be used.
+                       Unused if `devices` is set or `use_cuda` is False.
     :param multi_gpu: Whether to make use of all GPUs (if available).
+                      Unused if `devices` is set or `use_cuda` is False.
+    :param devices: an explicit list of which GPUs to use. Unused if `use_cuda` is False.
     """
-    if not use_cuda:
-        devices = [torch.device("cpu")]
+    if use_cuda is False:  # Note that it could be None, in which case we also want to just skip this step.
+        devices_to_use = [torch.device("cpu")]
         n_gpu = 0
+    elif devices:
+        devices_to_use = devices
+        n_gpu = sum([1 for device in devices if "cpu" not in device.type])
     elif local_rank == -1:
         if torch.cuda.is_available():
             if multi_gpu:
-                devices = [torch.device(device) for device in range(torch.cuda.device_count())]
+                devices_to_use = [torch.device(device) for device in range(torch.cuda.device_count())]
                 n_gpu = torch.cuda.device_count()
             else:
-                devices = [torch.device("cuda")]
+                devices_to_use = [torch.device("cuda")]
                 n_gpu = 1
         else:
-            devices = [torch.device("cpu")]
+            devices_to_use = [torch.device("cpu")]
             n_gpu = 0
     else:
-        devices = [torch.device("cuda", local_rank)]
-        torch.cuda.set_device(devices[0])
+        devices_to_use = [torch.device("cuda", local_rank)]
+        torch.cuda.set_device(devices_to_use[0])
         n_gpu = 1
         # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
         torch.distributed.init_process_group(backend="nccl")
-    logger.info(f"Using devices: {', '.join([str(device) for device in devices]).upper()}")
+    logger.info(f"Using devices: {', '.join([str(device) for device in devices_to_use]).upper()}")
     logger.info(f"Number of GPUs: {n_gpu}")
-    return devices, n_gpu
+    return devices_to_use, n_gpu
 
 
 def flatten_list(nested_list):
diff --git a/haystack/nodes/reader/farm.py b/haystack/nodes/reader/farm.py
--- a/haystack/nodes/reader/farm.py
+++ b/haystack/nodes/reader/farm.py
@@ -45,6 +45,7 @@ def __init__(
         context_window_size: int = 150,
         batch_size: int = 50,
         use_gpu: bool = True,
+        devices: List[torch.device] = [],
         no_ans_boost: float = 0.0,
         return_no_answer: bool = False,
         top_k: int = 10,
@@ -74,7 +75,9 @@ def __init__(
         :param batch_size: Number of samples the model receives in one batch for inference.
                            Memory consumption is much lower in inference mode. Recommendation: Increase the batch size
                            to a value so only a single batch is used.
-        :param use_gpu: Whether to use GPU (if available)
+        :param use_gpu: Whether to use GPUs or the CPU. Falls back on CPU if no GPU is available.
+        :param devices: List of GPU devices to limit inference to certain GPUs and not use all available ones (e.g. [torch.device('cuda:0')]).
+                        Unused if `use_gpu` is False.
         :param no_ans_boost: How much the no_answer logit is boosted/increased.
         If set to 0 (default), the no_answer logit is not changed.
         If a negative number, there is a lower chance of "no_answer" being predicted.
@@ -117,15 +120,14 @@ def __init__(
         """
         super().__init__()
 
-        self.devices, _ = initialize_device_settings(use_cuda=use_gpu, multi_gpu=False)
-
+        self.devices, self.n_gpu = initialize_device_settings(devices=devices, use_cuda=use_gpu, multi_gpu=True)
         self.return_no_answers = return_no_answer
         self.top_k = top_k
         self.top_k_per_candidate = top_k_per_candidate
         self.inferencer = QAInferencer.load(
             model_name_or_path,
             batch_size=batch_size,
-            gpu=use_gpu,
+            gpu=self.n_gpu > 0,
             task_type="question_answering",
             max_seq_len=max_seq_len,
             doc_stride=doc_stride,
@@ -152,7 +154,6 @@ def __init__(
         except:
             logger.warning("Could not set `duplicate_filtering` in FARM. Please update FARM version.")
         self.max_seq_len = max_seq_len
-        self.use_gpu = use_gpu
         self.progress_bar = progress_bar
         self.use_confidence_scores = use_confidence_scores
         self.confidence_threshold = confidence_threshold
@@ -165,6 +166,7 @@ def _training_procedure(
         dev_filename: Optional[str] = None,
         test_filename: Optional[str] = None,
         use_gpu: Optional[bool] = None,
+        devices: List[torch.device] = [],
         batch_size: int = 10,
         n_epochs: int = 2,
         learning_rate: float = 1e-5,
@@ -198,12 +200,12 @@ def _training_procedure(
 
         # For these variables, by default, we use the value set when initializing the FARMReader.
         # These can also be manually set when train() is called if you want a different value at train vs inference
-        if use_gpu is None:
-            use_gpu = self.use_gpu
+        if devices is None:
+            devices = self.devices
         if max_seq_len is None:
             max_seq_len = self.max_seq_len
 
-        devices, n_gpu = initialize_device_settings(use_cuda=use_gpu, multi_gpu=False)
+        devices, n_gpu = initialize_device_settings(devices=devices, use_cuda=use_gpu, multi_gpu=False)
 
         if not save_dir:
             save_dir = f"../../saved_models/{self.inferencer.model.language_model.name}"
@@ -334,6 +336,7 @@ def train(
         dev_filename: Optional[str] = None,
         test_filename: Optional[str] = None,
         use_gpu: Optional[bool] = None,
+        devices: List[torch.device] = [],
         batch_size: int = 10,
         n_epochs: int = 2,
         learning_rate: float = 1e-5,
@@ -365,6 +368,8 @@ def train(
         :param dev_split: Instead of specifying a dev_filename, you can also specify a ratio (e.g. 0.1) here
                           that gets split off from training data for eval.
         :param use_gpu: Whether to use GPU (if available)
+        :param devices: List of GPU devices to limit inference to certain GPUs and not use all available ones (e.g. [torch.device('cuda:0')]).
+                        Unused if `use_gpu` is False.
         :param batch_size: Number of samples the model receives in one batch for training
         :param n_epochs: Number of iterations on the whole training data set
         :param learning_rate: Learning rate of the optimizer
@@ -400,6 +405,7 @@ def train(
             dev_filename=dev_filename,
             test_filename=test_filename,
             use_gpu=use_gpu,
+            devices=devices,
             batch_size=batch_size,
             n_epochs=n_epochs,
             learning_rate=learning_rate,
@@ -425,6 +431,7 @@ def distil_prediction_layer_from(
         dev_filename: Optional[str] = None,
         test_filename: Optional[str] = None,
         use_gpu: Optional[bool] = None,
+        devices: List[torch.device] = [],
         student_batch_size: int = 10,
         teacher_batch_size: Optional[int] = None,
         n_epochs: int = 2,
@@ -471,6 +478,8 @@ def distil_prediction_layer_from(
         :param dev_split: Instead of specifying a dev_filename, you can also specify a ratio (e.g. 0.1) here
                           that gets split off from training data for eval.
         :param use_gpu: Whether to use GPU (if available)
+        :param devices: List of GPU devices to limit inference to certain GPUs and not use all available ones (e.g. [torch.device('cuda:0')]).
+                        Unused if `use_gpu` is False.
         :param student_batch_size: Number of samples the student model receives in one batch for training
         :param student_batch_size: Number of samples the teacher model receives in one batch for distillation
         :param n_epochs: Number of iterations on the whole training data set
@@ -514,6 +523,7 @@ def distil_prediction_layer_from(
             dev_filename=dev_filename,
             test_filename=test_filename,
             use_gpu=use_gpu,
+            devices=devices,
             batch_size=student_batch_size,
             n_epochs=n_epochs,
             learning_rate=learning_rate,
@@ -544,6 +554,7 @@ def distil_intermediate_layers_from(
         dev_filename: Optional[str] = None,
         test_filename: Optional[str] = None,
         use_gpu: Optional[bool] = None,
+        devices: List[torch.device] = [],
         batch_size: int = 10,
         n_epochs: int = 5,
         learning_rate: float = 5e-5,
@@ -585,6 +596,8 @@ def distil_intermediate_layers_from(
         :param dev_split: Instead of specifying a dev_filename, you can also specify a ratio (e.g. 0.1) here
                           that gets split off from training data for eval.
         :param use_gpu: Whether to use GPU (if available)
+        :param devices: List of GPU devices to limit inference to certain GPUs and not use all available ones (e.g. [torch.device('cuda:0')]).
+                        Unused if `use_gpu` is False.
         :param student_batch_size: Number of samples the student model receives in one batch for training
         :param student_batch_size: Number of samples the teacher model receives in one batch for distillation
         :param n_epochs: Number of iterations on the whole training data set
@@ -624,6 +637,7 @@ def distil_intermediate_layers_from(
             dev_filename=dev_filename,
             test_filename=test_filename,
             use_gpu=use_gpu,
+            devices=devices,
             batch_size=batch_size,
             n_epochs=n_epochs,
             learning_rate=learning_rate,


Test patch: