Coverage for ml_workbench/cli_experiment.py: 9%

1"""CLI tool for running YAML-defined ML experiments. 

2 

3This tool loads a YAML configuration file, selects one or more experiments, 

4and executes them using the Runner class. It also provides inspection options 

5to view configuration and dataset information without running experiments. 

6""" 

7 

8from __future__ import annotations 

9 

10import argparse 

11import json 

12from pathlib import Path 

13import sys 

14import traceback 

15from typing import Any 

16 

17from dotenv import load_dotenv 

18import yaml 

19 

20from .config import YamlConfig 

21from .dataset import Dataset 

22from .experiment import Experiment 

23from .runner import Runner 

24 

25load_dotenv() 

26 

27# Constants 

28MAX_COLUMNS_TO_DISPLAY = 10 

29 

30 

31def parse_kv_pairs(kv_list: list[str]) -> dict[str, str]: 

32 """Parse key=value pairs from command line arguments. 

33 

34 Parameters 

35 ---------- 

36 kv_list : list[str] 

37 List of "key=value" strings 

38 

39 Returns 

40 ------- 

41 dict[str, str] 

42 Dictionary of parsed key-value pairs 

43 

44 Raises 

45 ------ 

46 argparse.ArgumentTypeError 

47 If a pair is not in "key=value" format 

48 """ 

49 variables: dict[str, str] = {} 

50 for item in kv_list: 

51 if "=" not in item: 

52 raise argparse.ArgumentTypeError( # noqa: TRY003 

53 f"Invalid variable format {item!r}, expected key=value" 

54 ) 

55 key, value = item.split("=", 1) 

56 variables[key.strip()] = value.strip() 

57 return variables 

58 

59 

60def _basic_stats_for_dataset(name: str, cfg: YamlConfig) -> dict[str, Any]: 

61 """Get basic statistics for a dataset using the Dataset class. 

62 

63 Parameters 

64 ---------- 

65 name : str 

66 Dataset name 

67 cfg : YamlConfig 

68 Configuration object 

69 

70 Returns 

71 ------- 

72 dict[str, Any] 

73 Dictionary with dataset statistics including name, description, format, 

74 type, path, num_columns, num_rows, column_names, and is_combined flag 

75 """ 

76 try: 

77 ds = Dataset(name, cfg) 

78 

79 result: dict[str, Any] = { 

80 "name": name, 

81 "description": ds.description, 

82 "format": ds.format, 

83 "type": ds.type, 

84 "path": ds.path, 

85 "is_combined": ds.is_combined, 

86 } 

87 

88 # Try to read and get statistics 

89 try: 

90 stats = ds.get_statistics() 

91 result.update({ 

92 "num_columns": stats["num_columns"], 

93 "num_rows": stats["num_rows"], 

94 "column_names": stats["column_names"], 

95 }) 

96 except Exception as exc: 

97 # If reading fails, still return basic info with error 

98 result.update({ 

99 "num_columns": None, 

100 "num_rows": None, 

101 "column_names": None, 

102 "error": str(exc), 

103 }) 

104 

105 except Exception as exc: 

106 # If dataset creation fails, return minimal info 

107 return { 

108 "name": name, 

109 "error": f"Failed to create dataset: {exc}", 

110 } 

111 else: 

112 return result 

113 

114 

115def filter_sections(data: dict, sections: list[str]) -> dict: 

116 """Filter configuration to include only specified sections. 

117 

118 Parameters 

119 ---------- 

120 data : dict 

121 Full configuration dictionary 

122 sections : list[str] 

123 List of section names to include 

124 

125 Returns 

126 ------- 

127 dict 

128 Filtered configuration with only specified sections 

129 """ 

130 result = {} 

131 for section in sections: 

132 if section in data: 

133 result[section] = data[section] 

134 else: 

135 print(f"Warning: Section '{section}' not found", file=sys.stderr) # noqa: T201 

136 return result 

137 

138 

139def build_arg_parser() -> argparse.ArgumentParser: 

140 """Build the argument parser for the CLI. 

141 

142 Returns 

143 ------- 

144 argparse.ArgumentParser 

145 Configured argument parser 

146 """ 

147 parser = argparse.ArgumentParser( 

148 prog="cli-experiment", 

149 description="Run YAML-defined ML experiments or inspect configuration and datasets", 

150 formatter_class=argparse.RawDescriptionHelpFormatter, 

151 epilog=""" 

152Examples: 

153 # Run all experiments defined in the YAML file 

154 %(prog)s experiment.yaml 

155 

156 # Run specific experiment(s) 

157 %(prog)s experiment.yaml --experiments house_prices_prediction_simple 

158 

159 # Run multiple specific experiments 

160 %(prog)s experiment.yaml --experiments exp1 exp2 exp3 

161 

162 # Run with variable substitution 

163 %(prog)s experiment.yaml --var path=/data --var env=production 

164 

165 # Inspect configuration 

166 %(prog)s experiment.yaml --show-config 

167 

168 # Inspect configuration as JSON 

169 %(prog)s experiment.yaml --show-config --json 

170 

171 # Inspect specific sections 

172 %(prog)s experiment.yaml --show-config --section experiments --section datasets 

173 

174 # Inspect datasets 

175 %(prog)s experiment.yaml --show-datasets 

176 

177 # Inspect datasets as JSON 

178 %(prog)s experiment.yaml --show-datasets --json 

179 """, 

180 ) 

181 

182 parser.add_argument( 

183 "yaml", 

184 type=Path, 

185 help="Path to the YAML configuration file containing experiment definitions", 

186 ) 

187 

188 parser.add_argument( 

189 "--experiments", 

190 action="append", 

191 dest="experiments", 

192 metavar="EXPERIMENT", 

193 help="Name(s) of experiment(s) to run (can be specified multiple times). " 

194 "If not specified, all experiments in the YAML will be run. " 

195 "Ignored if --show-config or --show-datasets is used.", 

196 ) 

197 

198 parser.add_argument( 

199 "--var", 

200 action="append", 

201 dest="vars", 

202 metavar="KEY=VALUE", 

203 help="Variable for interpolation (can be specified multiple times)", 

204 ) 

205 

206 parser.add_argument( 

207 "--quiet", 

208 action="store_true", 

209 help="Suppress verbose output from Runner", 

210 ) 

211 

212 # Inspection options 

213 parser.add_argument( 

214 "--show-config", 

215 action="store_true", 

216 help="Show processed configuration . When used, experiments are not run.", 

217 ) 

218 

219 parser.add_argument( 

220 "--show-datasets", 

221 action="store_true", 

222 help="Show dataset statistics . When used, experiments are not run.", 

223 ) 

224 

225 parser.add_argument( 

226 "--json", 

227 action="store_true", 

228 help="Output as JSON instead of YAML/text (used with --show-config or --show-datasets)", 

229 ) 

230 

231 parser.add_argument( 

232 "--section", 

233 action="append", 

234 dest="sections", 

235 metavar="SECTION", 

236 help="Output only specific section(s) when using --show-config " 

237 "(can be specified multiple times)", 

238 ) 

239 

240 parser.add_argument( 

241 "--show-variables", 

242 action="store_true", 

243 help="Show resolved variables when using --show-config", 

244 ) 

245 

246 return parser 

247 

248 

249def main(argv: list[str] | None = None) -> int: 

250 """Main entry point for the CLI. 

251 

252 Parameters 

253 ---------- 

254 argv : Optional[list[str]] 

255 Command line arguments (default: sys.argv) 

256 

257 Returns 

258 ------- 

259 int 

260 Exit code (0 for success, non-zero for error) 

261 """ 

262 parser = build_arg_parser() 

263 args = parser.parse_args(argv) 

264 

265 exit_code = 0 # Initialize exit code 

266 try: 

267 # Parse variable substitutions 

268 variables = parse_kv_pairs(list(args.vars)) if args.vars else {} 

269 

270 # Load configuration 

271 if not args.yaml.exists(): 

272 print(f"YAML file not found: {args.yaml}", file=sys.stderr) # noqa: T201 

273 return 1 

274 

275 config = YamlConfig(args.yaml, **variables) 

276 

277 # Handle inspection modes (these don't run experiments) 

278 if args.show_config: 

279 # Show configuration 

280 data = config.to_dict() 

281 

282 # Filter sections if specified 

283 if args.sections: 

284 data = filter_sections(data, args.sections) 

285 

286 # Show variables if requested 

287 if args.show_variables: 

288 print("# Resolved Variables:", file=sys.stderr) # noqa: T201 

289 for key, value in sorted(config.variables.items()): 

290 print(f" {key}={value}", file=sys.stderr) # noqa: T201 

291 

292 # Format output 

293 if args.json: 

294 print(json.dumps(data, indent=2)) # noqa: T201 

295 else: 

296 print(yaml.dump( # noqa: T201 

297 data, 

298 default_flow_style=False, 

299 sort_keys=False, 

300 allow_unicode=True, 

301 )) 

302 

303 return 0 

304 

305 if args.show_datasets: 

306 # Show dataset statistics 

307 dataset_names = config.get_datasets_list() 

308 results = [_basic_stats_for_dataset(name, config) for name in dataset_names] 

309 

310 if args.json: 

311 print(json.dumps({"datasets": results}, indent=2)) # noqa: T201 

312 else: 

313 for item in results: 

314 if "error" in item: 

315 continue 

316 print(f"\nDataset: {item['name']}") # noqa: T201 

317 if item.get("is_combined"): 

318 print(" Type: Combined") # noqa: T201 

319 if item.get("description"): 

320 print(f" Description: {item['description']}") # noqa: T201 

321 if item.get("format"): 

322 print(f" format: {item['format']}") # noqa: T201 

323 if item.get("num_rows") is not None: 

324 print(f" Rows: {item['num_rows']}") # noqa: T201 

325 if item.get("num_columns") is not None: 

326 print(f" Columns: {item['num_columns']}") # noqa: T201 

327 # Print column names if available 

328 column_names = item.get("column_names") 

329 if column_names: 

330 print(f" Column names: {', '.join(column_names[:MAX_COLUMNS_TO_DISPLAY])}") # noqa: T201 

331 if len(column_names) > MAX_COLUMNS_TO_DISPLAY: 

332 print(f" ... and {len(column_names) - MAX_COLUMNS_TO_DISPLAY} more") # noqa: T201 

333 

334 return 0 

335 

336 # Normal mode: run experiments 

337 # Get list of available experiments 

338 all_experiment_names = Experiment.list_experiment_names(config) 

339 

340 if not all_experiment_names: 

341 print("No experiments found in configuration", file=sys.stderr) # noqa: T201 

342 return 1 

343 

344 # Determine which experiments to run 

345 if args.experiments: 

346 # User specified experiments to run 

347 experiments_to_run = args.experiments 

348 # Check if all specified experiments exist 

349 missing_experiments = [ 

350 exp for exp in experiments_to_run if exp not in all_experiment_names 

351 ] 

352 if missing_experiments: 

353 print(f"Experiment(s) not found: {', '.join(missing_experiments)}", file=sys.stderr) # noqa: T201 

354 return 1 

355 else: 

356 # Run all experiments by default 

357 experiments_to_run = all_experiment_names 

358 

359 # Verify configuration before running 

360 try: 

361 Experiment.verify_config(config) 

362 except Exception as e: 

363 if not args.quiet: 

364 print(f"Configuration validation failed: {e}", file=sys.stderr) # noqa: T201 

365 return 1 

366 

367 # Run each experiment 

368 verbose = not args.quiet 

369 exit_code = 0 

370 

371 for exp_name in experiments_to_run: 

372 try: 

373 if verbose: 

374 pass 

375 

376 # Create experiment object 

377 experiment = Experiment(config, exp_name) 

378 

379 # Create and run runner 

380 runner = Runner(experiment, verbose=verbose) 

381 results = runner.run() 

382 

383 if verbose: 

384 pass 

385 

386 except KeyboardInterrupt: 

387 return 130 

388 except Exception as e: 

389 print(f"Error running experiment: {e}", file=sys.stderr) # noqa: T201 

390 if verbose: 

391 traceback.print_exc() 

392 exit_code = 1 

393 

394 if verbose and len(experiments_to_run) > 1: 

395 pass 

396 

397 except FileNotFoundError as e: 

398 print(f"YAML file not found: {e}", file=sys.stderr) # noqa: T201 

399 return 1 

400 except ValueError as e: 

401 print(f"Error: {e}", file=sys.stderr) # noqa: T201 

402 return 1 

403 except Exception as e: 

404 if "--debug" in (argv or sys.argv): 

405 raise # noqa: TRY003 

406 print(f"Error: {e}", file=sys.stderr) # noqa: T201 

407 return 1 

408 else: 

409 return exit_code 

410 

411 

412if __name__ == "__main__": 

413 sys.exit(main())