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

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

@@ -1,5 +1,32 @@
 # Changelog
 
+## 0.2.0
+
+- **Run a program under a memory limit and see how much memory it used.** The new
+  launcher runs a script, test suite or other program in a child process, stops it
+  if it goes over a memory limit or a time limit, and reports why it ended and its
+  peak memory. The limit is a watchdog that checks memory every 10 ms (Linux,
+  macOS, Windows), or a kernel limit through a Linux cgroup. It also sets the
+  program's threads (`--threads`). It is an optional extra,
+  `pip install 'stream-limits[launcher]'`, which adds psutil; the rest of
+  stream-limits is unchanged and still needs only the standard library. See
+  [Running a program under a memory limit](https://leetucker.github.io/stream-limits/launcher.html).
+- **Pytest fixtures for memory tests.** `pytest_plugins =
+  ["stream_limits.pytest_plugin"]` gives a test suite `resource_runner`, which runs
+  a test's work in a child process under the watchdog on every platform, or under a
+  kernel limit when one is set up (`STREAM_LIMITS_TEST_CGROUP_PARENT`). Install
+  `stream-limits[testing]`.
+
+## 0.1.1
+
+- **Windows reports memory use.** `peak_rss_gb` and `current_rss_gb` returned `None` on Windows in
+  0.1.0; they now read the process's peak and current working set (the physical memory resident in it,
+  as on Linux) through the Windows API, with no new dependency. Windows cannot reset its peak, so
+  there it is the process's lifetime peak, as on macOS. `MemoryBudget` uses the real current memory
+  on Windows instead of assuming 0.5 GB.
+- Documentation: Polars 1.37 creates its thread pool at first use, while later versions (1.44 and
+  2.x) create it at import; the thread-limit pages say so.
+
 ## 0.1.0
 
 First release, coinciding with updates to hdfe-stream, logit-stream and

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

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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stream-limits
-Version: 0.1.0
+Version: 0.2.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
@@ -26,6 +26,11 @@ Classifier: Typing :: Typed
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE.md
+Provides-Extra: launcher
+Requires-Dist: psutil>=5.9; extra == "launcher"
+Provides-Extra: testing
+Requires-Dist: pytest>=7.0; extra == "testing"
+Requires-Dist: psutil>=5.9; extra == "testing"
 Provides-Extra: blas
 Requires-Dist: threadpoolctl>=3.1; extra == "blas"
 Provides-Extra: docs
@@ -35,6 +40,7 @@ Requires-Dist: furo>=2024.8.6; extra == "docs"
 Requires-Dist: sphinx-copybutton>=0.5; extra == "docs"
 Provides-Extra: test
 Requires-Dist: pytest>=7.0; extra == "test"
+Requires-Dist: psutil>=5.9; extra == "test"
 Requires-Dist: numpy; extra == "test"
 Requires-Dist: numba; extra == "test"
 Requires-Dist: polars; extra == "test"
@@ -85,7 +91,7 @@ table below (numexpr, PyTorch, OpenMP programs, TensorFlow, rayon, joblib, Julia
 
 | Library | Can it be capped after import? | What `limit_threads` does |
 |---|---|---|
-| Polars | **No**: the pool is sized once, at import | sets `POLARS_MAX_THREADS` (before import); warns if too late |
+| Polars | **No**: the pool is sized once, at import (Polars 1.37: at first use) | sets `POLARS_MAX_THREADS` (before import); warns if too late |
 | numba | down, never above `NUMBA_NUM_THREADS` at import | sets the variable; lowers `set_num_threads` |
 | BLAS (NumPy, SciPy) | yes, through `threadpoolctl` | sets `OMP_`/`OPENBLAS_`/`MKL_NUM_THREADS`; `threadpool_limits` |
 | pyarrow | yes | `set_cpu_count`, `set_io_thread_count` |
@@ -112,8 +118,8 @@ table below (numexpr, PyTorch, OpenMP programs, TensorFlow, rayon, joblib, Julia
   (default 15%, `STREAM_LIMITS_MEMORY_RESERVE`) for what the code does not model, less what the
   process holds now. Ask whether your exact allocation fits, not whether it is under a fixed
   fraction of the machine.
-- `peak_rss_gb()` / `reset_peak_rss()` read the kernel's own high-water mark (Linux): free, no
-  polling.
+- `peak_rss_gb()` / `reset_peak_rss()` read the system's own high-water mark (Linux, macOS and
+  Windows): free, no polling.
 
 `set_limits` is `limit_threads` and `limit_memory` in one call: values you give are applied, values you
 leave out are the job's allocation if one is found and nothing otherwise. The two remain, for their
@@ -152,8 +158,23 @@ pip install stream-limits            # nothing else is installed
 pip install "stream-limits[blas]"    # threadpoolctl, to limit BLAS after NumPy is loaded
 ```
 
-Linux and macOS are tested on Python 3.10 to 3.15, Windows on 3.10 and 3.14. On Windows the module imports
-and applies thread limits, but it detects no cgroup and cannot read peak memory. On macOS, NumPy's Apple
+Linux and macOS are tested on Python 3.10 to 3.15, Windows on 3.10 and 3.14. On Windows the module imports,
+applies thread limits and reads memory use (peak and current working set), but it detects no cgroup;
+macOS and Windows report a lifetime peak, as their counters cannot be reset. On macOS, NumPy's Apple
 Accelerate BLAS is limited only by `VECLIB_MAXIMUM_THREADS` before NumPy is imported (threadpoolctl cannot see it).
 
+## Running a program under a memory limit
+
+The optional launcher (`pip install 'stream-limits[launcher]'`, which adds psutil)
+runs a program in a child process, stops it if it goes over a memory or time
+limit, and reports why it ended and its peak memory:
+
+```bash
+python -m stream_limits.launcher --mode watchdog --memory-gb 4 --threads 2 -- python analysis.py
+```
+
+The watchdog works on Linux, macOS and Windows; a kernel-enforced limit needs a
+Linux cgroup set up for it. `stream-limits[testing]` adds pytest fixtures for
+memory tests. See [Running a program under a memory limit](docs/launcher.md).
+
 MIT licence.

{stream_limits-0.1.0 → stream_limits-0.2.0}/README.md

@@ -40,7 +40,7 @@ table below (numexpr, PyTorch, OpenMP programs, TensorFlow, rayon, joblib, Julia
 
 | Library | Can it be capped after import? | What `limit_threads` does |
 |---|---|---|
-| Polars | **No**: the pool is sized once, at import | sets `POLARS_MAX_THREADS` (before import); warns if too late |
+| Polars | **No**: the pool is sized once, at import (Polars 1.37: at first use) | sets `POLARS_MAX_THREADS` (before import); warns if too late |
 | numba | down, never above `NUMBA_NUM_THREADS` at import | sets the variable; lowers `set_num_threads` |
 | BLAS (NumPy, SciPy) | yes, through `threadpoolctl` | sets `OMP_`/`OPENBLAS_`/`MKL_NUM_THREADS`; `threadpool_limits` |
 | pyarrow | yes | `set_cpu_count`, `set_io_thread_count` |
@@ -67,8 +67,8 @@ table below (numexpr, PyTorch, OpenMP programs, TensorFlow, rayon, joblib, Julia
   (default 15%, `STREAM_LIMITS_MEMORY_RESERVE`) for what the code does not model, less what the
   process holds now. Ask whether your exact allocation fits, not whether it is under a fixed
   fraction of the machine.
-- `peak_rss_gb()` / `reset_peak_rss()` read the kernel's own high-water mark (Linux): free, no
-  polling.
+- `peak_rss_gb()` / `reset_peak_rss()` read the system's own high-water mark (Linux, macOS and
+  Windows): free, no polling.
 
 `set_limits` is `limit_threads` and `limit_memory` in one call: values you give are applied, values you
 leave out are the job's allocation if one is found and nothing otherwise. The two remain, for their
@@ -107,8 +107,23 @@ pip install stream-limits            # nothing else is installed
 pip install "stream-limits[blas]"    # threadpoolctl, to limit BLAS after NumPy is loaded
 ```
 
-Linux and macOS are tested on Python 3.10 to 3.15, Windows on 3.10 and 3.14. On Windows the module imports
-and applies thread limits, but it detects no cgroup and cannot read peak memory. On macOS, NumPy's Apple
+Linux and macOS are tested on Python 3.10 to 3.15, Windows on 3.10 and 3.14. On Windows the module imports,
+applies thread limits and reads memory use (peak and current working set), but it detects no cgroup;
+macOS and Windows report a lifetime peak, as their counters cannot be reset. On macOS, NumPy's Apple
 Accelerate BLAS is limited only by `VECLIB_MAXIMUM_THREADS` before NumPy is imported (threadpoolctl cannot see it).
 
+## Running a program under a memory limit
+
+The optional launcher (`pip install 'stream-limits[launcher]'`, which adds psutil)
+runs a program in a child process, stops it if it goes over a memory or time
+limit, and reports why it ended and its peak memory:
+
+```bash
+python -m stream_limits.launcher --mode watchdog --memory-gb 4 --threads 2 -- python analysis.py
+```
+
+The watchdog works on Linux, macOS and Windows; a kernel-enforced limit needs a
+Linux cgroup set up for it. `stream-limits[testing]` adds pytest fixtures for
+memory tests. See [Running a program under a memory limit](docs/launcher.md).
+
 MIT licence.

{stream_limits-0.1.0 → stream_limits-0.2.0}/docs/index.md

@@ -82,3 +82,15 @@ other-libraries
 reference
 changelog
 ```
+
+## Running a program under a memory limit
+
+The optional [launcher](launcher.md) runs a program in a child process, stops it
+if it goes over a memory or time limit, and reports its peak memory. It also
+gives other packages' test suites fixtures for memory tests.
+
+```{toctree}
+:hidden:
+
+launcher
+```

stream_limits-0.2.0/docs/launcher.md

@@ -0,0 +1,236 @@
+# Running a program under a memory limit
+
+The launcher runs a program (a Python or R script, a test suite, any executable)
+in a child process, watches its memory, and stops it if it goes over a limit you
+set. When the program ends it tells you why it ended and how much memory it used.
+
+It is an optional extra, which adds psutil and nothing else:
+
+```bash
+pip install 'stream-limits[launcher]'
+```
+
+Nothing else in stream-limits changes: `import stream_limits` still needs only
+the standard library, and the launcher imports no numerical library itself.
+
+## A first example
+
+```bash
+python -m stream_limits.launcher --mode watchdog --memory-gb 4 --threads 2 \
+    --timeout 600 -- python analysis.py
+```
+
+This runs `python analysis.py` with 2 threads for every numerical library, stops
+it if it and the processes it starts use more than 4 GB together or run longer
+than 10 minutes, and prints a JSON summary: the status, the exit code, the time,
+the peak memory, and the end of the program's output. The exit code is the
+program's own (124 on a timeout).
+
+The same from Python:
+
+```python
+import sys
+from stream_limits.launcher import run_limited
+
+result = run_limited([sys.executable, "analysis.py"],
+                     mode="watchdog", memory_gb=4, threads=2, timeout=600)
+print(result.status, result.os_peak_rss_bytes / 1e9, "GB")
+```
+
+The command is a list of arguments, run without a shell. GB are decimal
+(1 GB = 10^9 bytes).
+
+## Three ways to limit
+
+| Mode | Platforms | What it does |
+|---|---|---|
+| `watchdog` | Linux, macOS, Windows | Checks memory about every 10 ms and stops the program when it passes the limit |
+| `hard` | Linux, with a cgroup set up for it | The kernel enforces the limit |
+| `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,
+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
+process allocating 3 GB were each stopped at 1.51-1.52 GB. Memory shared
+between processes can be counted more than once. If the operating system denies
+it access to a process's memory, the watchdog stops the program and reports
+`monitoring_failed` rather than claim it enforced the limit.
+
+**Hard limits** leave nothing to sampling: the kernel will not let the program
+use more, and kills it if it tries. They need a Linux cgroup that you are
+allowed to manage (below). `hard` is the default mode, so that asking for a hard
+limit where none is possible fails before the program starts, instead of quietly
+running under the watchdog. macOS and Windows have no hard mode.
+
+## Results
+
+`run_limited` returns a `RunResult` (`as_dict()` turns it into JSON). Its
+`status` is one of:
+
+| Status | Meaning |
+|---|---|
+| `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 |
+| `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 |
+
+`kernel_oom` comes from the kernel's own counter, not from guessing at a killed
+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:
+
+- `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.
+
+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
+memory, so it differs from the two above), and its event counters.
+
+## Output, threads and the environment
+
+By default the program's output is saved to temporary files and the last
+megabyte of each stream is returned in the result (`max_output_bytes=`). To see
+it as it runs, as when wrapping a whole test run, pass `--no-capture`
+(`capture_output=False`). The program then writes to the terminal and the JSON
+summary goes to stderr; `--report FILE` also saves the summary to a file:
+
+```bash
+python -m stream_limits.launcher --mode watchdog --memory-gb 4 --threads 2 \
+    --no-capture --report memory.json -- python -m pytest -q
+```
+
+`--threads N` (`threads=N`) sets `STREAM_LIMITS_MAX_THREADS` and the thread
+variables of Polars, numba, OpenMP, the BLAS libraries, numexpr, rayon and Julia
+in the program's environment. `stream_limits.launcher.thread_env(N, env)`
+returns that environment. Without `threads`, no variable is added: `env`
+replaces the environment as in `subprocess`, and the launcher's own environment
+is never changed.
+
+On a timeout or when a limit is hit, the program and the processes it started
+are stopped. On Linux and macOS the launcher also stops anything left in the
+program's process group, and with a hard limit anything left in its cgroup. On
+Windows, a process that escapes before it is first sampled can survive.
+
+## Setting up hard limits
+
+Hard limits use a Linux cgroup with the memory controller, which an
+administrator (or you, with root) sets up once and lets you manage. The launcher
+never uses sudo and never changes a cgroup it did not create: for each run it
+creates a fresh group inside the one you give it, sets the limit there, starts
+the program inside it (the launcher itself stays outside the limit), and removes
+the group afterwards. It is a resource limit, not a security sandbox.
+
+`python -m stream_limits.launcher --capabilities` and `cat /proc/self/cgroup`
+show what the machine has.
+
+**cgroup v1** (WSL2, older distributions), once, as root:
+
+```bash
+sudo mkdir /sys/fs/cgroup/memory/stream-limits-test
+sudo chown -R "$USER" /sys/fs/cgroup/memory/stream-limits-test
+export STREAM_LIMITS_TEST_CGROUP_PARENT=/sys/fs/cgroup/memory/stream-limits-test
+```
+
+The group is gone after a reboot.
+
+**cgroup v2 with systemd**: start a shell in a user scope that you may manage:
+
+```bash
+systemd-run --user --scope -p Delegate=yes -p MemoryMax=8G bash
+export STREAM_LIMITS_TEST_CGROUP_PARENT=/sys/fs/cgroup$(cut -d: -f3 /proc/self/cgroup)
+```
+
+On v2 the group must have `memory` in its `cgroup.subtree_control` and no
+processes of its own, so the shell may first need to move into a child group.
+
+These recipes are standard, but stream-limits' tests do not run them (they need
+root or a systemd scope). Check a setup before relying on it:
+
+```bash
+python -c "from stream_limits.testing import verify_memory_limit as v; v('$STREAM_LIMITS_TEST_CGROUP_PARENT')"
+```
+
+This runs a small allocation that must succeed, checks that a program inside the
+group sees the limit, and runs a 256 MiB allocation in a 64 MiB group that the
+kernel must kill. It raises unless all three happen. (The group needs room for
+64 MiB.)
+
+Then:
+
+```python
+result = run_limited([sys.executable, "analysis.py"], mode="hard", memory_gb=4,
+                     cgroup_parent="/sys/fs/cgroup/memory/stream-limits-test")
+```
+
+Before running anything, the launcher checks that the group really is a kernel
+cgroup (not an ordinary directory with the same file names) and reads the limit
+back. It turns swap off for the group (v2: `memory.swap.max=0`; v1: a
+memory-plus-swap limit equal to the memory limit and `swappiness=0`), so the
+limit cannot be dodged by swapping. On v2 the kernel stops the whole group
+together when one process goes over. Older v2 kernels have no peak counter;
+their kernel peak is then `None`.
+
+The launcher's tests have exercised hard limits on a WSL2 host with cgroup v1.
+cgroup v2, macOS and Windows rely on CI and have not been checked by hand.
+
+## Memory tests in pytest
+
+The plugin gives other packages' test suites a common way to run a test's work in
+a child process under a memory limit. Install `stream-limits[testing]` and
+register it in `conftest.py` (it is not loaded automatically):
+
+```python
+pytest_plugins = ["stream_limits.pytest_plugin"]
+```
+
+A memory test then looks like this:
+
+```python
+import sys
+import pytest
+
+@pytest.mark.resource
+def test_fit_stays_within_memory(resource_runner, tmp_path):
+    result = resource_runner([sys.executable, "fit_large_case.py", str(tmp_path)],
+                             memory_gb=1.0, threads=2, timeout=120)
+    assert result.status == "ok", result.as_dict()
+```
+
+`resource_runner` takes `run_limited`'s arguments except the mode. It uses the
+watchdog on every platform, so these tests run on macOS and Windows too. When
+`STREAM_LIMITS_TEST_CGROUP_PARENT` names a cgroup set up as above, it checks it
+once per session with `verify_memory_limit` and uses a hard limit instead; a
+setting that does not work fails the session rather than falling back.
+
+```bash
+pytest -m resource                                                    # watchdog
+STREAM_LIMITS_TEST_CGROUP_PARENT=/sys/fs/cgroup/memory/stream-limits-test pytest -m resource
+```
+
+The other fixtures: `resource_enforcement` says which mode was chosen
+(`{"mode": "watchdog"}` or `{"mode": "hard", "cgroup_parent": ...}`);
+`memory_runner` is `run_limited` itself, for a test that picks its own mode; and
+`verified_memory_parent` is the checked cgroup, for tests that need a hard limit
+and skip without one. The plugin registers the markers `resource` and
+`kernel_memory`.
+
+**A memory test is only as good as its input.** A child process that imports
+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.
+
+Use the launcher only for tests about memory. Ordinary tests are better run in
+the test process, as before.

{stream_limits-0.1.0 → stream_limits-0.2.0}/docs/memory.md

@@ -91,5 +91,6 @@ print(stream_limits.peak_rss_gb())    # GB, the peak during the call
 ```
 
 On Linux the counter is reset by writing to `/proc/self/clear_refs`, so a second measurement in
-the same process reports its own peak; elsewhere it is the process's lifetime peak. Child
+the same process reports its own peak; on macOS and Windows the system's counter cannot be
+reset, so it is the process's lifetime peak (a second measurement is the largest of both). Child
 processes are not included. {py:func}`stream_limits.current_rss_gb` is the memory held now.

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

@@ -65,3 +65,13 @@
 `NUMEXPR_NUM_THREADS`, `NUMEXPR_MAX_THREADS`, `RAYON_NUM_THREADS`, `TF_NUM_INTRAOP_THREADS`,
 `TF_NUM_INTEROP_THREADS`, `LOKY_MAX_CPU_COUNT`, `JULIA_NUM_THREADS`. `limit_memory` sets
 `POLARS_OOC_MEMORY_BUDGET_MB` and, if asked, `POLARS_OOC_SPILL_DIR`.
+
+## Optional launcher
+
+```{eval-rst}
+.. automodule:: stream_limits.launcher
+   :members: RunResult, run_limited, capabilities, thread_env
+
+.. automodule:: stream_limits.testing
+   :members: verify_memory_limit, kernel_oom_proved
+```

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

@@ -8,7 +8,7 @@ threads and be killed for it.
 
 | Library | Can it be capped after import? | What `limit_threads` does |
 |---|---|---|
-| Polars | **No**: the pool is sized once, at import | sets `POLARS_MAX_THREADS` (before import); warns if too late |
+| Polars | **No**: the pool is sized once, at import (Polars 1.37: at first use) | sets `POLARS_MAX_THREADS` (before import); warns if too late |
 | numba | down, never above `NUMBA_NUM_THREADS` at import | sets the variable; lowers `set_num_threads` |
 | BLAS (NumPy, SciPy) | yes, through `threadpoolctl` | sets `OMP_`/`OPENBLAS_`/`MKL_NUM_THREADS`; `threadpool_limits` |
 | Apple Accelerate (NumPy on macOS) | **No**: `threadpoolctl` cannot see it | sets `VECLIB_MAXIMUM_THREADS` (before NumPy is imported) |
@@ -66,7 +66,7 @@ always wins over them.
 ## For one block of code
 
 {py:func}`stream_limits.thread_scope` holds numba and BLAS to `n` threads inside a block and
-restores them after, also when the block raises. Polars is not affected: its pool is fixed for
+restores them after, also when the block raises. Polars is not affected: its pool is fixed (at import, or at first use in Polars 1.37) for
 the process.
 
 ```python

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "stream-limits"
-version = "0.1.0"
+version = "0.2.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"
@@ -33,6 +33,9 @@ requires-python = ">=3.10"
 dependencies = []
 
 [project.optional-dependencies]
+# Opt-in subprocess launcher; core installs remain standard-library-only.
+launcher = ["psutil>=5.9"]
+testing = ["pytest>=7.0", "psutil>=5.9"]
 # threadpoolctl lets thread_scope()/limit_threads() limit BLAS after numpy is loaded
 blas = ["threadpoolctl>=3.1"]
 # docs/: the documentation site, built with Sphinx (see docs/conf.py)
@@ -42,7 +45,7 @@ docs = [
     "furo>=2024.8.6",
     "sphinx-copybutton>=0.5",
 ]
-test = ["pytest>=7.0", "numpy", "numba", "polars", "pyarrow", "threadpoolctl>=3.1",
+test = ["pytest>=7.0", "psutil>=5.9", "numpy", "numba", "polars", "pyarrow", "threadpoolctl>=3.1",
         'numexpr; python_version < "3.15"']
 
 [project.urls]

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

@@ -87,7 +87,7 @@ import warnings
 from pathlib import Path
 from typing import Iterator, Mapping
 
-__version__ = "0.1.0"
+__version__ = "0.2.0"
 
 __all__ = [
     "set_limits", "limit_threads", "limit_memory", "apply_scheduler_defaults", "thread_scope",
@@ -791,10 +791,49 @@ def thread_scope(n: int | None = None) -> Iterator[None]:
 
 
 # ----------------------------------------------------- peak memory used
+def _windows_memory() -> 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."""
+    if sys.platform != "win32":
+        return None
+    try:
+        import ctypes
+        from ctypes import wintypes
+
+        class _Counters(ctypes.Structure):          # PROCESS_MEMORY_COUNTERS
+            _fields_ = [("cb", wintypes.DWORD), ("PageFaultCount", wintypes.DWORD),
+                        ("PeakWorkingSetSize", ctypes.c_size_t),
+                        ("WorkingSetSize", ctypes.c_size_t),
+                        ("QuotaPeakPagedPoolUsage", ctypes.c_size_t),
+                        ("QuotaPagedPoolUsage", ctypes.c_size_t),
+                        ("QuotaPeakNonPagedPoolUsage", ctypes.c_size_t),
+                        ("QuotaNonPagedPoolUsage", ctypes.c_size_t),
+                        ("PagefileUsage", ctypes.c_size_t),
+                        ("PeakPagefileUsage", ctypes.c_size_t)]
+
+        kernel32 = ctypes.WinDLL("kernel32")
+        kernel32.GetCurrentProcess.restype = wintypes.HANDLE        # 64-bit safe
+        kernel32.K32GetProcessMemoryInfo.argtypes = [wintypes.HANDLE, ctypes.POINTER(_Counters),
+                                                     wintypes.DWORD]
+        kernel32.K32GetProcessMemoryInfo.restype = wintypes.BOOL
+        counters = _Counters()
+        counters.cb = ctypes.sizeof(_Counters)
+        if not kernel32.K32GetProcessMemoryInfo(kernel32.GetCurrentProcess(),
+                                                ctypes.byref(counters), counters.cb):
+            return None
+        return int(counters.WorkingSetSize), int(counters.PeakWorkingSetSize)
+    except Exception:                                          # noqa: BLE001
+        return None
+
+
 def reset_peak_rss() -> None:
-    """Restart the peak-memory counter at the current usage (Linux; elsewhere
-    it does nothing, and `peak_rss_gb` then reports the process's lifetime
-    peak)."""
+    """Restart the peak-memory counter at the current usage. Linux only: macOS
+    and Windows keep a counter that cannot be reset, so there `peak_rss_gb`
+    reports the process's lifetime peak (a second measurement in the same
+    process is the largest of both)."""
     try:
         with open("/proc/self/clear_refs", "w") as f:
             f.write("5")
@@ -804,8 +843,12 @@ def reset_peak_rss() -> None:
 
 def peak_rss_gb() -> float | None:
     """The peak resident memory, in GB, since `reset_peak_rss` (or since the
-    process started). The kernel keeps the counter, so reading it is free and
-    nothing is polled. Child processes are not included."""
+    process started). The system keeps the counter, so reading it is free and
+    nothing is polled: ``VmHWM`` on Linux, ``getrusage`` on macOS, the peak
+    working set on Windows. Child processes are not included."""
+    win = _windows_memory()
+    if win is not None:
+        return win[1] / 1e9
     try:
         for line in Path("/proc/self/status").read_text().splitlines():
             if line.startswith("VmHWM:"):
@@ -821,7 +864,11 @@ def peak_rss_gb() -> float | None:
 
 
 def current_rss_gb() -> float | None:
-    """The resident memory in GB right now (Linux; None elsewhere)."""
+    """The resident memory in GB right now (Linux and Windows; None on macOS,
+    where `MemoryBudget` then assumes 0.5 GB)."""
+    win = _windows_memory()
+    if win is not None:
+        return win[0] / 1e9
     try:
         for line in Path("/proc/self/status").read_text().splitlines():
             if line.startswith("VmRSS:"):

stream_limits-0.2.0/stream_limits/_exec_cgroup.py

@@ -0,0 +1,18 @@
+"""Internal pre-exec attachment helper, deliberately standard-library-only."""
+
+import os
+import sys
+from pathlib import Path
+
+
+def main():
+    if len(sys.argv) < 4 or sys.argv[2] != "--":
+        raise SystemExit("internal usage: PARENT -- COMMAND ...")
+    path = Path(sys.argv[1])
+    (path / "cgroup.procs").write_text(str(os.getpid()))
+    command = sys.argv[3:]
+    os.execvpe(command[0], command, os.environ)
+
+
+if __name__ == "__main__":
+    main()