Reference¶

Limiting¶

stream_limits.set_limits(threads=None, memory_gb=None, *, spill_dir=None, override=False, environ=None)¶

Cap the threads and Polars’ memory in one call: limit_threads and limit_memory, for most scripts the only line needed, first, before Polars is imported.

A value you give is also recorded in STREAM_LIMITS_MAX_THREADS or STREAM_LIMITS_MEMORY_GB, so detect_cpus, detect_memory_gb and MemoryBudget in this process (and in the packages built on this one, and in child processes) agree with it. set_limits() with no arguments does what apply_scheduler_defaults, the call the libraries make when imported, does.

Parameters:

threads (int, optional) – The thread cap. Given, it is applied (and overrides thread variables already set). Omitted, it is the job’s CPU allocation if one is found (see detect_cpus) and nothing otherwise.
memory_gb (float, optional) – The memory limit in GB. Given, it is applied. Omitted, it is the job’s allocation if one is found (see detect_memory_gb) and nothing otherwise: a workstation with no limit gets no Polars budget.
spill_dir (path, optional) – Where Polars spills; see limit_memory.
override (bool, default False) – Replace thread and Polars variables you have already set, instead of leaving them alone (and set loaded libraries to exactly threads instead of only lowering them); see limit_threads.

Returns:

dict – {"threads": ..., "memory": ...}, each the report of limit_threads or limit_memory, left out when there was no limit to apply.

stream_limits.limit_threads(n=None, *, override=False, environ=None, _auto=False, _stack=0)¶

Cap the threads that polars, numba, BLAS, pyarrow, numexpr and the other libraries named in the module docstring use, as far as each allows.

Parameters:

n (int, optional) – The cap. Default: detect_cpus(), the job’s allocation.
override (bool, default False) – By default a thread variable you have already set (OMP_NUM_THREADS=1, say) is left alone, and loaded libraries are only ever lowered: the cap is a ceiling, and a more specific choice of yours stays. A warning names any variable you set above the cap. override=True replaces the variables and sets the loaded libraries to exactly n.
it (Call it before importing polars. Before the libraries are imported)
variables (sets their environment)
Polars (which is the only way to limit)
lowers ((its pool is fixed at import). After they are imported it still)
numba's
BLAS's
pyarrow's
threads (numexpr's and PyTorch's)
cannot (but)
Polars' (touch)
to (and says so with a warning; a limit Polars would have)
POLARS_MAX_THREADS) (honor goes in STREAM_LIMITS_MAX_THREADS (or)
script. (in the job)

Returns:

dict – What it did: set_env (variables set), kept (variables you had set, with their values, left alone), lowered (libraries lowered at run time) and too_late (libraries it could not limit).

stream_limits.limit_memory(limit_gb=None, *, fraction=None, spill_dir=None, override=False, environ=None, _auto=False, _stack=0)¶

Give Polars a memory budget, past which its streaming engine spills to disk instead of growing.

A streaming sort of 1.9 GB of Parquet peaks 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 the engine’s operators hold, not the whole process: the Parquet readers hold about 1.5 GB more on 8 threads (less on fewer), and the packages’ own arrays are on top.

Parameters:

limit_gb (float, optional) – The memory the process may use. Default: detect_memory_gb().
fraction (float, optional) – The budget’s share of what is left after the readers’ floor: fraction x (limit - 1.5 GB), at least 0.25 GB. Default $STREAM_LIMITS_POLARS_FRACTION, else 0.5.
spill_dir (path, optional) – Where Polars spills (POLARS_OOC_SPILL_DIR). Default: its own temporary directory, which on a cluster should be the job’s scratch (check $TMPDIR).
override (bool, default False) – By default a Polars budget or spill directory you have already set in the environment is left alone. override=True replaces them.
warning (Like the thread limit this must come before Polars is imported (a)
2.0 (says so otherwise). It does nothing when Polars is older than)
a (when)
environment (budget is already set in the)
when (or)
2.x; (STREAM_LIMITS_POLARS_SPILL=0. Polars documents none of this in)
did. (the names are those of its binary. Returns what it)

stream_limits.apply_scheduler_defaults(environ=None)¶: Cap the threads and Polars’ memory at the scheduler’s allocation, if there is one. The estimator packages call this first thing when imported. It acts only on limits it finds (a variable, a scheduler’s environment, a cgroup, a restricted CPU affinity): the thread cap sets only the variables the user has not set; with no limit it does nothing, so a workstation is untouched. Returns {"threads": ..., "memory": ...} (each the report of limit_threads / limit_memory), or None.

stream_limits.thread_scope(n=None)¶: Hold numba and BLAS to n threads inside the block and restore them after. None (or 0) leaves everything as it is. Polars is not affected: its pool is fixed for the process.

stream_limits.set_openmp_wait_policy(policy='passive', *, override=False, environ=None)¶

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.

stream_limits.polars_spills()¶: Whether the installed Polars has the out-of-core memory budget that limit_memory sets: 2.0 and later (it is an environment variable there, undocumented, and read once at import).

Finding the allocation¶

stream_limits.detect_cpus(environ=None, *, cgroup_root=PosixPath('/sys/fs/cgroup'), proc_cgroup=PosixPath('/proc/self/cgroup'))¶: The number of CPUs this process may use: the limit a scheduler, container or STREAM_LIMITS_MAX_THREADS sets (see the module docstring for the order), otherwise the machine’s.

stream_limits.detect_memory_gb(environ=None, *, cgroup_root=PosixPath('/sys/fs/cgroup'), proc_cgroup=PosixPath('/proc/self/cgroup'))¶: The memory in GB this process may use: STREAM_LIMITS_MEMORY_GB, the cgroup limit a scheduler sets for the job, or else the machine’s physical memory.

stream_limits.describe(environ=None)¶: What was found, and what the libraries are doing now: for a log line, a bug report, or python -m stream_limits.

A memory budget¶

class stream_limits.MemoryBudget(limit_gb=None, reserve=None, *, environ=None)¶

How much memory is left to allocate, for guards that must choose between a fast path that needs a lot of memory and a slower one that needs less.

headroom_gb() is the limit, less the share reserve keeps back for what the code does not model (allocator slack, transient copies, pages the cgroup charges for files), less what the process holds now. A guard asks whether its exact allocation fits, instead of comparing it with a fixed fraction of the machine. Reading the process’s memory is one file read.

Parameters:

limit_gb (float, optional) – The memory the process may use. Default: detect_memory_gb().
reserve (float, optional) – The share of the limit kept back, 0 to 1. Default $STREAM_LIMITS_MEMORY_RESERVE, else 0.15: lower it to trust the estimates more, raise it for a job that must not be killed.

property usable_gb¶: The limit less the reserve.

held_gb()¶: What the process holds now (0.5 where it cannot be read).

headroom_gb(upcoming_gb=0.0)¶: What can still be allocated, after upcoming_gb that will be allocated later in the same phase; never negative.

fits(need_gb, upcoming_gb=0.0, share=1.0)¶: Whether an allocation of need_gb fits within share of the headroom.

Measuring memory¶

stream_limits.peak_rss_gb()¶: The peak resident memory, in GB, since reset_peak_rss (or since the 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.

stream_limits.reset_peak_rss()¶: 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).

stream_limits.current_rss_gb()¶: The resident memory in GB right now (Linux and Windows; None on macOS, where MemoryBudget then assumes 0.5 GB).

Environment variables¶

Variable	Meaning
`STREAM_LIMITS_MAX_THREADS`	the thread cap
`STREAM_LIMITS_MEMORY_GB`	the memory limit
`STREAM_LIMITS_CPU_VARS`	comma-separated names of variables holding a CPU count, checked before the built-in ones
`STREAM_LIMITS_MEMORY_VARS`	the same for memory; each holds a size (`16gb`, `512mb`; a bare number is bytes)
`STREAM_LIMITS_MEMORY_RESERVE`	`MemoryBudget`’s reserve (default 0.15)
`STREAM_LIMITS_POLARS_SPILL`	`0`: do not set Polars’ out-of-core budget
`STREAM_LIMITS_POLARS_FRACTION`	the budget’s share of the limit above the readers’ floor (default 0.5)
`STREAM_LIMITS_SPILL_DIR`	where Polars spills (default its temp dir: on a cluster, check `$TMPDIR`)
`STREAM_LIMITS_QSTAT`	`0`: do not run `qstat` to find a PBS job’s CPUs and memory

limit_threads also sets these, unless you already have (override=True replaces them): POLARS_MAX_THREADS, NUMBA_NUM_THREADS, OMP_NUM_THREADS, OMP_THREAD_LIMIT, OPENBLAS_NUM_THREADS, MKL_NUM_THREADS, BLIS_NUM_THREADS, VECLIB_MAXIMUM_THREADS, 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¶

Opt-in subprocess monitoring, sampled watchdogs and Linux kernel memory caps.

Install stream-limits[launcher] for psutil. Neither importing this module nor running a command imports a numerical library. Existing stream-limits defaults and the parent process’s environment are unchanged.

class stream_limits.launcher.RunResult(command: list[str], mode: str, status: str, returncode: int, elapsed_s: float, requested_memory_bytes: int | None, peak_rss_bytes: int, kernel: dict | None, stdout: str, stderr: str, stdout_truncated: bool, stderr_truncated: bool, descendants_terminated: bool, rss_sampling_complete: bool, os_peak_rss_bytes: int | None = None, peak_private_bytes: int | None = None, memory_measure: str = 'rss')¶

stream_limits.launcher.capabilities()¶: Backend availability, not a claim that a controller is delegated.

stream_limits.launcher.thread_env(threads, env=None)¶: A copy of env (default os.environ) with STREAM_LIMITS_MAX_THREADS and every thread variable stream-limits knows (Polars, numba, OpenMP, BLAS, numexpr, rayon, Julia) set to threads.

stream_limits.launcher.run_limited(command, *, memory_gb=None, mode='hard', cgroup_parent=None, timeout=None, env=None, cwd=None, sample_interval=0.01, max_output_bytes=1000000, threads=None, capture_output=True, memory_measure='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 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 output uses temporary disk files; only a bounded tail is retained. With capture_output=False the child writes to this process’s stdout/stderr and the result’s stdout/stderr are empty.

Reusable enforcement verification with no estimator or pytest dependency.

stream_limits.testing.kernel_oom_proved(result, metrics)¶: A signal or accepted launch alone is not evidence of kernel enforcement.

stream_limits.testing.verify_memory_limit(parent, *, python=None, env=None, timeout=15)¶

Verify a small successful allocation, inherited detection and kernel OOM.

The deliberate OOM affects only a fresh 64 MiB child group. The delegate must allow at least 64 MiB. Raises if enforcement/detection is not proved.