""" Usage logger that pushes alignment runs to a HF Dataset repo. Uses a ParquetScheduler (subclass of CommitScheduler) to buffer rows in memory and periodically write+upload parquet files with embedded audio to the Hub. Error logs use a separate CommitScheduler with JSONL files. Falls back to local-only logging if schedulers can't initialize. Scheduler creation is deferred to first use so that background threads don't interfere with ZeroGPU's startup function scan. """ import hashlib import json import threading from datetime import datetime from pathlib import Path from typing import Any, Dict, List, Optional, Union from uuid import uuid4 import numpy as np # ========================================================================= # Directory setup # ========================================================================= LOG_DIR = Path("usage_logs") LOG_DIR.mkdir(parents=True, exist_ok=True) ERROR_DIR = LOG_DIR / "errors" ERROR_DIR.mkdir(parents=True, exist_ok=True) ERROR_LOG_PATH = ERROR_DIR / f"error_log-{uuid4()}.jsonl" # ========================================================================= # ParquetScheduler class definition (no instances created at import time) # ========================================================================= _HAS_DEPS = False try: import pyarrow as pa import pyarrow.parquet as pq from huggingface_hub import CommitScheduler from config import USAGE_LOG_DATASET_REPO, USAGE_LOG_PUSH_INTERVAL_MINUTES _HAS_DEPS = True except Exception: pass # HF features schema V2 — flat columns for filtering, JSON columns for detail _ALIGNER_SCHEMA: Dict[str, Dict[str, str]] = { # Identity & request "audio": {"_type": "Audio"}, "audio_id": {"_type": "Value", "dtype": "string"}, "timestamp": {"_type": "Value", "dtype": "string"}, "user_id": {"_type": "Value", "dtype": "string"}, "endpoint": {"_type": "Value", "dtype": "string"}, # Top-level metrics (flat for filtering) "audio_duration_s": {"_type": "Value", "dtype": "float64"}, "total_time": {"_type": "Value", "dtype": "float64"}, # Session flags "resegmented": {"_type": "Value", "dtype": "bool"}, "retranscribed": {"_type": "Value", "dtype": "bool"}, # JSON detail columns "settings": {"_type": "Value", "dtype": "string"}, "profiling": {"_type": "Value", "dtype": "string"}, "gpu": {"_type": "Value", "dtype": "string"}, "results_summary": {"_type": "Value", "dtype": "string"}, "reciter_stats": {"_type": "Value", "dtype": "string"}, # Per-run details, timestamps & error "results_detailed": {"_type": "Value", "dtype": "string"}, "word_timestamps": {"_type": "Value", "dtype": "string"}, "char_timestamps": {"_type": "Value", "dtype": "string"}, "error": {"_type": "Value", "dtype": "string"}, } if _HAS_DEPS: class ParquetScheduler(CommitScheduler): """Buffers rows in memory and uploads a parquet file each interval. Audio values are stored as file paths in the row dict; on push they are read as bytes and embedded in the parquet using the HF Audio struct. """ def __init__( self, *, repo_id: str, schema: Optional[Dict[str, Dict[str, str]]] = None, every: Union[int, float] = 5, path_in_repo: Optional[str] = "data", repo_type: Optional[str] = "dataset", private: bool = False, ) -> None: super().__init__( repo_id=repo_id, folder_path="dummy", # not used — we upload directly every=every, path_in_repo=path_in_repo, repo_type=repo_type, private=private, ) self._rows: List[Dict[str, Any]] = [] self._schema = schema def append(self, row: Dict[str, Any]) -> None: with self.lock: self._rows.append(row) def push_to_hub(self) -> None: with self.lock: rows = self._rows self._rows = [] if not rows: return print(f"[USAGE_LOG] Pushing {len(rows)} alignment row(s) to Hub.") schema: Dict[str, Dict] = dict(self._schema) if self._schema else {} paths_to_cleanup: List[Path] = [] for row in rows: for key, value in row.items(): if key not in schema: schema[key] = _infer_schema(key, value) if value is not None and schema[key].get("_type") in ("Image", "Audio"): file_path = Path(value) if file_path.is_file(): row[key] = { "path": file_path.name, "bytes": file_path.read_bytes(), } paths_to_cleanup.append(file_path) else: row[key] = None for row in rows: for feature in schema: if feature not in row: row[feature] = None table = pa.Table.from_pylist(rows) # Cast null-typed columns to string so all parquet shards share # the same Arrow schema (prevents HF viewer concat errors). for i, field in enumerate(table.schema): if pa.types.is_null(field.type): table = table.set_column( i, field.name, pa.array([None] * len(table), type=pa.string()), ) table = table.replace_schema_metadata( {"huggingface": json.dumps({"info": {"features": schema}})} ) archive = None try: import tempfile archive = tempfile.NamedTemporaryFile(suffix=".parquet", delete=False) pq.write_table( table, archive.name, row_group_size=1, write_page_index=True, ) self.api.upload_file( repo_id=self.repo_id, repo_type=self.repo_type, revision=self.revision, path_in_repo=f"{self.path_in_repo}/{uuid4()}.parquet", path_or_fileobj=archive.name, ) print("[USAGE_LOG] Parquet commit completed.") except Exception as e: print(f"[USAGE_LOG] Failed to upload parquet: {e}") finally: if archive: archive.close() Path(archive.name).unlink(missing_ok=True) for path in paths_to_cleanup: path.unlink(missing_ok=True) def _infer_schema(key: str, value: Any) -> Dict[str, str]: if "image" in key: return {"_type": "Image"} if "audio" in key: return {"_type": "Audio"} if isinstance(value, bool): return {"_type": "Value", "dtype": "bool"} if isinstance(value, int): return {"_type": "Value", "dtype": "int64"} if isinstance(value, float): return {"_type": "Value", "dtype": "float64"} if isinstance(value, bytes): return {"_type": "Value", "dtype": "binary"} return {"_type": "Value", "dtype": "string"} # ========================================================================= # Lazy scheduler initialization (deferred to first use) # ========================================================================= _aligner_scheduler = None _error_scheduler = None _schedulers_initialized = False _init_lock = threading.Lock() _fallback_lock = threading.Lock() def _ensure_schedulers() -> None: global _aligner_scheduler, _error_scheduler, _schedulers_initialized if _schedulers_initialized: return with _init_lock: if _schedulers_initialized: return _schedulers_initialized = True if not _HAS_DEPS: print("[USAGE_LOG] Dependencies missing (local-only mode).") return try: _aligner_scheduler = ParquetScheduler( repo_id=USAGE_LOG_DATASET_REPO, schema=_ALIGNER_SCHEMA, every=USAGE_LOG_PUSH_INTERVAL_MINUTES, path_in_repo="data", repo_type="dataset", private=True, ) _error_scheduler = CommitScheduler( repo_id=USAGE_LOG_DATASET_REPO, repo_type="dataset", folder_path=ERROR_DIR, path_in_repo="data/errors", private=True, every=USAGE_LOG_PUSH_INTERVAL_MINUTES, ) except Exception as e: print(f"[USAGE_LOG] Scheduler init failed (local-only mode): {e}") # ========================================================================= # Helpers # ========================================================================= def _get_error_lock(): _ensure_schedulers() if _error_scheduler is not None: return _error_scheduler.lock return _fallback_lock def get_user_id(request) -> str: """SHA-256 hash (12-char) of client IP from a gr.Request, or 'unknown'.""" try: headers = request.headers ip = ( headers.get("x-forwarded-for", "").split(",")[0].strip() or headers.get("x-real-ip", "") or "" ) if not ip: return "unknown" return hashlib.sha256(ip.encode()).hexdigest()[:12] except Exception: return "unknown" def _compute_audio_id(audio: np.ndarray, ts: datetime) -> str: """Content hash (16-char) + compact timestamp.""" audio_hash = hashlib.sha256(audio.tobytes()).hexdigest()[:16] return f"{audio_hash}:{ts.strftime('%Y%m%dT%H%M%S')}" def _encode_audio_ogg(audio: np.ndarray, sample_rate: int, audio_id: str) -> str: """Encode audio to a temp OGG Vorbis file; returns the file path. Uses Vorbis instead of Opus because libsndfile (used by the HF dataset viewer) has buggy Opus support and crashes on many valid Opus files. """ import soundfile as sf import subprocess tmp_dir = LOG_DIR / "tmp_audio" tmp_dir.mkdir(parents=True, exist_ok=True) safe_id = audio_id.replace(":", "-") wav_path = tmp_dir / f"{safe_id}.wav" ogg_path = tmp_dir / f"{safe_id}.ogg" sf.write(str(wav_path), audio, sample_rate, format="WAV") try: subprocess.run( ["ffmpeg", "-y", "-i", str(wav_path), "-c:a", "libvorbis", "-q:a", "2", "-ar", "16000", "-ac", "1", str(ogg_path)], capture_output=True, check=True, ) finally: wav_path.unlink(missing_ok=True) return str(ogg_path) def _sync_row_to_scheduler(row: Dict[str, Any]) -> None: """Ensure *row* is represented in the scheduler buffer. gr.State may deserialize the dict (creating a copy), and push_to_hub detaches rows from the buffer. This helper finds the original row by audio_id and updates it, or re-appends if it was already pushed. """ if _aligner_scheduler is None: return audio_id = row.get("audio_id") if not audio_id: return with _aligner_scheduler.lock: for buffered in _aligner_scheduler._rows: if buffered.get("audio_id") == audio_id: # Update the buffered row in-place (handles gr.State copies) buffered.update(row) return # Row was already pushed — re-append (audio file may be gone, that's ok) _aligner_scheduler._rows.append(row) # ========================================================================= # Public logging API # ========================================================================= def log_alignment( *, audio: np.ndarray, sample_rate: int, request=None, # Flat fields audio_duration_s: float, endpoint: str, total_time: float, # JSON groups (caller passes dicts) settings: dict, profiling: dict, gpu: dict, results_summary: dict, reciter_stats: dict, # Segments log_segments: List[dict], _async: bool = False, ) -> Optional[Dict[str, Any]]: """Log an alignment run. Returns the row dict reference for in-place mutation. The returned dict can be stored in gr.State and mutated on resegment/retranscribe/timestamps before the scheduler pushes. When _async=True, the expensive OGG encoding and SHA256 hash run in a background daemon thread. The row is appended to the scheduler immediately (with audio=None) and updated in-place once the thread finishes. """ _ensure_schedulers() try: ts = datetime.now() user_id = get_user_id(request) if request else "unknown" # Build the segments JSON: array of run objects segments_runs = [{ "settings": settings, "segments": log_segments, }] if _async: audio_id = f"{uuid4().hex[:16]}:{ts.strftime('%Y%m%dT%H%M%S')}" else: audio_id = _compute_audio_id(audio, ts) row: Dict[str, Any] = { "audio": None, "audio_id": audio_id, "timestamp": ts.isoformat(timespec="seconds"), "user_id": user_id, "endpoint": endpoint, # Top-level metrics "audio_duration_s": audio_duration_s, "total_time": total_time, # Session flags "resegmented": False, "retranscribed": False, # JSON detail columns "settings": json.dumps(settings), "profiling": json.dumps(profiling), "gpu": json.dumps(gpu), "results_summary": json.dumps(results_summary), "reciter_stats": json.dumps(reciter_stats), # Per-run details & error "results_detailed": json.dumps(segments_runs), "word_timestamps": None, "char_timestamps": None, "error": None, } def _encode_and_append(): """Encode OGG and register row with scheduler/fallback.""" try: row["audio"] = _encode_audio_ogg(audio, sample_rate, audio_id) except Exception as e: print(f"[USAGE_LOG] OGG encoding failed: {e}") if _aligner_scheduler is None: _write_fallback(row) if _async: if _aligner_scheduler is not None: _aligner_scheduler.append(row) threading.Thread(target=_encode_and_append, daemon=True).start() else: _encode_and_append() if _aligner_scheduler is not None: _aligner_scheduler.append(row) return row except Exception as e: print(f"[USAGE_LOG] Failed to log alignment: {e}") return None def update_alignment_row( row: Dict[str, Any], *, action: str, # Flat fields audio_duration_s: float, endpoint: str, total_time: float, # JSON groups settings: dict, profiling: dict, gpu: dict, results_summary: dict, reciter_stats: dict, # Segments log_segments: List[dict], ) -> None: """Mutate an existing row dict in-place and ensure it's in the scheduler buffer. After mutation, syncs the row into the scheduler's buffer so the update is captured even if gr.State returned a deserialized copy or if the original row was already pushed to Hub. Args: row: The dict returned by log_alignment(), stored in gr.State. action: "resegment" or "retranscribe". """ try: # Overwrite flat fields row["audio_duration_s"] = audio_duration_s row["endpoint"] = endpoint row["total_time"] = total_time # Overwrite JSON detail columns row["settings"] = json.dumps(settings) row["profiling"] = json.dumps(profiling) row["gpu"] = json.dumps(gpu) row["results_summary"] = json.dumps(results_summary) row["reciter_stats"] = json.dumps(reciter_stats) # Set session flag if action == "resegment": row["resegmented"] = True elif action == "retranscribe": row["retranscribed"] = True # Append new run to results_detailed array segments_runs = json.loads(row.get("results_detailed") or "[]") segments_runs.append({ "settings": settings, "segments": log_segments, }) row["results_detailed"] = json.dumps(segments_runs) # Sync with scheduler buffer — the row from gr.State may be a # deserialized copy, or the original may have already been pushed. _sync_row_to_scheduler(row) except Exception as e: print(f"[USAGE_LOG] Failed to update alignment row: {e}") def update_word_timestamps( row: Dict[str, Any], word_timestamps_json: str, char_timestamps_json: Optional[str] = None, ) -> None: """Set word and char timestamps fields on an existing row and sync to scheduler.""" try: row["word_timestamps"] = word_timestamps_json if char_timestamps_json is not None: row["char_timestamps"] = char_timestamps_json _sync_row_to_scheduler(row) except Exception as e: print(f"[USAGE_LOG] Failed to update word timestamps: {e}") def _write_fallback(row: Dict[str, Any]) -> None: """Local-only fallback: write JSONL (without audio).""" fallback_path = LOG_DIR / "alignments_fallback.jsonl" with _fallback_lock: with fallback_path.open("a") as f: fallback_row = {k: v for k, v in row.items() if k != "audio"} json.dump(fallback_row, f) f.write("\n")