Coverage for packages/dqm-ml-job/src/dqm_ml_job/cli.py: 85%
163 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-07-21 08:27 +0000
1"""Command-line interface for DQM job execution.
3This module provides CLI functions for parsing arguments and running
4data quality assessment jobs from YAML configuration files.
5"""
7import argparse
8import logging
9from pathlib import Path
10from typing import Any, cast
12from dqm_ml_core import PluginLoadedRegistry
13from dqm_ml_core.models.config import JobConfig
14from dqm_ml_core.models.global_ import ComputeConfig, ErrorsConfig
15from dqm_ml_core.models.interfaces import FeaturesInterfaceConfig, GapInterfaceConfig, MetricsInterfaceConfig
16import pyarrow as pa
17import yaml
19from dqm_ml_job.job import DatasetJob
20from dqm_ml_job.outputwriter import OutputWriter
22logger = logging.getLogger(__name__)
25def _merge_errors(global_errors: ErrorsConfig | None, interface_errors: ErrorsConfig | None) -> ErrorsConfig:
26 """Merge global and interface-specific errors, with interface taking precedence."""
27 if interface_errors is None:
28 return global_errors or ErrorsConfig()
30 # Start with global defaults
31 merged = global_errors or ErrorsConfig()
33 # Override with interface-specific values (only where interface is not None)
34 if interface_errors.default is not None: 34 ↛ 36line 34 didn't jump to line 36 because the condition on line 34 was always true
35 merged.default = interface_errors.default
36 if interface_errors.images is not None: 36 ↛ 37line 36 didn't jump to line 37 because the condition on line 36 was never true
37 merged.images = interface_errors.images
38 if interface_errors.tabular is not None: 38 ↛ 39line 38 didn't jump to line 39 because the condition on line 38 was never true
39 merged.tabular = interface_errors.tabular
40 if interface_errors.max_failure_rate is not None: 40 ↛ 43line 40 didn't jump to line 43 because the condition on line 40 was always true
41 merged.max_failure_rate = interface_errors.max_failure_rate
43 return merged
46def parse_args(arg_list: list[str] | None) -> Any:
47 """
48 Parse command line arguments for the DQM job.
50 Args:
51 arg_list: List of arguments (default: sys.argv[1:]).
53 Returns:
54 The parsed Namespace object.
55 """
56 parser = argparse.ArgumentParser(
57 prog="dqm-ml",
58 description="DQM-ML Job client",
59 epilog="for more informations see README",
60 )
62 parser.add_argument(
63 "-p",
64 "--process-config",
65 type=str,
66 nargs="+",
67 required=True,
68 help="configuration files to execute",
69 )
71 parser.add_argument(
72 "--save-config",
73 type=str,
74 help="Path to save the resolved configuration",
75 )
77 # TODO add parameters to pass directly files / directory as inputs for loaders
78 args = parser.parse_args(arg_list)
80 return args
83# TODO get parameters, logs, ...
84def execute(arg_list: list[str] | None = None) -> None:
85 """
86 Main CLI entry point for executing DQM jobs from YAML configurations.
87 Args:
88 arg_list: List of command line arguments (default: sys.argv[1:]).
89 """
90 args = parse_args(arg_list)
91 config: dict[str, Any] = {}
93 for config_file in args.process_config:
94 logger.debug("Executing job from config file: %s", config_file)
96 config_path = Path(config_file).resolve()
97 if not config_path.is_file():
98 logger.error("Config file does not exist: %s", config_file)
99 return
101 with config_path.open() as stream:
102 try:
103 config_content = yaml.safe_load(stream)
104 config.update(config_content)
105 except yaml.YAMLError as exc:
106 logger.error("Fail to part job configuration: %s", config_file)
107 print(exc)
108 return
110 # if we succeed to load all config files, run the job
112 # Optionally save the resolved configuration
113 if args.save_config:
114 logger.debug("Saving resolved configuration to: %s", args.save_config)
115 save_path = Path(args.save_config).resolve()
116 save_path.parent.mkdir(parents=True, exist_ok=True)
117 with save_path.open("w") as stream:
118 yaml.safe_dump(config, stream)
120 run(config)
123def _init_components_from_list(
124 processor_list: list[dict[str, Any]], registry: dict[str, Any], component_name: str
125) -> dict[str, Any]:
126 """Initialize components from a list of processor configs (new format).
128 Each item in the list must have a 'name' and 'type' field.
130 Args:
131 processor_list: List of component configuration dicts.
132 registry: The registry containing the component classes.
133 component_name: The name of the component type (for error messages).
135 Returns:
136 A dictionary mapping component names to initialized instances.
137 """
138 components = {}
139 for comp_config in processor_list:
140 proc_name = comp_config.get("name")
141 if not proc_name:
142 raise ValueError(f"Configuration for {component_name} must contain 'name'")
143 if "type" not in comp_config:
144 raise ValueError(f"Configuration for {component_name} '{proc_name}' must contain 'type'")
145 comp_type = comp_config["type"]
146 if comp_type not in registry:
147 raise ValueError(f"{component_name.capitalize()} '{proc_name}' has invalid type '{comp_type}'")
148 components[proc_name] = registry[comp_type](name=proc_name, config=comp_config)
149 return components
152def _init_output_writer(
153 name: str,
154 path: str,
155 columns: list[str] | None,
156 outputs_registry: dict[str, Any],
157 exclude: list[str] | None = None,
158 storage: dict[str, Any] | None = None,
159) -> OutputWriter | None:
160 """Initialize a single output writer from interface outputs config.
162 Args:
163 name: Name for the writer instance.
164 path: Output path from the interface outputs config.
165 columns: Optional list of columns to include.
166 outputs_registry: Registry of available writer types.
167 exclude: Optional list of columns to exclude.
168 storage: Storage config dict to inject into the writer (optional).
170 Returns:
171 An initialized OutputWriter instance, or None if no writer type is available.
172 """
173 writer_type = "parquet"
174 if writer_type not in outputs_registry: 174 ↛ 175line 174 didn't jump to line 175 because the condition on line 174 was never true
175 logger.warning("Output writer type '%s' not found in registry", writer_type)
176 return None
177 writer_config: dict[str, Any] = {"path_pattern": path, "columns": columns or [], "exclude": exclude or []}
178 if storage: 178 ↛ 179line 178 didn't jump to line 179 because the condition on line 178 was never true
179 writer_config["storage"] = storage
180 return cast(OutputWriter, outputs_registry[writer_type](name=name, config=writer_config))
183def _init_interface_outputs(
184 interface_config: FeaturesInterfaceConfig | MetricsInterfaceConfig | GapInterfaceConfig | None,
185 outputs_registry: dict[str, Any],
186 kind: str,
187 storage: dict[str, Any] | None = None,
188) -> OutputWriter | None:
189 """Initialize the output writer for a given interface.
191 Args:
192 interface_config: The validated interface configuration.
193 outputs_registry: Registry of available writer types.
194 kind: The kind of interface ('features', 'metrics', or 'gap').
195 storage: Storage config dict to inject into the writer (optional).
197 Returns:
198 An initialized OutputWriter instance, or None.
199 """
200 if interface_config is None or interface_config.outputs is None:
201 return None
202 path = interface_config.outputs.path
203 columns: list[str] | None = None
204 if hasattr(interface_config.outputs, "include") and interface_config.outputs.include:
205 columns = interface_config.outputs.include
206 exclude: list[str] | None = None
207 if hasattr(interface_config.outputs, "exclude") and interface_config.outputs.exclude:
208 exclude = interface_config.outputs.exclude
209 return _init_output_writer(kind, path, columns, outputs_registry, exclude, storage)
212def _resolve_compute_config(validated: JobConfig) -> ComputeConfig:
213 """Resolve the compute config, providing defaults if not specified."""
214 if validated.compute:
215 return validated.compute
216 return ComputeConfig(
217 seed=42,
218 log_level="warning",
219 max_memory=None,
220 device="auto",
221 progress_bar=True,
222 threads=4,
223 )
226def _init_processors_from_interface(
227 interface: Any,
228 registry: dict[str, Any],
229 storage: dict[str, Any] | None = None,
230) -> dict[str, Any]:
231 """Initialize processors from an optional interface config.
233 Args:
234 interface: The validated interface config (or None).
235 registry: The component registry.
236 storage: Storage config dict to inject into each processor (optional).
238 Returns:
239 A dict of name-to-processor instances (empty if interface is None).
240 """
241 if interface is None:
242 return {}
243 proc_dicts = [p.model_dump() for p in interface.processors]
244 if storage: 244 ↛ 245line 244 didn't jump to line 245 because the condition on line 244 was never true
245 for proc_dict in proc_dicts:
246 proc_dict["storage"] = storage
247 return _init_components_from_list(proc_dicts, registry, "processor")
250def _safe_flush(writer: Any) -> None:
251 """Flush a writer if it has a flush method."""
252 if writer and hasattr(writer, "flush"): 252 ↛ exitline 252 didn't return from function '_safe_flush' because the condition on line 252 was always true
253 writer.flush()
256def _build_errors_by_interface(validated: JobConfig) -> dict[str, ErrorsConfig]:
257 """Build per-interface error configs by merging global and interface-specific errors.
259 Returns:
260 Dict mapping interface name to merged ErrorsConfig.
261 """
262 errors_by_interface: dict[str, ErrorsConfig] = {}
263 for name in ("features", "metrics", "gap"):
264 interface = getattr(validated, name, None)
265 interface_errors = interface.errors if interface else None
266 errors_by_interface[name] = _merge_errors(validated.errors, interface_errors)
267 return errors_by_interface
270def _enrich_delta_with_pairwise(
271 validated: JobConfig,
272 delta_data: dict[str, Any],
273 delta_metrics_table: pa.Table,
274) -> None:
275 """Add a ``source_target`` column to delta data if pairwise output is configured."""
276 if not (validated.gap and validated.gap.outputs and getattr(validated.gap.outputs, "pairwise", False)): 276 ↛ 277line 276 didn't jump to line 277 because the condition on line 276 was never true
277 return
278 source_target_values = [
279 f"{delta_metrics_table.column('selection_source')[i].as_py()}"
280 f"_{delta_metrics_table.column('selection_target')[i].as_py()}"
281 for i in range(delta_metrics_table.num_rows)
282 ]
283 delta_data["source_target"] = pa.array(source_target_values)
286def run(config: dict[str, Any]) -> None:
287 """
288 Execute a job from a validated configuration dictionary.
290 The config is validated against JobConfig and must follow the v2 structure:
291 - dataloaders: Contains loaders list and optional storage.
292 - features: Optional interface with outputs and processors list.
293 - metrics: Optional interface with outputs and processors list.
294 - gap: Optional interface with outputs and processors list.
295 """
296 if not config:
297 raise ValueError("Job requires a configuration dictionary.")
299 validated = JobConfig.model_validate(config)
301 dataloaders_registry = PluginLoadedRegistry.get_dataloaders_registry()
302 features_registry = PluginLoadedRegistry.get_features_registry()
303 metrics_registry = PluginLoadedRegistry.get_metrics_registry()
304 gap_registry = PluginLoadedRegistry.get_gap_registry()
305 outputs_registry = PluginLoadedRegistry.get_outputwriter_registry()
307 # Initialize dataloaders from list format
308 dataloader_dicts = [loader.model_dump() for loader in validated.dataloaders.loaders]
309 compute = _resolve_compute_config(validated)
310 dl_storage = validated.dataloaders.storage.model_dump() if validated.dataloaders.storage else None
311 for dl in dataloader_dicts:
312 dl["threads"] = compute.threads
313 if dl_storage and not dl.get("storage"): 313 ↛ 314line 313 didn't jump to line 314 because the condition on line 313 was never true
314 dl["storage"] = dl_storage
315 dataloaders = _init_components_from_list(dataloader_dicts, dataloaders_registry, "dataloader")
317 # Resolve storage config: interface override takes precedence over job-level
318 def _resolve_storage(interface: Any) -> dict[str, Any] | None:
319 if interface and interface.storage: 319 ↛ 320line 319 didn't jump to line 320 because the condition on line 319 was never true
320 result: dict[str, Any] = interface.storage.model_dump()
321 return result
322 if validated.storage: 322 ↛ 323line 322 didn't jump to line 323 because the condition on line 322 was never true
323 result = validated.storage.model_dump()
324 return result
325 return None
327 # Initialize processors from all interfaces
328 features_processors = _init_processors_from_interface(
329 validated.features,
330 features_registry,
331 _resolve_storage(validated.features),
332 )
333 metrics_processors = _init_processors_from_interface(
334 validated.metrics,
335 metrics_registry,
336 _resolve_storage(validated.metrics),
337 )
338 gap_processors = _init_processors_from_interface(
339 validated.gap,
340 gap_registry,
341 _resolve_storage(validated.gap),
342 )
344 # Initialize output writers from interfaces
345 features_output = _init_interface_outputs(
346 validated.features,
347 outputs_registry,
348 "features",
349 _resolve_storage(validated.features),
350 )
351 metrics_output = _init_interface_outputs(
352 validated.metrics,
353 outputs_registry,
354 "metrics",
355 _resolve_storage(validated.metrics),
356 )
357 delta_output = _init_interface_outputs(
358 validated.gap,
359 outputs_registry,
360 "delta",
361 _resolve_storage(validated.gap),
362 )
364 # Configure logging based on compute.log_level
365 if compute.log_level: 365 ↛ 370line 365 didn't jump to line 370 because the condition on line 365 was always true
366 log_level = compute.log_level.upper()
367 level = getattr(logging, log_level)
368 logging.basicConfig(level=level)
370 job = DatasetJob(
371 dataloaders=dataloaders,
372 features_processors=features_processors,
373 metrics_processors=metrics_processors,
374 gap_processors=gap_processors,
375 features_output=features_output,
376 progress_bar=compute.progress_bar,
377 threads=compute.threads,
378 errors_by_interface=_build_errors_by_interface(validated),
379 compute_seed=compute.seed,
380 compute_device=compute.device,
381 compute_max_memory=compute.max_memory,
382 )
384 dataselection_metrics_list, delta_metrics_table = job.run()
386 if metrics_output:
387 metrics_output.write_metrics_dict(dataselection_metrics_list)
388 _safe_flush(metrics_output)
390 if delta_output and delta_metrics_table:
391 delta_data = {col: delta_metrics_table.column(col) for col in delta_metrics_table.column_names}
392 _enrich_delta_with_pairwise(validated, delta_data, delta_metrics_table)
393 delta_output.write_table("delta", delta_data)
394 _safe_flush(delta_output)
397if __name__ == "__main__": 397 ↛ 398line 397 didn't jump to line 398 because the condition on line 397 was never true
398 execute()