stream-limits 0.2.0__tar.gz → 0.3.0__tar.gz

{stream_limits-0.2.0 → stream_limits-0.3.0}/CHANGELOG.md

@@ -1,5 +1,40 @@
 # Changelog
 
+## 0.3.0
+
+- **A new function, `set_openmp_wait_policy`, controls how idle OpenMP threads
+  wait.** Programs that make many short parallel calls, such as numba loops run
+  once per batch of data, can run noticeably slower when the idle threads spin
+  while they wait for the next call, especially on processors that mix fast and
+  efficient cores. Telling the threads to sleep instead (`OMP_WAIT_POLICY=passive`,
+  the default here) fixed this in diff-stream: fits ran 2% to 15% faster on a
+  cloud server, and on a laptop with mixed cores 8 threads went from slower than
+  one thread to faster. Programs with a few long parallel calls, such as
+  hdfe-stream and logit-stream, run the same either way. On Windows, where
+  numba's OpenMP runtime ignores the setting, the call changes nothing.
+- **A package has to ask for it.** Importing stream-limits, and
+  `apply_scheduler_defaults`, still change nothing, because the setting reaches
+  every OpenMP library in the process (BLAS builds on OpenMP, scikit-learn,
+  LightGBM). A value you have set yourself is kept unless you pass
+  `override=True`.
+- **It must run before OpenMP starts.** The setting is read once, when the
+  OpenMP runtime loads; numba loads its runtime at its first parallel call, so
+  calling the function after importing numba is still in time. The function
+  reports when a runtime was already loaded (Linux), in which case the setting
+  reaches only programs started afterwards.
+
+## 0.2.1
+
+- **The watchdog no longer counts the pages of memory-mapped files.** A program
+  that memory-maps large scratch files, as hdfe-stream does, has those pages in
+  its resident memory, but the kernel can reclaim them, so a watchdog that
+  counted them stopped healthy programs: a 20-million-row hdfe-stream fit with
+  1.2 GB of its own memory reached 2.3 GB of RSS. The watchdog now limits
+  private memory (Linux, Windows; macOS still uses RSS), reported as
+  `peak_private_bytes`. `--memory-measure rss` restores the old behaviour.
+- On Windows, `os_peak_rss_bytes` is now the program's exact peak working set,
+  read when it exits, instead of the last value seen while sampling.
+
 ## 0.2.0
 
 - **Run a program under a memory limit and see how much memory it used.** The new

{stream_limits-0.2.0 → stream_limits-0.3.0}/CITATION.cff

@@ -7,8 +7,8 @@ authors:
   - family-names: Tucker
     given-names: Lee
     email: lee.c.tucker@census.gov
-version: 0.2.0
-date-released: 2026-11-04
+version: 0.3.0
+date-released: 2026-11-08
 license: MIT
 repository-code: "https://github.com/leetucker/stream-limits"
 url: "https://github.com/leetucker/stream-limits"

{stream_limits-0.2.0 → stream_limits-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stream-limits
-Version: 0.2.0
+Version: 0.3.0
 Summary: Cap the threads and memory of numeric Python libraries (Polars, numba, BLAS, pyarrow) at a scheduler's allocation, and measure peak memory
 Author: Lee Tucker
 License-Expression: MIT

{stream_limits-0.2.0 → stream_limits-0.3.0}/docs/launcher.md

@@ -49,7 +49,18 @@ The command is a list of arguments, run without a shell. GB are decimal
 | `monitor` | Linux, macOS, Windows | No limit: measures, with an optional timeout |
 
 **The watchdog** adds up the memory of the program and every process it has
-started, and stops them all when the total passes the limit. Because it samples,
+started, and stops them all when the total passes the limit. It counts
+*private* memory: memory that is not the pages of a file. A program that
+memory-maps a large file (a scratch array on disk, a mapped Parquet file) has
+the file's pages in its resident memory (RSS), but the kernel can drop them and
+read them back when memory runs short, so they are not counted against the
+limit. hdfe-stream is such a program: on a 20-million-row panel its RSS reached
+2.3 GB while its private memory stayed at 1.2 GB, the same as on 2.6 million
+rows. On Linux private memory is anonymous plus shared memory; on Windows it is
+private bytes (committed memory, which can be more than RSS). macOS has no cheap
+equivalent, so there the watchdog counts RSS. `--memory-measure rss`
+(`memory_measure="rss"`) counts every resident page on every platform, and the
+result's `memory_measure` says which was used. Because it samples,
 it can overshoot a little before it acts, and a peak shorter than the interval
 can slip between samples. In practice it is tight: on Linux, with a 1.5 GB limit,
 a fast 3 GB NumPy allocation, a 2.5 GB spike lasting 20 ms and a grandchild
@@ -74,7 +85,7 @@ running under the watchdog. macOS and Windows have no hard mode.
 | `ok` | The program finished with exit code 0 |
 | `command_failed` | It finished with another exit code |
 | `timeout` | It ran past `timeout` and was stopped |
-| `rss_limit_exceeded` | The watchdog stopped it |
+| `rss_limit_exceeded` | The watchdog stopped it (by the measure in `memory_measure`, despite the name) |
 | `monitoring_failed` | The watchdog could not read its memory and stopped it |
 | `kernel_oom` | The kernel killed a process for going over a hard limit |
 
@@ -83,16 +94,19 @@ process. It is reported even when the process killed was a child the program
 started and the program itself carried on and exited with 0, so check `status`
 and not only the exit code.
 
-Two peaks are reported:
+Three peaks are reported:
 
-- `os_peak_rss_bytes`: the largest memory of any single process, as the
-  operating system records it. On Linux and macOS it is exact, and covers the
-  program and every process it started and waited for (not one it left running
-  in the background). On Windows it is read while sampling. Use it when a test
-  needs a reliable peak.
-- `peak_rss_bytes`: the largest sampled total over all the program's processes
-  at once. Use it for a program with several worker processes. A program that
-  ends before the first sample reports 0.
+- `peak_private_bytes`: the largest sampled total of private memory over all the
+  program's processes (`None` on macOS). This is the one to use for a program
+  that memory-maps files, and the one the watchdog limits.
+- `os_peak_rss_bytes`: the largest RSS of any single process, as the operating
+  system records it. On Linux and macOS it is exact, and covers the program and
+  every process it started and waited for (not one it left running in the
+  background). On Windows it is the program's own peak working set, read when
+  it exits (not its child processes'). It counts mapped file pages. Use it when a test needs a reliable peak of a program that maps no
+  large files, or on macOS.
+- `peak_rss_bytes`: the largest sampled total RSS over all the program's
+  processes at once. A program that ends before the first sample reports 0.
 
 With a hard limit, `result.kernel` adds the kernel's own figures: the limit, any
 smaller limit set higher up, its peak (which counts file cache and kernel
@@ -230,7 +244,10 @@ and skip without one. The plugin registers the markers `resource` and
 Polars, NumPy and numba already holds 0.2-0.4 GB, so a small input under a 1.5 GB
 limit passes whatever the code does with it. To catch code that holds the data in
 memory, give the test an input larger in memory than the limit less that
-baseline, or compare `os_peak_rss_bytes` at N rows and at several times N.
+baseline, or compare peaks at N rows and at several times N:
+`peak_private_bytes` where it is available, since a program that memory-maps
+its scratch files has an RSS that grows with the data even when its private
+memory does not.
 
 Use the launcher only for tests about memory. Ordinary tests are better run in
 the test process, as before.

{stream_limits-0.2.0 → stream_limits-0.3.0}/docs/reference.md

@@ -15,6 +15,8 @@
 
 .. autofunction:: thread_scope
 
+.. autofunction:: set_openmp_wait_policy
+
 .. autofunction:: polars_spills
 ```
 

{stream_limits-0.2.0 → stream_limits-0.3.0}/docs/threads.md

@@ -81,3 +81,23 @@ allocation, read from the environment, cgroups or the scheduler (see
 [Schedulers and containers](schedulers.md)), and otherwise the machine's. A job that runs several
 processes should give each its share, for example
 `STREAM_LIMITS_MAX_THREADS=$(( NCPUS / 4 ))` for four.
+
+
+## How OpenMP threads wait
+
+OpenMP threads that finish a parallel region either spin, waiting for the next one, or
+sleep. Spinning is quicker to resume, but a program that makes many short parallel calls
+(numba's `prange` loops, one per batch of data) can run several times slower with it,
+notably on CPUs that mix performance and efficiency cores, and a little slower on others.
+Programs with few, long parallel regions are unaffected: hdfe-stream's solvers and
+logit-stream's fits ran as fast either way.
+
+{py:func}`stream_limits.set_openmp_wait_policy` sets `OMP_WAIT_POLICY` (to `"passive"` unless you
+say `"active"`). A package opts in by calling it, because the setting reaches every
+OpenMP library in the process (BLAS builds on OpenMP, scikit-learn, LightGBM); importing
+stream-limits, and `apply_scheduler_defaults`, leave it alone. A value you have set is kept unless
+you pass `override=True`. The runtime reads the variable when it loads, which for numba is
+at its first parallel call, so call the function before that. The result says when a runtime
+was already loaded (Linux), in which case the variable is set for child processes only.
+On Windows, numba uses Microsoft's OpenMP runtime, which does not read the variable, so the
+call changes nothing there.

{stream_limits-0.2.0 → stream_limits-0.3.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "stream-limits"
-version = "0.2.0"
+version = "0.3.0"
 description = "Cap the threads and memory of numeric Python libraries (Polars, numba, BLAS, pyarrow) at a scheduler's allocation, and measure peak memory"
 readme = "README.md"
 license = "MIT"

{stream_limits-0.2.0 → stream_limits-0.3.0}/stream_limits/__init__.py

@@ -87,10 +87,11 @@ import warnings
 from pathlib import Path
 from typing import Iterator, Mapping
 
-__version__ = "0.2.0"
+__version__ = "0.3.0"
 
 __all__ = [
     "set_limits", "limit_threads", "limit_memory", "apply_scheduler_defaults", "thread_scope",
+    "set_openmp_wait_policy",
     "detect_cpus", "detect_memory_gb", "describe", "MemoryBudget",
     "peak_rss_gb", "reset_peak_rss", "current_rss_gb",
 ]
@@ -790,13 +791,59 @@ def thread_scope(n: int | None = None) -> Iterator[None]:
         numba.set_num_threads(previous)
 
 
+def _openmp_loaded() -> bool | None:
+    """Whether an OpenMP runtime is already in this process (Linux only; None
+    elsewhere). It reads ``OMP_WAIT_POLICY`` once, when it loads."""
+    try:
+        maps = Path("/proc/self/maps").read_text()
+    except OSError:
+        return None
+    return any(name in maps for name in ("libgomp", "libiomp5", "libomp"))
+
+
+def set_openmp_wait_policy(policy: str = "passive", *, override: bool = False,
+                           environ=None) -> dict:
+    """Set how OpenMP threads wait between parallel regions, by ``OMP_WAIT_POLICY``.
+
+    With ``"passive"`` an idle thread sleeps until the next region; with
+    ``"active"`` it spins. A package whose compiled code makes many short
+    parallel calls (numba's ``prange`` loops, one per batch of data) can run
+    several times slower with spinning threads, notably on CPUs with
+    performance and efficiency cores, and a little slower on others. Code with
+    few, long parallel regions (hdfe-stream's solvers, logit-stream's fits)
+    is unaffected, and this function is not called by importing the package:
+    a package opts in, since the setting reaches every OpenMP library in the
+    process (BLAS builds on OpenMP, scikit-learn, LightGBM). Microsoft's OpenMP
+    runtime, the one numba uses on Windows, does not read the variable, so there
+    the call changes nothing.
+
+    A value the user has set is left alone unless `override`. The runtime reads
+    the variable once, when it loads (numba loads its OpenMP layer at the first
+    parallel call), so call this before the first parallel call; the result's
+    ``"late"`` is True when a runtime was already loaded (Linux), in which case
+    the variable is still set, for child processes, but this process is
+    unaffected. Returns ``{"policy": ..., "set": bool, "late": bool | None}``,
+    where ``"policy"`` is the value in effect."""
+    if policy not in ("passive", "active"):
+        raise ValueError(f"policy must be 'passive' or 'active', not {policy!r}")
+    environ = os.environ if environ is None else environ
+    current = environ.get("OMP_WAIT_POLICY")
+    if current and not override:
+        return {"policy": current.lower(), "set": False, "late": None}
+    late = _openmp_loaded() if environ is os.environ else None
+    environ["OMP_WAIT_POLICY"] = policy
+    return {"policy": policy, "set": True, "late": late}
+
+
+
 # ----------------------------------------------------- peak memory used
-def _windows_memory() -> tuple[int, int] | None:
+def _windows_memory(handle=None) -> tuple[int, int] | None:
     """This process's ``(working set, peak working set)`` in bytes, from the
     Windows API (``GetProcessMemoryInfo``), or None: off Windows, or if the call
     fails. The working set is the physical memory resident in the process, which
     is what ``VmRSS`` is on Linux, and the peak is a high-water mark the system
-    keeps."""
+    keeps. `handle`: another process's handle instead (the launcher reads a
+    child's peak after it exits, while its handle is still open)."""
     if sys.platform != "win32":
         return None
     try:
@@ -821,8 +868,8 @@ def _windows_memory() -> tuple[int, int] | None:
         kernel32.K32GetProcessMemoryInfo.restype = wintypes.BOOL
         counters = _Counters()
         counters.cb = ctypes.sizeof(_Counters)
-        if not kernel32.K32GetProcessMemoryInfo(kernel32.GetCurrentProcess(),
-                                                ctypes.byref(counters), counters.cb):
+        target = kernel32.GetCurrentProcess() if handle is None else handle
+        if not kernel32.K32GetProcessMemoryInfo(target, ctypes.byref(counters), counters.cb):
             return None
         return int(counters.WorkingSetSize), int(counters.PeakWorkingSetSize)
     except Exception:                                          # noqa: BLE001

{stream_limits-0.2.0 → stream_limits-0.3.0}/stream_limits/launcher.py

@@ -18,7 +18,7 @@ import time
 from collections.abc import Mapping, Sequence
 from dataclasses import asdict, dataclass
 
-from . import _THREAD_ENV
+from . import _THREAD_ENV, _windows_memory
 from .enforcement import MemoryCgroup, UnsupportedEnforcement
 
 
@@ -39,13 +39,23 @@ class RunResult:
     descendants_terminated: bool
     rss_sampling_complete: bool
     os_peak_rss_bytes: int | None = None
+    peak_private_bytes: int | None = None
+    memory_measure: str = "rss"
 
     # os_peak_rss_bytes: the largest resident set any single process reached,
     # as the OS records it. POSIX: the child's ru_maxrss from wait4, which
     # covers the child and the descendants it waited for (exact, not sampled;
     # not summed over the tree; background processes it left running are not
-    # included). Windows: the largest peak working set read while sampling.
-    # None when unavailable.
+    # included). Windows: the child's peak working set, read from its handle
+    # after it exits (not its descendants'; the largest sampled value if that
+    # read fails). None when unavailable. Like RSS, it counts pages of
+    # memory-mapped files.
+    #
+    # peak_private_bytes: the largest sampled total over the process tree of
+    # memory that is not a file's pages (see _private_bytes). None on macOS.
+    #
+    # memory_measure: what the watchdog compared with the limit, "private" or
+    # "rss" (asked for, or private memory is unavailable on this platform).
 
     def as_dict(self):
         return asdict(self)
@@ -58,6 +68,7 @@ def capabilities():
         "psutil_available": importlib.util.find_spec("psutil") is not None,
         "monitor": True,
         "watchdog": True,
+        "private_memory": private_memory_available(),
         "hard_backends": ["cgroup-v1", "cgroup-v2"] if sys.platform == "linux" else [],
         "hard_limit_accounting": "cgroup-charged memory"
         if sys.platform == "linux"
@@ -78,6 +89,36 @@ def thread_env(threads, env=None):
     return out
 
 
+def _private_bytes(process, info):
+    """Memory of one process that is not a file's pages, which the kernel can
+    drop and read back (memory-mapped scratch files, mapped libraries).
+    Linux: anonymous plus shared memory (RssAnon + RssShmem; shared memory
+    is not reclaimable without swap), else RSS less psutil's shared pages.
+    Windows: private bytes (committed memory, so it can exceed RSS). None
+    elsewhere (macOS has no cheap equivalent)."""
+    if sys.platform.startswith("linux"):
+        try:
+            found = {}
+            with open(f"/proc/{process.pid}/status", encoding="ascii") as fh:
+                for line in fh:
+                    if line.startswith(("RssAnon:", "RssShmem:")):
+                        key, value = line.split()[:2]
+                        found[key] = int(value) * 1024
+            if "RssAnon:" in found:
+                return found["RssAnon:"] + found.get("RssShmem:", 0)
+        except (OSError, ValueError):
+            pass
+        return max(0, info.rss - getattr(info, "shared", 0))
+    if hasattr(info, "private"):
+        return info.private
+    return None
+
+
+def private_memory_available():
+    """Whether this platform reports private memory (Linux, Windows)."""
+    return sys.platform.startswith("linux") or os.name == "nt"
+
+
 def _positive(value, name):
     if (
         isinstance(value, bool)
@@ -101,13 +142,20 @@ def run_limited(
     max_output_bytes=1_000_000,
     threads: int | None = None,
     capture_output: bool = True,
+    memory_measure: str = "private",
 ):
     """Run a command without a shell and return measurements and termination cause.
 
     hard: verified Linux cgroup cap, including descendants; no automatic fallback.
-    watchdog: sampled process-tree RSS threshold, which can overshoot.
+    watchdog: sampled process-tree memory threshold, which can overshoot.
     monitor: measurements/timeouts only, without a memory ceiling.
 
+    memory_measure: what the watchdog compares with the limit. "private"
+    (default) leaves out pages of files, which the kernel can reclaim, so a
+    program that memory-maps large scratch files is not stopped for them;
+    where private memory is unavailable (macOS) RSS is used, and the result's
+    memory_measure says so. "rss" counts every resident page.
+
     GB is decimal. The supervisor stays outside the child's hard cap. env has
     subprocess replacement semantics; no budget/thread variables are injected
     unless `threads` is given, which sets them all (`thread_env`). Captured
@@ -125,6 +173,13 @@ def run_limited(
         )
     if mode not in ("hard", "watchdog", "monitor"):
         raise ValueError("mode must be hard, watchdog or monitor")
+    if memory_measure not in ("private", "rss"):
+        raise ValueError("memory_measure must be private or rss")
+    measure = (
+        "private"
+        if memory_measure == "private" and private_memory_available()
+        else "rss"
+    )
     if memory_gb is not None:
         _positive(memory_gb, "memory_gb")
     if mode != "monitor" and memory_gb is None:
@@ -158,6 +213,7 @@ def run_limited(
     group = MemoryCgroup(cgroup_parent, limit) if mode == "hard" else None
     proc, kernel, tracked = None, None, {}
     reason, peak, descendants = None, 0, False
+    private_peak = 0 if private_memory_available() else None
     os_peak, stopped = None, False
     sampling_complete = True
     started = time.monotonic()
@@ -216,6 +272,12 @@ def run_limited(
             if proc.poll() is None:
                 proc.kill()
             proc.wait()
+            # The exited process's counters stay readable while Popen holds
+            # its handle: the exact peak working set, not the last sample.
+            handle = getattr(proc, "_handle", None)
+            counters = _windows_memory(int(handle)) if handle is not None else None
+            if counters is not None:
+                os_peak = max(os_peak or 0, counters[1])
 
     def tail(handle):
         size = handle.seek(0, os.SEEK_END)
@@ -239,7 +301,7 @@ def run_limited(
             except psutil.NoSuchProcess:
                 process = None
             while not exited():
-                rss = 0
+                rss, private = 0, 0
                 try:
                     members = (
                         [process, *process.children(recursive=True)]
@@ -250,6 +312,8 @@ def run_limited(
                         try:
                             info = p.memory_info()
                             rss += info.rss
+                            if private_peak is not None:
+                                private += _private_bytes(p, info) or 0
                             if hasattr(info, "peak_wset"):  # Windows
                                 os_peak = max(os_peak or 0, info.peak_wset)
                             if p.pid != proc.pid:
@@ -263,10 +327,13 @@ def run_limited(
                 except psutil.AccessDenied:
                     sampling_complete = False
                 peak = max(peak, rss)
+                if private_peak is not None:
+                    private_peak = max(private_peak, private)
                 if mode == "watchdog" and not sampling_complete:
                     reason = "monitoring_failed"
                     break
-                if mode == "watchdog" and rss > limit:
+                watched = private if measure == "private" else rss
+                if mode == "watchdog" and watched > limit:
                     reason = "rss_limit_exceeded"
                     break
                 if timeout is not None and time.monotonic() - started >= timeout:
@@ -316,6 +383,8 @@ def run_limited(
                 descendants,
                 sampling_complete,
                 os_peak_rss_bytes=os_peak,
+                peak_private_bytes=private_peak,
+                memory_measure=measure,
             )
     finally:
         try:
@@ -343,6 +412,12 @@ def main():
         help="let the command write to this terminal; the JSON summary goes to stderr",
     )
     parser.add_argument("--report", help="also write the JSON summary to this file")
+    parser.add_argument(
+        "--memory-measure",
+        choices=["private", "rss"],
+        default="private",
+        help="what the watchdog limits: private memory (default; not file pages) or RSS",
+    )
     parser.add_argument("--capabilities", action="store_true")
     parser.add_argument("command", nargs=argparse.REMAINDER)
     args = parser.parse_args()
@@ -359,6 +434,7 @@ def main():
             timeout=args.timeout,
             threads=args.threads,
             capture_output=not args.no_capture,
+            memory_measure=args.memory_measure,
         )
     except (ValueError, OSError, RuntimeError, ImportError) as exc:
         parser.error(str(exc))

{stream_limits-0.2.0 → stream_limits-0.3.0}/stream_limits.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stream-limits
-Version: 0.2.0
+Version: 0.3.0
 Summary: Cap the threads and memory of numeric Python libraries (Polars, numba, BLAS, pyarrow) at a scheduler's allocation, and measure peak memory
 Author: Lee Tucker
 License-Expression: MIT

{stream_limits-0.2.0 → stream_limits-0.3.0}/tests/test_launcher.py

@@ -65,7 +65,11 @@ def test_watchdog_and_timeout_have_distinct_causes():
         timeout=10,
     )
     assert result.status == "rss_limit_exceeded" and result.returncode != 0
-    assert result.peak_rss_bytes > 64_000_000
+    # what the watchdog limited: on Windows private bytes count the bytearray
+    # once committed, before its pages are resident
+    watched = (result.peak_private_bytes if result.memory_measure == "private"
+               else result.peak_rss_bytes)
+    assert watched > 64_000_000
     result = run_limited(
         python("import time; time.sleep(5)"), mode="monitor", timeout=0.2
     )
@@ -247,8 +251,6 @@ def test_os_peak_is_exact_and_covers_short_commands():
     assert result.status == "ok"
     assert result.os_peak_rss_bytes is not None
     assert result.os_peak_rss_bytes >= 200_000_000
-    if os.name != "posix":   # Windows reads its peak while sampling
-        return
     # too short for any sample: the sampled peak can be 0, the OS's is not
     quick = run_limited(python("pass"), mode="monitor", timeout=30)
     assert quick.os_peak_rss_bytes is not None and quick.os_peak_rss_bytes > 0
@@ -328,3 +330,65 @@ def test_cli_without_capture_passes_output_through(tmp_path):
     summary = json.loads(result.stderr)
     assert summary["status"] == "ok" and summary["stdout"] == ""
     assert json.loads(report.read_text()) == summary
+
+
+# A program that writes a 600 MB file and memory-maps it, touching every page:
+# its RSS grows by the file's size, its private memory does not.
+MAPPED = """
+import mmap, os, sys
+path = sys.argv[1]
+with open(path, "wb") as fh:
+    fh.truncate(600_000_000)
+with open(path, "r+b") as fh:
+    m = mmap.mmap(fh.fileno(), 0)
+    for i in range(0, len(m), 4096):
+        m[i] = 1
+    m.flush()
+    total = sum(m[i] for i in range(0, len(m), 4096))
+    import time; time.sleep(0.3)
+print(total)
+"""
+
+
+@pytest.mark.skipif(not sys.platform.startswith("linux"), reason="Linux private memory")
+def test_watchdog_ignores_mapped_file_pages_by_default(tmp_path):
+    command = [sys.executable, "-c", MAPPED, str(tmp_path / "scratch.bin")]
+    result = run_limited(command, mode="watchdog", memory_gb=0.3, timeout=60)
+    assert result.status == "ok", result.as_dict()
+    assert result.memory_measure == "private"
+    assert result.peak_private_bytes < 200_000_000
+    assert result.peak_rss_bytes > 400_000_000      # the mapped pages are resident
+    rss = run_limited(command, mode="watchdog", memory_gb=0.3, timeout=60,
+                      memory_measure="rss")
+    assert rss.status == "rss_limit_exceeded" and rss.memory_measure == "rss"
+
+
+@pytest.mark.skipif(not sys.platform.startswith("linux"), reason="Linux private memory")
+def test_watchdog_stops_private_memory_over_the_limit():
+    code = "import time; x=bytearray(500_000_000); x[::4096]=b'1'*len(x[::4096]); time.sleep(5)"
+    result = run_limited(python(code), mode="watchdog", memory_gb=0.3, timeout=30)
+    assert result.status == "rss_limit_exceeded" and result.memory_measure == "private"
+    assert result.peak_private_bytes > 300_000_000
+
+
+@pytest.mark.skipif(not sys.platform.startswith("linux"), reason="Linux /proc")
+def test_shared_memory_counts_as_private():
+    # shared memory (/dev/shm, multiprocessing) cannot be reclaimed without swap
+    code = ("from multiprocessing import shared_memory as s; import time; "
+            "m=s.SharedMemory(create=True,size=300_000_000); m.buf[::4096]=b'1'*len(m.buf[::4096]); "
+            "time.sleep(0.5); m.close(); m.unlink()")
+    result = run_limited(python(code), mode="monitor", timeout=30)
+    assert result.status == "ok" and result.peak_private_bytes > 250_000_000
+
+
+def test_private_measure_is_reported_or_falls_back():
+    from stream_limits.launcher import private_memory_available
+
+    result = run_limited(python("import time; time.sleep(0.3)"), mode="watchdog",
+                         memory_gb=1, timeout=30)
+    if private_memory_available():
+        assert result.memory_measure == "private" and result.peak_private_bytes > 0
+    else:
+        assert result.memory_measure == "rss" and result.peak_private_bytes is None
+    with pytest.raises(ValueError):
+        run_limited(python("pass"), mode="monitor", memory_measure="uss")

{stream_limits-0.2.0 → stream_limits-0.3.0}/tests/test_override_and_libraries.py

@@ -213,3 +213,68 @@ def test_a_preset_spill_dir_is_kept_unless_override(tmp_path):
     a = _run(code.format(extra=""), POLARS_OOC_SPILL_DIR="/somewhere/else")
     b = _run(code.format(extra=", override=True"), POLARS_OOC_SPILL_DIR="/somewhere/else")
     assert a.stdout.strip() == "/somewhere/else" and b.stdout.strip() == str(tmp_path), (a.stderr, b.stderr)
+
+
+# --- set_openmp_wait_policy ---
+
+def test_wait_policy_sets_the_variable_and_respects_the_user():
+    env = {}
+    assert sl.set_openmp_wait_policy(environ=env) == {"policy": "passive", "set": True, "late": None}
+    assert env["OMP_WAIT_POLICY"] == "passive"
+    env = {"OMP_WAIT_POLICY": "ACTIVE"}
+    assert sl.set_openmp_wait_policy(environ=env) == {"policy": "active", "set": False, "late": None}
+    assert env["OMP_WAIT_POLICY"] == "ACTIVE"
+    assert sl.set_openmp_wait_policy(environ=env, override=True)["set"] is True
+    assert env["OMP_WAIT_POLICY"] == "passive"
+    with pytest.raises(ValueError):
+        sl.set_openmp_wait_policy("sleepy", environ={})
+
+
+def test_importing_the_module_leaves_the_policy_alone():
+    r = _run("import os, stream_limits; print(os.environ.get('OMP_WAIT_POLICY'))")
+    assert r.stdout.strip() == "None"
+    r = _run("import os, stream_limits; stream_limits.apply_scheduler_defaults(); "
+             "print(os.environ.get('OMP_WAIT_POLICY'))", NCPUS="2")
+    assert r.stdout.strip() == "None"
+
+
+NUMBA_PROBE = """
+import numpy as np, sys
+{before}
+import numba
+from numba import njit, prange
+import stream_limits
+print(stream_limits.set_openmp_wait_policy()["late"])
+@njit(parallel=True)
+def f(a):
+    s = 0.0
+    for i in prange(a.size):
+        s += a[i]
+    return s
+f(np.ones(100))
+{after}
+"""
+
+
+def test_wait_policy_takes_effect_when_numba_is_imported_first():
+    # numba loads its OpenMP runtime at the first parallel call, which is after
+    r = _run(NUMBA_PROBE.format(before="", after="print(numba.threading_layer())"),
+             OMP_DISPLAY_ENV="verbose")
+    if "omp" not in r.stdout.split():
+        pytest.skip("numba's OpenMP threading layer is not in use")
+    # "late" is only detectable on Linux (it reads /proc/self/maps)
+    assert r.stdout.split()[0] == ("False" if sys.platform.startswith("linux") else "None")
+    if sys.platform == "win32":
+        pytest.skip("Microsoft's OpenMP runtime does not read or display OMP_WAIT_POLICY")
+    assert "OMP_WAIT_POLICY = 'PASSIVE'" in r.stderr
+
+
+def test_wait_policy_reports_a_runtime_that_is_already_loaded():
+    r = _run(NUMBA_PROBE.format(before="", after="") +
+             "\nimport stream_limits as s\nprint(s.set_openmp_wait_policy(override=True)['late'])")
+    if not sys.platform.startswith("linux"):
+        pytest.skip("detection needs /proc/self/maps")
+    lines = r.stdout.split()
+    if lines[-1] == "False":
+        pytest.skip("numba's OpenMP layer is not in use")
+    assert lines[-1] == "True"