1  ---
  2  title: "Device Management"
  3  id: device-management
  4  slug: "/device-management"
  5  description: "This page discusses the concept of device management in the context of Haystack."
  6  ---
  7  
  8  # Device Management
  9  
 10  This page discusses the concept of device management in the context of Haystack.
 11  
 12  Many Haystack components, such as `HuggingFaceLocalGenerator` , `AzureOpenAIGenerator`, and others, allow users the ability to pick and choose which language model is to be queried and executed. For components that interface with cloud-based services, the service provider automatically takes care of the details of provisioning the requisite hardware (like GPUs). However, if you wish to use models on your local machine, you’ll need to figure out how to deploy them on your hardware. Further complicating things, different ML libraries have different APIs to launch models on specific devices.
 13  
 14  To make the process of running inference on local models as straightforward as possible, Haystack uses a framework-agnostic device management implementation. Exposing devices through this interface means you no longer need to worry about library-specific invocations and device representations.
 15  
 16  ## Concepts
 17  
 18  Haystack’s device management is built on the following abstractions:
 19  
 20  - `DeviceType`  - An enumeration that lists all the different types of supported devices.
 21  - `Device`  - A generic representation of a device composed of a `DeviceType` and a unique identifier. Together, it represents a single device in the group of all available devices.
 22  - `DeviceMap` - A mapping of strings to `Device` instances. The strings represent model-specific identifiers, usually model parameters. This allows us to map specific parts of a model to specific devices.
 23  - `ComponentDevice` - A tagged union of a single `Device` or a `DeviceMap` instance. Components that support local inference will expose an optional `device` parameter of this type in their constructor.
 24  
 25  With the above abstractions, Haystack can fully address any supported device that’s part of your local machine and can support the usage of multiple devices at the same time. Every component that supports local inference will internally handle the conversion of these generic representations to their backend-specific representations.
 26  
 27  :::info[Source Code]
 28  
 29  Find the full code for the abstractions above in the Haystack GitHub [repo](https://github.com/deepset-ai/haystack/blob/6a776e672fb69cc4ee42df9039066200f1baf24e/haystack/utils/device.py).
 30  :::
 31  
 32  ## Usage
 33  
 34  To use a single device for inference, use either the `ComponentDevice.from_single` or `ComponentDevice.from_str` class method:
 35  
 36  ```python
 37  from haystack.utils import ComponentDevice, Device
 38  
 39  device = ComponentDevice.from_single(Device.gpu(id=1))
 40  ## Alternatively, use a PyTorch device string
 41  device = ComponentDevice.from_str("cuda:1")
 42  generator = HuggingFaceLocalGenerator(model="llama2", device=device)
 43  ```
 44  
 45  To use multiple devices, use the `ComponentDevice.from_multiple` class method:
 46  
 47  ```python
 48  from haystack.utils import ComponentDevice, Device, DeviceMap
 49  
 50  device_map = DeviceMap(
 51      {
 52          "encoder.layer1": Device.gpu(id=0),
 53          "decoder.layer2": Device.gpu(id=1),
 54          "self_attention": Device.disk(),
 55          "lm_head": Device.cpu(),
 56      },
 57  )
 58  device = ComponentDevice.from_multiple(device_map)
 59  generator = HuggingFaceLocalGenerator(model="llama2", device=device)
 60  ```
 61  
 62  ### Integrating Devices in Custom Components
 63  
 64  Components should expose an optional `device` parameter of type `ComponentDevice`.  Once exposed, they can determine what to do with it:
 65  
 66  - If `device=None`, the component can pass that to the backend. In this case, the backend decides which device the model will be placed on.
 67  - Alternatively, the component can attempt to automatically pick an available device before passing it to the backend using the `ComponentDevice.resolve_device` class method.
 68  
 69  Once the device has been resolved, the component can use the `ComponentDevice.to_*` methods to get the backend-specific representation of the underlying device, which is then passed to the backend.
 70  
 71  The `ComponentDevice` instance should be serialized in the component’s `to_dict` and `from_dict` methods.
 72  
 73  ```python
 74  from haystack.utils import ComponentDevice, Device, DeviceMap
 75  
 76  class MyComponent(Component):
 77      def __init__(self, device: Optional[ComponentDevice] = None):
 78          # If device is None, automatically select a device.
 79          self.device = ComponentDevice.resolve_device(device)
 80  
 81      def warm_up(self):
 82          # Call the framework-specific conversion method.
 83          self.model = AutoModel.from_pretrained("deepset/bert-base-cased-squad2", device=self.device.to_hf())
 84  
 85  	  def to_dict(self):
 86  	    # Serialize the policy like any other (custom) data.
 87  	    return default_to_dict(self,
 88  							 device=self.device.to_dict() if self.device else None,
 89  							 ...)
 90  
 91  	  @classmethod
 92  	  def from_dict(cls, data):
 93  	    # Deserialize the device data inplace before passing
 94  			# it to the generic from_dict function.
 95  	    init_params = data["init_parameters"]
 96        init_params["device"] = ComponentDevice.from_dict(init_params["device"])
 97  			return default_from_dict(cls, data)
 98  
 99  ## Automatically selects a device.
100  c = MyComponent(device=None)
101  
102  ## Uses the first GPU available.
103  c = MyComponent(device=ComponentDevice.from_str("cuda:0"))
104  
105  ## Uses the CPU.
106  c = MyComponent(device=ComponentDevice.from_single(Device.cpu()))
107  
108  ## Allow the component to use multiple devices using a device map.
109  c = MyComponent(device=ComponentDevice.from_multiple(DeviceMap({
110        "layer1": Device.cpu(),
111        "layer2": Device.gpu(1),
112        "layer3": Device.disk()
113  })))
114  ```
115  
116  If the component’s backend provides a more specialized API to manage devices, it could add an additional init parameter that acts as a conduit. For instance, `HuggingFaceLocalGenerator` exposes a `huggingface_pipeline_kwargs` parameter through which Hugging Face-specific `device_map`  arguments can be passed:
117  
118  ```python
119  generator = HuggingFaceLocalGenerator(
120      model="llama2",
121      huggingface_pipeline_kwargs={"device_map": "balanced"},
122  )
123  ```
124  
125  In such cases, ensure that the parameter precedence and selection behavior is clearly documented. In the case of `HuggingFaceLocalGenerator`, the device map passed through the `huggingface_pipeline_kwargs` parameter overrides the explicit `device` parameter and is documented as such.