From 6c9b66d938647434f5c5fc67351f870ea1d0c365 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 --- source/python/omnitrace/__main__.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/source/python/omnitrace/__main__.py b/source/python/omnitrace/__main__.py index 71427d157..a4f43ec4e 100644 --- a/source/python/omnitrace/__main__.py +++ b/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"