From 9a3c5c6a96ef9b19957feace39794b0930da53aa Mon Sep 17 00:00:00 2001 From: Julian Risch Date: Tue, 20 Jan 2026 17:55:34 +0100 Subject: [PATCH 1/3] handle ComponentDevice in default serialization --- haystack/core/serialization.py | 31 +++++- test/core/test_serialization.py | 176 +++++++++++++++++++++++++++++++- 2 files changed, 202 insertions(+), 5 deletions(-) diff --git a/haystack/core/serialization.py b/haystack/core/serialization.py index 12514dda17..3ee6ff9715 100644 --- a/haystack/core/serialization.py +++ b/haystack/core/serialization.py @@ -11,6 +11,7 @@ from haystack.core.component.component import _hook_component_init from haystack.core.errors import DeserializationError, SerializationError from haystack.utils.auth import Secret, deserialize_secrets_inplace +from haystack.utils.device import ComponentDevice from haystack.utils.type_serialization import thread_safe_import logger = logging.getLogger(__name__) @@ -182,8 +183,8 @@ def default_to_dict(obj: Any, **init_parameters: Any) -> dict[str, Any]: instance of `obj` with `from_dict`. Omitting them might cause deserialisation errors or unexpected behaviours later, when calling `from_dict`. - Secret instances in `init_parameters` are automatically serialized by calling - their `to_dict()` method. + Secret and ComponentDevice instances in `init_parameters` are automatically + serialized by calling their `to_dict()` method. An example usage: @@ -213,10 +214,10 @@ def to_dict(self): :returns: A dictionary representation of the instance. """ - # Automatically serialize Secret instances + # Automatically serialize Secret and ComponentDevice instances serialized_params = {} for key, value in init_parameters.items(): - if isinstance(value, Secret): + if isinstance(value, (Secret, ComponentDevice)): serialized_params[key] = value.to_dict() else: serialized_params[key] = value @@ -224,6 +225,19 @@ def to_dict(self): return {"type": generate_qualified_class_name(type(obj)), "init_parameters": serialized_params} +def _is_serialized_component_device(value: Any) -> bool: + """ + Check if a value is a serialized ComponentDevice dictionary. + + A dictionary is considered a serialized ComponentDevice if it has a "type" key + with value "single" or "multiple". + """ + if not isinstance(value, dict): + return False + type_value = value.get("type") + return type_value in ("single", "multiple") + + def default_from_dict(cls: type[T], data: dict[str, Any]) -> T: """ Utility function to deserialize a dictionary to an object. @@ -240,6 +254,10 @@ def default_from_dict(cls: type[T], data: dict[str, Any]) -> T: deserialized. A dictionary is considered a serialized Secret if it has a "type" key with value "env_var". + Serialized ComponentDevice dictionaries in `init_parameters` are automatically detected + and deserialized. A dictionary is considered a serialized ComponentDevice if it has a + "type" key with value "single" or "multiple". + :param cls: The class to be used for deserialization. :param data: @@ -265,6 +283,11 @@ def default_from_dict(cls: type[T], data: dict[str, Any]) -> T: if secret_keys: deserialize_secrets_inplace(init_params, keys=secret_keys) + # Automatically detect and deserialize ComponentDevice instances + for key, value in init_params.items(): + if _is_serialized_component_device(value): + init_params[key] = ComponentDevice.from_dict(value) + return cls(**init_params) diff --git a/test/core/test_serialization.py b/test/core/test_serialization.py index 655ba649e3..d3221a7d58 100644 --- a/test/core/test_serialization.py +++ b/test/core/test_serialization.py @@ -19,7 +19,8 @@ import_class_by_name, ) from haystack.testing import factory -from haystack.utils import Secret +from haystack.utils import ComponentDevice, Secret +from haystack.utils.device import Device, DeviceMap def test_default_component_to_dict(): @@ -253,3 +254,176 @@ def test_component_to_dict_and_from_dict_roundtrip_with_secret(): assert deserialized_comp.regular_param == "test" assert deserialized_comp.api_key._env_vars == env_secret1._env_vars assert deserialized_comp.token._env_vars == env_secret2._env_vars + + +@component +class CustomComponentWithDevice: + def __init__( + self, + device: ComponentDevice | None = None, + other_device: ComponentDevice | None = None, + name: str | None = None, + ): + self.device = device + self.other_device = other_device + self.name = name + + @component.output_types(value=str) + def run(self, value: str): + return {"value": value} + + +def test_component_to_dict_with_component_device(): + """Test that ComponentDevice instances are automatically serialized in component_to_dict.""" + # Test with single device (CPU) + device = ComponentDevice.from_single(Device.cpu()) + comp = CustomComponentWithDevice(device=device) + res = component_to_dict(comp, "test_component") + assert res["init_parameters"]["device"] == {"type": "single", "device": "cpu"} + + # Test with single device (GPU with id) + device = ComponentDevice.from_single(Device.gpu(1)) + comp = CustomComponentWithDevice(device=device) + res = component_to_dict(comp, "test_component") + assert res["init_parameters"]["device"] == {"type": "single", "device": "cuda:1"} + + # Test with None + comp = CustomComponentWithDevice(device=None) + res = component_to_dict(comp, "test_component") + assert res["init_parameters"]["device"] is None + + # Test with multiple devices (device map) + device_map = DeviceMap({"layer1": Device.gpu(0), "layer2": Device.gpu(1)}) + device = ComponentDevice.from_multiple(device_map) + comp = CustomComponentWithDevice(device=device) + res = component_to_dict(comp, "test_component") + assert res["init_parameters"]["device"] == { + "type": "multiple", + "device_map": {"layer1": "cuda:0", "layer2": "cuda:1"}, + } + + # Test with multiple ComponentDevice params + device1 = ComponentDevice.from_single(Device.cpu()) + device2 = ComponentDevice.from_single(Device.gpu(0)) + comp = CustomComponentWithDevice(device=device1, other_device=device2, name="test") + res = component_to_dict(comp, "test_component") + assert res["init_parameters"]["device"] == {"type": "single", "device": "cpu"} + assert res["init_parameters"]["other_device"] == {"type": "single", "device": "cuda:0"} + assert res["init_parameters"]["name"] == "test" + + +def test_component_from_dict_with_component_device(): + """Test that serialized ComponentDevice dictionaries are automatically deserialized in component_from_dict.""" + # Test with single device (CPU) + data = { + "type": generate_qualified_class_name(CustomComponentWithDevice), + "init_parameters": {"device": {"type": "single", "device": "cpu"}, "name": "test"}, + } + comp = component_from_dict(CustomComponentWithDevice, data, "test_component") + assert isinstance(comp, CustomComponentWithDevice) + assert isinstance(comp.device, ComponentDevice) + assert comp.device.to_torch_str() == "cpu" + assert comp.name == "test" + + # Test with single device (GPU with id) + data = { + "type": generate_qualified_class_name(CustomComponentWithDevice), + "init_parameters": {"device": {"type": "single", "device": "cuda:1"}, "name": "test"}, + } + comp = component_from_dict(CustomComponentWithDevice, data, "test_component") + assert isinstance(comp.device, ComponentDevice) + assert comp.device.to_torch_str() == "cuda:1" + + # Test with None + data = { + "type": generate_qualified_class_name(CustomComponentWithDevice), + "init_parameters": {"device": None, "name": "test"}, + } + comp = component_from_dict(CustomComponentWithDevice, data, "test_component") + assert comp.device is None + assert comp.name == "test" + + # Test with multiple devices (device map) + data = { + "type": generate_qualified_class_name(CustomComponentWithDevice), + "init_parameters": {"device": {"type": "multiple", "device_map": {"layer1": "cuda:0", "layer2": "cuda:1"}}}, + } + comp = component_from_dict(CustomComponentWithDevice, data, "test_component") + assert isinstance(comp.device, ComponentDevice) + assert comp.device.has_multiple_devices + + # Test with regular dict (not a ComponentDevice - different structure) + data = { + "type": generate_qualified_class_name(CustomComponentWithDevice), + "init_parameters": {"device": {"some": "dict"}, "name": "test"}, + } + comp = component_from_dict(CustomComponentWithDevice, data, "test_component") + assert comp.device == {"some": "dict"} + assert comp.name == "test" + + # Test with multiple ComponentDevice params + data = { + "type": generate_qualified_class_name(CustomComponentWithDevice), + "init_parameters": { + "device": {"type": "single", "device": "cpu"}, + "other_device": {"type": "single", "device": "cuda:0"}, + "name": "test", + }, + } + comp = component_from_dict(CustomComponentWithDevice, data, "test_component") + assert isinstance(comp.device, ComponentDevice) + assert isinstance(comp.other_device, ComponentDevice) + assert comp.device.to_torch_str() == "cpu" + assert comp.other_device.to_torch_str() == "cuda:0" + assert comp.name == "test" + + +def test_component_to_dict_and_from_dict_roundtrip_with_component_device(): + """Test that serialization and deserialization work together for ComponentDevice.""" + # Test roundtrip with single device + original_device = ComponentDevice.from_single(Device.cpu()) + comp = CustomComponentWithDevice(device=original_device) + + serialized = component_to_dict(comp, "test_component") + assert serialized["init_parameters"]["device"]["type"] == "single" + + deserialized_comp = component_from_dict(CustomComponentWithDevice, serialized, "test_component") + assert isinstance(deserialized_comp.device, ComponentDevice) + assert deserialized_comp.device.to_torch_str() == original_device.to_torch_str() + + # Test roundtrip with GPU device + original_device = ComponentDevice.from_single(Device.gpu(2)) + comp = CustomComponentWithDevice(device=original_device) + + serialized = component_to_dict(comp, "test_component") + deserialized_comp = component_from_dict(CustomComponentWithDevice, serialized, "test_component") + assert deserialized_comp.device.to_torch_str() == "cuda:2" + + # Test roundtrip with device map + device_map = DeviceMap({"layer1": Device.gpu(0), "layer2": Device.cpu()}) + original_device = ComponentDevice.from_multiple(device_map) + comp = CustomComponentWithDevice(device=original_device) + + serialized = component_to_dict(comp, "test_component") + assert serialized["init_parameters"]["device"]["type"] == "multiple" + + deserialized_comp = component_from_dict(CustomComponentWithDevice, serialized, "test_component") + assert isinstance(deserialized_comp.device, ComponentDevice) + assert deserialized_comp.device.has_multiple_devices + + # Test roundtrip with multiple ComponentDevice params + device1 = ComponentDevice.from_single(Device.cpu()) + device2 = ComponentDevice.from_single(Device.gpu(0)) + comp = CustomComponentWithDevice(device=device1, other_device=device2, name="test") + + serialized = component_to_dict(comp, "test_component") + assert serialized["init_parameters"]["device"]["type"] == "single" + assert serialized["init_parameters"]["other_device"]["type"] == "single" + assert serialized["init_parameters"]["name"] == "test" + + deserialized_comp = component_from_dict(CustomComponentWithDevice, serialized, "test_component") + assert isinstance(deserialized_comp.device, ComponentDevice) + assert isinstance(deserialized_comp.other_device, ComponentDevice) + assert deserialized_comp.device.to_torch_str() == "cpu" + assert deserialized_comp.other_device.to_torch_str() == "cuda:0" + assert deserialized_comp.name == "test" From c0b353442872d64a60186f6d693013c5e7297d6b Mon Sep 17 00:00:00 2001 From: Julian Risch Date: Tue, 20 Jan 2026 18:02:13 +0100 Subject: [PATCH 2/3] add release note --- .../default-serde-component-device-889de2e1f885240e.yaml | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 releasenotes/notes/default-serde-component-device-889de2e1f885240e.yaml diff --git a/releasenotes/notes/default-serde-component-device-889de2e1f885240e.yaml b/releasenotes/notes/default-serde-component-device-889de2e1f885240e.yaml new file mode 100644 index 0000000000..243eee80ce --- /dev/null +++ b/releasenotes/notes/default-serde-component-device-889de2e1f885240e.yaml @@ -0,0 +1,6 @@ +--- +features: + - | + component_from_dict() and component_to_dict() now work with custom components out of the box also if the component has a ComponentDevice as an attribute. + Users no longer need to explicitly define to_dict() and from_dict() methods in their custom components to call ComponentDevice.from_dict() or device.to_dict(). + component_from_dict() and component_to_dict() now handle this automatically. From 270bfbd0f4cbac03c24a78dd23d503b240ca3092 Mon Sep 17 00:00:00 2001 From: Julian Risch Date: Wed, 21 Jan 2026 14:54:09 +0100 Subject: [PATCH 3/3] make component device check more robust --- haystack/core/serialization.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/haystack/core/serialization.py b/haystack/core/serialization.py index 3ee6ff9715..b92b15352a 100644 --- a/haystack/core/serialization.py +++ b/haystack/core/serialization.py @@ -229,13 +229,21 @@ def _is_serialized_component_device(value: Any) -> bool: """ Check if a value is a serialized ComponentDevice dictionary. - A dictionary is considered a serialized ComponentDevice if it has a "type" key - with value "single" or "multiple". + A dictionary is considered a serialized ComponentDevice if: + - It has "type": "single" and a "device" key with a string value, or + - It has "type": "multiple" and a "device_map" key with a dict value + + This matches the structure produced by ComponentDevice.to_dict(). """ if not isinstance(value, dict): return False + type_value = value.get("type") - return type_value in ("single", "multiple") + if type_value == "single": + return "device" in value and isinstance(value["device"], str) + elif type_value == "multiple": + return "device_map" in value and isinstance(value["device_map"], dict) + return False def default_from_dict(cls: type[T], data: dict[str, Any]) -> T: