pytest-isolate 0.0.4__tar.gz → 0.0.8__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.
- {pytest_isolate-0.0.4 → pytest_isolate-0.0.8}/PKG-INFO +40 -3
- {pytest_isolate-0.0.4 → pytest_isolate-0.0.8}/README.md +36 -2
- {pytest_isolate-0.0.4 → pytest_isolate-0.0.8}/pyproject.toml +4 -1
- pytest_isolate-0.0.8/src/pytest_isolate/__about__.py +1 -0
- {pytest_isolate-0.0.4 → pytest_isolate-0.0.8}/src/pytest_isolate/plugin.py +165 -17
- pytest_isolate-0.0.8/src/pytest_isolate/resource_management.py +148 -0
- pytest_isolate-0.0.8/tests/test_resource_management.py +81 -0
- {pytest_isolate-0.0.4 → pytest_isolate-0.0.8}/tox.ini +3 -1
- pytest_isolate-0.0.4/src/pytest_isolate/__about__.py +0 -1
- {pytest_isolate-0.0.4 → pytest_isolate-0.0.8}/.github/workflows/python-package.yml +0 -0
- {pytest_isolate-0.0.4 → pytest_isolate-0.0.8}/.github/workflows/python-publish.yml +0 -0
- {pytest_isolate-0.0.4 → pytest_isolate-0.0.8}/.gitignore +0 -0
- {pytest_isolate-0.0.4 → pytest_isolate-0.0.8}/LICENSE.txt +0 -0
- {pytest_isolate-0.0.4 → pytest_isolate-0.0.8}/src/pytest_isolate/__init__.py +0 -0
- {pytest_isolate-0.0.4 → pytest_isolate-0.0.8}/tests/conftest.py +0 -0
- {pytest_isolate-0.0.4 → pytest_isolate-0.0.8}/tests/test_isolate.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: pytest-isolate
|
|
3
|
-
Version: 0.0.
|
|
3
|
+
Version: 0.0.8
|
|
4
4
|
Summary: Run pytest tests in isolated subprocesses
|
|
5
5
|
Project-URL: Documentation, https://github.com/gilfree/pytest-isolate#readme
|
|
6
6
|
Project-URL: Issues, https://github.com/gilfree/pytest-isolate/issues
|
|
@@ -18,8 +18,11 @@ Classifier: Programming Language :: Python :: Implementation :: CPython
|
|
|
18
18
|
Classifier: Programming Language :: Python :: Implementation :: PyPy
|
|
19
19
|
Requires-Python: >=3.7
|
|
20
20
|
Requires-Dist: dill
|
|
21
|
+
Requires-Dist: filelock
|
|
21
22
|
Requires-Dist: pytest
|
|
22
23
|
Requires-Dist: tblib
|
|
24
|
+
Provides-Extra: gpu
|
|
25
|
+
Requires-Dist: pynvml; extra == 'gpu'
|
|
23
26
|
Description-Content-Type: text/markdown
|
|
24
27
|
|
|
25
28
|
# pytest-isolate
|
|
@@ -40,6 +43,7 @@ This pytest plugin was generated with Cookiecutter along with `@hackebrot`'s `co
|
|
|
40
43
|
* Add Timeout to a forked test
|
|
41
44
|
* Limit memory used by test
|
|
42
45
|
* Limit CPU time used by test
|
|
46
|
+
* Manage GPU resources with CUDA_VISIBLE_DEVICES
|
|
43
47
|
* Plays nice with pytest-xdist
|
|
44
48
|
* Shows warnings, even with xdist!
|
|
45
49
|
* *Create visual timeline of test execution (isolated or not)*
|
|
@@ -48,12 +52,22 @@ This pytest plugin was generated with Cookiecutter along with `@hackebrot`'s `co
|
|
|
48
52
|
|
|
49
53
|
* pytest
|
|
50
54
|
|
|
55
|
+
### Optional Dependencies
|
|
56
|
+
|
|
57
|
+
For GPU resource management:
|
|
58
|
+
|
|
59
|
+
* pynvml (optional, for automatic GPU detection)
|
|
60
|
+
|
|
51
61
|
## Installation
|
|
52
62
|
|
|
53
63
|
You can install "pytest-isolate" via `pip` from `PyPI`
|
|
54
64
|
|
|
55
65
|
pip install pytest-isolate
|
|
56
66
|
|
|
67
|
+
For GPU resource management support:
|
|
68
|
+
|
|
69
|
+
pip install pytest-isolate[gpu]
|
|
70
|
+
|
|
57
71
|
## Usage
|
|
58
72
|
|
|
59
73
|
pytest --isolate
|
|
@@ -84,7 +98,9 @@ they will take precedence. Uninstall them to use `pytest-isolate`.
|
|
|
84
98
|
|
|
85
99
|
Unlike `pytest-timeout`, timeout in `pytest-isolate` is implemented by forking the test to a separate subprocess, and setting timeout for that subprocess.
|
|
86
100
|
|
|
87
|
-
|
|
101
|
+
### Using the isolate marker
|
|
102
|
+
|
|
103
|
+
You can use a mark to isolate or time limit the memory and/or cpu usage test:
|
|
88
104
|
|
|
89
105
|
```python
|
|
90
106
|
@pytest.mark.isolate(timeout=10, mem_limit=10**6, cpu_limit=10)
|
|
@@ -92,7 +108,19 @@ def test_something():
|
|
|
92
108
|
pass
|
|
93
109
|
```
|
|
94
110
|
|
|
95
|
-
The
|
|
111
|
+
The `isolate` marker can also be used to request gpus for a test on a gpu machine:
|
|
112
|
+
|
|
113
|
+
```python
|
|
114
|
+
# Request 2 GPUs for this test
|
|
115
|
+
@pytest.mark.isolate(resources={'gpu': 2})
|
|
116
|
+
def test_with_gpus():
|
|
117
|
+
# The test will have CUDA_VISIBLE_DEVICES set to the allocated GPU IDs
|
|
118
|
+
assert len(os.environ.get("CUDA_VISIBLE_DEVICES").split(",")) == 2
|
|
119
|
+
```
|
|
120
|
+
|
|
121
|
+
### Configuration Options
|
|
122
|
+
|
|
123
|
+
The options can be set in an pytest configuration file, e.g:
|
|
96
124
|
|
|
97
125
|
```toml
|
|
98
126
|
[tool.pytest.ini_options]
|
|
@@ -101,6 +129,15 @@ isolate_mem_limit=1000000
|
|
|
101
129
|
isolate_cpu_limit=10
|
|
102
130
|
```
|
|
103
131
|
|
|
132
|
+
## CUDA_VISIBLE_DEVICES Handling
|
|
133
|
+
|
|
134
|
+
If `CUDA_VISIBLE_DEVICES` is already set when pytest starts, the plugin will respect this setting and only allocate from the GPUs specified there. This works even without pynvml installed.
|
|
135
|
+
|
|
136
|
+
For example:
|
|
137
|
+
|
|
138
|
+
* If `CUDA_VISIBLE_DEVICES=0,1,2` is set, tests will only use GPUs 0, 1, and 2.
|
|
139
|
+
* If `CUDA_VISIBLE_DEVICES=` is set (empty), no GPUs will be used.
|
|
140
|
+
|
|
104
141
|
## Contributing
|
|
105
142
|
|
|
106
143
|
Contributions are very welcome. Tests can be run with `tox`, please ensure
|
|
@@ -16,6 +16,7 @@ This pytest plugin was generated with Cookiecutter along with `@hackebrot`'s `co
|
|
|
16
16
|
* Add Timeout to a forked test
|
|
17
17
|
* Limit memory used by test
|
|
18
18
|
* Limit CPU time used by test
|
|
19
|
+
* Manage GPU resources with CUDA_VISIBLE_DEVICES
|
|
19
20
|
* Plays nice with pytest-xdist
|
|
20
21
|
* Shows warnings, even with xdist!
|
|
21
22
|
* *Create visual timeline of test execution (isolated or not)*
|
|
@@ -24,12 +25,22 @@ This pytest plugin was generated with Cookiecutter along with `@hackebrot`'s `co
|
|
|
24
25
|
|
|
25
26
|
* pytest
|
|
26
27
|
|
|
28
|
+
### Optional Dependencies
|
|
29
|
+
|
|
30
|
+
For GPU resource management:
|
|
31
|
+
|
|
32
|
+
* pynvml (optional, for automatic GPU detection)
|
|
33
|
+
|
|
27
34
|
## Installation
|
|
28
35
|
|
|
29
36
|
You can install "pytest-isolate" via `pip` from `PyPI`
|
|
30
37
|
|
|
31
38
|
pip install pytest-isolate
|
|
32
39
|
|
|
40
|
+
For GPU resource management support:
|
|
41
|
+
|
|
42
|
+
pip install pytest-isolate[gpu]
|
|
43
|
+
|
|
33
44
|
## Usage
|
|
34
45
|
|
|
35
46
|
pytest --isolate
|
|
@@ -60,7 +71,9 @@ they will take precedence. Uninstall them to use `pytest-isolate`.
|
|
|
60
71
|
|
|
61
72
|
Unlike `pytest-timeout`, timeout in `pytest-isolate` is implemented by forking the test to a separate subprocess, and setting timeout for that subprocess.
|
|
62
73
|
|
|
63
|
-
|
|
74
|
+
### Using the isolate marker
|
|
75
|
+
|
|
76
|
+
You can use a mark to isolate or time limit the memory and/or cpu usage test:
|
|
64
77
|
|
|
65
78
|
```python
|
|
66
79
|
@pytest.mark.isolate(timeout=10, mem_limit=10**6, cpu_limit=10)
|
|
@@ -68,7 +81,19 @@ def test_something():
|
|
|
68
81
|
pass
|
|
69
82
|
```
|
|
70
83
|
|
|
71
|
-
The
|
|
84
|
+
The `isolate` marker can also be used to request gpus for a test on a gpu machine:
|
|
85
|
+
|
|
86
|
+
```python
|
|
87
|
+
# Request 2 GPUs for this test
|
|
88
|
+
@pytest.mark.isolate(resources={'gpu': 2})
|
|
89
|
+
def test_with_gpus():
|
|
90
|
+
# The test will have CUDA_VISIBLE_DEVICES set to the allocated GPU IDs
|
|
91
|
+
assert len(os.environ.get("CUDA_VISIBLE_DEVICES").split(",")) == 2
|
|
92
|
+
```
|
|
93
|
+
|
|
94
|
+
### Configuration Options
|
|
95
|
+
|
|
96
|
+
The options can be set in an pytest configuration file, e.g:
|
|
72
97
|
|
|
73
98
|
```toml
|
|
74
99
|
[tool.pytest.ini_options]
|
|
@@ -77,6 +102,15 @@ isolate_mem_limit=1000000
|
|
|
77
102
|
isolate_cpu_limit=10
|
|
78
103
|
```
|
|
79
104
|
|
|
105
|
+
## CUDA_VISIBLE_DEVICES Handling
|
|
106
|
+
|
|
107
|
+
If `CUDA_VISIBLE_DEVICES` is already set when pytest starts, the plugin will respect this setting and only allocate from the GPUs specified there. This works even without pynvml installed.
|
|
108
|
+
|
|
109
|
+
For example:
|
|
110
|
+
|
|
111
|
+
* If `CUDA_VISIBLE_DEVICES=0,1,2` is set, tests will only use GPUs 0, 1, and 2.
|
|
112
|
+
* If `CUDA_VISIBLE_DEVICES=` is set (empty), no GPUs will be used.
|
|
113
|
+
|
|
80
114
|
## Contributing
|
|
81
115
|
|
|
82
116
|
Contributions are very welcome. Tests can be run with `tox`, please ensure
|
|
@@ -20,9 +20,12 @@ classifiers = [
|
|
|
20
20
|
"Programming Language :: Python :: Implementation :: CPython",
|
|
21
21
|
"Programming Language :: Python :: Implementation :: PyPy",
|
|
22
22
|
]
|
|
23
|
-
dependencies = ["dill", "tblib", "pytest"]
|
|
23
|
+
dependencies = ["dill", "tblib", "pytest","filelock"]
|
|
24
24
|
dynamic = ["version"]
|
|
25
25
|
|
|
26
|
+
[project.optional-dependencies]
|
|
27
|
+
gpu = ["pynvml"]
|
|
28
|
+
|
|
26
29
|
[project.urls]
|
|
27
30
|
Documentation = "https://github.com/gilfree/pytest-isolate#readme"
|
|
28
31
|
Issues = "https://github.com/gilfree/pytest-isolate/issues"
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
__version__ = "0.0.8"
|
|
@@ -40,13 +40,50 @@ try:
|
|
|
40
40
|
except ImportError:
|
|
41
41
|
pass
|
|
42
42
|
|
|
43
|
+
from pytest_isolate.resource_management import (
|
|
44
|
+
clean_resources,
|
|
45
|
+
cleanup_resource_environment,
|
|
46
|
+
parse_resource_list,
|
|
47
|
+
register_resource_provider,
|
|
48
|
+
setup_resource_environment,
|
|
49
|
+
)
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def get_available_gpus() -> List[int]:
|
|
53
|
+
# Check if CUDA_VISIBLE_DEVICES is already set
|
|
54
|
+
cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
|
|
55
|
+
if cuda_visible_devices is not None:
|
|
56
|
+
return parse_resource_list(cuda_visible_devices)
|
|
57
|
+
|
|
58
|
+
# If not set and pynvml is available, get all GPUs
|
|
59
|
+
try:
|
|
60
|
+
import pynvml
|
|
61
|
+
|
|
62
|
+
pynvml.nvmlInit()
|
|
63
|
+
count = pynvml.nvmlDeviceGetCount()
|
|
64
|
+
resources = list(range(count))
|
|
65
|
+
pynvml.nvmlShutdown()
|
|
66
|
+
return resources
|
|
67
|
+
except Exception:
|
|
68
|
+
warnings.warn("Failed to get GPU count using pynvml", RuntimeWarning)
|
|
69
|
+
|
|
70
|
+
return []
|
|
71
|
+
|
|
43
72
|
|
|
44
73
|
def forked_subprocess(
|
|
45
|
-
target,
|
|
74
|
+
target,
|
|
75
|
+
args=(),
|
|
76
|
+
timeout=None,
|
|
77
|
+
memlimt=None,
|
|
78
|
+
cpulimit=None,
|
|
79
|
+
wait_delta=0.05,
|
|
80
|
+
resource_dict=None,
|
|
81
|
+
test_id=None,
|
|
82
|
+
resource_timeout=None,
|
|
46
83
|
) -> Tuple[Any, int, bool]:
|
|
47
|
-
sub = ForkedSubprocess()
|
|
84
|
+
sub = ForkedSubprocess(wait_delta=wait_delta)
|
|
48
85
|
exitcode, timed_out, result = sub.run_in_subprocess(
|
|
49
|
-
timeout, memlimt, cpulimit, target, args
|
|
86
|
+
timeout, memlimt, cpulimit, resource_dict, test_id, resource_timeout, target, args
|
|
50
87
|
)
|
|
51
88
|
return result, exitcode, timed_out
|
|
52
89
|
|
|
@@ -73,7 +110,8 @@ def limits(max_mem, max_cpu):
|
|
|
73
110
|
|
|
74
111
|
|
|
75
112
|
class ForkedSubprocess:
|
|
76
|
-
def __init__(self) -> None:
|
|
113
|
+
def __init__(self, wait_delta=0.05) -> None:
|
|
114
|
+
self.wait_delta = wait_delta
|
|
77
115
|
self.r_out, self.w_out = os.pipe()
|
|
78
116
|
self.r_err, self.w_err = os.pipe()
|
|
79
117
|
fcntl.fcntl(self.r_err, fcntl.F_SETFL, os.O_NONBLOCK)
|
|
@@ -103,7 +141,17 @@ class ForkedSubprocess:
|
|
|
103
141
|
os.close(self.w_out)
|
|
104
142
|
os.close(self.w_err)
|
|
105
143
|
|
|
106
|
-
def run_in_subprocess(
|
|
144
|
+
def run_in_subprocess(
|
|
145
|
+
self,
|
|
146
|
+
timeout,
|
|
147
|
+
memlimit,
|
|
148
|
+
cpulimit,
|
|
149
|
+
resource_dict,
|
|
150
|
+
test_id,
|
|
151
|
+
resource_timeout,
|
|
152
|
+
target,
|
|
153
|
+
args=(),
|
|
154
|
+
):
|
|
107
155
|
ctx = mp.get_context("fork")
|
|
108
156
|
q = ctx.Queue()
|
|
109
157
|
timed_out = None
|
|
@@ -113,7 +161,8 @@ class ForkedSubprocess:
|
|
|
113
161
|
self.child_redirect_streams()
|
|
114
162
|
try:
|
|
115
163
|
with limits(memlimit, cpulimit):
|
|
116
|
-
|
|
164
|
+
with allocate_resources(resource_dict, test_id, resource_timeout):
|
|
165
|
+
q.put(dill.dumps(target(*args)))
|
|
117
166
|
except BaseException as e:
|
|
118
167
|
q.put(dill.dumps(e))
|
|
119
168
|
sys.stdout.flush()
|
|
@@ -123,16 +172,17 @@ class ForkedSubprocess:
|
|
|
123
172
|
p = ctx.Process(target=run_subprocess)
|
|
124
173
|
p.start()
|
|
125
174
|
self.parent_open_streams()
|
|
126
|
-
|
|
175
|
+
delta = self.wait_delta
|
|
176
|
+
time_left = timeout or delta
|
|
127
177
|
result = None
|
|
128
178
|
while time_left > 0:
|
|
129
179
|
try:
|
|
130
|
-
result = dill.loads(q.get(block=True, timeout=min(
|
|
180
|
+
result = dill.loads(q.get(block=True, timeout=min(delta, time_left)))
|
|
131
181
|
except Empty:
|
|
132
182
|
if not p.is_alive():
|
|
133
183
|
break
|
|
134
184
|
if timeout is not None:
|
|
135
|
-
time_left = time_left -
|
|
185
|
+
time_left = time_left - delta
|
|
136
186
|
except Exception as e:
|
|
137
187
|
result = e
|
|
138
188
|
break
|
|
@@ -163,7 +213,7 @@ def pytest_load_initial_conftests(early_config, parser, args):
|
|
|
163
213
|
)
|
|
164
214
|
early_config.addinivalue_line(
|
|
165
215
|
"markers",
|
|
166
|
-
"isolate: Always isolate this test.",
|
|
216
|
+
"isolate: Always isolate this test, possibly with resources.",
|
|
167
217
|
)
|
|
168
218
|
early_config.addinivalue_line(
|
|
169
219
|
"markers",
|
|
@@ -209,7 +259,18 @@ def pytest_addoption(parser):
|
|
|
209
259
|
parser.addini(
|
|
210
260
|
"isolate_cpu_limit", "Default cpu limit for isolated tests", type="string"
|
|
211
261
|
)
|
|
212
|
-
|
|
262
|
+
parser.addini(
|
|
263
|
+
"wait_delta",
|
|
264
|
+
"Delta between subprocess queue pollings",
|
|
265
|
+
type="string",
|
|
266
|
+
default="0.05",
|
|
267
|
+
)
|
|
268
|
+
parser.addini(
|
|
269
|
+
"resource_timeout",
|
|
270
|
+
"Timeout for resource allocation",
|
|
271
|
+
type="string",
|
|
272
|
+
default="300",
|
|
273
|
+
)
|
|
213
274
|
parser.addoption(
|
|
214
275
|
"--timeline",
|
|
215
276
|
dest="timeline",
|
|
@@ -228,6 +289,7 @@ def pytest_addoption(parser):
|
|
|
228
289
|
parser.addini("timeline_file", "timeline file", type="string")
|
|
229
290
|
|
|
230
291
|
|
|
292
|
+
|
|
231
293
|
@pytest.hookimpl(trylast=True)
|
|
232
294
|
def pytest_configure(config):
|
|
233
295
|
if config.pluginmanager.get_plugin("forked"):
|
|
@@ -253,6 +315,8 @@ def pytest_configure(config):
|
|
|
253
315
|
"markers",
|
|
254
316
|
"timeout: Run this tests in a separate, forked, process with timeout",
|
|
255
317
|
)
|
|
318
|
+
clean_resources()
|
|
319
|
+
register_resource_provider("gpu", "CUDA_VISIBLE_DEVICES", get_available_gpus)
|
|
256
320
|
|
|
257
321
|
|
|
258
322
|
# Taken from pytest 7.2:
|
|
@@ -314,23 +378,56 @@ def run_subprocess(item: pytest.Item):
|
|
|
314
378
|
config=item.config, report=report
|
|
315
379
|
)
|
|
316
380
|
)
|
|
381
|
+
|
|
317
382
|
return s_reports, warnings
|
|
318
383
|
except BaseException as e:
|
|
319
384
|
return e, None
|
|
320
385
|
|
|
321
386
|
|
|
387
|
+
@contextmanager
|
|
388
|
+
def allocate_resources(resource_dict, test_id, timeout):
|
|
389
|
+
resource_dict = resource_dict or {}
|
|
390
|
+
allocated_resources = {}
|
|
391
|
+
for resource_type, count in resource_dict.items():
|
|
392
|
+
# Allocate resources
|
|
393
|
+
resource_ids = setup_resource_environment(
|
|
394
|
+
test_id=f"{test_id}_{resource_type}",
|
|
395
|
+
count=count,
|
|
396
|
+
resource_type=resource_type,
|
|
397
|
+
wait_timeout=timeout,
|
|
398
|
+
)
|
|
399
|
+
allocated_resources[resource_type] = resource_ids
|
|
400
|
+
try:
|
|
401
|
+
yield allocated_resources
|
|
402
|
+
finally:
|
|
403
|
+
for res_type, _ in allocated_resources.items():
|
|
404
|
+
cleanup_resource_environment(f"{test_id}_{res_type}", res_type)
|
|
405
|
+
|
|
406
|
+
|
|
322
407
|
def run_in_subprocess(
|
|
323
408
|
item: pytest.Item,
|
|
324
409
|
timeout: Optional[float],
|
|
325
410
|
mem_limit: Optional[int],
|
|
326
411
|
cpu_limit: Optional[int],
|
|
412
|
+
wait_delta: float,
|
|
327
413
|
) -> List[pytest.TestReport]:
|
|
328
414
|
cap = item.config.pluginmanager.getplugin("capturemanager")
|
|
329
415
|
assert isinstance(cap, _pytest.capture.CaptureManager)
|
|
330
416
|
cap.resume_global_capture()
|
|
331
417
|
cap.activate_fixture()
|
|
418
|
+
resource_dict = get_resource_dict(item)
|
|
419
|
+
test_id = f"{item.nodeid}"
|
|
420
|
+
resource_timeout = float(item.config.getini("resource_timeout")) # type: ignore
|
|
332
421
|
result, exitcode, timed_out = forked_subprocess(
|
|
333
|
-
run_subprocess,
|
|
422
|
+
run_subprocess,
|
|
423
|
+
(item,),
|
|
424
|
+
timeout,
|
|
425
|
+
mem_limit,
|
|
426
|
+
cpu_limit,
|
|
427
|
+
wait_delta,
|
|
428
|
+
resource_dict,
|
|
429
|
+
test_id,
|
|
430
|
+
resource_timeout,
|
|
334
431
|
)
|
|
335
432
|
cap.deactivate_fixture()
|
|
336
433
|
cap.suspend_global_capture(in_=False)
|
|
@@ -414,7 +511,7 @@ def report_process_crash(
|
|
|
414
511
|
reason = xfail_marker.kwargs.get("reason", "")
|
|
415
512
|
rep.wasxfail = ""
|
|
416
513
|
if reason:
|
|
417
|
-
rep.wasxfail = f"reason: {xfail_marker.kwargs.get('reason','')}; "
|
|
514
|
+
rep.wasxfail = f"reason: {xfail_marker.kwargs.get('reason', '')}; "
|
|
418
515
|
rep.wasxfail += f"pytest-isolate reason: {call.excinfo}"
|
|
419
516
|
|
|
420
517
|
warnings.warn(
|
|
@@ -465,23 +562,66 @@ def get_marker(item, marker_name, argname=None, pos=None):
|
|
|
465
562
|
return marker
|
|
466
563
|
|
|
467
564
|
|
|
565
|
+
def get_resource_dict(item):
|
|
566
|
+
"""
|
|
567
|
+
Get resource requirements from the resources parameter of the isolate marker.
|
|
568
|
+
|
|
569
|
+
Example: @pytest.mark.isolate(resources={'gpu': 2})
|
|
570
|
+
"""
|
|
571
|
+
resources = {}
|
|
572
|
+
isolate_marker = item.get_closest_marker("isolate")
|
|
573
|
+
|
|
574
|
+
if isolate_marker is not None and "resources" in isolate_marker.kwargs:
|
|
575
|
+
resources_param = isolate_marker.kwargs.get("resources", {})
|
|
576
|
+
|
|
577
|
+
# Ensure the resources parameter is a dictionary
|
|
578
|
+
if not isinstance(resources_param, dict):
|
|
579
|
+
warnings.warn(
|
|
580
|
+
f"Invalid resources parameter {resources_param}. Must be a dictionary.",
|
|
581
|
+
RuntimeWarning,
|
|
582
|
+
)
|
|
583
|
+
return resources
|
|
584
|
+
|
|
585
|
+
# Process each resource requirement
|
|
586
|
+
for resource_type, count in resources_param.items():
|
|
587
|
+
try:
|
|
588
|
+
count = int(count)
|
|
589
|
+
if count < 0:
|
|
590
|
+
warnings.warn(
|
|
591
|
+
f"Invalid {resource_type} count: {count}. Must be >= 0.",
|
|
592
|
+
RuntimeWarning,
|
|
593
|
+
)
|
|
594
|
+
continue
|
|
595
|
+
resources[resource_type] = count
|
|
596
|
+
except (ValueError, TypeError):
|
|
597
|
+
warnings.warn(
|
|
598
|
+
f"Invalid {resource_type} count: {count}. Must be an integer.",
|
|
599
|
+
RuntimeWarning,
|
|
600
|
+
)
|
|
601
|
+
|
|
602
|
+
return resources
|
|
603
|
+
|
|
604
|
+
|
|
468
605
|
@pytest.hookimpl(tryfirst=True)
|
|
469
|
-
def pytest_runtest_protocol(item):
|
|
606
|
+
def pytest_runtest_protocol(item: pytest.Item):
|
|
470
607
|
if item.config.pluginmanager.get_plugin("forked"):
|
|
471
608
|
return
|
|
472
609
|
if item.config.pluginmanager.get_plugin("timeout"):
|
|
473
610
|
return
|
|
474
|
-
isolate, timeout, mem_limit, cpu_limit = get_isolation_options(item)
|
|
611
|
+
isolate, timeout, mem_limit, cpu_limit, resource_reqs = get_isolation_options(item)
|
|
612
|
+
|
|
613
|
+
wait_delta = float(item.config.getini("wait_delta")) # type: ignore
|
|
475
614
|
|
|
476
615
|
if (
|
|
477
616
|
isolate is not None
|
|
478
617
|
or mem_limit is not None
|
|
479
618
|
or timeout is not None
|
|
480
619
|
or cpu_limit is not None
|
|
620
|
+
or resource_reqs
|
|
481
621
|
):
|
|
482
622
|
ihook = item.ihook
|
|
483
623
|
ihook.pytest_runtest_logstart(nodeid=item.nodeid, location=item.location)
|
|
484
|
-
reports = run_in_subprocess(item, timeout, mem_limit, cpu_limit)
|
|
624
|
+
reports = run_in_subprocess(item, timeout, mem_limit, cpu_limit, wait_delta)
|
|
485
625
|
|
|
486
626
|
for rep in reports:
|
|
487
627
|
ihook.pytest_runtest_logreport(report=rep)
|
|
@@ -568,11 +708,19 @@ def get_isolation_options(item):
|
|
|
568
708
|
pass
|
|
569
709
|
else:
|
|
570
710
|
isolate = forked
|
|
711
|
+
|
|
571
712
|
timeout = get_timeout(item)
|
|
572
713
|
mem_limit = get_memory_limit(item)
|
|
573
714
|
cpu_limit = get_cpu_limit(item)
|
|
574
715
|
|
|
575
|
-
|
|
716
|
+
# Get resources from isolate(resources={'gpu': 2}) syntax
|
|
717
|
+
resource_dict = get_resource_dict(item)
|
|
718
|
+
|
|
719
|
+
# If resources are requested, ensure isolation
|
|
720
|
+
if resource_dict and not isolate:
|
|
721
|
+
isolate = True
|
|
722
|
+
|
|
723
|
+
return isolate, timeout, mem_limit, cpu_limit, resource_dict
|
|
576
724
|
|
|
577
725
|
|
|
578
726
|
@pytest.hookimpl(tryfirst=True)
|
|
@@ -0,0 +1,148 @@
|
|
|
1
|
+
import json
|
|
2
|
+
import os
|
|
3
|
+
import tempfile
|
|
4
|
+
import time
|
|
5
|
+
import warnings
|
|
6
|
+
from contextlib import contextmanager
|
|
7
|
+
from dataclasses import asdict, dataclass
|
|
8
|
+
from pathlib import Path
|
|
9
|
+
from typing import Callable, List, Optional
|
|
10
|
+
|
|
11
|
+
import filelock
|
|
12
|
+
|
|
13
|
+
# Constants
|
|
14
|
+
|
|
15
|
+
DEFAULT_STATE_FILE = Path(tempfile.gettempdir()) / "pytest_isolate_resources.json"
|
|
16
|
+
DEFAULT_WAIT_TIMEOUT = 300 # 5 minutes
|
|
17
|
+
POLL_INTERVAL = 0.1 # 1 second
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
@contextmanager
|
|
21
|
+
def lock_resource_file():
|
|
22
|
+
with filelock.FileLock(str(DEFAULT_STATE_FILE)+'.lock', timeout=4) as lock:
|
|
23
|
+
yield lock
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
@dataclass
|
|
27
|
+
class Resource:
|
|
28
|
+
env_variable: str
|
|
29
|
+
available: List[int]
|
|
30
|
+
allocated: dict[str, List[int]]
|
|
31
|
+
|
|
32
|
+
def allocate(self, test_id: str, count: int) -> List[int]:
|
|
33
|
+
"""Allocate resources for a test."""
|
|
34
|
+
if len(self.available) < count:
|
|
35
|
+
return []
|
|
36
|
+
|
|
37
|
+
allocated = self.available[:count]
|
|
38
|
+
self.allocated[test_id] = allocated
|
|
39
|
+
self.available = self.available[count:]
|
|
40
|
+
return allocated
|
|
41
|
+
|
|
42
|
+
def release(self, test_id: str) -> None:
|
|
43
|
+
"""Release resources for a test."""
|
|
44
|
+
if test_id in self.allocated:
|
|
45
|
+
released = self.allocated.pop(test_id)
|
|
46
|
+
self.available.extend(released)
|
|
47
|
+
else:
|
|
48
|
+
raise ValueError(f"Test ID {test_id} not found in allocated resources.")
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
@dataclass
|
|
52
|
+
class StateData:
|
|
53
|
+
resources: dict[str, Resource]
|
|
54
|
+
|
|
55
|
+
@classmethod
|
|
56
|
+
def get_instance(cls) -> "StateData":
|
|
57
|
+
if not DEFAULT_STATE_FILE.exists():
|
|
58
|
+
data = {}
|
|
59
|
+
try:
|
|
60
|
+
with open(DEFAULT_STATE_FILE, "r") as f:
|
|
61
|
+
data = json.load(f)
|
|
62
|
+
except (json.JSONDecodeError, IOError):
|
|
63
|
+
data = {}
|
|
64
|
+
return StateData(
|
|
65
|
+
resources={name: Resource(**res) for name, res in data.items()}
|
|
66
|
+
)
|
|
67
|
+
|
|
68
|
+
def save(self) -> None:
|
|
69
|
+
with open(DEFAULT_STATE_FILE, "w") as f:
|
|
70
|
+
json.dump({name: asdict(res) for name, res in self.resources.items()}, f)
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
def clean_resources() -> None:
|
|
74
|
+
"""Clean up resources by removing lock and state files."""
|
|
75
|
+
if DEFAULT_STATE_FILE.exists():
|
|
76
|
+
os.remove(DEFAULT_STATE_FILE)
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
def register_resource_provider(
|
|
80
|
+
resource_type: str, env_variable: str, provider_func: Callable[[], List[int]]
|
|
81
|
+
) -> None:
|
|
82
|
+
with lock_resource_file():
|
|
83
|
+
existing = provider_func()
|
|
84
|
+
state = StateData.get_instance()
|
|
85
|
+
if resource_type in state.resources:
|
|
86
|
+
warnings.warn(
|
|
87
|
+
f"Resource type {resource_type} already registered. Overwriting."
|
|
88
|
+
)
|
|
89
|
+
state.resources[resource_type] = Resource(
|
|
90
|
+
env_variable=env_variable,
|
|
91
|
+
available=existing,
|
|
92
|
+
allocated={},
|
|
93
|
+
)
|
|
94
|
+
state.save()
|
|
95
|
+
|
|
96
|
+
|
|
97
|
+
def parse_resource_list(env_value: Optional[str]) -> List[int]:
|
|
98
|
+
"""Parse a comma-separated list of resource IDs from an environment variable."""
|
|
99
|
+
if not env_value:
|
|
100
|
+
return []
|
|
101
|
+
try:
|
|
102
|
+
vals= [int(x.strip()) for x in env_value.split(",") if x.strip()]
|
|
103
|
+
vals = [x for x in vals if x >= 0]
|
|
104
|
+
return vals
|
|
105
|
+
except ValueError:
|
|
106
|
+
return []
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
def setup_resource_environment(
|
|
110
|
+
test_id: str,
|
|
111
|
+
count: int,
|
|
112
|
+
resource_type: str,
|
|
113
|
+
wait_timeout: Optional[float] = None,
|
|
114
|
+
) -> Optional[List[int]]:
|
|
115
|
+
def try_allocate() -> Optional[List[int]]:
|
|
116
|
+
with lock_resource_file():
|
|
117
|
+
state = StateData.get_instance()
|
|
118
|
+
if resource_type not in state.resources:
|
|
119
|
+
raise ValueError(f"Resource type {resource_type} not registered.")
|
|
120
|
+
|
|
121
|
+
resource = state.resources[resource_type]
|
|
122
|
+
allocated = resource.allocate(test_id, count)
|
|
123
|
+
if allocated:
|
|
124
|
+
os.environ[resource.env_variable] = ",".join(map(str, allocated))
|
|
125
|
+
state.save()
|
|
126
|
+
return allocated
|
|
127
|
+
return []
|
|
128
|
+
|
|
129
|
+
timeout = wait_timeout or DEFAULT_WAIT_TIMEOUT
|
|
130
|
+
start_time = time.time()
|
|
131
|
+
while time.time() - start_time < timeout:
|
|
132
|
+
allocated = try_allocate()
|
|
133
|
+
if allocated:
|
|
134
|
+
return allocated
|
|
135
|
+
time.sleep(POLL_INTERVAL)
|
|
136
|
+
raise TimeoutError(f"Timeout while waiting for resources of type {resource_type}.")
|
|
137
|
+
|
|
138
|
+
|
|
139
|
+
def cleanup_resource_environment(test_id: str, resource_type: str) -> None:
|
|
140
|
+
with lock_resource_file():
|
|
141
|
+
state = StateData.get_instance()
|
|
142
|
+
if resource_type not in state.resources:
|
|
143
|
+
raise ValueError(f"Resource type {resource_type} not registered.")
|
|
144
|
+
|
|
145
|
+
resource = state.resources[resource_type]
|
|
146
|
+
resource.release(test_id)
|
|
147
|
+
os.environ.pop(resource.env_variable, None)
|
|
148
|
+
state.save()
|
|
@@ -0,0 +1,81 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import random
|
|
3
|
+
import tempfile
|
|
4
|
+
from importlib import resources
|
|
5
|
+
from pathlib import Path
|
|
6
|
+
from time import sleep
|
|
7
|
+
from unittest import mock
|
|
8
|
+
|
|
9
|
+
import filelock
|
|
10
|
+
import pytest
|
|
11
|
+
|
|
12
|
+
from pytest_isolate.plugin import allocate_resources
|
|
13
|
+
from pytest_isolate.resource_management import (
|
|
14
|
+
parse_resource_list,
|
|
15
|
+
register_resource_provider,
|
|
16
|
+
)
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
@pytest.fixture
|
|
20
|
+
def temp_resource_files():
|
|
21
|
+
"""Create temporary files for resource manager testing"""
|
|
22
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
23
|
+
lock_file = Path(tmpdir) / "resources.lock"
|
|
24
|
+
state_file = Path(tmpdir) / "resources.json"
|
|
25
|
+
yield lock_file, state_file
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
@pytest.fixture
|
|
29
|
+
def mock_pynvml():
|
|
30
|
+
"""Mock pynvml module for testing"""
|
|
31
|
+
pynvml_mock = mock.MagicMock()
|
|
32
|
+
pynvml_mock.nvmlDeviceGetCount.return_value = 4
|
|
33
|
+
|
|
34
|
+
# Mock the import so it returns our mock when imported in plugin.py
|
|
35
|
+
with mock.patch.dict("sys.modules", {"pynvml": pynvml_mock}):
|
|
36
|
+
yield pynvml_mock
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def test_parse_resource_list():
|
|
40
|
+
"""Test parsing resources from environment variables."""
|
|
41
|
+
# Test empty input
|
|
42
|
+
assert parse_resource_list(None) == []
|
|
43
|
+
assert parse_resource_list("") == []
|
|
44
|
+
|
|
45
|
+
# Test simple cases
|
|
46
|
+
assert parse_resource_list("0") == [0]
|
|
47
|
+
assert parse_resource_list("0,1,2") == [0, 1, 2]
|
|
48
|
+
|
|
49
|
+
# Test with spaces
|
|
50
|
+
assert parse_resource_list(" 0, 1, 2 ") == [0, 1, 2]
|
|
51
|
+
|
|
52
|
+
# Test with invalid values - no warning in new implementation
|
|
53
|
+
assert parse_resource_list("0,a,2") == []
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
def test_allocate_resources():
|
|
57
|
+
with allocate_resources({"gpu": 1}, "foo", 100) as allocated:
|
|
58
|
+
assert allocated["gpu"] == [0]
|
|
59
|
+
assert os.environ.get("CUDA_VISIBLE_DEVICES") == "0"
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
@pytest.mark.isolate(resources={"gpu": 2}, timeout=10)
|
|
64
|
+
@pytest.mark.parametrize("dummy", [1] * 2)
|
|
65
|
+
def test_gpu_marker_2(dummy):
|
|
66
|
+
sleep(0.3)
|
|
67
|
+
assert len(os.environ.get("CUDA_VISIBLE_DEVICES", "").split(",")) == 2
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
@pytest.mark.isolate(resources={"gpu": 1}, timeout=10)
|
|
71
|
+
@pytest.mark.parametrize("dummy", [1] * 2)
|
|
72
|
+
def test_gpu_marker_1(dummy):
|
|
73
|
+
sleep(0.3)
|
|
74
|
+
assert len(os.environ.get("CUDA_VISIBLE_DEVICES", "").split(",")) == 1
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
@pytest.mark.xfail
|
|
78
|
+
@pytest.mark.isolate(resources={"gpu": 7}, timeout=4)
|
|
79
|
+
def test_gpu_marker_7():
|
|
80
|
+
sleep(0.3)
|
|
81
|
+
assert len(os.environ.get("CUDA_VISIBLE_DEVICES", "").split(",")) == 7
|
|
@@ -1 +0,0 @@
|
|
|
1
|
-
__version__ = "0.0.4"
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|