Coverage for src / bioimageio / core / _settings.py: 88%

1import os 

2import platform 

3from typing import Literal, Optional 

4 

5from loguru import logger 

6from pydantic import Field, field_validator 

7from typing_extensions import Annotated 

8 

9from bioimageio.spec._internal._settings import Settings as SpecSettings 

10 

11 

12class Settings(SpecSettings): 

13 """environment variables for bioimageio.spec and bioimageio.core""" 

14 

15 keras_backend: Annotated[ 

16 Literal["torch", "tensorflow", "jax"], Field(alias="KERAS_BACKEND") 

17 ] = "torch" 

18 

19 pytorch_enable_mps_fallback: Annotated[ 

20 Optional[bool], Field(alias="PYTORCH_ENABLE_MPS_FALLBACK") 

21 ] = None 

22 

23 @field_validator("pytorch_enable_mps_fallback", mode="after") 

24 @classmethod 

25 def _set_default_mps_fallback(cls, value: Optional[bool]): 

26 # pytorch versions up to the 2.6 don't support all operations (esp 3d) on MPS 

27 # this env variable allows falling back to CPU for those networks instead of failing 

28 # see for current status https://github.com/pytorch/pytorch/issues/141287 

29 if ( 

30 value is None 

31 and platform.system().lower() == "darwin" 

32 and platform.machine().lower() == "arm64" 

33 ): 

34 logger.info("Set environment variable 'PYTORCH_ENABLE_MPS_FALLBACK=1'.") 

35 os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" 

36 return True 

37 

38 return value 

39 

40 collection_index_url: str = "https://bioimage-io.github.io/collection/index.json" 

41 """URL to the bioimageio collection index""" 

42 

43 collection_config_url: str = ( 

44 "https://bioimage-io.github.io/collection/bioimageio_collection_config.json" 

45 ) 

46 """URL to the bioimageio collection config""" 

47 

48 

49settings = Settings() 

50"""parsed environment variables for bioimageio.spec and bioimageio.core"""