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.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: pytest-isolate
3
- Version: 0.0.4
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
- You can also use a mark to isolate or time limit the memory and/or cpu usage test:
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 options can also be set in an pytest configuration file, e.g:
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
- You can also use a mark to isolate or time limit the memory and/or cpu usage test:
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 options can also be set in an pytest configuration file, e.g:
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, args=(), timeout=None, memlimt=None, cpulimit=None
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(self, timeout, memlimit, cpulimit, target, args=()):
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
- q.put(dill.dumps(target(*args)))
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
- time_left = timeout or 1
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(1, time_left)))
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 - 1
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, (item,), timeout, mem_limit, cpu_limit
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
- return isolate, timeout, mem_limit, cpu_limit
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
@@ -11,7 +11,9 @@ deps =
11
11
  numpy
12
12
  pytest-xdist
13
13
  pytest-sugar
14
-
14
+ pytest-mock
15
+ setenv =
16
+ CUDA_VISIBLE_DEVICES = 0,1
15
17
  commands = pytest --isolate --import-mode importlib {posargs:tests} --timeline
16
18
 
17
19
  [testenv:lint]
@@ -1 +0,0 @@
1
- __version__ = "0.0.4"