Require functions at the top.

Add periodic memory logging

Author: jeromekelleher

pyproject.toml

@@ -138,7 +138,7 @@ line-length = 89
 indent-width = 4
 
 [tool.ruff.lint]
-select = ["E", "F", "B", "W", "I", "N", "UP", "A", "PT"]
+select = ["E", "F", "B", "W", "I", "N", "UP", "A", "PT", "PLC0415"]
 
 ignore = [
     "A001", "A002", "RUF",

tsinfer/pipeline.py

@@ -375,6 +375,8 @@ def _process_group(
             disable=not progress,
         )
         results = []
+        last_mem_log = time_.monotonic()
+        mem_log_interval = 30  # seconds
         for job, result in match_iter:
             if match_fh is not None:
                 doc = {
@@ -396,6 +398,18 @@ def _process_group(
                 match_fh.write(json.dumps(doc) + "\n")
             results.append((job, result))
             pbar.update(1)
+            now = time_.monotonic()
+            if now - last_mem_log >= mem_log_interval:
+                rss_mib = psutil.Process().memory_info().rss / (1024 * 1024)
+                reader.log_cache_state()
+                logger.info(
+                    "Group %d progress: %d/%d haplotypes, RSS=%.1f MiB",
+                    group_idx,
+                    len(results),
+                    len(job_list),
+                    rss_mib,
+                )
+                last_mem_log = now
         pbar.close()
         paired_results = sorted(results, key=lambda pair: pair[0].haplotype_index)
 

tsinfer/vcz.py

@@ -42,6 +42,7 @@
 from __future__ import annotations
 
 import dataclasses
+import functools
 import logging
 import math
 import pathlib
@@ -2325,8 +2326,6 @@ def __init__(
         sources,
         ancestral_state,
     ):
-        from functools import reduce
-
         if isinstance(sources, config.Source):
             sources = [sources]
 
@@ -2384,7 +2383,7 @@ def __init__(
         if len(valid_per_source) == 1:
             all_positions = valid_per_source[0]
         else:
-            all_positions = reduce(np.union1d, valid_per_source)
+            all_positions = functools.reduce(np.union1d, valid_per_source)
         self._positions = all_positions.astype(np.int32)
 
         # Build source_has_site mask via searchsorted