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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (35) hide show
  1. {stream_limits-0.1.1 → stream_limits-0.2.0}/CHANGELOG.md +17 -0
  2. {stream_limits-0.1.1 → stream_limits-0.2.0}/CITATION.cff +2 -2
  3. {stream_limits-0.1.1 → stream_limits-0.2.0}/PKG-INFO +21 -1
  4. {stream_limits-0.1.1 → stream_limits-0.2.0}/README.md +14 -0
  5. {stream_limits-0.1.1 → stream_limits-0.2.0}/docs/index.md +12 -0
  6. stream_limits-0.2.0/docs/launcher.md +236 -0
  7. {stream_limits-0.1.1 → stream_limits-0.2.0}/docs/reference.md +10 -0
  8. {stream_limits-0.1.1 → stream_limits-0.2.0}/pyproject.toml +5 -2
  9. {stream_limits-0.1.1 → stream_limits-0.2.0}/stream_limits/__init__.py +1 -1
  10. stream_limits-0.2.0/stream_limits/_exec_cgroup.py +18 -0
  11. stream_limits-0.2.0/stream_limits/enforcement.py +199 -0
  12. stream_limits-0.2.0/stream_limits/launcher.py +376 -0
  13. stream_limits-0.2.0/stream_limits/pytest_plugin.py +76 -0
  14. stream_limits-0.2.0/stream_limits/testing.py +64 -0
  15. {stream_limits-0.1.1 → stream_limits-0.2.0}/stream_limits.egg-info/PKG-INFO +21 -1
  16. {stream_limits-0.1.1 → stream_limits-0.2.0}/stream_limits.egg-info/SOURCES.txt +8 -0
  17. {stream_limits-0.1.1 → stream_limits-0.2.0}/stream_limits.egg-info/requires.txt +8 -0
  18. stream_limits-0.2.0/tests/conftest.py +1 -0
  19. stream_limits-0.2.0/tests/test_launcher.py +330 -0
  20. {stream_limits-0.1.1 → stream_limits-0.2.0}/LICENSE.md +0 -0
  21. {stream_limits-0.1.1 → stream_limits-0.2.0}/MANIFEST.in +0 -0
  22. {stream_limits-0.1.1 → stream_limits-0.2.0}/docs/_static/.gitkeep +0 -0
  23. {stream_limits-0.1.1 → stream_limits-0.2.0}/docs/changelog.md +0 -0
  24. {stream_limits-0.1.1 → stream_limits-0.2.0}/docs/conf.py +0 -0
  25. {stream_limits-0.1.1 → stream_limits-0.2.0}/docs/memory.md +0 -0
  26. {stream_limits-0.1.1 → stream_limits-0.2.0}/docs/other-libraries.md +0 -0
  27. {stream_limits-0.1.1 → stream_limits-0.2.0}/docs/schedulers.md +0 -0
  28. {stream_limits-0.1.1 → stream_limits-0.2.0}/docs/threads.md +0 -0
  29. {stream_limits-0.1.1 → stream_limits-0.2.0}/setup.cfg +0 -0
  30. {stream_limits-0.1.1 → stream_limits-0.2.0}/stream_limits/__main__.py +0 -0
  31. {stream_limits-0.1.1 → stream_limits-0.2.0}/stream_limits/py.typed +0 -0
  32. {stream_limits-0.1.1 → stream_limits-0.2.0}/stream_limits.egg-info/dependency_links.txt +0 -0
  33. {stream_limits-0.1.1 → stream_limits-0.2.0}/stream_limits.egg-info/top_level.txt +0 -0
  34. {stream_limits-0.1.1 → stream_limits-0.2.0}/tests/test_limits.py +0 -0
  35. {stream_limits-0.1.1 → stream_limits-0.2.0}/tests/test_override_and_libraries.py +0 -0
@@ -1,5 +1,22 @@
1
1
  # Changelog
2
2
 
3
+ ## 0.2.0
4
+
5
+ - **Run a program under a memory limit and see how much memory it used.** The new
6
+ launcher runs a script, test suite or other program in a child process, stops it
7
+ if it goes over a memory limit or a time limit, and reports why it ended and its
8
+ peak memory. The limit is a watchdog that checks memory every 10 ms (Linux,
9
+ macOS, Windows), or a kernel limit through a Linux cgroup. It also sets the
10
+ program's threads (`--threads`). It is an optional extra,
11
+ `pip install 'stream-limits[launcher]'`, which adds psutil; the rest of
12
+ stream-limits is unchanged and still needs only the standard library. See
13
+ [Running a program under a memory limit](https://leetucker.github.io/stream-limits/launcher.html).
14
+ - **Pytest fixtures for memory tests.** `pytest_plugins =
15
+ ["stream_limits.pytest_plugin"]` gives a test suite `resource_runner`, which runs
16
+ a test's work in a child process under the watchdog on every platform, or under a
17
+ kernel limit when one is set up (`STREAM_LIMITS_TEST_CGROUP_PARENT`). Install
18
+ `stream-limits[testing]`.
19
+
3
20
  ## 0.1.1
4
21
 
5
22
  - **Windows reports memory use.** `peak_rss_gb` and `current_rss_gb` returned `None` on Windows in
@@ -7,8 +7,8 @@ authors:
7
7
  - family-names: Tucker
8
8
  given-names: Lee
9
9
  email: lee.c.tucker@census.gov
10
- version: 0.1.1
11
- date-released: 2026-10-31
10
+ version: 0.2.0
11
+ date-released: 2026-11-04
12
12
  license: MIT
13
13
  repository-code: "https://github.com/leetucker/stream-limits"
14
14
  url: "https://github.com/leetucker/stream-limits"
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: stream-limits
3
- Version: 0.1.1
3
+ Version: 0.2.0
4
4
  Summary: Cap the threads and memory of numeric Python libraries (Polars, numba, BLAS, pyarrow) at a scheduler's allocation, and measure peak memory
5
5
  Author: Lee Tucker
6
6
  License-Expression: MIT
@@ -26,6 +26,11 @@ Classifier: Typing :: Typed
26
26
  Requires-Python: >=3.10
27
27
  Description-Content-Type: text/markdown
28
28
  License-File: LICENSE.md
29
+ Provides-Extra: launcher
30
+ Requires-Dist: psutil>=5.9; extra == "launcher"
31
+ Provides-Extra: testing
32
+ Requires-Dist: pytest>=7.0; extra == "testing"
33
+ Requires-Dist: psutil>=5.9; extra == "testing"
29
34
  Provides-Extra: blas
30
35
  Requires-Dist: threadpoolctl>=3.1; extra == "blas"
31
36
  Provides-Extra: docs
@@ -35,6 +40,7 @@ Requires-Dist: furo>=2024.8.6; extra == "docs"
35
40
  Requires-Dist: sphinx-copybutton>=0.5; extra == "docs"
36
41
  Provides-Extra: test
37
42
  Requires-Dist: pytest>=7.0; extra == "test"
43
+ Requires-Dist: psutil>=5.9; extra == "test"
38
44
  Requires-Dist: numpy; extra == "test"
39
45
  Requires-Dist: numba; extra == "test"
40
46
  Requires-Dist: polars; extra == "test"
@@ -157,4 +163,18 @@ applies thread limits and reads memory use (peak and current working set), but i
157
163
  macOS and Windows report a lifetime peak, as their counters cannot be reset. On macOS, NumPy's Apple
158
164
  Accelerate BLAS is limited only by `VECLIB_MAXIMUM_THREADS` before NumPy is imported (threadpoolctl cannot see it).
159
165
 
166
+ ## Running a program under a memory limit
167
+
168
+ The optional launcher (`pip install 'stream-limits[launcher]'`, which adds psutil)
169
+ runs a program in a child process, stops it if it goes over a memory or time
170
+ limit, and reports why it ended and its peak memory:
171
+
172
+ ```bash
173
+ python -m stream_limits.launcher --mode watchdog --memory-gb 4 --threads 2 -- python analysis.py
174
+ ```
175
+
176
+ The watchdog works on Linux, macOS and Windows; a kernel-enforced limit needs a
177
+ Linux cgroup set up for it. `stream-limits[testing]` adds pytest fixtures for
178
+ memory tests. See [Running a program under a memory limit](docs/launcher.md).
179
+
160
180
  MIT licence.
@@ -112,4 +112,18 @@ applies thread limits and reads memory use (peak and current working set), but i
112
112
  macOS and Windows report a lifetime peak, as their counters cannot be reset. On macOS, NumPy's Apple
113
113
  Accelerate BLAS is limited only by `VECLIB_MAXIMUM_THREADS` before NumPy is imported (threadpoolctl cannot see it).
114
114
 
115
+ ## Running a program under a memory limit
116
+
117
+ The optional launcher (`pip install 'stream-limits[launcher]'`, which adds psutil)
118
+ runs a program in a child process, stops it if it goes over a memory or time
119
+ limit, and reports why it ended and its peak memory:
120
+
121
+ ```bash
122
+ python -m stream_limits.launcher --mode watchdog --memory-gb 4 --threads 2 -- python analysis.py
123
+ ```
124
+
125
+ The watchdog works on Linux, macOS and Windows; a kernel-enforced limit needs a
126
+ Linux cgroup set up for it. `stream-limits[testing]` adds pytest fixtures for
127
+ memory tests. See [Running a program under a memory limit](docs/launcher.md).
128
+
115
129
  MIT licence.
@@ -82,3 +82,15 @@ other-libraries
82
82
  reference
83
83
  changelog
84
84
  ```
85
+
86
+ ## Running a program under a memory limit
87
+
88
+ The optional [launcher](launcher.md) runs a program in a child process, stops it
89
+ if it goes over a memory or time limit, and reports its peak memory. It also
90
+ gives other packages' test suites fixtures for memory tests.
91
+
92
+ ```{toctree}
93
+ :hidden:
94
+
95
+ launcher
96
+ ```
@@ -0,0 +1,236 @@
1
+ # Running a program under a memory limit
2
+
3
+ The launcher runs a program (a Python or R script, a test suite, any executable)
4
+ in a child process, watches its memory, and stops it if it goes over a limit you
5
+ set. When the program ends it tells you why it ended and how much memory it used.
6
+
7
+ It is an optional extra, which adds psutil and nothing else:
8
+
9
+ ```bash
10
+ pip install 'stream-limits[launcher]'
11
+ ```
12
+
13
+ Nothing else in stream-limits changes: `import stream_limits` still needs only
14
+ the standard library, and the launcher imports no numerical library itself.
15
+
16
+ ## A first example
17
+
18
+ ```bash
19
+ python -m stream_limits.launcher --mode watchdog --memory-gb 4 --threads 2 \
20
+ --timeout 600 -- python analysis.py
21
+ ```
22
+
23
+ This runs `python analysis.py` with 2 threads for every numerical library, stops
24
+ it if it and the processes it starts use more than 4 GB together or run longer
25
+ than 10 minutes, and prints a JSON summary: the status, the exit code, the time,
26
+ the peak memory, and the end of the program's output. The exit code is the
27
+ program's own (124 on a timeout).
28
+
29
+ The same from Python:
30
+
31
+ ```python
32
+ import sys
33
+ from stream_limits.launcher import run_limited
34
+
35
+ result = run_limited([sys.executable, "analysis.py"],
36
+ mode="watchdog", memory_gb=4, threads=2, timeout=600)
37
+ print(result.status, result.os_peak_rss_bytes / 1e9, "GB")
38
+ ```
39
+
40
+ The command is a list of arguments, run without a shell. GB are decimal
41
+ (1 GB = 10^9 bytes).
42
+
43
+ ## Three ways to limit
44
+
45
+ | Mode | Platforms | What it does |
46
+ |---|---|---|
47
+ | `watchdog` | Linux, macOS, Windows | Checks memory about every 10 ms and stops the program when it passes the limit |
48
+ | `hard` | Linux, with a cgroup set up for it | The kernel enforces the limit |
49
+ | `monitor` | Linux, macOS, Windows | No limit: measures, with an optional timeout |
50
+
51
+ **The watchdog** adds up the memory of the program and every process it has
52
+ started, and stops them all when the total passes the limit. Because it samples,
53
+ it can overshoot a little before it acts, and a peak shorter than the interval
54
+ can slip between samples. In practice it is tight: on Linux, with a 1.5 GB limit,
55
+ a fast 3 GB NumPy allocation, a 2.5 GB spike lasting 20 ms and a grandchild
56
+ process allocating 3 GB were each stopped at 1.51-1.52 GB. Memory shared
57
+ between processes can be counted more than once. If the operating system denies
58
+ it access to a process's memory, the watchdog stops the program and reports
59
+ `monitoring_failed` rather than claim it enforced the limit.
60
+
61
+ **Hard limits** leave nothing to sampling: the kernel will not let the program
62
+ use more, and kills it if it tries. They need a Linux cgroup that you are
63
+ allowed to manage (below). `hard` is the default mode, so that asking for a hard
64
+ limit where none is possible fails before the program starts, instead of quietly
65
+ running under the watchdog. macOS and Windows have no hard mode.
66
+
67
+ ## Results
68
+
69
+ `run_limited` returns a `RunResult` (`as_dict()` turns it into JSON). Its
70
+ `status` is one of:
71
+
72
+ | Status | Meaning |
73
+ |---|---|
74
+ | `ok` | The program finished with exit code 0 |
75
+ | `command_failed` | It finished with another exit code |
76
+ | `timeout` | It ran past `timeout` and was stopped |
77
+ | `rss_limit_exceeded` | The watchdog stopped it |
78
+ | `monitoring_failed` | The watchdog could not read its memory and stopped it |
79
+ | `kernel_oom` | The kernel killed a process for going over a hard limit |
80
+
81
+ `kernel_oom` comes from the kernel's own counter, not from guessing at a killed
82
+ process. It is reported even when the process killed was a child the program
83
+ started and the program itself carried on and exited with 0, so check `status`
84
+ and not only the exit code.
85
+
86
+ Two peaks are reported:
87
+
88
+ - `os_peak_rss_bytes`: the largest memory of any single process, as the
89
+ operating system records it. On Linux and macOS it is exact, and covers the
90
+ program and every process it started and waited for (not one it left running
91
+ in the background). On Windows it is read while sampling. Use it when a test
92
+ needs a reliable peak.
93
+ - `peak_rss_bytes`: the largest sampled total over all the program's processes
94
+ at once. Use it for a program with several worker processes. A program that
95
+ ends before the first sample reports 0.
96
+
97
+ With a hard limit, `result.kernel` adds the kernel's own figures: the limit, any
98
+ smaller limit set higher up, its peak (which counts file cache and kernel
99
+ memory, so it differs from the two above), and its event counters.
100
+
101
+ ## Output, threads and the environment
102
+
103
+ By default the program's output is saved to temporary files and the last
104
+ megabyte of each stream is returned in the result (`max_output_bytes=`). To see
105
+ it as it runs, as when wrapping a whole test run, pass `--no-capture`
106
+ (`capture_output=False`). The program then writes to the terminal and the JSON
107
+ summary goes to stderr; `--report FILE` also saves the summary to a file:
108
+
109
+ ```bash
110
+ python -m stream_limits.launcher --mode watchdog --memory-gb 4 --threads 2 \
111
+ --no-capture --report memory.json -- python -m pytest -q
112
+ ```
113
+
114
+ `--threads N` (`threads=N`) sets `STREAM_LIMITS_MAX_THREADS` and the thread
115
+ variables of Polars, numba, OpenMP, the BLAS libraries, numexpr, rayon and Julia
116
+ in the program's environment. `stream_limits.launcher.thread_env(N, env)`
117
+ returns that environment. Without `threads`, no variable is added: `env`
118
+ replaces the environment as in `subprocess`, and the launcher's own environment
119
+ is never changed.
120
+
121
+ On a timeout or when a limit is hit, the program and the processes it started
122
+ are stopped. On Linux and macOS the launcher also stops anything left in the
123
+ program's process group, and with a hard limit anything left in its cgroup. On
124
+ Windows, a process that escapes before it is first sampled can survive.
125
+
126
+ ## Setting up hard limits
127
+
128
+ Hard limits use a Linux cgroup with the memory controller, which an
129
+ administrator (or you, with root) sets up once and lets you manage. The launcher
130
+ never uses sudo and never changes a cgroup it did not create: for each run it
131
+ creates a fresh group inside the one you give it, sets the limit there, starts
132
+ the program inside it (the launcher itself stays outside the limit), and removes
133
+ the group afterwards. It is a resource limit, not a security sandbox.
134
+
135
+ `python -m stream_limits.launcher --capabilities` and `cat /proc/self/cgroup`
136
+ show what the machine has.
137
+
138
+ **cgroup v1** (WSL2, older distributions), once, as root:
139
+
140
+ ```bash
141
+ sudo mkdir /sys/fs/cgroup/memory/stream-limits-test
142
+ sudo chown -R "$USER" /sys/fs/cgroup/memory/stream-limits-test
143
+ export STREAM_LIMITS_TEST_CGROUP_PARENT=/sys/fs/cgroup/memory/stream-limits-test
144
+ ```
145
+
146
+ The group is gone after a reboot.
147
+
148
+ **cgroup v2 with systemd**: start a shell in a user scope that you may manage:
149
+
150
+ ```bash
151
+ systemd-run --user --scope -p Delegate=yes -p MemoryMax=8G bash
152
+ export STREAM_LIMITS_TEST_CGROUP_PARENT=/sys/fs/cgroup$(cut -d: -f3 /proc/self/cgroup)
153
+ ```
154
+
155
+ On v2 the group must have `memory` in its `cgroup.subtree_control` and no
156
+ processes of its own, so the shell may first need to move into a child group.
157
+
158
+ These recipes are standard, but stream-limits' tests do not run them (they need
159
+ root or a systemd scope). Check a setup before relying on it:
160
+
161
+ ```bash
162
+ python -c "from stream_limits.testing import verify_memory_limit as v; v('$STREAM_LIMITS_TEST_CGROUP_PARENT')"
163
+ ```
164
+
165
+ This runs a small allocation that must succeed, checks that a program inside the
166
+ group sees the limit, and runs a 256 MiB allocation in a 64 MiB group that the
167
+ kernel must kill. It raises unless all three happen. (The group needs room for
168
+ 64 MiB.)
169
+
170
+ Then:
171
+
172
+ ```python
173
+ result = run_limited([sys.executable, "analysis.py"], mode="hard", memory_gb=4,
174
+ cgroup_parent="/sys/fs/cgroup/memory/stream-limits-test")
175
+ ```
176
+
177
+ Before running anything, the launcher checks that the group really is a kernel
178
+ cgroup (not an ordinary directory with the same file names) and reads the limit
179
+ back. It turns swap off for the group (v2: `memory.swap.max=0`; v1: a
180
+ memory-plus-swap limit equal to the memory limit and `swappiness=0`), so the
181
+ limit cannot be dodged by swapping. On v2 the kernel stops the whole group
182
+ together when one process goes over. Older v2 kernels have no peak counter;
183
+ their kernel peak is then `None`.
184
+
185
+ The launcher's tests have exercised hard limits on a WSL2 host with cgroup v1.
186
+ cgroup v2, macOS and Windows rely on CI and have not been checked by hand.
187
+
188
+ ## Memory tests in pytest
189
+
190
+ The plugin gives other packages' test suites a common way to run a test's work in
191
+ a child process under a memory limit. Install `stream-limits[testing]` and
192
+ register it in `conftest.py` (it is not loaded automatically):
193
+
194
+ ```python
195
+ pytest_plugins = ["stream_limits.pytest_plugin"]
196
+ ```
197
+
198
+ A memory test then looks like this:
199
+
200
+ ```python
201
+ import sys
202
+ import pytest
203
+
204
+ @pytest.mark.resource
205
+ def test_fit_stays_within_memory(resource_runner, tmp_path):
206
+ result = resource_runner([sys.executable, "fit_large_case.py", str(tmp_path)],
207
+ memory_gb=1.0, threads=2, timeout=120)
208
+ assert result.status == "ok", result.as_dict()
209
+ ```
210
+
211
+ `resource_runner` takes `run_limited`'s arguments except the mode. It uses the
212
+ watchdog on every platform, so these tests run on macOS and Windows too. When
213
+ `STREAM_LIMITS_TEST_CGROUP_PARENT` names a cgroup set up as above, it checks it
214
+ once per session with `verify_memory_limit` and uses a hard limit instead; a
215
+ setting that does not work fails the session rather than falling back.
216
+
217
+ ```bash
218
+ pytest -m resource # watchdog
219
+ STREAM_LIMITS_TEST_CGROUP_PARENT=/sys/fs/cgroup/memory/stream-limits-test pytest -m resource
220
+ ```
221
+
222
+ The other fixtures: `resource_enforcement` says which mode was chosen
223
+ (`{"mode": "watchdog"}` or `{"mode": "hard", "cgroup_parent": ...}`);
224
+ `memory_runner` is `run_limited` itself, for a test that picks its own mode; and
225
+ `verified_memory_parent` is the checked cgroup, for tests that need a hard limit
226
+ and skip without one. The plugin registers the markers `resource` and
227
+ `kernel_memory`.
228
+
229
+ **A memory test is only as good as its input.** A child process that imports
230
+ Polars, NumPy and numba already holds 0.2-0.4 GB, so a small input under a 1.5 GB
231
+ limit passes whatever the code does with it. To catch code that holds the data in
232
+ memory, give the test an input larger in memory than the limit less that
233
+ baseline, or compare `os_peak_rss_bytes` at N rows and at several times N.
234
+
235
+ Use the launcher only for tests about memory. Ordinary tests are better run in
236
+ the test process, as before.
@@ -65,3 +65,13 @@
65
65
  `NUMEXPR_NUM_THREADS`, `NUMEXPR_MAX_THREADS`, `RAYON_NUM_THREADS`, `TF_NUM_INTRAOP_THREADS`,
66
66
  `TF_NUM_INTEROP_THREADS`, `LOKY_MAX_CPU_COUNT`, `JULIA_NUM_THREADS`. `limit_memory` sets
67
67
  `POLARS_OOC_MEMORY_BUDGET_MB` and, if asked, `POLARS_OOC_SPILL_DIR`.
68
+
69
+ ## Optional launcher
70
+
71
+ ```{eval-rst}
72
+ .. automodule:: stream_limits.launcher
73
+ :members: RunResult, run_limited, capabilities, thread_env
74
+
75
+ .. automodule:: stream_limits.testing
76
+ :members: verify_memory_limit, kernel_oom_proved
77
+ ```
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "stream-limits"
7
- version = "0.1.1"
7
+ version = "0.2.0"
8
8
  description = "Cap the threads and memory of numeric Python libraries (Polars, numba, BLAS, pyarrow) at a scheduler's allocation, and measure peak memory"
9
9
  readme = "README.md"
10
10
  license = "MIT"
@@ -33,6 +33,9 @@ requires-python = ">=3.10"
33
33
  dependencies = []
34
34
 
35
35
  [project.optional-dependencies]
36
+ # Opt-in subprocess launcher; core installs remain standard-library-only.
37
+ launcher = ["psutil>=5.9"]
38
+ testing = ["pytest>=7.0", "psutil>=5.9"]
36
39
  # threadpoolctl lets thread_scope()/limit_threads() limit BLAS after numpy is loaded
37
40
  blas = ["threadpoolctl>=3.1"]
38
41
  # docs/: the documentation site, built with Sphinx (see docs/conf.py)
@@ -42,7 +45,7 @@ docs = [
42
45
  "furo>=2024.8.6",
43
46
  "sphinx-copybutton>=0.5",
44
47
  ]
45
- test = ["pytest>=7.0", "numpy", "numba", "polars", "pyarrow", "threadpoolctl>=3.1",
48
+ test = ["pytest>=7.0", "psutil>=5.9", "numpy", "numba", "polars", "pyarrow", "threadpoolctl>=3.1",
46
49
  'numexpr; python_version < "3.15"']
47
50
 
48
51
  [project.urls]
@@ -87,7 +87,7 @@ import warnings
87
87
  from pathlib import Path
88
88
  from typing import Iterator, Mapping
89
89
 
90
- __version__ = "0.1.1"
90
+ __version__ = "0.2.0"
91
91
 
92
92
  __all__ = [
93
93
  "set_limits", "limit_threads", "limit_memory", "apply_scheduler_defaults", "thread_scope",
@@ -0,0 +1,18 @@
1
+ """Internal pre-exec attachment helper, deliberately standard-library-only."""
2
+
3
+ import os
4
+ import sys
5
+ from pathlib import Path
6
+
7
+
8
+ def main():
9
+ if len(sys.argv) < 4 or sys.argv[2] != "--":
10
+ raise SystemExit("internal usage: PARENT -- COMMAND ...")
11
+ path = Path(sys.argv[1])
12
+ (path / "cgroup.procs").write_text(str(os.getpid()))
13
+ command = sys.argv[3:]
14
+ os.execvpe(command[0], command, os.environ)
15
+
16
+
17
+ if __name__ == "__main__":
18
+ main()
@@ -0,0 +1,199 @@
1
+ """Explicit Linux cgroup memory enforcement; no numerical dependencies.
2
+
3
+ The caller supplies an existing delegated parent. This module never invokes
4
+ sudo, mounts controllers, or changes limits on a pre-existing group.
5
+ """
6
+
7
+ import math
8
+ import os
9
+ import signal
10
+ import sys
11
+ import time
12
+ from pathlib import Path
13
+
14
+
15
+ class UnsupportedEnforcement(RuntimeError):
16
+ """The requested hard-limit backend is unavailable; no fallback was used."""
17
+
18
+
19
+ def cgroup_capability(parent=None):
20
+ path = Path(parent) if parent else Path("/sys/fs/cgroup/unified")
21
+ if not path.exists():
22
+ path = Path("/sys/fs/cgroup")
23
+ return {
24
+ "path": str(path),
25
+ "v2": (path / "cgroup.controllers").exists(),
26
+ "v1": (path / "memory.limit_in_bytes").exists(),
27
+ "writable": os.access(path, os.W_OK),
28
+ "memory_controller": (path / "memory.max").exists()
29
+ or (path / "memory.limit_in_bytes").exists(),
30
+ "mounts": [
31
+ line
32
+ for line in Path("/proc/mounts").read_text().splitlines()
33
+ if " cgroup" in line
34
+ ]
35
+ if Path("/proc/mounts").exists()
36
+ else [],
37
+ "note": "No cgroup is created unless --cgroup-parent is supplied explicitly.",
38
+ }
39
+
40
+
41
+ class MemoryCgroup:
42
+ """Own one child group; the supervising driver stays outside its limit."""
43
+
44
+ def __init__(self, parent, limit):
45
+ if sys.platform != "linux":
46
+ raise UnsupportedEnforcement("cgroup memory enforcement requires Linux")
47
+ if isinstance(limit, bool) or not math.isfinite(limit) or limit < 1:
48
+ raise ValueError("memory limit must be finite and positive")
49
+ parent = Path(parent)
50
+ if (parent / "memory.limit_in_bytes").exists():
51
+ self.version = 1
52
+ elif (parent / "memory.max").exists() or (
53
+ parent / "cgroup.controllers"
54
+ ).exists():
55
+ self.version = 2
56
+ else:
57
+ raise ValueError("parent has no supported memory controller")
58
+ # File names/readback in an ordinary directory do not prove that a
59
+ # kernel controller is present. Verify the actual filesystem first.
60
+ import ctypes
61
+
62
+ libc = ctypes.CDLL(None, use_errno=True)
63
+ statfs = libc.statfs
64
+ statfs.argtypes = [ctypes.c_char_p, ctypes.c_void_p]
65
+ statfs.restype = ctypes.c_int
66
+ data = ctypes.create_string_buffer(256)
67
+ if statfs(os.fsencode(parent), data) != 0:
68
+ raise OSError(
69
+ ctypes.get_errno(), "cannot inspect cgroup filesystem", str(parent)
70
+ )
71
+ magic = ctypes.c_long.from_buffer(data).value
72
+ expected_magic = 0x27E0EB if self.version == 1 else 0x63677270
73
+ if magic != expected_magic:
74
+ raise UnsupportedEnforcement(
75
+ "parent is not on a genuine kernel memory cgroup filesystem"
76
+ )
77
+ self.limit = int(limit)
78
+ self.path = Path(parent) / f"stream-limits-{os.getpid()}-{time.time_ns()}"
79
+ self.path.mkdir()
80
+ try:
81
+ if self.version == 2:
82
+ (self.path / "memory.max").write_text(str(self.limit))
83
+ (self.path / "memory.swap.max").write_text("0")
84
+ (self.path / "memory.oom.group").write_text("1")
85
+ else:
86
+ (self.path / "memory.limit_in_bytes").write_text(str(self.limit))
87
+ (self.path / "memory.memsw.limit_in_bytes").write_text(str(self.limit))
88
+ (self.path / "memory.swappiness").write_text("0")
89
+ (self.path / "memory.oom_control").write_text("0")
90
+ metrics = self.metrics()
91
+ actual = metrics["memory_max_bytes"]
92
+ if self.version == 1 and (
93
+ metrics["memory_memsw_max_bytes"] != actual
94
+ or metrics["swappiness"] != 0
95
+ ):
96
+ raise RuntimeError(
97
+ "kernel swap-policy readback differs from requested policy"
98
+ )
99
+ if self.version == 2 and metrics["swap_max_bytes"] != 0:
100
+ raise RuntimeError(
101
+ "kernel swap-policy readback differs from requested policy"
102
+ )
103
+ if not 0 < actual <= self.limit or self.limit - actual >= os.sysconf(
104
+ "SC_PAGE_SIZE"
105
+ ):
106
+ raise RuntimeError(
107
+ "kernel memory-limit readback differs from requested cap"
108
+ )
109
+ except BaseException:
110
+ self.path.rmdir()
111
+ raise
112
+
113
+ def command(self, command):
114
+ return [
115
+ sys.executable,
116
+ "-m",
117
+ "stream_limits._exec_cgroup",
118
+ str(self.path),
119
+ "--",
120
+ *map(str, command),
121
+ ]
122
+
123
+ def _effective_limit(self):
124
+ best = self.limit
125
+ filename = "memory.limit_in_bytes" if self.version == 1 else "memory.max"
126
+ for path in [self.path, *self.path.parents]:
127
+ control = path / filename
128
+ if not control.exists():
129
+ break
130
+ value = control.read_text().strip()
131
+ if value.isdigit() and 0 < int(value) < (1 << 60):
132
+ best = min(best, int(value))
133
+ return best
134
+
135
+ def metrics(self):
136
+ if self.version == 1:
137
+ oom = dict(
138
+ line.split()
139
+ for line in (self.path / "memory.oom_control").read_text().splitlines()
140
+ )
141
+ return {
142
+ "version": 1,
143
+ "effective_memory_max_bytes": self._effective_limit(),
144
+ "memory_max_bytes": int(
145
+ (self.path / "memory.limit_in_bytes").read_text()
146
+ ),
147
+ "memory_memsw_max_bytes": int(
148
+ (self.path / "memory.memsw.limit_in_bytes").read_text()
149
+ ),
150
+ "swappiness": int((self.path / "memory.swappiness").read_text()),
151
+ "memory_peak_bytes": int(
152
+ (self.path / "memory.max_usage_in_bytes").read_text()
153
+ ),
154
+ "events": {
155
+ "oom_kill": int(oom.get("oom_kill", 0)),
156
+ "limit_failures": int((self.path / "memory.failcnt").read_text()),
157
+ },
158
+ "swap_policy": "memory+swap limit equals memory limit; swappiness=0",
159
+ }
160
+ values = {
161
+ "version": 2,
162
+ "effective_memory_max_bytes": self._effective_limit(),
163
+ "memory_max_bytes": int((self.path / "memory.max").read_text()),
164
+ "swap_max_bytes": int((self.path / "memory.swap.max").read_text()),
165
+ "events": {
166
+ k: int(v)
167
+ for k, v in (
168
+ line.split()
169
+ for line in (self.path / "memory.events").read_text().splitlines()
170
+ )
171
+ },
172
+ }
173
+ peak = self.path / "memory.peak"
174
+ values["memory_peak_bytes"] = int(peak.read_text()) if peak.exists() else None
175
+ return values
176
+
177
+ def close(self):
178
+ # A command can leave background descendants. Only this owned child
179
+ # group is touched, never the delegated parent or a system service.
180
+ deadline = time.monotonic() + 3
181
+ while True:
182
+ pids = (self.path / "cgroup.procs").read_text().split()
183
+ if not pids:
184
+ break
185
+ for pid in pids:
186
+ try:
187
+ os.kill(int(pid), signal.SIGKILL)
188
+ except ProcessLookupError:
189
+ pass
190
+ if time.monotonic() >= deadline:
191
+ raise RuntimeError("processes remain in the owned memory cgroup")
192
+ time.sleep(0.01)
193
+ self.path.rmdir()
194
+
195
+ def __enter__(self):
196
+ return self
197
+
198
+ def __exit__(self, *exc):
199
+ self.close()