Coverage for src/arraybridge/decorators.py: 95%

1""" 

2Memory type declaration decorators. 

3 

4This module provides decorators for explicitly declaring the memory interface 

5of pure functions and supporting memory-type-aware dispatching and orchestration. 

6 

7These decorators annotate functions with input_memory_type and output_memory_type 

8attributes and provide automatic thread-local CUDA stream management for GPU 

9frameworks to enable true parallelization across multiple threads. 

10 

11Framework-specific runtime capabilities are delegated to ``MemoryType`` members. 

12""" 

13 

14import functools 

15import inspect 

16import logging 

17import threading 

18from abc import ABC, abstractmethod 

19from collections.abc import Callable 

20from dataclasses import dataclass 

21from enum import Enum 

22from typing import Any, ClassVar, TypeVar, cast 

23 

24import numpy as np 

25from metaclass_registry import AutoRegisterMeta, RegistryFamily, RegistryKeyAttribute 

26from python_introspect import RuntimeParameterDeclarationABC 

27 

28from arraybridge.array_payload import ArrayPayload 

29from arraybridge.oom_recovery import _execute_with_oom_recovery 

30from arraybridge.slice_processing import process_slices 

31from arraybridge.types import MemoryContractAttribute, MemoryType 

32 

33logger = logging.getLogger(__name__) 

34 

35F = TypeVar("F", bound=Callable[..., Any]) 

36 

37 

38class DtypeConversion(Enum): 

39 """Data type conversion modes for all memory type functions.""" 

40 

41 PRESERVE_INPUT = "preserve" # Keep input dtype (default) 

42 NATIVE_OUTPUT = "native" # Use framework's native output 

43 UINT8 = "uint8" # Force uint8 (0-255 range) 

44 UINT16 = "uint16" # Force uint16 (microscopy standard) 

45 INT16 = "int16" # Force int16 (signed microscopy data) 

46 INT32 = "int32" # Force int32 (large integer values) 

47 FLOAT32 = "float32" # Force float32 (GPU performance) 

48 FLOAT64 = "float64" # Force float64 (maximum precision) 

49 

50 @property 

51 def numpy_dtype(self): 

52 """Get the corresponding numpy dtype.""" 

53 dtype_map = { 

54 self.UINT8: np.uint8, 

55 self.UINT16: np.uint16, 

56 self.INT16: np.int16, 

57 self.INT32: np.int32, 

58 self.FLOAT32: np.float32, 

59 self.FLOAT64: np.float64, 

60 } 

61 return dtype_map.get(self, None) 

62 

63 

64class DtypeConversionConfig(RuntimeParameterDeclarationABC): 

65 """Nominal dtype conversion config surface consumed by decorators.""" 

66 

67 @property 

68 @abstractmethod 

69 def default_dtype_conversion(self) -> DtypeConversion: 

70 """Return the dtype conversion mode for decorated function output.""" 

71 

72 @classmethod 

73 def require_parameter_name(cls) -> str: 

74 return "dtype_config" 

75 

76 @classmethod 

77 def default_value(cls): 

78 return PRESERVE_INPUT_DTYPE_CONFIG 

79 

80 @classmethod 

81 def annotation_type(cls): 

82 return DtypeConversionConfig 

83 

84 @classmethod 

85 def parameter(cls) -> inspect.Parameter: 

86 return inspect.Parameter( 

87 cls.require_parameter_name(), 

88 inspect.Parameter.KEYWORD_ONLY, 

89 default=cls.default_value(), 

90 annotation=cls.annotation_type(), 

91 ) 

92 

93 

94class SliceBySliceRuntimeParameter(RuntimeParameterDeclarationABC): 

95 """Nominal slice-by-slice execution parameter consumed by decorators.""" 

96 

97 preserve_for_execution = True 

98 is_semantic_control = True 

99 

100 @classmethod 

101 def require_parameter_name(cls) -> str: 

102 return "slice_by_slice" 

103 

104 @classmethod 

105 def default_value(cls) -> bool: 

106 return False 

107 

108 @classmethod 

109 def annotation_type(cls) -> type[bool]: 

110 return bool 

111 

112 @classmethod 

113 def parameter(cls, *, default_value: bool | None = None) -> inspect.Parameter: 

114 return inspect.Parameter( 

115 cls.require_parameter_name(), 

116 inspect.Parameter.KEYWORD_ONLY, 

117 default=(cls.default_value() if default_value is None else default_value), 

118 annotation=cls.annotation_type(), 

119 ) 

120 

121 

122@dataclass(frozen=True, slots=True) 

123class PreserveInputDtypeConfig(DtypeConversionConfig): 

124 """Direct-call dtype config for wrappers executed outside a pipeline runtime.""" 

125 

126 default_dtype_conversion: DtypeConversion = DtypeConversion.PRESERVE_INPUT 

127 

128 

129PRESERVE_INPUT_DTYPE_CONFIG = PreserveInputDtypeConfig() 

130 

131 

132class EnumValueRegistryKeyMixin: 

133 """Derive AutoRegisterMeta strategy labels from enum-valued class members.""" 

134 

135 strategy_label: ClassVar[str | None] = None 

136 

137 def __init_subclass__(cls, **kwargs: Any) -> None: 

138 super().__init_subclass__(**kwargs) 

139 member = cls.registry_enum_member() 

140 if isinstance(member, Enum) and cls.__dict__.get("strategy_label") is None: 

141 cls.strategy_label = member.value 

142 

143 @classmethod 

144 @abstractmethod 

145 def registry_enum_member(cls) -> Enum | None: 

146 """Return the enum member that should key this concrete strategy.""" 

147 

148 

149@dataclass(frozen=True, slots=True) 

150class DtypeConversionRequest: 

151 """Runtime data needed to convert one decorated function output.""" 

152 

153 array: Any 

154 original_dtype_name: str | None 

155 array_dtype_name: str | None 

156 scale_func: Callable[[Any, Any], Any] 

157 

158 

159class DtypeConversionRunner( 

160 EnumValueRegistryKeyMixin, 

161 ABC, 

162 metaclass=AutoRegisterMeta, 

163): 

164 """Registered dtype conversion behavior selected by DtypeConversion.""" 

165 

166 __registry_family__ = RegistryFamily(RegistryKeyAttribute.STRATEGY_LABEL) 

167 

168 dtype_conversion: ClassVar[DtypeConversion | None] = None 

169 

170 @classmethod 

171 def registry_enum_member(cls) -> Enum | None: 

172 return cls.dtype_conversion 

173 

174 @classmethod 

175 def for_dtype_conversion( 

176 cls, 

177 dtype_conversion: DtypeConversion, 

178 ) -> "DtypeConversionRunner": 

179 return cast(DtypeConversionRunner, cls.__registry__[dtype_conversion.value]()) 

180 

181 @abstractmethod 

182 def apply(self, request: DtypeConversionRequest) -> Any: 

183 """Return output converted according to the configured dtype policy.""" 

184 

185 

186class PreserveInputDtypeConversionRunner(DtypeConversionRunner): 

187 """Scale output back to the input dtype when the wrapped function changed it.""" 

188 

189 dtype_conversion = DtypeConversion.PRESERVE_INPUT 

190 

191 def apply(self, request: DtypeConversionRequest) -> Any: 

192 if ( 

193 request.original_dtype_name is not None 

194 and request.array_dtype_name != request.original_dtype_name 

195 ): 

196 return request.scale_func(request.array, request.original_dtype_name) 

197 return request.array 

198 

199 

200class NativeOutputDtypeConversionRunner(DtypeConversionRunner): 

201 """Keep the wrapped framework function's native output dtype.""" 

202 

203 dtype_conversion = DtypeConversion.NATIVE_OUTPUT 

204 

205 def apply(self, request: DtypeConversionRequest) -> Any: 

206 return request.array 

207 

208 

209class FixedDtypeConversionRunner(DtypeConversionRunner): 

210 """Scale output to the dtype declared by a fixed DtypeConversion member.""" 

211 

212 def apply(self, request: DtypeConversionRequest) -> Any: 

213 if self.dtype_conversion is None: 

214 raise TypeError("FixedDtypeConversionRunner requires dtype_conversion.") 

215 target_dtype = self.dtype_conversion.numpy_dtype 

216 if target_dtype is None: 

217 return request.array 

218 return request.scale_func(request.array, target_dtype) 

219 

220 

221class Uint8DtypeConversionRunner(FixedDtypeConversionRunner): 

222 dtype_conversion = DtypeConversion.UINT8 

223 

224 

225class Uint16DtypeConversionRunner(FixedDtypeConversionRunner): 

226 dtype_conversion = DtypeConversion.UINT16 

227 

228 

229class Int16DtypeConversionRunner(FixedDtypeConversionRunner): 

230 dtype_conversion = DtypeConversion.INT16 

231 

232 

233class Int32DtypeConversionRunner(FixedDtypeConversionRunner): 

234 dtype_conversion = DtypeConversion.INT32 

235 

236 

237class Float32DtypeConversionRunner(FixedDtypeConversionRunner): 

238 dtype_conversion = DtypeConversion.FLOAT32 

239 

240 

241class Float64DtypeConversionRunner(FixedDtypeConversionRunner): 

242 dtype_conversion = DtypeConversion.FLOAT64 

243 

244 

245class KeywordOnlySignatureExtension: 

246 """Insert decorator-owned keyword-only parameters in valid signature order.""" 

247 

248 def __init__(self, signature: inspect.Signature): 

249 self.signature = signature 

250 

251 def with_parameter(self, parameter: inspect.Parameter) -> inspect.Signature: 

252 parameters = list(self.signature.parameters.values()) 

253 if parameter.name in self.signature.parameters: 

254 return self.signature 

255 insertion_index = self._insertion_index(parameters) 

256 parameters.insert(insertion_index, parameter) 

257 return self.signature.replace(parameters=parameters) 

258 

259 @staticmethod 

260 def _insertion_index(parameters: list[inspect.Parameter]) -> int: 

261 for index, candidate in enumerate(parameters): 

262 if candidate.kind is inspect.Parameter.VAR_KEYWORD: 

263 return index 

264 return len(parameters) 

265 

266 

267# Thread-local storage for GPU streams and contexts 

268_thread_gpu_contexts = threading.local() 

269 

270 

271class ThreadGPUContext: 

272 """Thread-local streams keyed by framework-local device identity.""" 

273 

274 def __init__(self): 

275 self._streams: dict[tuple[MemoryType, int], Any] = {} 

276 

277 def stream_for( 

278 self, 

279 memory_type: MemoryType, 

280 module: Any, 

281 ) -> tuple[int | None, Any | None]: 

282 """Return the current device and its stable thread-local stream.""" 

283 

284 device_id = memory_type.current_device_id(module) 

285 if device_id is None: 

286 return None, None 

287 memory_type.require_device(device_id, module) 

288 key = (memory_type, device_id) 

289 if key not in self._streams: 

290 with memory_type.device_scope(device_id, module): 

291 stream = memory_type.create_stream(module) 

292 if stream is None: 

293 return device_id, None 

294 self._streams[key] = stream 

295 logger.debug( 

296 "Created %s stream for device %d in thread %s", 

297 memory_type.display_name, 

298 device_id, 

299 threading.current_thread().name, 

300 ) 

301 return device_id, self._streams[key] 

302 

303 

304def _get_thread_gpu_context(): 

305 """Get or create thread-local GPU context.""" 

306 if not hasattr(_thread_gpu_contexts, "context"): 

307 _thread_gpu_contexts.context = ThreadGPUContext() 

308 return _thread_gpu_contexts.context 

309 

310 

311def memory_types( 

312 input_type: str | MemoryType, 

313 output_type: str | MemoryType, 

314 contract: Any | None = None, 

315) -> Callable[[F], F]: 

316 """ 

317 Base decorator for declaring memory types of a function. 

318 

319 This is the foundation decorator that all memory-type-specific decorators build upon. 

320 """ 

321 

322 input_member = input_type if isinstance(input_type, MemoryType) else MemoryType(input_type) 

323 output_member = output_type if isinstance(output_type, MemoryType) else MemoryType(output_type) 

324 input_memory_type = input_member.value 

325 output_memory_type = output_member.value 

326 

327 def decorator(func: F) -> F: 

328 @functools.wraps(func) 

329 def wrapper(*args, **kwargs): 

330 result = func(*args, **kwargs) 

331 

332 # Apply output validation only when a callable contract was provided. 

333 # Non-callable contracts are declarative metadata consumed by runtimes. 

334 if callable(contract) and not contract(result): 

335 raise ValueError(f"Function {func.__name__} violated its output contract") 

336 

337 return result 

338 

339 # Attach memory type metadata 

340 MemoryContractAttribute.INPUT.write(wrapper, input_memory_type) 

341 MemoryContractAttribute.OUTPUT.write(wrapper, output_memory_type) 

342 if contract is not None and not callable(contract): 

343 setattr(wrapper, "__processing_contract__", contract) 

344 

345 return cast(F, wrapper) 

346 

347 return decorator 

348 

349 

350def wrap_dtype_preserving_callable( 

351 func, 

352 mem_type: MemoryType, 

353 *, 

354 slice_by_slice_default: bool = False, 

355): 

356 """ 

357 Return a callable with ArrayBridge dtype and slice controls. 

358 

359 Host registries can use this public boundary without depending on the 

360 complete framework decorator or ArrayBridge internals. 

361 """ 

362 func_name = func.__name__ 

363 input_memory_type = MemoryType(MemoryContractAttribute.INPUT.read(func, mem_type.value)) 

364 output_memory_type = MemoryType(MemoryContractAttribute.OUTPUT.read(func, mem_type.value)) 

365 scale_func = output_memory_type.scale_dtype 

366 

367 @functools.wraps(func) 

368 def dtype_wrapper(image, *args, **kwargs): 

369 # Pipeline runtimes may inject dtype_config; direct calls use the same 

370 # preserve-input default explicitly. 

371 slice_by_slice = kwargs.pop( 

372 SliceBySliceRuntimeParameter.require_parameter_name(), 

373 slice_by_slice_default, 

374 ) 

375 dtype_config: DtypeConversionConfig = kwargs.pop( 

376 DtypeConversionConfig.require_parameter_name(), 

377 DtypeConversionConfig.default_value(), 

378 ) 

379 dtype_conversion = dtype_config.default_dtype_conversion 

380 

381 # Store original dtype 

382 original_dtype = getattr(image, "dtype", None) 

383 original_dtype_name = ( 

384 None 

385 if original_dtype is None 

386 else input_memory_type.canonical_dtype_name(original_dtype) 

387 ) 

388 

389 # Handle slice_by_slice processing for 3D arrays 

390 if slice_by_slice and hasattr(image, "ndim") and image.ndim == 3: 

391 result = process_slices(image, func, args, kwargs) 

392 else: 

393 # Call the original function normally 

394 result = func(image, *args, **kwargs) 

395 

396 def _apply_dtype_conversion(array): 

397 if isinstance(array, ArrayPayload): 

398 return array.map_array_payload(_apply_dtype_conversion) 

399 if not hasattr(array, "dtype"): 

400 return array 

401 return DtypeConversionRunner.for_dtype_conversion(dtype_conversion).apply( 

402 DtypeConversionRequest( 

403 array=array, 

404 original_dtype_name=original_dtype_name, 

405 array_dtype_name=output_memory_type.canonical_dtype_name(array.dtype), 

406 scale_func=scale_func, 

407 ) 

408 ) 

409 

410 # Apply dtype conversion to the main output. Conversion errors are 

411 # contract violations and must remain visible to the caller. 

412 if isinstance(result, tuple): 

413 if not result: 

414 return result 

415 converted_main = _apply_dtype_conversion(result[0]) 

416 return (converted_main, *result[1:]) 

417 return _apply_dtype_conversion(result) 

418 

419 # Update function signature to include new parameters 

420 try: 

421 dtype_signature = KeywordOnlySignatureExtension(inspect.signature(func)).with_parameter( 

422 SliceBySliceRuntimeParameter.parameter( 

423 default_value=slice_by_slice_default, 

424 ) 

425 ) 

426 dtype_signature = KeywordOnlySignatureExtension(dtype_signature).with_parameter( 

427 DtypeConversionConfig.parameter() 

428 ) 

429 setattr(dtype_wrapper, "__signature__", dtype_signature) 

430 

431 # Update docstring 

432 if dtype_wrapper.__doc__: 

433 dtype_wrapper.__doc__ += "\n\n Additional Parameters\n ---------------------\n" 

434 dtype_wrapper.__doc__ += ( 

435 " slice_by_slice : bool, optional\n" 

436 f" Added by the {mem_type.value} memory decorator. " 

437 "Process 3D arrays slice-by-slice.\n" 

438 ) 

439 dtype_wrapper.__doc__ += ( 

440 f" Defaults to {slice_by_slice_default}. " 

441 "Prevents cross-slice contamination when enabled.\n" 

442 ) 

443 

444 except Exception as e: 

445 logger.warning(f"Could not update signature for {func_name}: {e}") 

446 

447 return dtype_wrapper 

448 

449 

450def _create_gpu_wrapper(func, mem_type: MemoryType, oom_recovery: bool): 

451 """ 

452 Auto-generate GPU stream/device wrapper for any GPU memory type. 

453 

454 This function creates the GPU-specific wrapper with stream management and OOM recovery. 

455 """ 

456 

457 @functools.wraps(func) 

458 def gpu_wrapper(*args, **kwargs): 

459 framework = mem_type.import_if_installed() 

460 

461 # Check if GPU is available for this framework 

462 if framework is not None and mem_type.available_device_ids(framework): 

463 # Get thread-local context 

464 ctx = _get_thread_gpu_context() 

465 

466 device_id, stream = ctx.stream_for(mem_type, framework) 

467 

468 # Define execution function that captures args/kwargs 

469 def execute_with_stream(): 

470 with mem_type.stream_scope(stream, framework): 

471 return func(*args, **kwargs) 

472 

473 # Execute with OOM recovery if enabled 

474 if oom_recovery: 

475 return _execute_with_oom_recovery( 

476 execute_with_stream, 

477 mem_type.value, 

478 device_id=device_id, 

479 ) 

480 return execute_with_stream() 

481 

482 # CPU fallback or framework not available 

483 return func(*args, **kwargs) 

484 

485 # Preserve memory type attributes 

486 MemoryContractAttribute.INPUT.write( 

487 gpu_wrapper, 

488 MemoryContractAttribute.INPUT.read(func), 

489 ) 

490 MemoryContractAttribute.OUTPUT.write( 

491 gpu_wrapper, 

492 MemoryContractAttribute.OUTPUT.read(func), 

493 ) 

494 

495 return gpu_wrapper 

496 

497 

498def _create_memory_decorator(mem_type: MemoryType): 

499 """ 

500 Factory function that creates a decorator for a specific memory type. 

501 

502 This single factory replaces 6 nearly-identical decorator functions. 

503 """ 

504 

505 def decorator( 

506 func=None, 

507 *, 

508 input_type=mem_type.value, 

509 output_type=mem_type.value, 

510 oom_recovery=True, 

511 contract=None, 

512 slice_by_slice_default=False, 

513 ): 

514 """ 

515 Decorator for {mem_type} memory type functions. 

516 

517 Args: 

518 func: Function to decorate (when used as @decorator) 

519 input_type: Expected input memory type (default: {mem_type}) 

520 output_type: Expected output memory type (default: {mem_type}) 

521 oom_recovery: Enable automatic OOM recovery (default: True) 

522 contract: Optional validation function for outputs 

523 slice_by_slice_default: Default for the decorator-owned slice control 

524 

525 Returns: 

526 Decorated function with memory type metadata and dtype preservation 

527 """ 

528 

529 def inner_decorator(func): 

530 # Apply base memory_types decorator 

531 memory_decorator = memory_types( 

532 input_type=input_type, output_type=output_type, contract=contract 

533 ) 

534 func = memory_decorator(func) 

535 

536 # Apply dtype preservation wrapper 

537 func = wrap_dtype_preserving_callable( 

538 func, 

539 mem_type, 

540 slice_by_slice_default=slice_by_slice_default, 

541 ) 

542 

543 # Apply GPU wrapper if this is a GPU memory type 

544 if mem_type.is_gpu: 

545 func = _create_gpu_wrapper(func, mem_type, oom_recovery) 

546 

547 MemoryContractAttribute.EXECUTION.write(func, mem_type.value) 

548 

549 return func 

550 

551 # Handle both @decorator and @decorator() forms 

552 if func is None: 

553 return inner_decorator 

554 return inner_decorator(func) 

555 

556 # Set proper function name and docstring 

557 decorator.__name__ = mem_type.value 

558 decorator.__doc__ = (decorator.__doc__ or "").format(mem_type=mem_type.display_name) 

559 

560 return decorator 

561 

562 

563# Auto-generate all 6 memory type decorators 

564for mem_type in MemoryType: 

565 decorator_func = _create_memory_decorator(mem_type) 

566 globals()[mem_type.value] = decorator_func 

567 

568 

569# Export the fixed decorator infrastructure plus declaration-derived helpers. 

570__all__ = [ 

571 "memory_types", 

572 "DtypeConversion", 

573 "DtypeConversionConfig", 

574 "PreserveInputDtypeConfig", 

575 "PRESERVE_INPUT_DTYPE_CONFIG", 

576 "SliceBySliceRuntimeParameter", 

577 "wrap_dtype_preserving_callable", 

578] + [memory_type.value for memory_type in MemoryType]