pytest-isolate 0.0.4__tar.gz → 0.0.9__tar.gz

{pytest_isolate-0.0.4 → pytest_isolate-0.0.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytest-isolate
-Version: 0.0.4
+Version: 0.0.9
 Summary: Run pytest tests in isolated subprocesses 
 Project-URL: Documentation, https://github.com/gilfree/pytest-isolate#readme
 Project-URL: Issues, https://github.com/gilfree/pytest-isolate/issues
@@ -18,8 +18,11 @@ Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
 Requires-Python: >=3.7
 Requires-Dist: dill
+Requires-Dist: filelock
 Requires-Dist: pytest
 Requires-Dist: tblib
+Provides-Extra: gpu
+Requires-Dist: pynvml; extra == 'gpu'
 Description-Content-Type: text/markdown
 
 # pytest-isolate
@@ -40,6 +43,7 @@ This pytest plugin was generated with Cookiecutter along with `@hackebrot`'s `co
 * Add Timeout to a forked test
 * Limit memory used by test
 * Limit CPU time used by test
+* Manage GPU resources with CUDA_VISIBLE_DEVICES
 * Plays nice with pytest-xdist
 * Shows warnings, even with xdist!
 * *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
 
 * pytest
 
+### Optional Dependencies
+
+For GPU resource management:
+
+* pynvml (optional, for automatic GPU detection)
+
 ## Installation
 
 You can install "pytest-isolate" via `pip` from `PyPI`
 
     pip install pytest-isolate
 
+For GPU resource management support:
+
+    pip install pytest-isolate[gpu]
+
 ## Usage
 
     pytest --isolate
@@ -84,7 +98,9 @@ they will take precedence. Uninstall them to use `pytest-isolate`.
 
 Unlike `pytest-timeout`, timeout in `pytest-isolate` is implemented by forking the test to a separate subprocess, and setting timeout for that subprocess.
 
-You can also use a mark to isolate or time limit the memory and/or cpu usage test:
+### Using the isolate marker
+
+You can use a mark to isolate or time limit the memory and/or cpu usage test:
 
 ```python
 @pytest.mark.isolate(timeout=10, mem_limit=10**6, cpu_limit=10)
@@ -92,7 +108,19 @@ def test_something():
     pass
 ```
 
-The options can also be set in an pytest configuration file, e.g:
+The `isolate` marker can also be used to request gpus for a test on a gpu machine:
+
+```python
+# Request 2 GPUs for this test
+@pytest.mark.isolate(resources={'gpu': 2})
+def test_with_gpus():
+    # The test will have CUDA_VISIBLE_DEVICES set to the allocated GPU IDs
+    assert len(os.environ.get("CUDA_VISIBLE_DEVICES").split(",")) == 2
+```
+
+### Configuration Options
+
+The options can be set in an pytest configuration file, e.g:
 
 ```toml
 [tool.pytest.ini_options]
@@ -101,6 +129,15 @@ isolate_mem_limit=1000000
 isolate_cpu_limit=10
 ```
 
+## CUDA_VISIBLE_DEVICES Handling
+
+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.
+
+For example:
+
+* If `CUDA_VISIBLE_DEVICES=0,1,2` is set, tests will only use GPUs 0, 1, and 2.
+* If `CUDA_VISIBLE_DEVICES=` is set (empty), no GPUs will be used.
+
 ## Contributing
 
 Contributions are very welcome. Tests can be run with `tox`, please ensure

{pytest_isolate-0.0.4 → pytest_isolate-0.0.9}/README.md

@@ -16,6 +16,7 @@ This pytest plugin was generated with Cookiecutter along with `@hackebrot`'s `co
 * Add Timeout to a forked test
 * Limit memory used by test
 * Limit CPU time used by test
+* Manage GPU resources with CUDA_VISIBLE_DEVICES
 * Plays nice with pytest-xdist
 * Shows warnings, even with xdist!
 * *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
 
 * pytest
 
+### Optional Dependencies
+
+For GPU resource management:
+
+* pynvml (optional, for automatic GPU detection)
+
 ## Installation
 
 You can install "pytest-isolate" via `pip` from `PyPI`
 
     pip install pytest-isolate
 
+For GPU resource management support:
+
+    pip install pytest-isolate[gpu]
+
 ## Usage
 
     pytest --isolate
@@ -60,7 +71,9 @@ they will take precedence. Uninstall them to use `pytest-isolate`.
 
 Unlike `pytest-timeout`, timeout in `pytest-isolate` is implemented by forking the test to a separate subprocess, and setting timeout for that subprocess.
 
-You can also use a mark to isolate or time limit the memory and/or cpu usage test:
+### Using the isolate marker
+
+You can use a mark to isolate or time limit the memory and/or cpu usage test:
 
 ```python
 @pytest.mark.isolate(timeout=10, mem_limit=10**6, cpu_limit=10)
@@ -68,7 +81,19 @@ def test_something():
     pass
 ```
 
-The options can also be set in an pytest configuration file, e.g:
+The `isolate` marker can also be used to request gpus for a test on a gpu machine:
+
+```python
+# Request 2 GPUs for this test
+@pytest.mark.isolate(resources={'gpu': 2})
+def test_with_gpus():
+    # The test will have CUDA_VISIBLE_DEVICES set to the allocated GPU IDs
+    assert len(os.environ.get("CUDA_VISIBLE_DEVICES").split(",")) == 2
+```
+
+### Configuration Options
+
+The options can be set in an pytest configuration file, e.g:
 
 ```toml
 [tool.pytest.ini_options]
@@ -77,6 +102,15 @@ isolate_mem_limit=1000000
 isolate_cpu_limit=10
 ```
 
+## CUDA_VISIBLE_DEVICES Handling
+
+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.
+
+For example:
+
+* If `CUDA_VISIBLE_DEVICES=0,1,2` is set, tests will only use GPUs 0, 1, and 2.
+* If `CUDA_VISIBLE_DEVICES=` is set (empty), no GPUs will be used.
+
 ## Contributing
 
 Contributions are very welcome. Tests can be run with `tox`, please ensure

{pytest_isolate-0.0.4 → pytest_isolate-0.0.9}/pyproject.toml

@@ -20,9 +20,12 @@ classifiers = [
   "Programming Language :: Python :: Implementation :: CPython",
   "Programming Language :: Python :: Implementation :: PyPy",
 ]
-dependencies = ["dill", "tblib", "pytest"]
+dependencies = ["dill", "tblib", "pytest","filelock"]
 dynamic = ["version"]
 
+[project.optional-dependencies]
+gpu = ["pynvml"]
+
 [project.urls]
 Documentation = "https://github.com/gilfree/pytest-isolate#readme"
 Issues = "https://github.com/gilfree/pytest-isolate/issues"

pytest_isolate-0.0.9/src/pytest_isolate/__about__.py

@@ -0,0 +1 @@
+__version__ = "0.0.9"

{pytest_isolate-0.0.4 → pytest_isolate-0.0.9}/src/pytest_isolate/plugin.py

@@ -40,13 +40,52 @@ try:
 except ImportError:
     pass
 
+from pytest_isolate.tracing import create_event
+from pytest_isolate.resource_management import (
+    clean_resources,
+    cleanup_resource_environment,
+    get_resource_events,
+    parse_resource_list,
+    register_resource_provider,
+    setup_resource_environment,
+)
+
+
+def get_available_gpus() -> List[int]:
+    # Check if CUDA_VISIBLE_DEVICES is already set
+    cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
+    if cuda_visible_devices is not None:
+        return parse_resource_list(cuda_visible_devices)
+
+    # If not set and pynvml is available, get all GPUs
+    try:
+        import pynvml
+
+        pynvml.nvmlInit()
+        count = pynvml.nvmlDeviceGetCount()
+        resources = list(range(count))
+        pynvml.nvmlShutdown()
+        return resources
+    except Exception:
+        warnings.warn("Failed to get GPU count using pynvml", RuntimeWarning)
+
+    return []
+
 
 def forked_subprocess(
-    target, args=(), timeout=None, memlimt=None, cpulimit=None
+    target,
+    args=(),
+    timeout=None,
+    memlimt=None,
+    cpulimit=None,
+    wait_delta=0.05,
+    resource_dict=None,
+    test_id=None,
+    resource_timeout=None,
 ) -> Tuple[Any, int, bool]:
-    sub = ForkedSubprocess()
+    sub = ForkedSubprocess(wait_delta=wait_delta)
     exitcode, timed_out, result = sub.run_in_subprocess(
-        timeout, memlimt, cpulimit, target, args
+        timeout, memlimt, cpulimit, resource_dict, test_id, resource_timeout, target, args
     )
     return result, exitcode, timed_out
 
@@ -73,7 +112,8 @@ def limits(max_mem, max_cpu):
 
 
 class ForkedSubprocess:
-    def __init__(self) -> None:
+    def __init__(self, wait_delta=0.05) -> None:
+        self.wait_delta = wait_delta
         self.r_out, self.w_out = os.pipe()
         self.r_err, self.w_err = os.pipe()
         fcntl.fcntl(self.r_err, fcntl.F_SETFL, os.O_NONBLOCK)
@@ -103,7 +143,17 @@ class ForkedSubprocess:
         os.close(self.w_out)
         os.close(self.w_err)
 
-    def run_in_subprocess(self, timeout, memlimit, cpulimit, target, args=()):
+    def run_in_subprocess(
+        self,
+        timeout,
+        memlimit,
+        cpulimit,
+        resource_dict,
+        test_id,
+        resource_timeout,
+        target,
+        args=(),
+    ):
         ctx = mp.get_context("fork")
         q = ctx.Queue()
         timed_out = None
@@ -113,7 +163,8 @@ class ForkedSubprocess:
             self.child_redirect_streams()
             try:
                 with limits(memlimit, cpulimit):
-                    q.put(dill.dumps(target(*args)))
+                    with allocate_resources(resource_dict, test_id, resource_timeout):
+                        q.put(dill.dumps(target(*args)))
             except BaseException as e:
                 q.put(dill.dumps(e))
             sys.stdout.flush()
@@ -123,16 +174,17 @@ class ForkedSubprocess:
         p = ctx.Process(target=run_subprocess)
         p.start()
         self.parent_open_streams()
-        time_left = timeout or 1
+        delta = self.wait_delta
+        time_left = timeout or delta
         result = None
         while time_left > 0:
             try:
-                result = dill.loads(q.get(block=True, timeout=min(1, time_left)))
+                result = dill.loads(q.get(block=True, timeout=min(delta, time_left)))
             except Empty:
                 if not p.is_alive():
                     break
                 if timeout is not None:
-                    time_left = time_left - 1
+                    time_left = time_left - delta
             except Exception as e:
                 result = e
                 break
@@ -163,7 +215,7 @@ def pytest_load_initial_conftests(early_config, parser, args):
     )
     early_config.addinivalue_line(
         "markers",
-        "isolate: Always isolate this test.",
+        "isolate: Always isolate this test, possibly with resources.",
     )
     early_config.addinivalue_line(
         "markers",
@@ -209,7 +261,18 @@ def pytest_addoption(parser):
     parser.addini(
         "isolate_cpu_limit", "Default cpu limit for isolated tests", type="string"
     )
-
+    parser.addini(
+        "wait_delta",
+        "Delta between subprocess queue pollings",
+        type="string",
+        default="0.05",
+    )
+    parser.addini(
+        "resource_timeout",
+        "Timeout for resource allocation",
+        type="string",
+        default="300",
+    )
     parser.addoption(
         "--timeline",
         dest="timeline",
@@ -228,6 +291,7 @@ def pytest_addoption(parser):
     parser.addini("timeline_file", "timeline file", type="string")
 
 
+
 @pytest.hookimpl(trylast=True)
 def pytest_configure(config):
     if config.pluginmanager.get_plugin("forked"):
@@ -253,6 +317,8 @@ def pytest_configure(config):
             "markers",
             "timeout: Run this tests in a separate, forked, process with timeout",
         )
+    clean_resources()
+    register_resource_provider("gpu", "CUDA_VISIBLE_DEVICES", get_available_gpus)
 
 
 # Taken from pytest 7.2:
@@ -314,23 +380,56 @@ def run_subprocess(item: pytest.Item):
                     config=item.config, report=report
                 )
             )
+
         return s_reports, warnings
     except BaseException as e:
         return e, None
 
 
+@contextmanager
+def allocate_resources(resource_dict, test_id, timeout):
+    resource_dict = resource_dict or {}
+    allocated_resources = {}
+    for resource_type, count in resource_dict.items():
+        # Allocate resources
+        resource_ids = setup_resource_environment(
+            test_id=f"{test_id}_{resource_type}",
+            count=count,
+            resource_type=resource_type,
+            wait_timeout=timeout,
+        )
+        allocated_resources[resource_type] = resource_ids
+    try:
+        yield allocated_resources
+    finally:
+        for res_type, _ in allocated_resources.items():
+            cleanup_resource_environment(f"{test_id}_{res_type}", res_type)
+
+
 def run_in_subprocess(
     item: pytest.Item,
     timeout: Optional[float],
     mem_limit: Optional[int],
     cpu_limit: Optional[int],
+    wait_delta: float,
 ) -> List[pytest.TestReport]:
     cap = item.config.pluginmanager.getplugin("capturemanager")
     assert isinstance(cap, _pytest.capture.CaptureManager)
     cap.resume_global_capture()
     cap.activate_fixture()
+    resource_dict = get_resource_dict(item)
+    test_id = f"{item.nodeid}"
+    resource_timeout = float(item.config.getini("resource_timeout"))  # type: ignore
     result, exitcode, timed_out = forked_subprocess(
-        run_subprocess, (item,), timeout, mem_limit, cpu_limit
+        run_subprocess,
+        (item,),
+        timeout,
+        mem_limit,
+        cpu_limit,
+        wait_delta,
+        resource_dict,
+        test_id,
+        resource_timeout,
     )
     cap.deactivate_fixture()
     cap.suspend_global_capture(in_=False)
@@ -414,7 +513,7 @@ def report_process_crash(
     reason = xfail_marker.kwargs.get("reason", "")
     rep.wasxfail = ""
     if reason:
-        rep.wasxfail = f"reason: {xfail_marker.kwargs.get('reason','')}; "
+        rep.wasxfail = f"reason: {xfail_marker.kwargs.get('reason', '')}; "
     rep.wasxfail += f"pytest-isolate reason: {call.excinfo}"
 
     warnings.warn(
@@ -465,23 +564,66 @@ def get_marker(item, marker_name, argname=None, pos=None):
     return marker
 
 
+def get_resource_dict(item):
+    """
+    Get resource requirements from the resources parameter of the isolate marker.
+
+    Example: @pytest.mark.isolate(resources={'gpu': 2})
+    """
+    resources = {}
+    isolate_marker = item.get_closest_marker("isolate")
+
+    if isolate_marker is not None and "resources" in isolate_marker.kwargs:
+        resources_param = isolate_marker.kwargs.get("resources", {})
+
+        # Ensure the resources parameter is a dictionary
+        if not isinstance(resources_param, dict):
+            warnings.warn(
+                f"Invalid resources parameter {resources_param}. Must be a dictionary.",
+                RuntimeWarning,
+            )
+            return resources
+
+        # Process each resource requirement
+        for resource_type, count in resources_param.items():
+            try:
+                count = int(count)
+                if count < 0:
+                    warnings.warn(
+                        f"Invalid {resource_type} count: {count}. Must be >= 0.",
+                        RuntimeWarning,
+                    )
+                    continue
+                resources[resource_type] = count
+            except (ValueError, TypeError):
+                warnings.warn(
+                    f"Invalid {resource_type} count: {count}. Must be an integer.",
+                    RuntimeWarning,
+                )
+
+    return resources
+
+
 @pytest.hookimpl(tryfirst=True)
-def pytest_runtest_protocol(item):
+def pytest_runtest_protocol(item: pytest.Item):
     if item.config.pluginmanager.get_plugin("forked"):
         return
     if item.config.pluginmanager.get_plugin("timeout"):
         return
-    isolate, timeout, mem_limit, cpu_limit = get_isolation_options(item)
+    isolate, timeout, mem_limit, cpu_limit, resource_reqs = get_isolation_options(item)
+
+    wait_delta = float(item.config.getini("wait_delta"))  # type: ignore
 
     if (
         isolate is not None
         or mem_limit is not None
         or timeout is not None
         or cpu_limit is not None
+        or resource_reqs
     ):
         ihook = item.ihook
         ihook.pytest_runtest_logstart(nodeid=item.nodeid, location=item.location)
-        reports = run_in_subprocess(item, timeout, mem_limit, cpu_limit)
+        reports = run_in_subprocess(item, timeout, mem_limit, cpu_limit, wait_delta)
 
         for rep in reports:
             ihook.pytest_runtest_logreport(report=rep)
@@ -568,11 +710,19 @@ def get_isolation_options(item):
         pass
     else:
         isolate = forked
+
     timeout = get_timeout(item)
     mem_limit = get_memory_limit(item)
     cpu_limit = get_cpu_limit(item)
 
-    return isolate, timeout, mem_limit, cpu_limit
+    # Get resources from isolate(resources={'gpu': 2}) syntax
+    resource_dict = get_resource_dict(item)
+
+    # If resources are requested, ensure isolation
+    if resource_dict and not isolate:
+        isolate = True
+
+    return isolate, timeout, mem_limit, cpu_limit, resource_dict
 
 
 @pytest.hookimpl(tryfirst=True)
@@ -628,7 +778,8 @@ def pytest_terminal_summary(terminalreporter, exitstatus, config) -> None:
         filename = config.getoption("timeline_file")
         os.makedirs(os.path.dirname(os.path.abspath(filename)), exist_ok=True)
         json.dump(events, open(filename, "w"))
-
+        filename= filename.replace('.json', '.resources.json')
+        json.dump(get_resource_events(), open(filename, "w"))
 
 @pytest.hookimpl(tryfirst=False, hookwrapper=True)
 def pytest_runtest_call(item: pytest.Item):
@@ -638,22 +789,3 @@ def pytest_runtest_call(item: pytest.Item):
     )
 
 
-def create_event(worker_name, test_name, category, start_time, end_time, **kwargs):
-    start_event = {
-        "name": test_name.rsplit("/")[-1],
-        "cat": category,
-        "ph": "B",
-        "pid": worker_name,
-        "tid": 0,
-        "ts": start_time * (1000**2),
-        "args": kwargs,
-    }
-    end_event = {
-        "name": test_name.rsplit("/")[-1],
-        "cat": "pipeline",
-        "ph": "E",
-        "pid": worker_name,
-        "tid": 0,
-        "ts": end_time * (1000**2),
-    }
-    return start_event, end_event

pytest_isolate-0.0.9/src/pytest_isolate/resource_management.py

@@ -0,0 +1,229 @@
+import json
+import os
+import tempfile
+import time
+import warnings
+from contextlib import contextmanager
+from copy import deepcopy
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Callable, List, Optional
+
+import filelock
+
+from pytest_isolate.tracing import create_event
+
+# Constants
+
+DEFAULT_STATE_FILE = Path(tempfile.gettempdir()) / "pytest_isolate_resources.json"
+DEFAULT_EVENTS_FILE = Path(str(DEFAULT_STATE_FILE).replace(".json", "_events.json"))
+DEFAULT_WAIT_TIMEOUT = 300  # 5 minutes
+POLL_INTERVAL = 0.1  # 1 second
+
+
+@contextmanager
+def lock_resource_file():
+    with filelock.FileLock(str(DEFAULT_STATE_FILE) + ".lock", timeout=4) as lock:
+        yield lock
+
+
+@dataclass
+class Resource:
+    env_variable: str
+    available: List[int]
+    allocated: dict[str, List[int]]
+
+    def allocate(self, test_id: str, count: int) -> List[int]:
+        """Allocate resources for a test."""
+        if len(self.available) < count:
+            return []
+
+        allocated = self.available[:count]
+        self.allocated[test_id] = allocated
+        self.available = self.available[count:]
+        return allocated
+
+    def release(self, test_id: str) -> None:
+        """Release resources for a test."""
+        if test_id in self.allocated:
+            released = self.allocated.pop(test_id)
+            self.available.extend(released)
+        else:
+            raise ValueError(f"Test ID {test_id} not found in allocated resources.")
+
+
+@dataclass
+class StateData:
+    resources: dict[str, Resource]
+
+    @classmethod
+    def get_instance(cls) -> "StateData":
+        if not DEFAULT_STATE_FILE.exists():
+            data = {}
+        try:
+            with open(DEFAULT_STATE_FILE, "r") as f:
+                data = json.load(f)
+        except (json.JSONDecodeError, IOError):
+            data = {}
+        return StateData(
+            resources={name: Resource(**res) for name, res in data.items()}
+        )
+
+    def save(self) -> None:
+        with open(DEFAULT_STATE_FILE, "w") as f:
+            json.dump({name: asdict(res) for name, res in self.resources.items()}, f)
+
+
+def clean_resources() -> None:
+    """Clean up resources by removing lock and state files."""
+    if DEFAULT_STATE_FILE.exists():
+        os.remove(DEFAULT_STATE_FILE)
+    if DEFAULT_EVENTS_FILE.exists():
+        os.remove(DEFAULT_EVENTS_FILE)
+    if DEFAULT_STATE_FILE.with_suffix(".lock").exists():
+        os.remove(DEFAULT_STATE_FILE.with_suffix(".lock"))
+
+
+def register_resource_provider(
+    resource_type: str, env_variable: str, provider_func: Callable[[], List[int]]
+) -> None:
+    with lock_resource_file():
+        existing = provider_func()
+        state = StateData.get_instance()
+        if resource_type in state.resources:
+            warnings.warn(
+                f"Resource type {resource_type} already registered. Overwriting."
+            )
+        state.resources[resource_type] = Resource(
+            env_variable=env_variable,
+            available=existing,
+            allocated={},
+        )
+        state.save()
+
+
+def parse_resource_list(env_value: Optional[str]) -> List[int]:
+    """Parse a comma-separated list of resource IDs from an environment variable."""
+    if not env_value:
+        return []
+    try:
+        vals = [int(x.strip()) for x in env_value.split(",") if x.strip()]
+        vals = [x for x in vals if x >= 0]
+        return vals
+    except ValueError:
+        return []
+
+
+def log_resource_allocation(
+    allocated,
+    resource_type: str,
+    test_id: str,
+    start_time=None,
+    end_time=None,
+    file=None,
+) -> None:
+    file = str(file or DEFAULT_EVENTS_FILE)
+    folder = os.path.dirname(file)
+    if folder and not os.path.exists(folder):
+        os.makedirs(folder, exist_ok=True)
+    if not os.path.exists(file):
+        resource_events = {}
+    else:
+        with open(file, "r") as f:
+            resource_events = json.load(f)
+    if test_id not in resource_events:
+        resource_events[test_id] = dict(
+            allocated=allocated,
+            worker_name=resource_type,
+            test_name=test_id,
+            category="resource_allocation",
+        )
+    if start_time is not None:
+        resource_events[test_id]["start_time"] = start_time
+    if end_time is not None:
+        resource_events[test_id]["end_time"] = end_time
+    with open(file, "w") as f:
+        json.dump(resource_events, f)
+
+
+def setup_resource_environment(
+    test_id: str,
+    count: int,
+    resource_type: str,
+    wait_timeout: Optional[float] = None,
+) -> Optional[List[int]]:
+    def try_allocate() -> Optional[List[int]]:
+        with lock_resource_file():
+            state = StateData.get_instance()
+            if resource_type not in state.resources:
+                raise ValueError(f"Resource type {resource_type} not registered.")
+
+            resource = state.resources[resource_type]
+            allocated = resource.allocate(test_id, count)
+            if allocated:
+                os.environ[resource.env_variable] = ",".join(map(str, allocated))
+                state.save()
+                log_resource_allocation(
+                    resource.allocated.get(test_id),
+                    resource_type,
+                    test_id,
+                    start_time=time.time(),
+                )
+                return allocated
+        return []
+
+    timeout = wait_timeout or DEFAULT_WAIT_TIMEOUT
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        allocated = try_allocate()
+        if allocated:
+            return allocated
+        time.sleep(POLL_INTERVAL)
+    raise TimeoutError(f"Timeout while waiting for resources of type {resource_type}.")
+
+
+def cleanup_resource_environment(test_id: str, resource_type: str) -> None:
+    with lock_resource_file():
+        state = StateData.get_instance()
+        if resource_type not in state.resources:
+            raise ValueError(f"Resource type {resource_type} not registered.")
+
+        resource = state.resources[resource_type]
+        allocated = resource.allocated.get(test_id)
+        if allocated:
+            log_resource_allocation(
+                resource.allocated.get(test_id),
+                resource_type,
+                test_id,
+                end_time=time.time(),
+            )
+        resource.release(test_id)
+
+        os.environ.pop(resource.env_variable, None)
+        state.save()
+
+
+def get_resource_events(file=None) -> list[dict]:
+    # Convert the resource events to the format used in tracing.py,
+    # by merging the start and end events.
+    file = Path(file or str(DEFAULT_EVENTS_FILE))
+    if not file.exists():
+        return []
+    resource_events = json.load(open(file))
+    events = []
+    for test_id, event in resource_events.items():
+        event_args = dict(**event)
+        allocated = event_args.pop("allocated", None)
+        for resource_id in allocated:
+            resource_event = deepcopy(event_args)
+            resource_event["worker_name"] = (
+                f"{resource_event['worker_name']}_{resource_id}"
+            )
+            if "end_time" not in resource_event:
+                resource_event["end_time"] = time.time()
+            if "start_time" not in resource_event:
+                raise ValueError(
+                    f"Missing start_time for resource event {resource_event}"
+                )
+            events.extend(create_event(**resource_event))
+    return events

pytest_isolate-0.0.9/src/pytest_isolate/tracing.py

@@ -0,0 +1,19 @@
+def create_event(worker_name, test_name, category, start_time, end_time, **kwargs):
+    start_event = {
+        "name": test_name.rsplit("/")[-1],
+        "cat": category,
+        "ph": "B",
+        "pid": worker_name,
+        "tid": 0,
+        "ts": start_time * (1000**2),
+        "args": kwargs,
+    }
+    end_event = {
+        "name": test_name.rsplit("/")[-1],
+        "cat": "pipeline",
+        "ph": "E",
+        "pid": worker_name,
+        "tid": 0,
+        "ts": end_time * (1000**2),
+    }
+    return start_event, end_event

pytest_isolate-0.0.9/tests/test_resource_management.py

@@ -0,0 +1,90 @@
+import json
+import os
+import random
+import tempfile
+from importlib import resources
+from pathlib import Path
+from time import sleep
+from unittest import mock
+
+import filelock
+import pytest
+
+from pytest_isolate.plugin import allocate_resources, get_resource_events
+from pytest_isolate.resource_management import (
+    log_resource_allocation,
+    parse_resource_list,
+    register_resource_provider,
+)
+
+
+@pytest.fixture
+def temp_resource_files():
+    """Create temporary files for resource manager testing"""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        lock_file = Path(tmpdir) / "resources.lock"
+        state_file = Path(tmpdir) / "resources.json"
+        yield lock_file, state_file
+
+
+@pytest.fixture
+def mock_pynvml():
+    """Mock pynvml module for testing"""
+    pynvml_mock = mock.MagicMock()
+    pynvml_mock.nvmlDeviceGetCount.return_value = 4
+
+    # Mock the import so it returns our mock when imported in plugin.py
+    with mock.patch.dict("sys.modules", {"pynvml": pynvml_mock}):
+        yield pynvml_mock
+
+
+def test_parse_resource_list():
+    """Test parsing resources from environment variables."""
+    # Test empty input
+    assert parse_resource_list(None) == []
+    assert parse_resource_list("") == []
+
+    # Test simple cases
+    assert parse_resource_list("0") == [0]
+    assert parse_resource_list("0,1,2") == [0, 1, 2]
+
+    # Test with spaces
+    assert parse_resource_list(" 0, 1, 2 ") == [0, 1, 2]
+
+    # Test with invalid values - no warning in new implementation
+    assert parse_resource_list("0,a,2") == []
+
+def test_log_resource_allocation():
+    log_resource_allocation([0,1],'gpu','foo', start_time=100,file="timeline.events.test.json")
+    
+    log_resource_allocation([0,1],'gpu','foo', end_time=200,file="timeline.events.test.json")
+    events = get_resource_events(file="timeline.events.test.json")
+    open("timeline.resources.test.json", "w").write(json.dumps(events))
+    
+
+def test_allocate_resources():
+    with allocate_resources({"gpu": 1}, "foo", 100) as allocated:
+        assert len(allocated["gpu"]) == 1
+        assert len(os.environ.get("CUDA_VISIBLE_DEVICES","").split(","))
+
+
+
+@pytest.mark.isolate(resources={"gpu": 2}, timeout=10)
+@pytest.mark.parametrize("dummy", [1, 2])
+def test_gpu_marker_2(dummy):
+    sleep(0.3)
+    assert len(os.environ.get("CUDA_VISIBLE_DEVICES", "").split(",")) == 2
+
+
+@pytest.mark.isolate(resources={"gpu": 1}, timeout=10)
+@pytest.mark.parametrize("dummy", [1, 2])
+def test_gpu_marker_1(dummy):
+    sleep(0.3)
+    assert len(os.environ.get("CUDA_VISIBLE_DEVICES", "").split(",")) == 1
+
+
+@pytest.mark.xfail
+@pytest.mark.isolate(resources={"gpu": 7}, timeout=4)
+def test_gpu_marker_7():
+    sleep(0.3)
+    assert len(os.environ.get("CUDA_VISIBLE_DEVICES", "").split(",")) == 7

{pytest_isolate-0.0.4 → pytest_isolate-0.0.9}/tox.ini

@@ -11,7 +11,9 @@ deps =
     numpy
     pytest-xdist
     pytest-sugar
-
+    pytest-mock
+setenv = 
+    CUDA_VISIBLE_DEVICES = 0,1
 commands = pytest --isolate  --import-mode importlib {posargs:tests} --timeline 
 
 [testenv:lint]

pytest_isolate-0.0.4/src/pytest_isolate/__about__.py

@@ -1 +0,0 @@
-__version__ = "0.0.4"