Blablablab/audio-classification
0
1"""2Directory Watcher Module3 4This module provides functionality for loading annotation instances from a directory5and optionally watching that directory for new or modified files. When watching is6enabled, a background thread periodically scans the directory and dynamically loads7new instances or updates existing ones.8 9The module supports the same file formats as the standard data_files configuration:10JSON, JSONL, CSV, and TSV.11 12Configuration:13 data_directory: str - Path to the directory containing data files14 watch_data_directory: bool - Whether to watch for changes (default: False)15 watch_poll_interval: float - Seconds between directory scans (default: 5.0)16 data_directory_encoding: str - File encoding for directory files (default: "utf-8")17 18Example config:19 data_directory: "./data/incoming"20 watch_data_directory: true21 watch_poll_interval: 10.022"""23 24from __future__ import annotations25 26import json27import logging28import os29import threading30import glob31from dataclasses import dataclass, field32from typing import Dict, List, Optional, Set, Tuple, TYPE_CHECKING33 34try:35 import pandas as pd36 HAS_PANDAS = True37except ImportError:38 HAS_PANDAS = False39 40if TYPE_CHECKING:41 from potato.item_state_management import ItemStateManager42 43logger = logging.getLogger(__name__)44 45# Singleton instance with thread-safe initialization46DIRECTORY_WATCHER: Optional['DirectoryWatcher'] = None47_DIRECTORY_WATCHER_LOCK = threading.Lock()48 49 50@dataclass51class FileState:52 """53 Tracks the state of a watched file.54 55 Attributes:56 file_path: Absolute path to the file57 last_modified: Last modification time (os.path.getmtime)58 file_size: File size in bytes59 instance_ids: Set of instance IDs loaded from this file60 last_error: Last error message if processing failed, None otherwise61 last_processed: Timestamp of last successful processing62 """63 file_path: str64 last_modified: float = 0.065 file_size: int = 066 instance_ids: Set[str] = field(default_factory=set)67 last_error: Optional[str] = None68 last_processed: Optional[float] = None69 70 71class DirectoryWatcher:72 """73 Watches a directory for new or modified data files and loads them as annotation instances.74 75 This class provides two modes of operation:76 1. Static loading: Load all files from a directory at startup (load_directory())77 2. Dynamic watching: Continuously monitor for changes (start_watching())78 79 The watcher tracks which instances came from which file, enabling proper handling80 of file modifications (updating existing instances rather than creating duplicates).81 82 Thread Safety:83 All public methods are thread-safe. The internal state is protected by84 a reentrant lock (_lock) to allow safe concurrent access from the main85 application thread and the background watching thread.86 87 Attributes:88 data_directory: Path to the directory to watch89 poll_interval: Seconds between directory scans90 id_key: Key in data items containing the unique instance ID91 text_key: Key in data items containing the text to annotate92 """93 94 # Supported file extensions95 SUPPORTED_EXTENSIONS = ('.json', '.jsonl', '.csv', '.tsv')96 97 def __init__(self, config: dict, item_state_manager: 'ItemStateManager'):98 """99 Initialize the directory watcher.100 101 Args:102 config: Configuration dictionary containing:103 - data_directory: Path to watch104 - watch_poll_interval: Seconds between scans (default: 5.0)105 - data_directory_encoding: File encoding (default: "utf-8")106 - item_properties.id_key: Key for instance IDs107 - item_properties.text_key: Key for text content108 item_state_manager: The ItemStateManager instance to add items to109 110 Raises:111 ValueError: If data_directory is not configured or doesn't exist112 """113 self.data_directory = config.get("data_directory")114 if not self.data_directory:115 raise ValueError("data_directory must be configured")116 117 # Resolve relative paths based on task_dir if available118 if not os.path.isabs(self.data_directory):119 task_dir = config.get("task_dir", "")120 if task_dir:121 self.data_directory = os.path.join(task_dir, self.data_directory)122 self.data_directory = os.path.abspath(self.data_directory)123 124 if not os.path.isdir(self.data_directory):125 raise ValueError(f"data_directory does not exist or is not a directory: {self.data_directory}")126 127 self.poll_interval = config.get("watch_poll_interval", 5.0)128 self.encoding = config.get("data_directory_encoding", "utf-8")129 self.id_key = config["item_properties"]["id_key"]130 self.text_key = config["item_properties"]["text_key"]131 132 self._item_state_manager = item_state_manager133 134 # File tracking state135 self._file_states: Dict[str, FileState] = {}136 self._instance_to_file: Dict[str, str] = {} # instance_id -> file_path137 138 # Threading139 self._lock = threading.RLock()140 self._stop_event = threading.Event()141 self._watch_thread: Optional[threading.Thread] = None142 143 logger.info(f"DirectoryWatcher initialized for: {self.data_directory}")144 145 def load_directory(self) -> int:146 """147 Load all supported files from the data directory.148 149 This method performs an initial scan of the directory and loads all150 instances from supported file formats. It should be called once at151 startup before start_watching().152 153 Returns:154 int: Total number of instances loaded155 156 Side Effects:157 - Populates ItemStateManager with loaded instances158 - Updates internal file tracking state159 """160 total_added = 0161 162 with self._lock:163 files = self._scan_directory()164 logger.info(f"Found {len(files)} supported files in {self.data_directory}")165 166 for file_path in files:167 try:168 added, updated = self._process_file(file_path)169 total_added += added170 if added > 0 or updated > 0:171 logger.info(f"Loaded {file_path}: {added} added, {updated} updated")172 except Exception as e:173 logger.error(f"Error loading {file_path}: {e}")174 175 logger.info(f"Directory load complete: {total_added} total instances loaded")176 return total_added177 178 def start_watching(self) -> None:179 """180 Start the background directory watching thread.181 182 The watching thread will periodically scan the directory for new or183 modified files and process them. Use stop() to terminate the thread.184 185 Note:186 This method is idempotent - calling it multiple times has no effect187 if the thread is already running.188 """189 with self._lock:190 if self._watch_thread is not None and self._watch_thread.is_alive():191 logger.warning("Directory watcher thread is already running")192 return193 194 self._stop_event.clear()195 self._watch_thread = threading.Thread(196 target=self._watch_loop,197 name="DirectoryWatcher",198 daemon=True199 )200 self._watch_thread.start()201 logger.info(f"Directory watching started (poll interval: {self.poll_interval}s)")202 203 def stop(self) -> None:204 """205 Stop the directory watching thread gracefully.206 207 This method signals the watching thread to stop and waits for it208 to terminate (up to 5 seconds). It's safe to call this method209 even if watching was never started.210 """211 self._stop_event.set()212 213 if self._watch_thread is not None and self._watch_thread.is_alive():214 self._watch_thread.join(timeout=5.0)215 if self._watch_thread.is_alive():216 logger.warning("Directory watcher thread did not stop gracefully")217 else:218 logger.info("Directory watcher stopped")219 220 self._watch_thread = None221 222 def get_stats(self) -> dict:223 """224 Get statistics about the directory watcher state.225 226 Returns:227 dict: Statistics including:228 - data_directory: Path being watched229 - is_watching: Whether the watch thread is running230 - poll_interval: Seconds between scans231 - files_tracked: Number of files being tracked232 - total_instances: Total instances loaded from this directory233 - files: List of file states with details234 """235 with self._lock:236 return {237 "data_directory": self.data_directory,238 "is_watching": self._watch_thread is not None and self._watch_thread.is_alive(),239 "poll_interval": self.poll_interval,240 "files_tracked": len(self._file_states),241 "total_instances": len(self._instance_to_file),242 "files": [243 {244 "path": fs.file_path,245 "last_modified": fs.last_modified,246 "instance_count": len(fs.instance_ids),247 "last_error": fs.last_error248 }249 for fs in self._file_states.values()250 ]251 }252 253 def force_rescan(self) -> Tuple[int, int]:254 """255 Force an immediate rescan of the directory.256 257 This method can be called to trigger an immediate check for changes258 without waiting for the next poll interval.259 260 Returns:261 Tuple[int, int]: (total_added, total_updated) counts262 """263 return self._scan_and_process()264 265 def _watch_loop(self) -> None:266 """267 Main watching loop that runs in the background thread.268 269 This loop periodically scans the directory for changes and processes270 any new or modified files. It continues until stop() is called.271 """272 logger.debug("Directory watch loop started")273 274 while not self._stop_event.is_set():275 try:276 added, updated = self._scan_and_process()277 if added > 0 or updated > 0:278 logger.info(f"Directory scan: {added} instances added, {updated} updated")279 except Exception as e:280 logger.error(f"Error in directory watch loop: {e}", exc_info=True)281 282 # Wait for the poll interval or until stopped283 self._stop_event.wait(timeout=self.poll_interval)284 285 logger.debug("Directory watch loop ended")286 287 def _scan_and_process(self) -> Tuple[int, int]:288 """289 Scan for changed files and process them.290 291 Returns:292 Tuple[int, int]: (total_added, total_updated) counts293 """294 total_added = 0295 total_updated = 0296 297 with self._lock:298 current_files = set(self._scan_directory())299 tracked_files = set(self._file_states.keys())300 301 # Find new and potentially modified files302 for file_path in current_files:303 try:304 stat = os.stat(file_path)305 current_mtime = stat.st_mtime306 current_size = stat.st_size307 except OSError as e:308 logger.warning(f"Cannot stat file {file_path}: {e}")309 continue310 311 # Check if file is new or modified312 if file_path not in self._file_states:313 # New file314 added, updated = self._process_file(file_path)315 total_added += added316 total_updated += updated317 else:318 # Check if modified319 fs = self._file_states[file_path]320 if current_mtime > fs.last_modified or current_size != fs.file_size:321 logger.debug(f"File modified: {file_path}")322 added, updated = self._process_file(file_path)323 total_added += added324 total_updated += updated325 326 # Note: We don't remove instances when files are deleted - this preserves327 # annotations that may have been made on those instances.328 removed_files = tracked_files - current_files329 for file_path in removed_files:330 logger.info(f"File removed (instances preserved): {file_path}")331 # Keep the file state but mark that the file is gone332 if file_path in self._file_states:333 self._file_states[file_path].last_error = "File removed from directory"334 335 return total_added, total_updated336 337 def _scan_directory(self) -> List[str]:338 """339 Scan the data directory for supported files.340 341 Returns:342 List[str]: List of absolute paths to supported files343 """344 files = []345 for ext in self.SUPPORTED_EXTENSIONS:346 pattern = os.path.join(self.data_directory, f"*{ext}")347 files.extend(glob.glob(pattern))348 return sorted(files)349 350 def _process_file(self, file_path: str) -> Tuple[int, int]:351 """352 Process a single data file, adding or updating instances.353 354 Args:355 file_path: Absolute path to the file to process356 357 Returns:358 Tuple[int, int]: (added_count, updated_count)359 360 Side Effects:361 - Updates ItemStateManager with new/updated instances362 - Updates file tracking state363 """364 added_count = 0365 updated_count = 0366 367 try:368 instances = self._parse_file(file_path)369 stat = os.stat(file_path)370 371 # Get or create file state372 if file_path not in self._file_states:373 self._file_states[file_path] = FileState(file_path=file_path)374 375 fs = self._file_states[file_path]376 new_instance_ids: Set[str] = set()377 378 for instance_data in instances:379 # Validate ID key exists380 if self.id_key not in instance_data:381 logger.warning(f"Missing id_key '{self.id_key}' in {file_path}, skipping instance")382 continue383 384 instance_id = str(instance_data[self.id_key])385 new_instance_ids.add(instance_id)386 387 # Check if text_key is missing (warning only)388 if self.text_key not in instance_data:389 logger.warning(f"Missing text_key '{self.text_key}' for instance {instance_id}")390 391 # Add or update the instance392 if self._item_state_manager.has_item(instance_id):393 # Update existing instance394 if self._item_state_manager.update_item(instance_id, instance_data):395 updated_count += 1396 logger.debug(f"Updated instance: {instance_id}")397 else:398 # Add new instance399 try:400 self._item_state_manager.add_item(instance_id, instance_data)401 self._instance_to_file[instance_id] = file_path402 added_count += 1403 logger.debug(f"Added instance: {instance_id}")404 except ValueError as e:405 logger.error(f"Failed to add instance {instance_id}: {e}")406 407 # Update file state408 fs.last_modified = stat.st_mtime409 fs.file_size = stat.st_size410 fs.instance_ids = new_instance_ids411 fs.last_error = None412 fs.last_processed = stat.st_mtime413 414 except Exception as e:415 logger.error(f"Error processing file {file_path}: {e}")416 if file_path in self._file_states:417 self._file_states[file_path].last_error = str(e)418 else:419 self._file_states[file_path] = FileState(420 file_path=file_path,421 last_error=str(e)422 )423 424 return added_count, updated_count425 426 def _parse_file(self, file_path: str) -> List[dict]:427 """428 Parse a data file and return a list of instance dictionaries.429 430 Args:431 file_path: Absolute path to the file432 433 Returns:434 List[dict]: List of instance data dictionaries435 436 Raises:437 ValueError: If file format is unsupported or parsing fails438 """439 ext = os.path.splitext(file_path)[1].lower()440 441 if ext in ('.json', '.jsonl'):442 return self._parse_json_file(file_path)443 elif ext == '.csv':444 return self._parse_csv_file(file_path, separator=',')445 elif ext == '.tsv':446 return self._parse_csv_file(file_path, separator='\t')447 else:448 raise ValueError(f"Unsupported file format: {ext}")449 450 def _parse_json_file(self, file_path: str) -> List[dict]:451 """452 Parse a JSON or JSONL file.453 454 Supports both:455 - JSONL format: One JSON object per line456 - JSON format: Single JSON array or object per line457 458 Args:459 file_path: Path to the JSON/JSONL file460 461 Returns:462 List[dict]: List of parsed instance dictionaries463 """464 instances = []465 466 with open(file_path, 'rt', encoding=self.encoding) as f:467 for line_no, line in enumerate(f, 1):468 line = line.strip()469 if not line:470 continue471 472 try:473 item = json.loads(line)474 475 # Handle both single objects and arrays476 if isinstance(item, list):477 instances.extend(item)478 else:479 instances.append(item)480 481 except json.JSONDecodeError as e:482 raise ValueError(483 f"Invalid JSON at line {line_no} in {file_path}: {e}"484 ) from e485 486 return instances487 488 def _parse_csv_file(self, file_path: str, separator: str) -> List[dict]:489 """490 Parse a CSV or TSV file.491 492 Args:493 file_path: Path to the CSV/TSV file494 separator: Column separator (',' for CSV, '\t' for TSV)495 496 Returns:497 List[dict]: List of row dictionaries498 499 Raises:500 ImportError: If pandas is not available501 ValueError: If required columns are missing502 """503 if not HAS_PANDAS:504 raise ImportError(505 "pandas is required for CSV/TSV file support. "506 "Install it with: pip install pandas"507 )508 509 df = pd.read_csv(file_path, sep=separator, encoding=self.encoding)510 511 # Validate ID column exists512 if self.id_key not in df.columns:513 raise ValueError(f"ID column '{self.id_key}' not found in {file_path}")514 515 # Convert ID column to string516 df[self.id_key] = df[self.id_key].astype(str)517 518 # Convert text column to string if present519 if self.text_key in df.columns:520 df[self.text_key] = df[self.text_key].astype(str)521 522 return df.to_dict('records')523 524 525def init_directory_watcher(config: dict) -> Optional[DirectoryWatcher]:526 """527 Initialize the global DirectoryWatcher singleton if data_directory is configured.528 529 This function creates a DirectoryWatcher instance if the configuration includes530 a data_directory setting. The watcher is initialized but not started - call531 load_directory() and optionally start_watching() after initialization.532 533 Args:534 config: Configuration dictionary535 536 Returns:537 DirectoryWatcher: The initialized watcher, or None if not configured538 539 Note:540 Thread-safe initialization using double-checked locking pattern.541 """542 global DIRECTORY_WATCHER543 544 # Check if data_directory is configured545 if "data_directory" not in config:546 return None547 548 # Double-checked locking for thread safety549 if DIRECTORY_WATCHER is None:550 with _DIRECTORY_WATCHER_LOCK:551 if DIRECTORY_WATCHER is None:552 from potato.item_state_management import get_item_state_manager553 ism = get_item_state_manager()554 DIRECTORY_WATCHER = DirectoryWatcher(config, ism)555 556 return DIRECTORY_WATCHER557 558 559def get_directory_watcher() -> Optional[DirectoryWatcher]:560 """561 Get the global DirectoryWatcher singleton instance.562 563 Returns:564 DirectoryWatcher: The singleton instance, or None if not initialized565 """566 return DIRECTORY_WATCHER567 568 569def clear_directory_watcher() -> None:570 """571 Clear the global DirectoryWatcher singleton (for testing).572 573 This function stops any running watch thread and clears the global instance.574 Thread-safe.575 """576 global DIRECTORY_WATCHER577 578 with _DIRECTORY_WATCHER_LOCK:579 if DIRECTORY_WATCHER is not None:580 DIRECTORY_WATCHER.stop()581 DIRECTORY_WATCHER = None582