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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytest-isolate
-Version: 0.0.8
+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

pytest_isolate-0.0.9/src/pytest_isolate/__about__.py

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

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

@@ -40,9 +40,11 @@ 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,
@@ -776,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):
@@ -786,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.8 → pytest_isolate-0.0.9}/src/pytest_isolate/resource_management.py

@@ -4,22 +4,26 @@ 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:
+    with filelock.FileLock(str(DEFAULT_STATE_FILE) + ".lock", timeout=4) as lock:
         yield lock
 
 
@@ -74,6 +78,10 @@ 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(
@@ -99,13 +107,45 @@ def parse_resource_list(env_value: Optional[str]) -> List[int]:
     if not env_value:
         return []
     try:
-        vals= [int(x.strip()) for x in env_value.split(",") if x.strip()]
+        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,
@@ -123,6 +163,12 @@ def setup_resource_environment(
             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 []
 
@@ -143,6 +189,41 @@ def cleanup_resource_environment(test_id: str, resource_type: str) -> None:
             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.8 → pytest_isolate-0.0.9}/tests/test_resource_management.py

@@ -1,3 +1,4 @@
+import json
 import os
 import random
 import tempfile
@@ -9,8 +10,9 @@ from unittest import mock
 import filelock
 import pytest
 
-from pytest_isolate.plugin import allocate_resources
+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,
 )
@@ -52,23 +54,30 @@ def test_parse_resource_list():
     # 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 allocated["gpu"] == [0]
-        assert os.environ.get("CUDA_VISIBLE_DEVICES") == "0"
+        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)
+@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)
+@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_isolate-0.0.8/src/pytest_isolate/__about__.py

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