feat!: Framework-agnostic device management

[shadeMe]

Related Issues

• fixes GPU/devices management in Haystack 2.x #5977

Proposed Changes:

This PR introduces the concept of framework-agnostic device representations. The main impetus behind this change is to move away from stringified representations of devices that are not portable between different frameworks. It also enables support for multi-device inference in a generic manner.

device.py contains the following new classes:

• DeviceType - An enum representing the types of devices we support.
• Device - A tuple of a DeviceType and an integer identifier that represents a single device.
• DeviceMap - An arbitrary mapping of model parameters to single devices, similar to HF's device_map in their accelerate library.
• ComponentDevice - Essentially a tagged union of Device and DeviceMap. This class consumed by downstream components.

Going forward, components can expose a single, optional device parameter in their constructor (Optional[ComponentDevice]). The component can then decide what to do with it:

• If None, the component can either do nothing or automatically pick the best device to load the model. In the case of the latter, it is done by simply passing the optional parameter to ComponentDevice.resolve_device(device).
• If not None, the component can still call ComponentDevice.resolve_device(device). In this case, the function will immediately return the passed device without any modification.
• After device resolution, the ComponentDevice instance can optionally be persisted in the component itself. Since the class itself is not trivially serializable to JSON, the Component.to_dict and Component.from_dict functions must explicitly call ComponentDevice.to_dict/from_dict during serde.
• Finally, the ComponentDevice can be converted to the component's backend representation by calling its to_xxx methods.

The following components have been updated to support the above workflow: NamedEntityExtractor, HuggingFaceLocalGenerator, TransformersSimilarityRanker, ExtractiveReader.

How did you test it?

Unit and e2e tests.

Notes for the reviewer

• I'm not happy with the ComponentDevice name. Suggestions are welcome.

Checklist

• I have read the contributors guidelines and the code of conduct
• I have updated the related issue with new insights and changes
• I added unit tests and updated the docstrings
• I've used one of the conventional commit types for my PR title: fix:, feat:, build:, chore:, ci:, docs:, style:, refactor:, perf:, test:.
• I documented my code
• I ran pre-commit hooks and fixed any issue

[coveralls]

Pull Request Test Coverage Report for Build 7545294479

• 0 of 0 changed or added relevant lines in 0 files are covered.
• No unchanged relevant lines lost coverage.
• Overall coverage increased (+0.1%) to 88.01%

---

Totals
Change from base Build 7543114073:	0.1%
Covered Lines:	4529
Relevant Lines:	5146

---

💛 - Coveralls

[sjrl]

Hey @shadeMe this is looking really good! I'm still looking through the PR, but I thought I would add a few questions that have already popped into mind:

• With the current setup if users want to specify their own device are we requiring them to have to pass a ComponentDevice which means if they want to pass "cuda:0" they would have to use ComponentDevice.from_str("cuda:0")? Do you think it would make sense for convenience to allow users to pass in a string representation as well instead of needing it to be a ComponentDevice? Or are there some problems with this approach?
• Could you provide an example of how a user could create a DeviceMap version of ComponentDevice? I'm curious to figure this out for the PR feat: Add support for device_map #6679

[shadeMe]

* With the current setup if users want to specify their own device are we requiring them to have to pass a `ComponentDevice` which means if they want to pass `"cuda:0"` they would have to use `ComponentDevice.from_str("cuda:0")`? Do you think it would make sense for convenience to allow users to pass in a string representation as well instead of needing it to be a `ComponentDevice`? Or are there some problems with this approach?

There's no inherent limitation that prevents us from converting that parameter into a Union[str, ComponentDevice] and have ComponentDevice.resolve_device call from_str internally. That being said, I think there's an argument to be made for keeping the API boundaries as explicit as possible and avoid untagged Unions.

* Could you provide an example of how a user could create a `DeviceMap` version of `ComponentDevice`? I'm curious to figure this out for the PR [feat: Add support for `device_map` #6679](https://github.com/deepset-ai/haystack/pull/6679)

device_map = ComponentDevice.from_multiple(DeviceMap({
	"classifier": Device.gpu(1),
	"layer_1": Device.cpu(),
	"lm_head": Device.disk()
}))

device_map = ComponentDevice.from_multiple(DeviceMap.from_hf({
	"classifier": 1,
	"layer_1": "cpu",
	"lm_head": "disk"
}))

It's not super ergonomic and there's admittedly a lot of ceremony, but that's the tradeoff for being more explicit/generic.

[ZanSara]

Left a few comments, in general looks good!

[sjrl]

Oh one other thing that would be very useful is that when using device_map with multiple devices it would be great if we could have an easy way to access the first device in the list. I'm not entirely sure if that is setup out of the box here. Ofc we can access this by doing a lookup on the underlying DeviceMap and converting the first Device in that dict.

I can see how this looks and feels when finishing up the device_map PR where the first device in the list is needed to know where to pass the output of the tokenizer to get it into the model.

[shadeMe]

Oh one other thing that would be very useful is that when using device_map with multiple devices it would be great if we could have an easy way to access the first device in the list. I'm not entirely sure if that is setup out of the box here. Ofc we can access this by doing a lookup on the underlying DeviceMap and converting the first Device in that dict.

Ah, good point. I can add a property for that.

[sjrl]

This looks good to me!

[ZanSara]

LGTM 😊