Coverage for haystack/utils/device.py: 92%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>

3# SPDX-License-Identifier: Apache-2.0

5import os

6from collections.abc import Iterator

7from dataclasses import dataclass, field

8from enum import Enum

9from typing import Any, Optional, Union

11from haystack.lazy_imports import LazyImport

13with LazyImport(

14 message="PyTorch must be installed to use torch.device or use GPU support in HuggingFace transformers. "

15 "Run 'pip install \"transformers[torch]\"'"

16) as torch_import:

17 import torch

20class DeviceType(Enum):

21 """

22 Represents device types supported by Haystack.

24 This also includes devices that are not directly used by models - for example, the disk device is exclusively used

25 in device maps for frameworks that support offloading model weights to disk.

26 """

28 CPU = "cpu"

29 GPU = "cuda"

30 DISK = "disk"

31 MPS = "mps"

32 XPU = "xpu"

34 def __str__(self) -> str:

35 return self.value

37 @staticmethod

38 def from_str(string: str) -> "DeviceType":

39 """

40 Create a device type from a string.

42 :param string:

43 The string to convert.

44 :returns:

45 The device type.

46 """

47 mapping = {e.value: e for e in DeviceType}

48 _type = mapping.get(string)

49 if _type is None:

50 raise ValueError(f"Unknown device type string '{string}'")

51 return _type

54@dataclass

55class Device:

56 """

57 A generic representation of a device.

59 :param type:

60 The device type.

61 :param id:

62 The optional device id.

63 """

65 type: DeviceType

66 id: int | None = field(default=None)

68 def __init__(self, type: DeviceType, id: int | None = None) -> None: # noqa:A002

69 """

70 Create a generic device.

72 :param type:

73 The device type.

74 :param id:

75 The device id.

76 """

77 if id is not None and id < 0:

78 raise ValueError(f"Device id must be >= 0, got {id}")

80 self.type = type

81 self.id = id

83 def __str__(self) -> str:

84 if self.id is None:

85 return str(self.type)

86 return f"{self.type}:{self.id}"

88 @staticmethod

89 def cpu() -> "Device":

90 """

91 Create a generic CPU device.

93 :returns:

94 The CPU device.

95 """

96 return Device(DeviceType.CPU)

98 @staticmethod

99 def gpu(id: int = 0) -> "Device": # noqa:A002

100 """

101 Create a generic GPU device.

102

103 :param id:

104 The GPU id.

105 :returns:

106 The GPU device.

107 """

108 return Device(DeviceType.GPU, id)

109

110 @staticmethod

111 def disk() -> "Device":

112 """

113 Create a generic disk device.

114

115 :returns:

116 The disk device.

117 """

118 return Device(DeviceType.DISK)

119

120 @staticmethod

121 def mps() -> "Device":

122 """

123 Create a generic Apple Metal Performance Shader device.

124

125 :returns:

126 The MPS device.

127 """

128 return Device(DeviceType.MPS)

129

130 @staticmethod

131 def xpu() -> "Device":

132 """

133 Create a generic Intel GPU Optimization device.

134

135 :returns:

136 The XPU device.

137 """

138 return Device(DeviceType.XPU)

139

140 @staticmethod

141 def from_str(string: str) -> "Device":

142 """

143 Create a generic device from a string.

144

145 :returns:

146 The device.

147

148 """

149 device_type_str, device_id = _split_device_string(string)

150 return Device(DeviceType.from_str(device_type_str), device_id)

151

152

153@dataclass

154class DeviceMap:

155 """

156 A generic mapping from strings to devices.

157

158 The semantics of the strings are dependent on target framework. Primarily used to deploy HuggingFace models to

159 multiple devices.

160

161 :param mapping:

162 Dictionary mapping strings to devices.

163 """

164

165 mapping: dict[str, Device] = field(default_factory=dict, hash=False)

166

167 def __getitem__(self, key: str) -> Device:

168 return self.mapping[key]

169

170 def __setitem__(self, key: str, value: Device) -> None:

171 self.mapping[key] = value

172

173 def __contains__(self, key: str) -> bool:

174 return key in self.mapping

175

176 def __len__(self) -> int:

177 return len(self.mapping)

178

179 def __iter__(self) -> Iterator[tuple[str, Device]]:

180 return iter(self.mapping.items())

181

182 def to_dict(self) -> dict[str, str]:

183 """

184 Serialize the mapping to a JSON-serializable dictionary.

185

186 :returns:

187 The serialized mapping.

188 """

189 return {key: str(device) for key, device in self.mapping.items()}

190

191 @property

192 def first_device(self) -> Device | None:

193 """

194 Return the first device in the mapping, if any.

195

196 :returns:

197 The first device.

198 """

199 if not self.mapping:

200 return None

201 return next(iter(self.mapping.values()))

202

203 @staticmethod

204 def from_dict(dict: dict[str, str]) -> "DeviceMap": # noqa:A002

205 """

206 Create a generic device map from a JSON-serialized dictionary.

207

208 :param dict:

209 The serialized mapping.

210 :returns:

211 The generic device map.

212 """

213 mapping = {}

214 for key, device_str in dict.items():

215 mapping[key] = Device.from_str(device_str)

216 return DeviceMap(mapping)

217

218 @staticmethod

219 def from_hf(hf_device_map: dict[str, Union[int, str, "torch.device"]]) -> "DeviceMap":

220 """

221 Create a generic device map from a HuggingFace device map.

222

223 :param hf_device_map:

224 The HuggingFace device map.

225 :returns:

226 The deserialized device map.

227 :raises TypeError: If a device value in the map is not an int, str, or torch.device.

228 """

229 mapping = {}

230 for key, device in hf_device_map.items():

231 if isinstance(device, int):

232 mapping[key] = Device(DeviceType.GPU, device)

233 elif isinstance(device, str):

234 device_type, device_id = _split_device_string(device)

235 mapping[key] = Device(DeviceType.from_str(device_type), device_id)

236 elif isinstance(device, torch.device):

237 device_type = device.type

238 device_id = device.index

239 mapping[key] = Device(DeviceType.from_str(device_type), device_id)

240 else:

241 raise TypeError(

242 f"Couldn't convert HuggingFace device map - unexpected device '{str(device)}' for '{key}'"

243 )

244 return DeviceMap(mapping)

245

246

247@dataclass(frozen=True)

248class ComponentDevice:

249 """

250 A representation of a device for a component.

251

252 This can be either a single device or a device map.

253 """

254

255 _single_device: Device | None = field(default=None)

256 _multiple_devices: DeviceMap | None = field(default=None)

257

258 @classmethod

259 def from_str(cls, device_str: str) -> "ComponentDevice":

260 """

261 Create a component device representation from a device string.

262

263 The device string can only represent a single device.

264

265 :param device_str:

266 The device string.

267 :returns:

268 The component device representation.

269 """

270 device = Device.from_str(device_str)

271 return cls.from_single(device)

272

273 @classmethod

274 def from_single(cls, device: Device) -> "ComponentDevice":

275 """

276 Create a component device representation from a single device.

277

278 Disks cannot be used as single devices.

279

280 :param device:

281 The device.

282 :returns:

283 The component device representation.

284 """

285 if device.type == DeviceType.DISK:

286 raise ValueError("The disk device can only be used as a part of device maps")

287

288 return cls(_single_device=device)

289

290 @classmethod

291 def from_multiple(cls, device_map: DeviceMap) -> "ComponentDevice":

292 """

293 Create a component device representation from a device map.

294

295 :param device_map:

296 The device map.

297 :returns:

298 The component device representation.

299 """

300 return cls(_multiple_devices=device_map)

301

302 def _validate(self) -> None:

303 """

304 Validate the component device representation.

305 """

306 if not (self._single_device is not None) ^ (self._multiple_devices is not None):

307 raise ValueError(

308 "The component device can neither be empty nor contain both a single device and a device map"

309 )

310

311 def to_torch(self) -> "torch.device":

312 """

313 Convert the component device representation to PyTorch format.

314

315 Device maps are not supported.

316

317 :returns:

318 The PyTorch device representation.

319 """

320 self._validate()

321

322 if self._single_device is None:

323 raise ValueError("Only single devices can be converted to PyTorch format")

324

325 torch_import.check()

326 assert self._single_device is not None

327 return torch.device(str(self._single_device))

328

329 def to_torch_str(self) -> str:

330 """

331 Convert the component device representation to PyTorch string format.

332

333 Device maps are not supported.

334

335 :returns:

336 The PyTorch device string representation.

337 """

338 self._validate()

339

340 if self._single_device is None:

341 raise ValueError("Only single devices can be converted to PyTorch format")

342

343 assert self._single_device is not None

344 return str(self._single_device)

345

346 def to_spacy(self) -> int:

347 """

348 Convert the component device representation to spaCy format.

349

350 Device maps are not supported.

351

352 :returns:

353 The spaCy device representation.

354 """

355 self._validate()

356

357 if self._single_device is None:

358 raise ValueError("Only single devices can be converted to spaCy format")

359

360 assert self._single_device is not None

361 if self._single_device.type == DeviceType.GPU:

362 assert self._single_device.id is not None

363 return self._single_device.id

364 return -1

365

366 def to_hf(self) -> int | str | dict[str, int | str]:

367 """

368 Convert the component device representation to HuggingFace format.

369

370 :returns:

371 The HuggingFace device representation.

372 """

373 self._validate()

374

375 def convert_device(device: Device, *, gpu_id_only: bool = False) -> int | str:

376 if gpu_id_only and device.type == DeviceType.GPU:

377 assert device.id is not None

378 return device.id

379 return str(device)

380

381 if self._single_device is not None:

382 return convert_device(self._single_device)

383

384 assert self._multiple_devices is not None

385 return {key: convert_device(device, gpu_id_only=True) for key, device in self._multiple_devices.mapping.items()}

386

387 def update_hf_kwargs(self, hf_kwargs: dict[str, Any], *, overwrite: bool) -> dict[str, Any]:

388 """

389 Convert the component device representation to HuggingFace format.

390

391 Add them as canonical keyword arguments to the keyword arguments dictionary.

392

393 :param hf_kwargs:

394 The HuggingFace keyword arguments dictionary.

395 :param overwrite:

396 Whether to overwrite existing device arguments.

397 :returns:

398 The HuggingFace keyword arguments dictionary.

399 """

400 self._validate()

401

402 if not overwrite and any(x in hf_kwargs for x in ("device", "device_map")):

403 return hf_kwargs

404

405 converted = self.to_hf()

406 key = "device_map" if self.has_multiple_devices else "device"

407 hf_kwargs[key] = converted

408 return hf_kwargs

409

410 @property

411 def has_multiple_devices(self) -> bool:

412 """

413 Whether this component device representation contains multiple devices.

414 """

415 self._validate()

416

417 return self._multiple_devices is not None

418

419 @property

420 def first_device(self) -> Optional["ComponentDevice"]:

421 """

422 Return either the single device or the first usable device in the device map, if any.

423

424 Disk devices are skipped because they can only be used as part of a device map and not as a

425 single device. If the device map is empty or contains only disk devices, a ``ValueError`` is

426 raised so callers that do ``first_device.to_torch()`` still get a clear error instead of

427 ``AttributeError: 'NoneType' object has no attribute 'to_torch'``.

428

429 :returns:

430 The first usable device.

431 """

432 self._validate()

433

434 if self._single_device is not None:

435 return self.from_single(self._single_device)

436

437 assert self._multiple_devices is not None

438 first_usable_device = next(

439 (device for device in self._multiple_devices.mapping.values() if device.type != DeviceType.DISK), None

440 )

441 if first_usable_device is None:

442 raise ValueError("No usable device found in the device map; disk devices cannot be used as a single device")

443 return self.from_single(first_usable_device)

444

445 @staticmethod

446 def resolve_device(device: Optional["ComponentDevice"] = None) -> "ComponentDevice":

447 """

448 Select a device for a component. If a device is specified, it's used. Otherwise, the default device is used.

449

450 :param device:

451 The provided device, if any.

452 :returns:

453 The resolved device.

454 """

455 if not isinstance(device, ComponentDevice) and device is not None:

456 raise ValueError(

457 f"Invalid component device type '{type(device).__name__}'. Must either be None or ComponentDevice."

458 )

459

460 if device is None:

461 device = ComponentDevice.from_single(_get_default_device())

462

463 return device

464

465 def to_dict(self) -> dict[str, Any]:

466 """

467 Convert the component device representation to a JSON-serializable dictionary.

468

469 :returns:

470 The dictionary representation.

471 """

472 if self._single_device is not None:

473 return {"type": "single", "device": str(self._single_device)}

474 if self._multiple_devices is not None:

475 return {"type": "multiple", "device_map": self._multiple_devices.to_dict()}

476 # Unreachable

477 raise AssertionError()

478

479 @classmethod

480 def from_dict(cls, dict: dict[str, Any]) -> "ComponentDevice": # noqa:A002

481 """

482 Create a component device representation from a JSON-serialized dictionary.

483

484 :param dict:

485 The serialized representation.

486 :returns:

487 The deserialized component device.

488 """

489 if dict["type"] == "single":

490 return cls.from_str(dict["device"])

491 if dict["type"] == "multiple":

492 return cls.from_multiple(DeviceMap.from_dict(dict["device_map"]))

493 raise ValueError(f"Unknown component device type '{dict['type']}' in serialized data")

494

495

496def _get_default_device() -> Device:

497 """

498 Return the default device for Haystack.

499

500 Precedence:

501 GPU > XPU > MPS > CPU. If PyTorch is not installed, only CPU is available.

502

503 :returns:

504 The default device.

505 """

506 try:

507 torch_import.check()

508

509 has_mps = (

510 hasattr(torch.backends, "mps")

511 and torch.backends.mps.is_available()

512 and os.getenv("HAYSTACK_MPS_ENABLED", "true") != "false"

513 )

514 has_cuda = torch.cuda.is_available()

515 has_xpu = (

516 hasattr(torch, "xpu")

517 and hasattr(torch.xpu, "is_available")

518 and torch.xpu.is_available()

519 and os.getenv("HAYSTACK_XPU_ENABLED", "true") != "false"

520 )

521 except ImportError:

522 has_mps = False

523 has_cuda = False

524 has_xpu = False

525

526 if has_cuda:

527 return Device.gpu()

528 if has_xpu:

529 return Device.xpu()

530 if has_mps:

531 return Device.mps()

532 return Device.cpu()

533

534

535def _split_device_string(string: str) -> tuple[str, int | None]:

536 """

537 Split a device string into device type and device id.

538

539 :param string:

540 The device string to split.

541 :returns:

542 The device type and device id, if any.

543 """

544 if ":" in string:

545 device_type, device_id_str = string.split(":")

546 try:

547 device_id = int(device_id_str)

548 except ValueError as e:

549 raise ValueError(f"Device id must be an integer, got {device_id_str}") from e

550 else:

551 device_type = string

552 device_id = None

553 return device_type, device_id