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(
 84              "deepset/bert-base-cased-squad2", device=self.device.to_hf()
 85              )
 86  
 87      def to_dict(self):
 88          # Serialize the policy like any other (custom) data.
 89          return default_to_dict(
 90              self, device=self.device.to_dict() if self.device else None, ...
 91          )
 92  
 93      @classmethod
 94      def from_dict(cls, data):
 95          # Deserialize the device data inplace before passing
 96          # it to the generic from_dict function.
 97  	    init_params = data["init_parameters"]
 98          init_params["device"] = ComponentDevice.from_dict(init_params["device"])
 99          return default_from_dict(cls, data)
100  
101  ## Automatically selects a device.
102  c = MyComponent(device=None)
103  
104  ## Uses the first GPU available.
105  c = MyComponent(device=ComponentDevice.from_str("cuda:0"))
106  
107  ## Uses the CPU.
108  c = MyComponent(device=ComponentDevice.from_single(Device.cpu()))
109  
110  ## Allow the component to use multiple devices using a device map.
111  c = MyComponent(device=ComponentDevice.from_multiple(DeviceMap({
112      "layer1": Device.cpu(),
113      "layer2": Device.gpu(1),
114      "layer3": Device.disk()
115  })))
116  ```
117  
118  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:
119  
120  ```python
121  generator = HuggingFaceLocalGenerator(
122      model="llama2",
123      huggingface_pipeline_kwargs={"device_map": "balanced"},
124  )
125  ```
126  
127  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.