BatchalignHK 0.8.0__tar.gz → 0.8.0.post2__tar.gz

{batchalignhk-0.8.0 → batchalignhk-0.8.0.post2}/BatchalignHK.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: BatchalignHK
-Version: 0.8.0
+Version: 0.8.0.post2
 Summary: Python Speech Language Sample Analysis
 Author: Brian MacWhinney, Houjun Liu
 Author-email: macw@cmu.edu, houjun@cmu.edu

{batchalignhk-0.8.0 → batchalignhk-0.8.0.post2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: BatchalignHK
-Version: 0.8.0
+Version: 0.8.0.post2
 Summary: Python Speech Language Sample Analysis
 Author: Brian MacWhinney, Houjun Liu
 Author-email: macw@cmu.edu, houjun@cmu.edu

{batchalignhk-0.8.0 → batchalignhk-0.8.0.post2}/batchalign/cli/dispatch.py

@@ -33,6 +33,7 @@ import time
 import traceback
 import logging as L
 baL = L.getLogger('batchalign')
+import psutil
 
 warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')
 
@@ -47,103 +48,169 @@ def _get_worker_pipeline(command, lang, num_speakers, **kwargs):
                                                 lang=lang, num_speakers=num_speakers, **kwargs)
     return _worker_pipeline
 
-def _worker_task(file_info, command, lang, num_speakers, loader_info, writer_info, progress_queue=None, **kwargs):
+def _worker_task(file_info, command, lang, num_speakers, loader_info, writer_info, progress_queue=None, verbose=0, **kwargs):
     """The task executed in each worker process."""
     import sys
     import os
     import tempfile
-    
+    import logging
+
     file, output = file_info
     pid = os.getpid()
-    
-    # Use a temporary file to capture ALL output at the FD level
-    # This is the most robust way to prevent interleaved output
-    with tempfile.TemporaryFile(mode='w+') as log_file:
+    rss_start = None
+    rss_end = None
+    rss_peak = None
+
+    def _safe_rss():
+        try:
+            import psutil
+            return psutil.Process(pid).memory_info().rss
+        except Exception:
+            return None
+
+    def _safe_peak_rss():
+        try:
+            import resource
+            peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+            if peak is None:
+                return None
+            # ru_maxrss is KB on Linux, bytes on macOS; normalize to bytes.
+            return int(peak * 1024) if peak < 1024 * 1024 * 1024 else int(peak)
+        except Exception:
+            return None
+
+    rss_start = _safe_rss()
+
+    # Configure logging in this worker process
+    if verbose >= 1:
+        # Ensure basicConfig is called so logging works
+        logging.basicConfig(format="%(levelname)s - %(name)s - %(message)s", level=logging.ERROR)
+
+    # Configure batchalign logger level in this worker process
+    baL = logging.getLogger('batchalign')
+    if verbose == 0:
+        baL.setLevel(logging.WARN)
+    elif verbose == 1:
+        baL.setLevel(logging.INFO)
+    else:
+        baL.setLevel(logging.DEBUG)
+
+    # Always capture output to avoid interleaving with progress rendering.
+    should_capture = True
+
+    if should_capture:
+        # Use a temporary file to capture ALL output at the FD level
+        # This is the most robust way to prevent interleaved output
+        log_file = tempfile.TemporaryFile(mode='w+')
         old_stdout_fd = os.dup(sys.stdout.fileno())
         old_stderr_fd = os.dup(sys.stderr.fileno())
-        
-        try:
-            # Redirect FD 1 and 2 to our temp file
-            os.dup2(log_file.fileno(), sys.stdout.fileno())
-            os.dup2(log_file.fileno(), sys.stderr.fileno())
-            
-            pipeline = _get_worker_pipeline(command, lang, num_speakers, **kwargs)
 
-            def progress_callback(completed, total, tasks):
-                if not progress_queue:
-                    return
-                try:
-                    progress_queue.put((file, completed, total, tasks))
-                except Exception:
-                    pass
-            
-            # For now, we'll re-import what we need
-            from batchalign.formats.chat import CHATFile
-            
-            # Morphosyntax specific loader/writer logic moved here for picklability
-            if command == "morphotag":
-                # Extract morphotag-specific arguments from kwargs
-                mwt = kwargs.pop("mwt", {})
-                retokenize = kwargs.pop("retokenize", False)
-                skipmultilang = kwargs.pop("skipmultilang", False)
-                
-                cf = CHATFile(path=os.path.abspath(file), special_mor_=True)
-                doc = cf.doc
-                if str(cf).count("%mor") > 0:
-                    doc.ba_special_["special_mor_notation"] = True
-                
-                # Prepare arguments for the pipeline
-                pipeline_kwargs = {
-                    "retokenize": retokenize,
-                    "skipmultilang": skipmultilang,
-                    "mwt": mwt
-                }
-                # Add any remaining kwargs
-                pipeline_kwargs.update(kwargs)
-                
-                # Process
-                doc = pipeline(doc, callback=progress_callback, **pipeline_kwargs)
-                
-                # Write
-                CHATFile(doc=doc, special_mor_=doc.ba_special_.get("special_mor_notation", False)).write(output)
-            
-            # Add other commands as needed, or use a more generic registry
-            elif command == "align":
-                cf = CHATFile(path=os.path.abspath(file))
-                doc = cf.doc
-                kw = {"pauses": kwargs.get("pauses", False)}
-                doc = pipeline(doc, callback=progress_callback, **kw)
-                CHATFile(doc=doc).write(output, write_wor=kwargs.get("wor", True))
-            
-            else:
-                loader, writer = loader_info, writer_info
-                doc = loader(os.path.abspath(file))
-                kw = {}
-                if isinstance(doc, tuple) and len(doc) > 1:
-                    doc, kw = doc
-                doc = pipeline(doc, callback=progress_callback, **kw)
-                writer(doc, output)
-            
-            # Flush everything before reading back
+        # Redirect FD 1 and 2 to our temp file
+        os.dup2(log_file.fileno(), sys.stdout.fileno())
+        os.dup2(log_file.fileno(), sys.stderr.fileno())
+
+    try:
+        pipeline = _get_worker_pipeline(command, lang, num_speakers, **kwargs)
+
+        def progress_callback(completed, total, tasks):
+            if not progress_queue:
+                return
+            try:
+                progress_queue.put((file, completed, total, tasks))
+            except Exception:
+                pass
+
+        # For now, we'll re-import what we need
+        from batchalign.formats.chat import CHATFile
+
+        # Morphosyntax specific loader/writer logic moved here for picklability
+        if command == "morphotag":
+            # Extract morphotag-specific arguments from kwargs
+            mwt = kwargs.pop("mwt", {})
+            retokenize = kwargs.pop("retokenize", False)
+            skipmultilang = kwargs.pop("skipmultilang", False)
+
+            cf = CHATFile(path=os.path.abspath(file), special_mor_=True)
+            doc = cf.doc
+            if str(cf).count("%mor") > 0:
+                doc.ba_special_["special_mor_notation"] = True
+
+            # Prepare arguments for the pipeline
+            pipeline_kwargs = {
+                "retokenize": retokenize,
+                "skipmultilang": skipmultilang,
+                "mwt": mwt
+            }
+            # Add any remaining kwargs
+            pipeline_kwargs.update(kwargs)
+
+            # Process
+            doc = pipeline(doc, callback=progress_callback, **pipeline_kwargs)
+
+            # Write
+            CHATFile(doc=doc, special_mor_=doc.ba_special_.get("special_mor_notation", False)).write(output)
+
+        # Add other commands as needed, or use a more generic registry
+        elif command == "align":
+            cf = CHATFile(path=os.path.abspath(file))
+            doc = cf.doc
+            kw = {"pauses": kwargs.get("pauses", False)}
+            doc = pipeline(doc, callback=progress_callback, **kw)
+            CHATFile(doc=doc).write(output, write_wor=kwargs.get("wor", True))
+
+        else:
+            loader, writer = loader_info, writer_info
+            doc = loader(os.path.abspath(file))
+            kw = {}
+            if isinstance(doc, tuple) and len(doc) > 1:
+                doc, kw = doc
+            doc = pipeline(doc, callback=progress_callback, **kw)
+            writer(doc, output)
+
+        # Flush and read captured output if we were capturing
+        if should_capture:
             sys.stdout.flush()
             sys.stderr.flush()
             log_file.seek(0)
             captured = log_file.read()
-            
-            return file, None, None, captured
-        except Exception as e:
-            # Flush everything before reading back
+        else:
+            captured = ""
+
+        rss_end = _safe_rss()
+        rss_peak = _safe_peak_rss()
+        mem_info = {
+            "pid": pid,
+            "rss_start": rss_start,
+            "rss_end": rss_end,
+            "rss_peak": rss_peak,
+        }
+        return file, None, None, captured, mem_info
+    except Exception as e:
+        # Flush and read captured output if we were capturing
+        if should_capture:
             sys.stdout.flush()
             sys.stderr.flush()
             log_file.seek(0)
             captured = log_file.read()
-            return file, traceback.format_exc(), e, captured
-        finally:
-            # Restore original FDs
+        else:
+            captured = ""
+        rss_end = _safe_rss()
+        rss_peak = _safe_peak_rss()
+        mem_info = {
+            "pid": pid,
+            "rss_start": rss_start,
+            "rss_end": rss_end,
+            "rss_peak": rss_peak,
+        }
+        return file, traceback.format_exc(), e, captured, mem_info
+    finally:
+        # Restore original FDs only if we redirected them
+        if should_capture:
             os.dup2(old_stdout_fd, sys.stdout.fileno())
             os.dup2(old_stderr_fd, sys.stderr.fileno())
             os.close(old_stdout_fd)
             os.close(old_stderr_fd)
+            log_file.close()
 
 # this dictionary maps what commands are executed
 # against what BatchalignPipeline tasks are actually ran 
@@ -253,6 +320,7 @@ def _dispatch(command, lang, num_speakers,
     file_pairs = list(zip(files, outputs))
     file_pairs.sort(key=lambda fo: os.path.getsize(fo[0]) if os.path.exists(fo[0]) else 0, reverse=True)
     files, outputs = zip(*file_pairs) if file_pairs else ([], [])
+    file_sizes = {f: os.path.getsize(f) if os.path.exists(f) else 0 for f in files}
 
     C.print(f"\nMode: [blue]{command}[/blue]; got [bold cyan]{len(files)}[/bold cyan] transcript{'s' if len(files) > 1 else ''} to process from {in_dir}:\n")
 
@@ -287,8 +355,66 @@ def _dispatch(command, lang, num_speakers,
     # create the spinner
     prog = Progress(SpinnerColumn(), *Progress.get_default_columns()[:-1],
                     TimeElapsedColumn(),
-                    TextColumn("[cyan]{task.fields[processor]}[/cyan]"), console=C)
+                    TextColumn("[magenta]{task.fields[mem]}[/magenta]"),
+                    TextColumn("[cyan]{task.fields[processor]}[/cyan]"),
+                    console=C, refresh_per_second=5)
     errors = []
+    mem_records = {}
+    mem_samples = []
+    last_low_mem_warn = 0.0
+
+    def _format_bytes(count, precision=2):
+        if count is None:
+            return "unknown"
+        units = ["B", "KB", "MB", "GB", "TB"]
+        idx = 0
+        size = float(count)
+        while size >= 1024 and idx < len(units) - 1:
+            size /= 1024
+            idx += 1
+        if idx == 0:
+            return f"{int(size)}{units[idx]}"
+        return f"{size:.{precision}f}{units[idx]}"
+
+    def _mem_label(base, available=None, low_mem=False):
+        parts = [base]
+        if available is not None:
+            parts.append(f"avail {_format_bytes(available, precision=1)}")
+        if low_mem:
+            parts.append("LOW MEM")
+        return " | ".join(parts)
+
+    def _system_memory():
+        try:
+            vm = psutil.virtual_memory()
+            return vm.total, vm.available
+        except Exception:
+            return None, None
+
+    def _memory_reserve(total):
+        if total is None:
+            return None
+        return max(int(total * 0.10), 2 * 1024 * 1024 * 1024)
+
+    def _estimate_worker_bytes(file_size):
+        if not mem_samples:
+            return 512 * 1024 * 1024
+        ratios = [mem / size for size, mem in mem_samples if size and mem]
+        if not ratios:
+            return 512 * 1024 * 1024
+        ratios.sort()
+        median_ratio = ratios[len(ratios) // 2]
+        est = int(median_ratio * file_size)
+        return max(512 * 1024 * 1024, min(est, 6 * 1024 * 1024 * 1024))
+
+    def _should_throttle(est_bytes):
+        total, available = _system_memory()
+        if total is None or available is None:
+            return False, total, available
+        reserve = _memory_reserve(total)
+        if reserve is None:
+            return False, total, available
+        return (available - est_bytes) < reserve, total, available
 
     try:
         with prog as prog:
@@ -296,8 +422,9 @@ def _dispatch(command, lang, num_speakers,
             task_totals = {}
 
             for f in files:
-                tasks[f] = prog.add_task(Path(f).name, start=False, total=1, processor="Waiting...")
+                tasks[f] = prog.add_task(Path(f).name, start=False, total=1, processor="Waiting...", mem="queued")
                 task_totals[f] = 1
+                prog.start_task(tasks[f])
 
             def drain_progress_queue():
                 if not progress_queue:
@@ -313,10 +440,16 @@ def _dispatch(command, lang, num_speakers,
                         continue
                     task_total = max(int(total) if total else task_totals.get(file, 1), 1)
                     task_totals[file] = task_total
+                    total_mem, available_mem = _system_memory()
+                    reserve = _memory_reserve(total_mem)
+                    low_mem = False
+                    if reserve is not None and available_mem is not None:
+                        low_mem = available_mem < reserve
                     prog.update(tasks[file],
                                 total=task_total,
                                 completed=min(int(completed), task_total),
-                                processor=render_stage(stage_tasks))
+                                processor=render_stage(stage_tasks),
+                                mem=_mem_label("running", available_mem, low_mem))
 
             with concurrent.futures.ProcessPoolExecutor(max_workers=num_workers) as executor:
                 worker_func = partial(_worker_task,
@@ -326,13 +459,57 @@ def _dispatch(command, lang, num_speakers,
                                       loader_info=None,
                                       writer_info=None,
                                       progress_queue=progress_queue,
+                                      verbose=ctx.obj["verbose"],
                                       **kwargs)
 
-                future_to_file = {executor.submit(worker_func, (f, o)): f for f, o in zip(files, outputs)}
+                file_iter = iter(zip(files, outputs))
+                future_to_file = {}
+
+                def submit_one(file_path, output_path):
+                    future = executor.submit(worker_func, (file_path, output_path))
+                    future_to_file[future] = file_path
+                    est_bytes = _estimate_worker_bytes(file_sizes.get(file_path, 0))
+                    total_mem, available_mem = _system_memory()
+                    reserve = _memory_reserve(total_mem)
+                    low_mem = False
+                    if reserve is not None and available_mem is not None:
+                        low_mem = available_mem < reserve
+                    prog.update(
+                        tasks[file_path],
+                        processor="Processing...",
+                        mem=_mem_label(f"est {_format_bytes(est_bytes)}", available_mem, low_mem),
+                    )
 
-                for f in files:
-                    prog.start_task(tasks[f])
-                    prog.update(tasks[f], processor="Processing...")
+                def schedule_available():
+                    nonlocal last_low_mem_warn
+                    while len(future_to_file) < num_workers:
+                        try:
+                            next_file, next_output = next(file_iter)
+                        except StopIteration:
+                            break
+                        est_bytes = _estimate_worker_bytes(file_sizes.get(next_file, 0))
+                        throttle, total, available = _should_throttle(est_bytes)
+                        if throttle and future_to_file:
+                            now = time.time()
+                            if now - last_low_mem_warn > 10:
+                                reserve = _memory_reserve(total)
+                                prog.console.print(
+                                    f"[bold yellow]Low memory[/bold yellow]: "
+                                    f"{_format_bytes(available)} free, "
+                                    f"{_format_bytes(reserve)} reserve. "
+                                    f"Throttling new workers."
+                                )
+                                last_low_mem_warn = now
+                            break
+                        if throttle and not future_to_file:
+                            prog.console.print(
+                                f"[bold yellow]Low memory[/bold yellow]: "
+                                f"{_format_bytes(available)} free. "
+                                "Continuing with a single worker."
+                            )
+                        submit_one(next_file, next_output)
+
+                schedule_available()
 
                 pending = set(future_to_file.keys())
                 while pending:
@@ -345,8 +522,9 @@ def _dispatch(command, lang, num_speakers,
 
                     for future in done:
                         file = future_to_file[future]
+                        future_to_file.pop(future, None)
                         try:
-                            res_file, trcbk, e, captured = future.result()
+                            res_file, trcbk, e, captured, mem_info = future.result()
                             final_total = max(task_totals.get(file, 1), 1)
                             if e:
                                 prog.update(tasks[file], total=final_total, completed=final_total, processor="[bold red]FAIL[/bold red]")
@@ -354,12 +532,25 @@ def _dispatch(command, lang, num_speakers,
                             else:
                                 prog.update(tasks[file], total=final_total, completed=final_total, processor="[bold green]DONE[/bold green]")
                                 if ctx.obj["verbose"] >= 1 and captured.strip():
-                                    errors.append((res_file, "Logs only (Success)", None, captured))
+                                    prog.console.print(f"[bold blue]INFO[/bold blue] on file [italic]{Path(file).name}[/italic]:\n{escape(captured.strip())}\n")
+                            if mem_info:
+                                mem_records[file] = mem_info
+                                peak = mem_info.get("rss_peak") or mem_info.get("rss_end")
+                                if peak:
+                                    mem_samples.append((file_sizes.get(file, 0), peak))
+                                    total_mem, available_mem = _system_memory()
+                                    reserve = _memory_reserve(total_mem)
+                                    low_mem = False
+                                    if reserve is not None and available_mem is not None:
+                                        low_mem = available_mem < reserve
+                                    prog.update(tasks[file], mem=_mem_label(_format_bytes(peak), available_mem, low_mem))
                         except Exception as e:
                             final_total = max(task_totals.get(file, 1), 1)
                             prog.update(tasks[file], total=final_total, completed=final_total, processor="[bold red]FAIL[/bold red]")
                             errors.append((file, traceback.format_exc(), e, ""))
 
+                    schedule_available()
+                    pending = set(future_to_file.keys())
                 drain_progress_queue()
     finally:
         if manager:
@@ -383,6 +574,16 @@ def _dispatch(command, lang, num_speakers,
     else:
         C.print(f"\nAll done. Results saved to {out_dir}!\n")
 
+    if mem_records and ctx.obj["verbose"] >= 1:
+        C.print("\nMemory usage per file (worker RSS peak):")
+        for file, info in mem_records.items():
+            rel_path = os.path.relpath(str(Path(file).absolute()), in_dir)
+            peak = info.get("rss_peak") or info.get("rss_end")
+            C.print(f"- {rel_path}: {_format_bytes(peak)}")
+        total, available = _system_memory()
+        if total is not None and available is not None:
+            C.print(f"\nSystem memory available: {_format_bytes(available)} / {_format_bytes(total)}")
+
     if ctx.obj["verbose"] > 1:
         C.end_capture()
 

batchalignhk-0.8.0.post2/batchalign/version

@@ -0,0 +1,3 @@
+0.8.0-post.2
+Jan 15th, 2025
+Memory Safegaurds

batchalignhk-0.8.0/batchalign/version

@@ -1,3 +0,0 @@
-0.8.0
-Jan 13th, 2025
-Speed