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

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5import os 

6from collections.abc import Iterator 

7from dataclasses import dataclass, field 

8from enum import Enum 

9from typing import Any, Optional, Union 

10 

11from haystack.lazy_imports import LazyImport 

12 

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 

18 

19 

20class DeviceType(Enum): 

21 """ 

22 Represents device types supported by Haystack. 

23 

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 """ 

27 

28 CPU = "cpu" 

29 GPU = "cuda" 

30 DISK = "disk" 

31 MPS = "mps" 

32 XPU = "xpu" 

33 

34 def __str__(self) -> str: 

35 return self.value 

36 

37 @staticmethod 

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

39 """ 

40 Create a device type from a string. 

41 

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 

52 

53 

54@dataclass 

55class Device: 

56 """ 

57 A generic representation of a device. 

58 

59 :param type: 

60 The device type. 

61 :param id: 

62 The optional device id. 

63 """ 

64 

65 type: DeviceType 

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

67 

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

69 """ 

70 Create a generic device. 

71 

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}") 

79 

80 self.type = type 

81 self.id = id 

82 

83 def __str__(self) -> str: 

84 if self.id is None: 

85 return str(self.type) 

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

87 

88 @staticmethod 

89 def cpu() -> "Device": 

90 """ 

91 Create a generic CPU device. 

92 

93 :returns: 

94 The CPU device. 

95 """ 

96 return Device(DeviceType.CPU) 

97 

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 device in the device map, if any. 

423 

424 :returns: 

425 The first device. 

426 """ 

427 self._validate() 

428 

429 if self._single_device is not None: 

430 return self.from_single(self._single_device) 

431 

432 assert self._multiple_devices is not None 

433 assert self._multiple_devices.first_device is not None 

434 return self.from_single(self._multiple_devices.first_device) 

435 

436 @staticmethod 

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

438 """ 

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

440 

441 :param device: 

442 The provided device, if any. 

443 :returns: 

444 The resolved device. 

445 """ 

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

447 raise ValueError( 

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

449 ) 

450 

451 if device is None: 

452 device = ComponentDevice.from_single(_get_default_device()) 

453 

454 return device 

455 

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

457 """ 

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

459 

460 :returns: 

461 The dictionary representation. 

462 """ 

463 if self._single_device is not None: 

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

465 if self._multiple_devices is not None: 

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

467 # Unreachable 

468 raise AssertionError() 

469 

470 @classmethod 

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

472 """ 

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

474 

475 :param dict: 

476 The serialized representation. 

477 :returns: 

478 The deserialized component device. 

479 """ 

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

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

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

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

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

485 

486 

487def _get_default_device() -> Device: 

488 """ 

489 Return the default device for Haystack. 

490 

491 Precedence: 

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

493 

494 :returns: 

495 The default device. 

496 """ 

497 try: 

498 torch_import.check() 

499 

500 has_mps = ( 

501 hasattr(torch.backends, "mps") 

502 and torch.backends.mps.is_available() 

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

504 ) 

505 has_cuda = torch.cuda.is_available() 

506 has_xpu = ( 

507 hasattr(torch, "xpu") 

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

509 and torch.xpu.is_available() 

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

511 ) 

512 except ImportError: 

513 has_mps = False 

514 has_cuda = False 

515 has_xpu = False 

516 

517 if has_cuda: 

518 return Device.gpu() 

519 if has_xpu: 

520 return Device.xpu() 

521 if has_mps: 

522 return Device.mps() 

523 return Device.cpu() 

524 

525 

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

527 """ 

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

529 

530 :param string: 

531 The device string to split. 

532 :returns: 

533 The device type and device id, if any. 

534 """ 

535 if ":" in string: 

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

537 try: 

538 device_id = int(device_id_str) 

539 except ValueError as e: 

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

541 else: 

542 device_type = string 

543 device_id = None 

544 return device_type, device_id