Coverage for ml_workbench/mlflow_conf.py: 20%

1from __future__ import annotations 

2 

3from collections.abc import Mapping 

4from dataclasses import dataclass 

5import os 

6from typing import TYPE_CHECKING, Any 

7 

8if TYPE_CHECKING: 8 ↛ 9line 8 didn't jump to line 9 because the condition on line 8 was never true

9 from .config import YamlConfig 

10 

11 

12@dataclass(frozen=True) 

13class MlflowSpec: 

14 enabled: bool 

15 type: str # Always "local" or "databricks", inferred from MLFLOW_TRACKING_URI if not specified 

16 experiment_name_prefix: str 

17 tags: dict[str, Any] 

18 

19 

20class MlflowConf: 

21 """Encapsulates MLflow configuration from the YAML config. 

22 

23 Expected YAML structure: 

24 

25 mlflow: 

26 enabled: true # optional, defaults to True 

27 type: "local" # optional, "local" or "databricks", inferred from MLFLOW_TRACKING_URI if not specified 

28 experiment_name_prefix: "/Shared/" # optional, defaults to "" 

29 tags: # optional, defaults to {} 

30 environment: "development" 

31 data_version: "v1" 

32 

33 If the `mlflow` section is not present in the config, default values are used: 

34 - enabled: True 

35 - type: inferred from MLFLOW_TRACKING_URI env var (if starts with "databricks" -> "databricks", else "local") 

36 - experiment_name_prefix: "" 

37 - tags: {} 

38 """ 

39 

40 def __init__(self, config: YamlConfig) -> None: 

41 self.config = config 

42 

43 mlflow_section = self._get_mlflow_section(config) 

44 

45 # Extract values with defaults 

46 enabled = mlflow_section.get("enabled", True) 

47 if not isinstance(enabled, bool): 

48 raise TypeError("mlflow.enabled must be a boolean") # noqa: TRY003 

49 

50 mlflow_type = mlflow_section.get("type") 

51 if mlflow_type is not None and not isinstance(mlflow_type, str): 

52 raise TypeError("mlflow.type must be a string or None") # noqa: TRY003 

53 # Validate type if provided 

54 if mlflow_type is not None: 

55 mlflow_type_lower = mlflow_type.lower() 

56 if mlflow_type_lower not in ("local", "databricks"): 

57 raise ValueError( # noqa: TRY003 

58 f"mlflow.type must be 'local' or 'databricks', got '{mlflow_type}'" 

59 ) 

60 mlflow_type = mlflow_type_lower 

61 else: 

62 # If type not specified, infer from MLFLOW_TRACKING_URI environment variable 

63 tracking_uri = os.environ.get("MLFLOW_TRACKING_URI", "") 

64 if tracking_uri.startswith("databricks"): 

65 mlflow_type = "databricks" 

66 else: 

67 mlflow_type = "local" 

68 

69 experiment_name_prefix = mlflow_section.get("experiment_name_prefix", "") 

70 if not isinstance(experiment_name_prefix, str): 

71 raise TypeError("mlflow.experiment_name_prefix must be a string") # noqa: TRY003 

72 

73 tags = mlflow_section.get("tags", {}) 

74 if not isinstance(tags, Mapping): 

75 raise TypeError("mlflow.tags must be a mapping") # noqa: TRY003 

76 

77 self.spec = MlflowSpec( 

78 enabled=enabled, 

79 type=mlflow_type, 

80 experiment_name_prefix=experiment_name_prefix, 

81 tags=dict(tags), 

82 ) 

83 

84 @staticmethod 

85 def _get_mlflow_section(config: YamlConfig) -> Mapping[str, Any]: 

86 """Get mlflow section from config, returning empty dict if not present.""" 

87 mlflow = config.get_data().get("mlflow") 

88 if mlflow is None: 

89 return {} 

90 if not isinstance(mlflow, Mapping): 

91 raise TypeError("'mlflow' section must be a mapping") # noqa: TRY003 

92 return mlflow 

93 

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

95 """Return dictionary representation of MLflow configuration.""" 

96 return { 

97 "enabled": self.spec.enabled, 

98 "type": self.spec.type, 

99 "experiment_name_prefix": self.spec.experiment_name_prefix, 

100 "tags": dict(self.spec.tags), 

101 } 

102 

103 def is_enabled(self) -> bool: 

104 """Check if MLflow tracking is enabled. 

105 

106 Returns 

107 ------- 

108 bool 

109 True if MLflow tracking is enabled, False otherwise 

110 """ 

111 return self.spec.enabled 

112 

113 def get_type(self) -> str: 

114 """Get the MLflow tracking type. 

115 

116 Returns 

117 ------- 

118 str 

119 The tracking type: "local" or "databricks" 

120 """ 

121 return self.spec.type 

122 

123 def get_tags(self) -> dict[str, Any]: 

124 """Get all tags. 

125 

126 Returns 

127 ------- 

128 dict[str, Any] 

129 All tags 

130 """ 

131 return self.spec.tags 

132 

133 def get_name(self, experiment_name: str) -> str: 

134 """Get the full experiment name by combining experiment name prefix with experiment name. 

135 

136 Parameters 

137 ---------- 

138 experiment_name : str 

139 The experiment name 

140 

141 Returns 

142 ------- 

143 str 

144 The full experiment name (prefix + experiment_name) 

145 """ 

146 prefix = self.spec.experiment_name_prefix 

147 # Ensure prefix ends with / if it's not empty and doesn't already end with / 

148 if prefix and not prefix.endswith("/"): 

149 prefix = prefix + "/" 

150 # Remove leading / from experiment_name if prefix already has it 

151 if prefix.startswith("/") and experiment_name.startswith("/"): 

152 experiment_name = experiment_name.lstrip("/") 

153 return prefix + experiment_name 

154 

155 @classmethod 

156 def verify_config(cls, config: YamlConfig) -> None: 

157 """Validate that mlflow section is well-formed if present. 

158 

159 Checks 

160 ------ 

161 - mlflow section, if present, is a mapping 

162 - enabled is a boolean (if provided) 

163 - type is "local" or "databricks" (if provided) 

164 - experiment_name_prefix is a string (if provided) 

165 - tags is a mapping (if provided) 

166 

167 Notes 

168 ----- 

169 If mlflow section is not present, validation passes (defaults will be used). 

170 """ 

171 mlflow = config.get_data().get("mlflow") 

172 if mlflow is None: 

173 return # Section not present, defaults will be used 

174 

175 if not isinstance(mlflow, Mapping): 

176 raise TypeError("'mlflow' section must be a mapping") # noqa: TRY003 

177 

178 # Validate enabled if present 

179 enabled = mlflow.get("enabled") 

180 if enabled is not None and not isinstance(enabled, bool): 

181 raise TypeError("mlflow.enabled must be a boolean") # noqa: TRY003 

182 

183 # Validate type if present 

184 mlflow_type = mlflow.get("type") 

185 if mlflow_type is not None: 

186 if not isinstance(mlflow_type, str): 

187 raise TypeError("mlflow.type must be a string") # noqa: TRY003 

188 mlflow_type_lower = mlflow_type.lower() 

189 if mlflow_type_lower not in ("local", "databricks"): 

190 raise ValueError( # noqa: TRY003 

191 f"mlflow.type must be 'local' or 'databricks', got '{mlflow_type}'" 

192 ) 

193 

194 # Validate experiment_name_prefix if present 

195 experiment_name_prefix = mlflow.get("experiment_name_prefix") 

196 if experiment_name_prefix is not None and not isinstance( 

197 experiment_name_prefix, str 

198 ): 

199 raise TypeError("mlflow.experiment_name_prefix must be a string") # noqa: TRY003 

200 

201 # Validate tags if present 

202 tags = mlflow.get("tags") 

203 if tags is not None and not isinstance(tags, Mapping): 

204 raise TypeError("mlflow.tags must be a mapping") # noqa: TRY003 

205 

206 

207__all__ = ["MlflowConf", "MlflowSpec"]