Coverage for src/bioimageio/spec/_get_conda_env.py: 56%

1from __future__ import annotations 

2 

3from typing import Literal, Union 

4 

5from typing_extensions import assert_never 

6 

7from ._internal.gh_utils import set_github_warning 

8from ._internal.io import FileDescr, get_reader 

9from ._internal.io_utils import read_yaml 

10from .conda_env import BioimageioCondaEnv, PipDeps 

11from .model import v0_4, v0_5 

12from .model.v0_5 import Version 

13 

14SupportedWeightsEntry = Union[ 

15 v0_4.KerasHdf5WeightsDescr, 

16 v0_4.OnnxWeightsDescr, 

17 v0_4.PytorchStateDictWeightsDescr, 

18 v0_4.TensorflowSavedModelBundleWeightsDescr, 

19 v0_4.TorchscriptWeightsDescr, 

20 v0_5.KerasHdf5WeightsDescr, 

21 v0_5.KerasV3WeightsDescr, 

22 v0_5.OnnxWeightsDescr, 

23 v0_5.PytorchStateDictWeightsDescr, 

24 v0_5.TensorflowSavedModelBundleWeightsDescr, 

25 v0_5.TorchscriptWeightsDescr, 

26] 

27 

28 

29def get_conda_env( 

30 *, 

31 entry: SupportedWeightsEntry, 

32 env_name: Literal["DROP"] | str | None = None, 

33) -> BioimageioCondaEnv: 

34 """get the recommended Conda environment for a given weights entry description""" 

35 if isinstance(entry, (v0_4.OnnxWeightsDescr, v0_5.OnnxWeightsDescr)): 

36 conda_env = _get_default_onnx_env(opset_version=entry.opset_version) 

37 elif isinstance( 

38 entry, 

39 ( 

40 v0_4.PytorchStateDictWeightsDescr, 

41 v0_5.PytorchStateDictWeightsDescr, 

42 v0_4.TorchscriptWeightsDescr, 

43 v0_5.TorchscriptWeightsDescr, 

44 ), 

45 ): 

46 if ( 

47 isinstance(entry, v0_5.TorchscriptWeightsDescr) 

48 or entry.dependencies is None 

49 ): 

50 conda_env = _get_default_pytorch_env(pytorch_version=entry.pytorch_version) 

51 else: 

52 conda_env = _get_env_from_deps(entry.dependencies) 

53 

54 elif isinstance( 

55 entry, 

56 ( 

57 v0_4.TensorflowSavedModelBundleWeightsDescr, 

58 v0_5.TensorflowSavedModelBundleWeightsDescr, 

59 ), 

60 ): 

61 if entry.dependencies is None: 

62 conda_env = _get_default_tf_env(tensorflow_version=entry.tensorflow_version) 

63 else: 

64 conda_env = _get_env_from_deps(entry.dependencies) 

65 elif isinstance( 

66 entry, 

67 (v0_4.KerasHdf5WeightsDescr, v0_5.KerasHdf5WeightsDescr), 

68 ): 

69 conda_env = _get_default_tf_env(tensorflow_version=entry.tensorflow_version) 

70 elif isinstance(entry, v0_5.KerasV3WeightsDescr): 

71 conda_env = _get_default_keras3_env(entry.backend) 

72 else: 

73 assert_never(entry) 

74 

75 if env_name == "DROP": 

76 conda_env.name = None 

77 elif env_name is not None: 

78 conda_env.name = env_name 

79 

80 return conda_env 

81 

82 

83def _get_default_keras3_env( 

84 backend: tuple[Literal["tensorflow", "jax", "torch"], Version], 

85) -> BioimageioCondaEnv: 

86 if backend[0] == "tensorflow": 

87 env = _get_default_tf_env(tensorflow_version=backend[1]) 

88 elif backend[0] == "torch": 

89 env = _get_default_pytorch_env(pytorch_version=backend[1]) 

90 elif backend[0] == "jax": 

91 env = BioimageioCondaEnv( 

92 dependencies=[ 

93 f"jax=={backend[1]}", 

94 ] 

95 ) 

96 else: 

97 assert_never(backend[0]) 

98 

99 env.dependencies.append("keras >=3.0, <4") 

100 return env 

101 

102 

103def _get_default_pytorch_env( 

104 *, 

105 pytorch_version: Version | None = None, 

106) -> BioimageioCondaEnv: 

107 if pytorch_version is None: 

108 pytorch_version = Version("1.10.1") 

109 

110 channels = ["conda-forge", "nodefaults"] 

111 

112 # dependencies to install pytorch according to 

113 # https://pytorch.org/get-started/previous-versions/ 

114 v = pytorch_version.base_version 

115 if v.count(".") == 0: 

116 v += ".0.0" 

117 elif v.count(".") == 1: 

118 v += ".0" 

119 

120 deps: list[str | PipDeps] = [f"pytorch=={v}"] 

121 additional_deps = { 

122 "1.5.1": "torchvision==0.6.1", 

123 "1.6.0": "torchvision==0.7.0", 

124 "1.7.0": "torchvision==0.8.0", 

125 "1.7.1": "torchvision==0.8.2", 

126 "1.8.0": "torchvision==0.9.0", 

127 "1.8.1": "torchvision==0.9.1", 

128 "1.9.0": "torchvision==0.10.0", 

129 "1.9.1": "torchvision==0.10.1", 

130 "1.10.0": "torchvision==0.11.0", 

131 "1.10.1": "torchvision==0.11.2", 

132 "1.11.0": "torchvision==0.12.0", 

133 "1.12.0": "torchvision==0.13.0", 

134 "1.12.1": "torchvision==0.13.1", 

135 "1.13.0": "torchvision==0.14.0", 

136 "1.13.1": "torchvision==0.14.1", 

137 "2.0.0": "torchvision==0.15.0", 

138 "2.0.1": "torchvision==0.15.2", 

139 "2.1.0": "torchvision==0.16.0", 

140 "2.1.1": "torchvision==0.16.1", 

141 "2.1.2": "torchvision==0.16.2", 

142 "2.2.0": "torchvision==0.17.0", 

143 "2.2.1": "torchvision==0.17.1", 

144 "2.2.2": "torchvision==0.17.2", 

145 "2.3.0": "torchvision==0.18.0", 

146 "2.3.1": "torchvision==0.18.1", 

147 "2.4.0": "torchvision==0.19.0", 

148 "2.4.1": "torchvision==0.19.1", 

149 "2.5.0": "torchvision==0.20.0", 

150 "2.5.1": "torchvision==0.20.1", 

151 "2.6.0": "torchvision==0.21.0", 

152 "2.7.0": "torchvision==0.22.0", 

153 "2.7.1": "torchvision==0.22.1", 

154 "2.8.0": "torchvision==0.23.0", 

155 "2.9.0": "torchvision==0.24.0", 

156 "2.9.1": "torchvision==0.24.1", 

157 }.get(v) 

158 if additional_deps is None: 

159 set_github_warning( 

160 "UPDATE NEEDED", 

161 f"Leaving torchvision unpinned for pytorch=={v}", 

162 ) 

163 additional_deps = "torchvision" 

164 

165 deps.append(additional_deps) 

166 

167 # avoid `undefined symbol: iJIT_NotifyEvent` from `torch/lib/libtorch_cpu.so` 

168 # see https://github.com/pytorch/pytorch/issues/123097 

169 if ( 

170 pytorch_version 

171 < Version( 

172 "2.1.0" # TODO: check if this is the correct cutoff where the fix is not longer needed 

173 ) 

174 ): 

175 deps.append("mkl ==2024.0.0") 

176 

177 if pytorch_version < Version("2.2"): 

178 # avoid ImportError: cannot import name 'packaging' from 'pkg_resources' 

179 # see https://github.com/pypa/setuptools/issues/4376#issuecomment-2126162839 

180 deps.append("setuptools <70.0.0") 

181 

182 if pytorch_version < Version("2.3"): 

183 # see https://github.com/pytorch/pytorch/issues/107302 

184 deps.append("numpy <2") 

185 

186 return BioimageioCondaEnv(channels=channels, dependencies=deps) 

187 

188 

189def _get_default_onnx_env(*, opset_version: int | None) -> BioimageioCondaEnv: 

190 if opset_version is None: 

191 opset_version = 15 

192 

193 # note: we should not need to worry about the opset version, 

194 # see https://github.com/microsoft/onnxruntime/blob/master/docs/Versioning.md 

195 return BioimageioCondaEnv(dependencies=["onnxruntime"]) 

196 

197 

198def _get_default_tf_env(tensorflow_version: Version | None) -> BioimageioCondaEnv: 

199 if tensorflow_version is None or tensorflow_version.major < 2: 

200 tensorflow_version = Version("2.17") 

201 

202 return BioimageioCondaEnv( 

203 dependencies=[f"tensorflow =={tensorflow_version}"], 

204 ) 

205 

206 

207def _get_env_from_deps( 

208 deps: v0_4.Dependencies | FileDescr, 

209) -> BioimageioCondaEnv: 

210 if isinstance(deps, v0_4.Dependencies): 

211 deps_reader = get_reader(deps.file) 

212 if deps.manager == "pip": 

213 pip_deps_str = deps_reader.read_text() 

214 pip_deps = [d.strip() for d in pip_deps_str.split("\n")] 

215 if "bioimageio.core" not in pip_deps: 

216 pip_deps.append("bioimageio.core>=0.9.4") 

217 

218 return BioimageioCondaEnv( 

219 dependencies=[PipDeps(pip=pip_deps)], 

220 ) 

221 elif deps.manager in ("conda", "mamba"): 

222 return BioimageioCondaEnv.model_validate(read_yaml(deps_reader)) 

223 else: 

224 raise ValueError(f"Dependency manager {deps.manager} not supported") 

225 

226 elif isinstance(deps, FileDescr): 

227 deps_reader = deps.get_reader() 

228 return BioimageioCondaEnv.model_validate(read_yaml(deps_reader)) 

229 else: 

230 assert_never(deps)