Skip to content

Commit 5993fb1

Browse files
Added test for onetrace_enabled() context manager. Fixed docstring example
1 parent f216df6 commit 5993fb1

File tree

2 files changed

+14
-2
lines changed

2 files changed

+14
-2
lines changed

dpctl/tests/test_utils.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,3 +111,13 @@ def validate_usm_type_arg():
111111
dpctl.utils.validate_usm_type("inv", allow_none=True)
112112
with pytest.raises(ValueError):
113113
dpctl.utils.validate_usm_type("inv", allow_none=False)
114+
115+
116+
def test_onetrace_enabled():
117+
import os
118+
119+
v_name = "PTI_ENABLE_COLLECTION"
120+
v_v = os.getenv(v_name, None)
121+
with dpctl.utils.onetrace_enabled():
122+
assert os.getenv(v_name, None) == "1"
123+
assert os.getenv(v_name, None) == v_v

dpctl/utils/_onetrace_context.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,12 +25,14 @@ def onetrace_enabled():
2525
N.B.: Proper working of this utility assumes that Python interpreter
2626
has been launched by `onetrace` tool from intel/pti-gpu project.
2727
28-
.. Example:: bash
28+
:Example:
29+
Launch the Python interpreter using `onetrace` tool: ::
2930
3031
$ onetrace --conditional-collection -v -t --demangle \
3132
python script.py
3233
33-
.. Example:: bash
34+
Now using the context manager in the Python sessions enables
35+
data collection and its output for every offloaded kernel ::
3436
3537
import dpctl.tensor as dpt
3638
from dpctl.utils import onetrace_enabled

0 commit comments

Comments
 (0)