Merge pull request #49

Enhance logging, error reporting, and add tests for benchmarks

Author: yokotoka

src/palabra_ai/benchmark/__main__.py

@@ -4,6 +4,7 @@
 import asyncio
 import bisect
 import re
+import sys
 import traceback
 import wave
 from base64 import b64decode
@@ -248,6 +249,27 @@ def truncate_text(text: str, max_len: int = 25) -> str:
     remaining = len(text) - max_len
     return f"{text[:max_len]}...(+{remaining})"
 
+def flatten_dict_to_paths(d: dict | list, prefix: str = "") -> list[tuple[str, Any]]:
+    """Flatten nested dict/list to materialized paths like ('a.b.c', value)"""
+    result = []
+
+    if isinstance(d, dict):
+        for key, value in d.items():
+            new_prefix = f"{prefix}.{key}" if prefix else key
+            if isinstance(value, (dict, list)):
+                result.extend(flatten_dict_to_paths(value, new_prefix))
+            else:
+                result.append((new_prefix, value))
+    elif isinstance(d, list):
+        for i, value in enumerate(d):
+            new_prefix = f"{prefix}.{i}"
+            if isinstance(value, (dict, list)):
+                result.extend(flatten_dict_to_paths(value, new_prefix))
+            else:
+                result.append((new_prefix, value))
+
+    return result
+
 def format_report(report: Report, io_data: IoData, source_lang: str, target_lang: str, in_file: str, out_file: str, config: Config) -> str:
     """Format report as text with tables and histogram"""
     lines = []
@@ -273,6 +295,23 @@ def format_report(report: Report, io_data: IoData, source_lang: str, target_lang
             lines.append(f"TTS autotempo: ✅ on ({queue_config.min_tempo}-{queue_config.max_tempo})")
         else:
             lines.append(f"TTS autotempo: ❌ off")
+
+    # CONFIG - exact same dict that goes to SetTaskMessage
+    lines.append("")
+    lines.append("CONFIG (sent to set_task)")
+    lines.append("-" * 80)
+    config_dict = config.to_dict()  # Same as SetTaskMessage.from_config uses!
+    config_paths = flatten_dict_to_paths(config_dict)
+
+    table = PrettyTable()
+    table.field_names = ["Key", "Value"]
+    table.align["Key"] = "l"
+    table.align["Value"] = "l"
+
+    for key, value in config_paths:
+        table.add_row([key, value])
+
+    lines.append(str(table))
     lines.append("")
 
     # Metrics summary table
@@ -374,8 +413,12 @@ def main():
             timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
 
             # Save sysinfo immediately at startup
+            sysinfo = get_system_info()
+            sysinfo["command"] = " ".join(sys.argv)
+            sysinfo["argv"] = sys.argv
+            sysinfo["cwd"] = str(Path.cwd())
             sysinfo_path = output_dir / f"{timestamp}_bench_sysinfo.json"
-            sysinfo_path.write_bytes(to_json(get_system_info(), True))
+            sysinfo_path.write_bytes(to_json(sysinfo, True))
 
         # Get audio duration for progress tracking
         with av.open(str(audio_path)) as container:
@@ -426,6 +469,11 @@ def on_transcription(msg):
             config.debug = True
             config.log_file = str(output_dir / f"{timestamp}_bench.log")
 
+            # Save exact config that goes to set_task (SetTaskMessage.from_config uses to_dict)
+            config_dict = config.to_dict()
+            config_path = output_dir / f"{timestamp}_bench_config.json"
+            config_path.write_bytes(to_json(config_dict, True))
+
         # Create progress bar with language info
         progress_bar[0] = tqdm(
             total=100,
@@ -475,18 +523,26 @@ def on_transcription(msg):
                     print("This usually means:")
                     print("  - User interrupted with Ctrl+C")
                     print("  - Task was cancelled by timeout")
-                    print("  - Internal cancellation due to error\n")
+                    print("  - Internal cancellation due to error")
+                    print("  - One of the subtasks failed and caused cascade cancellation\n")
+
+                    # For CancelledError, show ALL logs to understand what happened
+                    if result.log_data and result.log_data.logs:
+                        print(f"Full logs (all {len(result.log_data.logs)} entries):")
+                        for log_line in result.log_data.logs:
+                            print(log_line, end='')
+                        print()
                 else:
                     print(f"\n{'='*80}")
                     print(f"BENCHMARK FAILED: {exc_type}: {exc_msg}")
                     print(f"{'='*80}\n")
 
-                # Print more logs (100 instead of 50)
-                if result.log_data and result.log_data.logs:
-                    print("Last 100 log entries:")
-                    for log_line in result.log_data.logs[-100:]:
-                        print(log_line, end='')
-                    print()
+                    # For other errors, show last 100
+                    if result.log_data and result.log_data.logs:
+                        print("Last 100 log entries:")
+                        for log_line in result.log_data.logs[-100:]:
+                            print(log_line, end='')
+                        print()
 
                 # Print traceback from exception if available
                 if hasattr(result.exc, '__traceback__') and result.exc.__traceback__:
@@ -542,21 +598,47 @@ def on_transcription(msg):
         print("\n" + report_text)
 
     except Exception as e:
+        # Capture traceback IMMEDIATELY - must be done in except block!
+        tb_string = traceback.format_exc()
+
         # Print full traceback to console
         print(f"\n{'='*80}")
         print("BENCHMARK CRASHED - FULL TRACEBACK:")
         print(f"{'='*80}\n")
-        traceback.print_exc()
+        print(tb_string)
 
         # Save error to file if output directory exists
         if output_dir and timestamp:
             try:
                 error_file = output_dir / f"{timestamp}_bench_error.txt"
-                error_file.write_text(f"Benchmark Error:\n\n{traceback.format_exc()}")
+                error_file.write_text(f"Benchmark Error:\n\n{tb_string}")
                 print(f"\nError details saved to: {error_file}")
             except Exception as save_error:
                 print(f"Failed to save error file: {save_error}")
 
+        # Try to save partial report/audio even on error (for debugging)
+        if output_dir and timestamp and result and result.io_data:
+            try:
+                print("\nAttempting to save partial results for debugging...")
+
+                # Try to parse report
+                report, in_audio, out_audio = Report.parse(result.io_data)
+
+                # Save report files
+                report_path = output_dir / f"{timestamp}_bench_report_partial.json"
+                report_path.write_bytes(to_json(report, True))
+                print(f"✓ Partial report saved to: {report_path}")
+
+                # Save audio (always when --out is specified)
+                in_wav = output_dir / f"{timestamp}_bench_in_partial.wav"
+                out_wav = output_dir / f"{timestamp}_bench_out_partial.wav"
+                save_wav(in_audio, in_wav, result.io_data.in_sr, result.io_data.channels)
+                save_wav(out_audio, out_wav, result.io_data.out_sr, result.io_data.channels)
+                print(f"✓ Partial audio saved: {in_wav.name}, {out_wav.name}")
+
+            except Exception as save_err:
+                print(f"Could not save partial results: {save_err}")
+
         # Re-raise the exception
         raise
 

src/palabra_ai/client.py

@@ -115,7 +115,7 @@ async def _run_with_result(manager: Manager) -> RunResult:
                 await manager.task
                 ok = True
             except asyncio.CancelledError as e:
-                debug("Manager task was cancelled")
+                exception("Manager task was cancelled")
                 exc = e
             except BaseException as e:
                 exception("Error in manager task")
@@ -234,9 +234,14 @@ async def process(
             success("🎉🎉🎉 Translation completed 🎉🎉🎉")
 
         except* asyncio.CancelledError as eg:
-            debug("TaskGroup received CancelledError")
-            # Re-raise CancelledError
-            raise eg.exceptions[0] from eg
+            # Check if this was graceful completion (internal shutdown)
+            if manager._graceful_completion:
+                debug("Graceful completion detected - not propagating CancelledError")
+                # Don't re-raise - this is expected graceful shutdown
+            else:
+                # External cancellation (Ctrl+C, timeout, error)
+                debug("External cancellation - propagating CancelledError")
+                raise eg.exceptions[0] from eg
         except* Exception as eg:
             excs = unwrap_exceptions(eg)
             excs_wo_cancel = [

src/palabra_ai/task/io/base.py

@@ -7,6 +7,7 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
+from websockets.exceptions import ConnectionClosed
 
 from palabra_ai.audio import AudioFrame
 from palabra_ai.constant import BOOT_TIMEOUT, BYTES_PER_SAMPLE, SLEEP_INTERVAL_LONG
@@ -120,8 +121,17 @@ async def in_msg_sender(self):
                     return
                 raw = to_json(msg)
                 debug(f"<- {raw[0:30]} / {msg.dbg_delta=}")
-                await self.send_message(raw)
-                self.io_events.append(IoEvent(msg._dbg, raw))
+                try:
+                    await self.send_message(raw)
+                    self.io_events.append(IoEvent(msg._dbg, raw))
+                except Exception as e:
+                    # Connection closed during shutdown is OK
+                    if isinstance(e, ConnectionClosed):
+                        debug(
+                            f"Connection closed while sending message (OK during shutdown): {e}"
+                        )
+                        return
+                    raise
 
     async def do(self):
         """Main processing loop - read audio chunks and push them."""

src/palabra_ai/task/io/ws.py

@@ -5,7 +5,7 @@
 
 from palabra_ai.audio import AudioFrame
 from palabra_ai.enum import Channel, Direction, Kind
-from palabra_ai.message import Dbg, IoEvent
+from palabra_ai.message import Dbg, EosMessage, IoEvent, Message
 from palabra_ai.task.io.base import Io
 from palabra_ai.util.logger import debug, trace
 from palabra_ai.util.timing import get_perf_ts, get_utc_ts
@@ -34,17 +34,20 @@ async def send_frame(self, frame: AudioFrame, raw: bytes | None = None) -> None:
         debug(f"<- {frame} / {frame.dbg_delta=}")
         if self.global_start_perf_ts is None:
             self.init_global_start_ts()
-            frame._dbg.dawn_ts = 0.0
+            if hasattr(frame, "_dbg") and frame._dbg:
+                frame._dbg.dawn_ts = 0.0
         await self.ws.send(raw)
 
     def new_input_frame(self) -> AudioFrame:
         return AudioFrame.create(*self.cfg.mode.for_input_audio_frame)
 
     async def ws_receiver(self):
-        from palabra_ai.message import EosMessage, Message
-
         try:
             async for raw_msg in self.ws:
+                # Check for cancellation/stopper eagerly
+                if self.stopper:
+                    debug("ws_receiver: stopper detected, exiting")
+                    break
                 perf_ts = get_perf_ts()
                 utc_ts = get_utc_ts()
                 if self.stopper or raw_msg is None:
@@ -115,6 +118,12 @@ async def boot(self):
 
     async def exit(self):
         """Clean up WebSocket connection"""
-        if self._ws_cm and self.ws:
-            await self._ws_cm.__aexit__(None, None, None)
+        # Websocket should already be closed in do(), just cleanup context manager
+        if self._ws_cm:
+            try:
+                await self._ws_cm.__aexit__(None, None, None)
+                debug(f"{self.name}: WebSocket context manager cleaned up")
+            except Exception as e:
+                debug(f"{self.name}: Error exiting websocket context: {e}")
+
         self.ws = None

src/palabra_ai/task/manager.py

@@ -22,7 +22,7 @@
 from palabra_ai.task.logger import Logger
 from palabra_ai.task.stat import Stat
 from palabra_ai.task.transcription import Transcription
-from palabra_ai.util.logger import debug, success, warning
+from palabra_ai.util.logger import debug, exception, success, warning
 
 
 @dataclass
@@ -50,6 +50,7 @@ class Manager(Task):
     _debug_mode: bool = field(default=True, init=False)
     _transcriptions_shown: set = field(default_factory=set, init=False)
     _show_banner_loop: asyncio.Task | None = field(default=None, init=False)
+    _graceful_completion: bool = field(default=False, init=False)
 
     def __post_init__(self):
         self.stat = Stat(manager=self, cfg=self.cfg)
@@ -215,6 +216,7 @@ async def do(self):
                     debug(f"🔚 {self.name}.do() sleep cancelled, exiting...")
                 debug(f"🔚 {self.name}.do() received EOF or stopper, exiting...")
                 success("🏁 Done! ⏻ Shutting down...")
+                self._graceful_completion = True
                 break
         +self.stopper  # noqa
         await self.graceful_exit()
@@ -247,7 +249,9 @@ async def shutdown_task(self, task, timeout=SHUTDOWN_TIMEOUT):
         try:
             await asyncio.wait_for(task._task, timeout=timeout)
         except TimeoutError:
-            debug(f"🔧 {self.name}.shutdown_task() {task.name} shutdown timeout!")
+            # TimeoutError during graceful shutdown is expected for some tasks (e.g. WebSocket)
+            # Only log as warning, not error
+            warning(f"{task.name} shutdown timeout ({timeout}s) - cancelling task")
             task._task.cancel()
             try:
                 await task._task
@@ -256,7 +260,7 @@ async def shutdown_task(self, task, timeout=SHUTDOWN_TIMEOUT):
         except asyncio.CancelledError:
             debug(f"🔧 {self.name}.shutdown_task() {task.name} shutdown cancelled!")
         except Exception as e:
-            debug(f"🔧 {self.name}.shutdown_task() {task.name} shutdown error: {e}")
+            exception(f"Cancelling {task.name} due to shutdown error: {e}")
             task._task.cancel()
             try:
                 await task._task
@@ -317,8 +321,8 @@ async def writer_mercy(self):
                 debug(f"🔧 {self.name}.writer_mercy() writer shutdown timeout!")
                 attempt += 1
                 if attempt >= max_attempts:
-                    debug(
-                        f"🔧 {self.name}.writer_mercy() max attempts reached, cancelling writer!"
+                    exception(
+                        f"Cancelling writer {self.writer.name} due to max shutdown attempts ({max_attempts}) reached"
                     )
                     self.writer._task.cancel()
                     try: