From 077a844cb644828c57a708fe9aa4a9c98f74417d Mon Sep 17 00:00:00 2001 From: "Jonathan R. Madsen" Date: Thu, 22 Jun 2023 17:14:40 -0500 Subject: [PATCH] PyTorch Python fork fix part 2 (#292) PyTorch Python fork fix part 2 - store script file in environment for robustness against restart after fork [ROCm/rocprofiler-systems commit: 6c9b66d938647434f5c5fc67351f870ea1d0c365] --- .../rocprofiler-systems/source/python/omnitrace/__main__.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/projects/rocprofiler-systems/source/python/omnitrace/__main__.py b/projects/rocprofiler-systems/source/python/omnitrace/__main__.py index 71427d157e..a4f43ec4ef 100644 --- a/projects/rocprofiler-systems/source/python/omnitrace/__main__.py +++ b/projects/rocprofiler-systems/source/python/omnitrace/__main__.py @@ -358,6 +358,7 @@ def main(main_args=sys.argv): sys.path.insert(0, os.path.dirname(script_file)) _OMNITRACE_PYTHON_SCRIPT_FILE = script_file + os.environ["OMNITRACE_PYTHON_SCRIPT_FILE"] = script_file prof = Profiler() fake = FakeProfiler() @@ -397,6 +398,11 @@ def main(main_args=sys.argv): if __name__ == "__main__": args = sys.argv + if _OMNITRACE_PYTHON_SCRIPT_FILE is None: + _OMNITRACE_PYTHON_SCRIPT_FILE = os.environ.get( + "OMNITRACE_PYTHON_SCRIPT_FILE", None + ) + if "--" not in args and _OMNITRACE_PYTHON_SCRIPT_FILE is not None: args = [args[0]] + ["--", _OMNITRACE_PYTHON_SCRIPT_FILE] + args[1:] os.environ["OMNITRACE_USE_PID"] = "ON"