Added test for onetrace_enabled() context manager. Fixed docstring example

Author: oleksandr-pavlyk

dpctl/tests/test_utils.py

@@ -111,3 +111,13 @@ def validate_usm_type_arg():
         dpctl.utils.validate_usm_type("inv", allow_none=True)
     with pytest.raises(ValueError):
         dpctl.utils.validate_usm_type("inv", allow_none=False)
+
+
+def test_onetrace_enabled():
+    import os
+
+    v_name = "PTI_ENABLE_COLLECTION"
+    v_v = os.getenv(v_name, None)
+    with dpctl.utils.onetrace_enabled():
+        assert os.getenv(v_name, None) == "1"
+    assert os.getenv(v_name, None) == v_v

dpctl/utils/_onetrace_context.py

@@ -20,17 +20,18 @@
 
 @contextmanager
 def onetrace_enabled():
-    """ Enable `onetrace` collection for kernels executed in this context.
+    """Enable `onetrace` collection for kernels executed in this context.
 
     N.B.: Proper working of this utility assumes that Python interpreter
     has been launched by `onetrace` tool from intel/pti-gpu project.
 
-        .. Example:: bash
+    :Example:
+        Launch the Python interpreter using `onetrace` tool: ::
 
-            $ onetrace --conditional-collection -v -t --demangle \
-                python script.py
+            $ onetrace --conditional-collection -v -t --demangle python app.py
 
-        .. Example:: bash
+        Now using the context manager in the Python sessions enables
+        data collection and its output for every offloaded kernel ::
 
             import dpctl.tensor as dpt
             from dpctl.utils import onetrace_enabled