diff --git a/dpctl/tests/test_utils.py b/dpctl/tests/test_utils.py index 5960d3fbe5..df4a9f503f 100644 --- a/dpctl/tests/test_utils.py +++ b/dpctl/tests/test_utils.py @@ -111,3 +111,14 @@ 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) + + +@pytest.mark.filterwarnings("ignore:.*:RuntimeWarning") +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 diff --git a/dpctl/utils/__init__.py b/dpctl/utils/__init__.py index 5601290af6..7589f9de9f 100644 --- a/dpctl/utils/__init__.py +++ b/dpctl/utils/__init__.py @@ -23,9 +23,11 @@ get_execution_queue, validate_usm_type, ) +from ._onetrace_context import onetrace_enabled __all__ = [ "get_execution_queue", "get_coerced_usm_type", "validate_usm_type", + "onetrace_enabled", ] diff --git a/dpctl/utils/_onetrace_context.py b/dpctl/utils/_onetrace_context.py new file mode 100644 index 0000000000..b7e68e69bf --- /dev/null +++ b/dpctl/utils/_onetrace_context.py @@ -0,0 +1,76 @@ +# Data Parallel Control (dpctl) +# +# Copyright 2020-2022 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from contextlib import contextmanager +from os import environ, getenv +from platform import system as sys_platform + +_UNCHECKED = sys_platform() == "Linux" +del sys_platform + + +@contextmanager +def onetrace_enabled(): + """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: + Launch the Python interpreter using `onetrace` tool: :: + + $ onetrace --conditional-collection -v -t --demangle python app.py + + 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 + + # onetrace output reporting on execution of the kernel + # should be seen, starting with "Device Timeline" + with onetrace_enabled(): + dpt.arange(100, dtype='int16') + + """ + global _UNCHECKED + + if _UNCHECKED: + _UNCHECKED = False + if not ( + getenv("PTI_ENABLE", None) == "1" + and "onetrace_tool" in getenv("LD_PRELOAD", "") + ): + import warnings + + warnings.warn( + "It looks like Python interpreter was not started using " + "`onetrace` utility. Using `onetrace_enabled` may have " + "no effect. See `onetrace_enabled.__doc__` for usage.", + RuntimeWarning, + stacklevel=2, + ) + + _env_var_name = "PTI_ENABLE_COLLECTION" + saved = getenv(_env_var_name, None) + try: + environ[_env_var_name] = "1" + yield + finally: + if saved is None: + del environ[_env_var_name] + else: + environ[_env_var_name] = saved