stream-limits 0.1.0__tar.gz

stream_limits-0.1.0/CHANGELOG.md

@@ -0,0 +1,44 @@
+# Changelog
+
+## 0.1.0
+
+First release, coinciding with updates to hdfe-stream, logit-stream and
+dpd-stream to use it. Anyone running other numeric code on a scheduler can use
+it on its own.
+
+**Limits**
+
+- `set_limits(threads=None, memory_gb=None)` applies thread and memory limits in one call. Values you
+  give are used; values you omit are the job's allocation if one is found, and nothing otherwise.
+- `limit_threads` caps Polars, numba, NumPy/SciPy's BLAS (OpenBLAS, MKL, BLIS, Accelerate), OpenMP
+  programs (LightGBM, XGBoost, scikit-learn), pyarrow, numexpr, PyTorch, TensorFlow, rayon-based Rust
+  libraries, joblib and Julia. It warns when Polars was imported first and can no longer be capped.
+- `limit_memory` gives Polars 2.x an out-of-core memory budget, so its streaming engine spills to disk
+  instead of growing (a streaming sort of 1.9 GB peaked at 6.2 GB; with a budget, 1.8 GB).
+- `apply_scheduler_defaults` does both, only when a limit is found. The packages built on this one call
+  it when imported.
+- Variables you have set are left alone, and loaded libraries are only lowered; a variable set above the
+  cap is named in a warning. `override=True` replaces them.
+- `thread_scope` holds numba and BLAS to a thread count for one block of code.
+- Every setting can be given as a `STREAM_LIMITS_*` environment variable (`STREAM_LIMITS_MAX_THREADS`,
+  `STREAM_LIMITS_MEMORY_GB`, ...), which is how a job script sets it. Values passed to `set_limits` are
+  recorded there too, so the rest of the process, child processes and the packages built on this one agree.
+
+**Finding the allocation**
+
+- `detect_cpus` and `detect_memory_gb` read the environment, cgroup v1/v2 limits (the lowest on the path
+  to the root), PBS (`NCPUS`, `PBS_NUM_PPN`, ...; `qstat -f $PBS_JOBID` when no variable or cgroup says)
+  and Slurm.
+- `STREAM_LIMITS_CPU_VARS` and `STREAM_LIMITS_MEMORY_VARS` name a site's own variables.
+- `describe` and `python -m stream_limits` show what was found and where it came from.
+
+**Memory**
+
+- `MemoryBudget` gives the headroom left, for code that chooses between a fast path that needs memory and
+  a slower one that needs less.
+- `peak_rss_gb`, `reset_peak_rss` and `current_rss_gb` read the kernel's own counters.
+
+**Documentation**
+
+- A Sphinx site, built in CI: thread limits, memory limits, schedulers and containers, and recipes for
+  libraries with only per-object settings (DuckDB, Dask, Ray, JAX, Spark).

stream_limits-0.1.0/CITATION.cff

@@ -0,0 +1,21 @@
+cff-version: 1.2.0
+message: "If you use this software in your research, please cite it as below."
+type: software
+title: "stream-limits"
+abstract: "Cap the threads and memory of numeric Python libraries (Polars, numba, BLAS, pyarrow, ...) at a scheduler's allocation, and measure peak memory."
+authors:
+  - family-names: Tucker
+    given-names: Lee
+    email: lee.c.tucker@census.gov
+version: 0.1.0
+date-released: 2026-10-31
+license: MIT
+repository-code: "https://github.com/leetucker/stream-limits"
+url: "https://github.com/leetucker/stream-limits"
+keywords:
+  - threads
+  - memory
+  - PBS
+  - Slurm
+  - cgroups
+  - HPC

stream_limits-0.1.0/LICENSE.md

@@ -0,0 +1,21 @@
+# MIT License
+
+Copyright (c) 2026 Lee Tucker
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

stream_limits-0.1.0/MANIFEST.in

@@ -0,0 +1,5 @@
+include README.md CHANGELOG.md CITATION.cff LICENSE.md
+graft tests
+graft docs
+prune docs/_build
+global-exclude __pycache__ *.py[cod]

stream_limits-0.1.0/PKG-INFO

@@ -0,0 +1,159 @@
+Metadata-Version: 2.4
+Name: stream-limits
+Version: 0.1.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
+Project-URL: Homepage, https://github.com/leetucker/stream-limits
+Project-URL: Documentation, https://leetucker.github.io/stream-limits/
+Project-URL: Changelog, https://github.com/leetucker/stream-limits/blob/main/CHANGELOG.md
+Project-URL: Source, https://github.com/leetucker/stream-limits
+Keywords: threads,memory,PBS,Slurm,cgroup,polars,numba,HPC
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Operating System :: MacOS
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Programming Language :: Python :: 3.15
+Classifier: Topic :: System :: Hardware
+Classifier: Typing :: Typed
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE.md
+Provides-Extra: blas
+Requires-Dist: threadpoolctl>=3.1; extra == "blas"
+Provides-Extra: docs
+Requires-Dist: sphinx>=8.0; extra == "docs"
+Requires-Dist: myst-parser>=4.0; extra == "docs"
+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: numpy; extra == "test"
+Requires-Dist: numba; extra == "test"
+Requires-Dist: polars; extra == "test"
+Requires-Dist: pyarrow; extra == "test"
+Requires-Dist: threadpoolctl>=3.1; extra == "test"
+Requires-Dist: numexpr; python_version < "3.15" and extra == "test"
+Dynamic: license-file
+
+# stream-limits
+
+Cap the threads and memory of numeric Python libraries at a job's allocation, with one
+interface, and measure the memory a block of code used.
+
+On a cluster (PBS, Slurm, a container) a job gets a fixed number of CPUs and a fixed
+amount of memory, and is killed if it exceeds either. Polars, numba, BLAS and pyarrow
+each size their thread pools separately and default to every core of the *node*; code
+that asks how much memory there is sees the node's too. `stream-limits` finds the job's
+allocation, applies it to each library as far as the library allows, and says so when it
+cannot. It depends on nothing outside the standard library and imports none of the
+libraries it limits, so it can run before they are imported, which is when some of them
+(Polars) need to be told.
+
+It is what [hdfe-stream](https://github.com/leetucker/hdfe-stream),
+[logit-stream](https://github.com/leetucker/logit-stream) and
+[dpd-stream](https://github.com/leetucker/dpd-stream) use, and it works on its own.
+Documentation: <https://leetucker.github.io/stream-limits/> (build locally with
+`pip install -e ".[docs]"` and `sphinx-build -W -b html docs docs/_build/html`).
+
+```python
+import stream_limits
+stream_limits.set_limits()         # the job's allocation; or set_limits(threads=8, memory_gb=16)
+
+import polars as pl                # now capped
+import numpy as np
+```
+
+```
+$ python -m stream_limits          # what was found, from inside a job
+{"cpus": 8, "cpus_source": "PBS_NCPUS", "memory_gb": 32.0, "memory_source": "cgroup memory limit", ...}
+```
+
+Variables you have set are **not overridden** by default (a warning names any you set above the cap);
+loaded libraries are only lowered. `override=True` replaces them. More libraries are covered than the
+table below (numexpr, PyTorch, OpenMP programs, TensorFlow, rayon, joblib, Julia), and the documentation's
+"Other libraries" has recipes for those with only per-object settings (DuckDB, Dask, Ray, JAX, Spark).
+
+## What each library allows
+
+| 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 |
+| 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` |
+
+`thread_scope(n)` holds numba and BLAS to `n` threads for a block and restores them.
+
+## Memory
+
+- `detect_memory_gb()` is the memory the process may use: `STREAM_LIMITS_MEMORY_GB`, else the
+  cgroup limit (the lowest on the path to the root, so a scheduler's job limit is found), else
+  PBS's own record of the job, else `SLURM_MEM_PER_NODE`, else the machine's. PBS sites often set
+  no memory variable and enforce no cgroup; there `qstat -f $PBS_JOBID` is run once (no shell, a
+  10 s timeout, cached) and `Resource_List.mem` (else `vmem`) is read, with PBS's size suffixes.
+  `STREAM_LIMITS_QSTAT=0` turns this off.
+- `limit_memory()` gives Polars 2.x an out-of-core budget (`POLARS_OOC_MEMORY_BUDGET_MB`),
+  past which its streaming engine spills to disk instead of growing. A streaming sort of 1.9 GB
+  of Parquet peaked at 6.2 GB with the defaults and 1.8 GB with a 400 MB budget, in twice the
+  time. The budget bounds what Polars' operators hold, not the whole process (its Parquet
+  readers hold about 1.5 GB more on 8 threads). This is undocumented in Polars 2.x and read at
+  import, so it is set only for Polars 2.0 or later, only before Polars is imported, never over
+  a budget you set, and `STREAM_LIMITS_POLARS_SPILL=0` turns it off.
+- `MemoryBudget(limit_gb, reserve)` is for code that chooses between a fast path that needs
+  memory and a slower one that needs less: `budget.headroom_gb()` is the limit, less a reserve
+  (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.
+
+`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
+own options (`spill_dir`, `fraction`, ...).
+
+## Automatically
+
+`apply_scheduler_defaults()` (which is `set_limits()` with no arguments) does `limit_threads` and `limit_memory`, only for limits it finds
+(a variable, a scheduler's environment, a cgroup, a restricted CPU affinity), and sets only
+variables you have not set. Libraries built on this call it when they are imported; importing
+`stream_limits` itself changes nothing.
+
+## Environment variables
+
+| Variable | Meaning |
+|---|---|
+| `STREAM_LIMITS_MAX_THREADS` | the thread cap |
+| `STREAM_LIMITS_MEMORY_GB` | the memory limit |
+| `STREAM_LIMITS_MEMORY_RESERVE` | `MemoryBudget`'s reserve (default 0.15) |
+| `STREAM_LIMITS_POLARS_SPILL` | `0`: do not set Polars' budget |
+| `STREAM_LIMITS_POLARS_FRACTION` | the budget's share of the limit above the readers' floor (default 0.5) |
+| `STREAM_LIMITS_QSTAT` | `0`: do not run `qstat` to find a PBS job's CPUs and memory |
+| `STREAM_LIMITS_CPU_VARS` | names of variables holding a CPU count, checked first (comma-separated) |
+| `STREAM_LIMITS_MEMORY_VARS` | the same for memory; each holds a size (`16gb`) |
+| `STREAM_LIMITS_SPILL_DIR` | where Polars spills (default its temp dir: on a cluster, check `$TMPDIR`) |
+
+Thread limits are read from, in order: the argument; `STREAM_LIMITS_MAX_THREADS`; the variables named
+in `STREAM_LIMITS_CPU_VARS`; `NCPUS`, `PBS_NCPUS`, `PBS_NUM_PPN`, `PBS_NP`; `SLURM_CPUS_PER_TASK`,
+`SLURM_CPUS_ON_NODE`; the cgroup CPU quota; `qstat -f $PBS_JOBID`; the CPU affinity; `os.cpu_count()`.
+See the documentation's "Schedulers and containers" for each source and for sites that do it differently.
+
+## Install
+
+```
+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
+Accelerate BLAS is limited only by `VECLIB_MAXIMUM_THREADS` before NumPy is imported (threadpoolctl cannot see it).
+
+MIT licence.

stream_limits-0.1.0/README.md

@@ -0,0 +1,114 @@
+# stream-limits
+
+Cap the threads and memory of numeric Python libraries at a job's allocation, with one
+interface, and measure the memory a block of code used.
+
+On a cluster (PBS, Slurm, a container) a job gets a fixed number of CPUs and a fixed
+amount of memory, and is killed if it exceeds either. Polars, numba, BLAS and pyarrow
+each size their thread pools separately and default to every core of the *node*; code
+that asks how much memory there is sees the node's too. `stream-limits` finds the job's
+allocation, applies it to each library as far as the library allows, and says so when it
+cannot. It depends on nothing outside the standard library and imports none of the
+libraries it limits, so it can run before they are imported, which is when some of them
+(Polars) need to be told.
+
+It is what [hdfe-stream](https://github.com/leetucker/hdfe-stream),
+[logit-stream](https://github.com/leetucker/logit-stream) and
+[dpd-stream](https://github.com/leetucker/dpd-stream) use, and it works on its own.
+Documentation: <https://leetucker.github.io/stream-limits/> (build locally with
+`pip install -e ".[docs]"` and `sphinx-build -W -b html docs docs/_build/html`).
+
+```python
+import stream_limits
+stream_limits.set_limits()         # the job's allocation; or set_limits(threads=8, memory_gb=16)
+
+import polars as pl                # now capped
+import numpy as np
+```
+
+```
+$ python -m stream_limits          # what was found, from inside a job
+{"cpus": 8, "cpus_source": "PBS_NCPUS", "memory_gb": 32.0, "memory_source": "cgroup memory limit", ...}
+```
+
+Variables you have set are **not overridden** by default (a warning names any you set above the cap);
+loaded libraries are only lowered. `override=True` replaces them. More libraries are covered than the
+table below (numexpr, PyTorch, OpenMP programs, TensorFlow, rayon, joblib, Julia), and the documentation's
+"Other libraries" has recipes for those with only per-object settings (DuckDB, Dask, Ray, JAX, Spark).
+
+## What each library allows
+
+| 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 |
+| 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` |
+
+`thread_scope(n)` holds numba and BLAS to `n` threads for a block and restores them.
+
+## Memory
+
+- `detect_memory_gb()` is the memory the process may use: `STREAM_LIMITS_MEMORY_GB`, else the
+  cgroup limit (the lowest on the path to the root, so a scheduler's job limit is found), else
+  PBS's own record of the job, else `SLURM_MEM_PER_NODE`, else the machine's. PBS sites often set
+  no memory variable and enforce no cgroup; there `qstat -f $PBS_JOBID` is run once (no shell, a
+  10 s timeout, cached) and `Resource_List.mem` (else `vmem`) is read, with PBS's size suffixes.
+  `STREAM_LIMITS_QSTAT=0` turns this off.
+- `limit_memory()` gives Polars 2.x an out-of-core budget (`POLARS_OOC_MEMORY_BUDGET_MB`),
+  past which its streaming engine spills to disk instead of growing. A streaming sort of 1.9 GB
+  of Parquet peaked at 6.2 GB with the defaults and 1.8 GB with a 400 MB budget, in twice the
+  time. The budget bounds what Polars' operators hold, not the whole process (its Parquet
+  readers hold about 1.5 GB more on 8 threads). This is undocumented in Polars 2.x and read at
+  import, so it is set only for Polars 2.0 or later, only before Polars is imported, never over
+  a budget you set, and `STREAM_LIMITS_POLARS_SPILL=0` turns it off.
+- `MemoryBudget(limit_gb, reserve)` is for code that chooses between a fast path that needs
+  memory and a slower one that needs less: `budget.headroom_gb()` is the limit, less a reserve
+  (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.
+
+`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
+own options (`spill_dir`, `fraction`, ...).
+
+## Automatically
+
+`apply_scheduler_defaults()` (which is `set_limits()` with no arguments) does `limit_threads` and `limit_memory`, only for limits it finds
+(a variable, a scheduler's environment, a cgroup, a restricted CPU affinity), and sets only
+variables you have not set. Libraries built on this call it when they are imported; importing
+`stream_limits` itself changes nothing.
+
+## Environment variables
+
+| Variable | Meaning |
+|---|---|
+| `STREAM_LIMITS_MAX_THREADS` | the thread cap |
+| `STREAM_LIMITS_MEMORY_GB` | the memory limit |
+| `STREAM_LIMITS_MEMORY_RESERVE` | `MemoryBudget`'s reserve (default 0.15) |
+| `STREAM_LIMITS_POLARS_SPILL` | `0`: do not set Polars' budget |
+| `STREAM_LIMITS_POLARS_FRACTION` | the budget's share of the limit above the readers' floor (default 0.5) |
+| `STREAM_LIMITS_QSTAT` | `0`: do not run `qstat` to find a PBS job's CPUs and memory |
+| `STREAM_LIMITS_CPU_VARS` | names of variables holding a CPU count, checked first (comma-separated) |
+| `STREAM_LIMITS_MEMORY_VARS` | the same for memory; each holds a size (`16gb`) |
+| `STREAM_LIMITS_SPILL_DIR` | where Polars spills (default its temp dir: on a cluster, check `$TMPDIR`) |
+
+Thread limits are read from, in order: the argument; `STREAM_LIMITS_MAX_THREADS`; the variables named
+in `STREAM_LIMITS_CPU_VARS`; `NCPUS`, `PBS_NCPUS`, `PBS_NUM_PPN`, `PBS_NP`; `SLURM_CPUS_PER_TASK`,
+`SLURM_CPUS_ON_NODE`; the cgroup CPU quota; `qstat -f $PBS_JOBID`; the CPU affinity; `os.cpu_count()`.
+See the documentation's "Schedulers and containers" for each source and for sites that do it differently.
+
+## Install
+
+```
+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
+Accelerate BLAS is limited only by `VECLIB_MAXIMUM_THREADS` before NumPy is imported (threadpoolctl cannot see it).
+
+MIT licence.

stream_limits-0.1.0/docs/changelog.md

@@ -0,0 +1,2 @@
+```{include} ../CHANGELOG.md
+```

stream_limits-0.1.0/docs/conf.py

@@ -0,0 +1,59 @@
+"""Sphinx configuration for the stream-limits documentation.
+
+Build locally with
+
+    pip install -e ".[docs]"
+    sphinx-build -W --keep-going -b html docs docs/_build/html
+"""
+
+import stream_limits
+
+project = "stream-limits"
+author = "Lee Tucker"
+copyright = "2026, Lee Tucker"
+release = stream_limits.__version__
+version = release
+
+extensions = [
+    "myst_parser",
+    "sphinx.ext.autodoc",
+    "sphinx.ext.napoleon",
+    "sphinx.ext.intersphinx",
+    "sphinx_copybutton",
+]
+
+source_suffix = {".md": "markdown", ".rst": "restructuredtext"}
+exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
+# fenced blocks without a language are output and tables, not code
+highlight_language = "text"
+
+myst_enable_extensions = ["colon_fence", "deflist", "attrs_inline"]
+myst_heading_anchors = 4
+
+autodoc_typehints = "none"
+autodoc_member_order = "bysource"
+napoleon_google_docstring = False
+napoleon_numpy_docstring = True
+napoleon_use_rtype = False
+default_role = "py:obj"
+
+intersphinx_mapping = {
+    "python": ("https://docs.python.org/3", None),
+    "polars": ("https://docs.pola.rs/api/python/stable", None),
+    "hdfe_stream": ("https://leetucker.github.io/hdfe-stream", None),
+}
+# the inventories are fetched at build time; a build without network access
+# still succeeds, with those links left as plain text
+suppress_warnings = ["intersphinx.external", "config.cache"]
+
+html_theme = "furo"
+html_title = f"stream-limits {release}"
+html_static_path = ["_static"]
+html_theme_options = {
+    "source_repository": "https://github.com/leetucker/stream-limits/",
+    "source_branch": "main",
+    "source_directory": "docs/",
+}
+
+copybutton_prompt_text = r">>> |\.\.\. |\$ "
+copybutton_prompt_is_regexp = True

stream_limits-0.1.0/docs/index.md

@@ -0,0 +1,84 @@
+# stream-limits
+
+Cap the threads and memory of numeric Python libraries at a job's allocation, with one
+interface, and measure the memory a block of code used.
+
+On a cluster (PBS, Slurm, a container) a job gets a fixed number of CPUs and a fixed
+amount of memory, and is killed if it exceeds either. Polars, numba, BLAS and pyarrow
+each size their thread pools separately and default to every core of the *node*; code
+that asks how much memory there is sees the node's too. `stream-limits` finds the job's
+allocation, applies it to each library as far as the library allows, and says so when it
+cannot. It depends on nothing outside the standard library and imports none of the
+libraries it limits, so it can run before they are imported, which is when some of them
+(Polars) need to be told.
+
+It is what [hdfe-stream](https://leetucker.github.io/hdfe-stream/),
+[logit-stream](https://leetucker.github.io/logit-stream/) and
+[dpd-stream](https://leetucker.github.io/dpd-stream/) use, and it works on its own.
+
+```python
+import stream_limits
+stream_limits.set_limits()         # the job's allocation; or set_limits(threads=8, memory_gb=16)
+
+import polars as pl                # now capped
+import numpy as np
+```
+
+`set_limits` is {py:func}`~stream_limits.limit_threads` and {py:func}`~stream_limits.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, so the same line is right on a cluster and on a laptop. The two
+functions it calls remain, for the options the one call leaves out.
+
+```
+$ python -m stream_limits          # what was found, from inside a job
+{
+  "cpus": 8,
+  "cpus_source": "NCPUS",
+  "memory_gb": 17.18,
+  "memory_source": "qstat Resource_List.mem",
+  ...
+}
+```
+
+## Where to start
+
+- **On a cluster, before anything else**: run `python -m stream_limits` inside a job and check
+  that the CPUs and memory it reports are the ones you asked for. If they are not, see
+  [Schedulers and containers](schedulers.md), which says where each number comes from and how
+  to tell the library about a site that does it differently.
+- **To cap threads**: [Thread limits](threads.md).
+- **For a library that is not covered** (DuckDB, Dask, Ray, JAX, Spark, ...): [Other libraries](other-libraries.md).
+- **To bound memory**: [Memory limits](memory.md), which covers the memory a job has,
+  Polars' spill-to-disk budget, and `MemoryBudget`, for code that chooses between a fast path
+  that needs memory and a slower one that needs less.
+- **If you use hdfe-stream, logit-stream or dpd-stream**: they call
+  `apply_scheduler_defaults()` when imported, so a job is capped without a line of your code,
+  and their `max_memory_gb` and `n_threads` options use what is described here. Their
+  `profile=True` option, which estimates a fit's memory before you run it, is documented in
+  [hdfe-stream](https://leetucker.github.io/hdfe-stream/resources.html).
+
+## Automatically, or explicitly
+
+`apply_scheduler_defaults()` (which is `set_limits()` with no arguments) does `limit_threads` and
+`limit_memory`, **only for limits it finds** (a variable, a scheduler's environment, a cgroup, a restricted CPU affinity), and sets
+only variables you have not set yourself; with no limit it does nothing, so a workstation is
+untouched. Libraries built on `stream-limits` call it when they are imported. Importing
+`stream_limits` itself changes nothing.
+
+## Install
+
+```
+pip install stream-limits            # nothing else is installed
+pip install "stream-limits[blas]"    # threadpoolctl, to limit BLAS after NumPy is loaded
+```
+
+```{toctree}
+:hidden:
+
+threads
+memory
+schedulers
+other-libraries
+reference
+changelog
+```

stream_limits-0.1.0/docs/memory.md

@@ -0,0 +1,95 @@
+# Memory limits
+
+## The memory a job has
+
+{py:func}`stream_limits.detect_memory_gb` is the memory the process may use. It is read from,
+first match winning:
+
+1. `STREAM_LIMITS_MEMORY_GB`;
+2. the variables named in `STREAM_LIMITS_MEMORY_VARS` (a size such as `16gb`);
+3. the cgroup memory limit, the lowest one on the path to the root, so a scheduler's limit on the
+   job is found even though the process sits in a child group;
+4. PBS's own record of the job, `qstat -f $PBS_JOBID` (`Resource_List.mem`, else `vmem`), for
+   sites that set no memory variable and enforce no cgroup;
+5. `SLURM_MEM_PER_NODE`, or `SLURM_MEM_PER_CPU` times the CPUs;
+6. the machine's physical memory.
+
+Use it instead of asking the operating system how much memory there is: on a shared node that
+number is the node's, and a guard of "a quarter of the memory" then allows many times what the
+job has. [Schedulers and containers](schedulers.md) describes each source.
+
+## Polars' memory budget
+
+Polars 2.x has an out-of-core memory budget past which its streaming engine spills to disk
+instead of growing. {py:func}`stream_limits.limit_memory` sets it from the memory limit.
+
+A streaming sort of 1.9 GB of Parquet peaked at 6.2 GB with Polars' defaults, a multiple of the
+data. With a 400 MB budget it peaked at 1.8 GB, in twice the time. The budget bounds what
+Polars' operators hold, not the whole process: its Parquet readers hold about 1.5 GB more on
+8 threads (less on fewer), and your own arrays are on top. The default budget is
+`0.5 x (limit - 1.5 GB)`, at least 0.25 GB; change the share with `fraction=` or
+`STREAM_LIMITS_POLARS_FRACTION`.
+
+Some things to know:
+
+- **It is read once, at import**, like the thread pool, so `limit_memory` must run before Polars
+  is imported (a warning says so otherwise).
+- **It is undocumented in Polars 2.x.** The names are those of Polars' binary
+  (`POLARS_OOC_MEMORY_BUDGET_MB`, `POLARS_OOC_SPILL_DIR`, ...). So it is set only for Polars 2.0
+  or later (read from the package metadata, without importing Polars), never over a budget you
+  set yourself, and `STREAM_LIMITS_POLARS_SPILL=0` turns it off.
+- **Where it spills** is Polars' temporary directory. On a cluster, make sure `$TMPDIR` is the
+  job's scratch space, or give `spill_dir=` or `STREAM_LIMITS_SPILL_DIR`.
+- **How much it helps depends on the code.** Code that sorts or joins a large Parquet scan
+  through Polars' streaming engine gains the most. Code that already works in bounded pieces
+  gains little: in hdfe-stream the peak of a 3-million-row fit did not change, and in dpd-stream
+  on 4.8 million rows it fell from 2.0 to 1.6 GB, with identical results and no slowdown.
+
+## Choosing between a fast path and a smaller one
+
+{py:class}`stream_limits.MemoryBudget` is for code that has a fast path that needs a lot of
+memory and a slower one that needs less, or an allocation that cannot shrink and should be
+refused up front instead of at the operating system's kill.
+
+> headroom = limit x (1 - reserve) - memory the process holds now
+
+```python
+from stream_limits import MemoryBudget
+
+budget = MemoryBudget(limit_gb=None, reserve=None)   # defaults: the job's allocation, 0.15
+need_gb = nnz * 12 / 1e9                              # the exact size of the allocation
+if budget.fits(need_gb):
+    build_matrix()
+else:
+    stream_instead()
+```
+
+- The **limit** defaults to {py:func}`stream_limits.detect_memory_gb`.
+- The **reserve** is the share of the limit held back for what the code does not model:
+  allocator slack, transient copies, pages a cgroup charges for files. The default is 0.15
+  (`STREAM_LIMITS_MEMORY_RESERVE`). Lower it to run closer to the limit; raise it for a job that
+  must not be killed.
+- `headroom_gb(upcoming_gb)` subtracts allocations you know are coming later in the same phase,
+  and `fits(need, upcoming, share)` asks about a share of the headroom (half, for per-thread
+  copies that must leave room for others).
+- It reads the process's resident memory (one file read on Linux; 0.5 GB is assumed where it
+  cannot be read), so the answer reflects what the process holds *at that moment*.
+
+Ask whether your **exact** allocation fits, not whether an estimate's upper bound does:
+estimates are pessimistic by design, and refusing a job for an estimate's pessimism helps no one.
+
+## Measuring what was used
+
+{py:func}`stream_limits.peak_rss_gb` and {py:func}`stream_limits.reset_peak_rss` read the
+kernel's own high-water mark of the process's resident memory. Nothing is polled, so it costs
+nothing.
+
+```python
+stream_limits.reset_peak_rss()
+fit_the_model()
+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
+processes are not included. {py:func}`stream_limits.current_rss_gb` is the memory held now.