Coverage for ml_workbench/mlflow.py: 28%

1from __future__ import annotations 

2 

3import os 

4from collections.abc import Mapping 

5from dataclasses import dataclass 

6from typing import Any 

7 

8from .config import YamlConfig 

9 

10 

11@dataclass(frozen=True) 

12class MlflowSpec: 

13 enabled: bool 

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

15 experiment_name_prefix: str 

16 tags: dict[str, Any] 

17 

18 

19class Mlflow: 

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

21 

22 Expected YAML structure: 

23 

24 mlflow: 

25 enabled: true # optional, defaults to True 

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

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

28 tags: # optional, defaults to {} 

29 environment: "development" 

30 data_version: "v1" 

31 

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

33 - enabled: True 

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

35 - experiment_name_prefix: "" 

36 - tags: {} 

37 """ 

38 

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

40 self.config = config 

41 

42 mlflow_section = self._get_mlflow_section(config) 

43 

44 # Extract values with defaults 

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

46 if not isinstance(enabled, bool): 

47 raise ValueError("mlflow.enabled must be a boolean") 

48 

49 mlflow_type = mlflow_section.get("type") 

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

51 raise ValueError("mlflow.type must be a string or None") 

52 # Validate type if provided 

53 if mlflow_type is not None: 

54 mlflow_type_lower = mlflow_type.lower() 

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

56 raise ValueError( 

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

58 ) 

59 mlflow_type = mlflow_type_lower 

60 else: 

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

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

63 if tracking_uri.startswith("databricks"): 

64 mlflow_type = "databricks" 

65 else: 

66 mlflow_type = "local" 

67 

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

69 if not isinstance(experiment_name_prefix, str): 

70 raise ValueError("mlflow.experiment_name_prefix must be a string") 

71 

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

73 if not isinstance(tags, Mapping): 

74 raise ValueError("mlflow.tags must be a mapping") 

75 

76 self.spec = MlflowSpec( 

77 enabled=enabled, 

78 type=mlflow_type, 

79 experiment_name_prefix=experiment_name_prefix, 

80 tags=dict(tags), 

81 ) 

82 

83 @staticmethod 

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

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

86 mlflow = config._data.get("mlflow") # type: ignore[attr-defined] 

87 if mlflow is None: 

88 return {} 

89 if not isinstance(mlflow, Mapping): 

90 raise ValueError("'mlflow' section must be a mapping") 

91 return mlflow 

92 

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

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

95 return { 

96 "enabled": self.spec.enabled, 

97 "type": self.spec.type, 

98 "experiment_name_prefix": self.spec.experiment_name_prefix, 

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

100 } 

101 

102 def is_enabled(self) -> bool: 

103 """Check if MLflow tracking is enabled. 

104 

105 Returns 

106 ------- 

107 bool 

108 True if MLflow tracking is enabled, False otherwise 

109 """ 

110 return self.spec.enabled 

111 

112 def get_type(self) -> str: 

113 """Get the MLflow tracking type. 

114 

115 Returns 

116 ------- 

117 str 

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

119 """ 

120 return self.spec.type 

121 

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

123 """Get all tags. 

124 

125 Returns 

126 ------- 

127 dict[str, Any] 

128 All tags 

129 """ 

130 return self.spec.tags 

131 

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

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

134 

135 Parameters 

136 ---------- 

137 experiment_name : str 

138 The experiment name 

139 

140 Returns 

141 ------- 

142 str 

143 The full experiment name (prefix + experiment_name) 

144 """ 

145 prefix = self.spec.experiment_name_prefix 

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

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

148 prefix = prefix + "/" 

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

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

151 experiment_name = experiment_name.lstrip("/") 

152 return prefix + experiment_name 

153 

154 @classmethod 

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

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

157 

158 Checks 

159 ------ 

160 - mlflow section, if present, is a mapping 

161 - enabled is a boolean (if provided) 

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

163 - experiment_name_prefix is a string (if provided) 

164 - tags is a mapping (if provided) 

165 

166 Notes 

167 ----- 

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

169 """ 

170 mlflow = config._data.get("mlflow") # type: ignore[attr-defined] 

171 if mlflow is None: 

172 return # Section not present, defaults will be used 

173 

174 if not isinstance(mlflow, Mapping): 

175 raise ValueError("'mlflow' section must be a mapping") 

176 

177 # Validate enabled if present 

178 enabled = mlflow.get("enabled") 

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

180 raise ValueError("mlflow.enabled must be a boolean") 

181 

182 # Validate type if present 

183 mlflow_type = mlflow.get("type") 

184 if mlflow_type is not None: 

185 if not isinstance(mlflow_type, str): 

186 raise ValueError("mlflow.type must be a string") 

187 mlflow_type_lower = mlflow_type.lower() 

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

189 raise ValueError( 

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

191 ) 

192 

193 # Validate experiment_name_prefix if present 

194 experiment_name_prefix = mlflow.get("experiment_name_prefix") 

195 if experiment_name_prefix is not None and not isinstance( 

196 experiment_name_prefix, str 

197 ): 

198 raise ValueError("mlflow.experiment_name_prefix must be a string") 

199 

200 # Validate tags if present 

201 tags = mlflow.get("tags") 

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

203 raise ValueError("mlflow.tags must be a mapping") 

204 

205 

206__all__ = ["Mlflow", "MlflowSpec"] 

207